Compare commits
2 Commits
5a2844c7cb
...
87ecaec3f3
| Author | SHA1 | Date | |
|---|---|---|---|
| 87ecaec3f3 | |||
| 150cc57932 |
+37
-9
@@ -24,8 +24,17 @@ app = typer.Typer(
|
|||||||
"Storage location:\n"
|
"Storage location:\n"
|
||||||
"- By default, Chromy uses Chroma's default persistent location behavior.\n"
|
"- By default, Chromy uses Chroma's default persistent location behavior.\n"
|
||||||
f"- Set {CHROMA_FOLDER_ENV_VAR} to a parent directory to override it.\n"
|
f"- Set {CHROMA_FOLDER_ENV_VAR} to a parent directory to override it.\n"
|
||||||
f"- Chromy stores data in <{CHROMA_FOLDER_ENV_VAR}>/chroma."
|
f"- Chromy stores data in <{CHROMA_FOLDER_ENV_VAR}>/chroma.\n\n"
|
||||||
)
|
"Command aliases:\n"
|
||||||
|
"- list-collections: lc\n"
|
||||||
|
"- create-collection: cc\n"
|
||||||
|
"- delete-collection: dc\n"
|
||||||
|
"- count: c\n"
|
||||||
|
"- import: i\n"
|
||||||
|
"- query: q\n"
|
||||||
|
"- delete: del"
|
||||||
|
),
|
||||||
|
invoke_without_command=True,
|
||||||
)
|
)
|
||||||
|
|
||||||
ExitCodeHandler = Callable[[], int]
|
ExitCodeHandler = Callable[[], int]
|
||||||
@@ -46,12 +55,22 @@ def _fail(message: str) -> None:
|
|||||||
raise typer.Exit(1)
|
raise typer.Exit(1)
|
||||||
|
|
||||||
|
|
||||||
|
@app.callback()
|
||||||
|
def main(ctx: typer.Context) -> None:
|
||||||
|
"""Run the CLI and show help when no command is provided."""
|
||||||
|
if ctx.invoked_subcommand is None:
|
||||||
|
print("[bold red]Error[/]: Missing command.")
|
||||||
|
typer.echo(ctx.get_help())
|
||||||
|
raise typer.Exit(1)
|
||||||
|
|
||||||
|
|
||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
# LIST COLLECTIONS
|
# LIST COLLECTIONS
|
||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
|
@app.command("lc", help="Alias for list-collections.")
|
||||||
@app.command(
|
@app.command(
|
||||||
"list-collections",
|
"list-collections",
|
||||||
help="List all collections stored in the local Chroma database.",
|
help="List all collections stored in the local Chroma database. Alias: lc.",
|
||||||
)
|
)
|
||||||
def list_collections() -> None:
|
def list_collections() -> None:
|
||||||
_run(handle_list_collections)
|
_run(handle_list_collections)
|
||||||
@@ -60,9 +79,10 @@ def list_collections() -> None:
|
|||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
# CREATE A COLLECTION
|
# CREATE A COLLECTION
|
||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
|
@app.command("cc", help="Alias for create-collection.")
|
||||||
@app.command(
|
@app.command(
|
||||||
"create-collection",
|
"create-collection",
|
||||||
help="Create a collection in the local Chroma database.",
|
help="Create a collection in the local Chroma database. Alias: cc.",
|
||||||
)
|
)
|
||||||
def create_collection(
|
def create_collection(
|
||||||
collection: Annotated[
|
collection: Annotated[
|
||||||
@@ -79,9 +99,10 @@ def create_collection(
|
|||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
# DELETE A COLLECTION
|
# DELETE A COLLECTION
|
||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
|
@app.command("dc", help="Alias for delete-collection.")
|
||||||
@app.command(
|
@app.command(
|
||||||
"delete-collection",
|
"delete-collection",
|
||||||
help="Delete a collection from the local Chroma database.",
|
help="Delete a collection from the local Chroma database. Alias: dc.",
|
||||||
)
|
)
|
||||||
def delete_collection(
|
def delete_collection(
|
||||||
collection: Annotated[
|
collection: Annotated[
|
||||||
@@ -98,9 +119,10 @@ def delete_collection(
|
|||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
# COUNT RECORDS
|
# COUNT RECORDS
|
||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
|
@app.command("c", help="Alias for count.")
|
||||||
@app.command(
|
@app.command(
|
||||||
"count",
|
"count",
|
||||||
help="Count records in a collection from the local Chroma database.",
|
help="Count records in a collection from the local Chroma database. Alias: c.",
|
||||||
)
|
)
|
||||||
def count(
|
def count(
|
||||||
collection: Annotated[
|
collection: Annotated[
|
||||||
@@ -117,11 +139,12 @@ def count(
|
|||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
# IMPORT DATA
|
# IMPORT DATA
|
||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
|
@app.command("i", help="Alias for import.")
|
||||||
@app.command(
|
@app.command(
|
||||||
"import",
|
"import",
|
||||||
help=(
|
help=(
|
||||||
"Chunk, embed, and add one or more files to a collection in the "
|
"Chunk, embed, and add one or more files to a collection in the "
|
||||||
"local Chroma database."
|
"local Chroma database. Alias: i."
|
||||||
),
|
),
|
||||||
)
|
)
|
||||||
def import_data(
|
def import_data(
|
||||||
@@ -145,7 +168,8 @@ def import_data(
|
|||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
# QUERY
|
# QUERY
|
||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
@app.command("query", help="Query a collection with the provided text.")
|
@app.command("q", help="Alias for query.")
|
||||||
|
@app.command("query", help="Query a collection with the provided text. Alias: q.")
|
||||||
def query(
|
def query(
|
||||||
collection: Annotated[
|
collection: Annotated[
|
||||||
str,
|
str,
|
||||||
@@ -165,7 +189,11 @@ def query(
|
|||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
# DELETE DATA
|
# DELETE DATA
|
||||||
# ------------------------------------------------------------------------------
|
# ------------------------------------------------------------------------------
|
||||||
@app.command("delete", help="Delete records from a collection using a metadata filter.")
|
@app.command("del", help="Alias for delete.")
|
||||||
|
@app.command(
|
||||||
|
"delete",
|
||||||
|
help="Delete records from a collection using a metadata filter. Alias: del.",
|
||||||
|
)
|
||||||
def delete_records(
|
def delete_records(
|
||||||
collection: Annotated[
|
collection: Annotated[
|
||||||
str,
|
str,
|
||||||
|
|||||||
@@ -95,6 +95,106 @@ class CliTests(unittest.TestCase):
|
|||||||
"The 'notes' collection contains 7 records.\n",
|
"The 'notes' collection contains 7 records.\n",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_command_aliases(self) -> None:
|
||||||
|
with self.subTest(alias="lc"):
|
||||||
|
with patch(
|
||||||
|
"chromy.handlers.list_collections.list_collections",
|
||||||
|
return_value=[],
|
||||||
|
) as mocked:
|
||||||
|
result = _invoke(["lc"])
|
||||||
|
|
||||||
|
mocked.assert_called_once_with()
|
||||||
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
self.assertEqual(result.stdout, "No collections found.\n")
|
||||||
|
|
||||||
|
with self.subTest(alias="cc"):
|
||||||
|
with patch(
|
||||||
|
"chromy.handlers.create_collection.create_collection",
|
||||||
|
return_value="notes",
|
||||||
|
) as mocked:
|
||||||
|
result = _invoke(["cc", "notes"])
|
||||||
|
|
||||||
|
mocked.assert_called_once_with("notes")
|
||||||
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
self.assertEqual(result.stdout, "Created: collection 'notes'.\n")
|
||||||
|
|
||||||
|
with self.subTest(alias="dc"):
|
||||||
|
with patch(
|
||||||
|
"chromy.handlers.delete_collection.delete_collection",
|
||||||
|
) as mocked:
|
||||||
|
result = _invoke(["dc", "notes"])
|
||||||
|
|
||||||
|
mocked.assert_called_once_with("notes")
|
||||||
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
self.assertEqual(result.stdout, "Deleted collection 'notes'.\n")
|
||||||
|
|
||||||
|
with self.subTest(alias="c"):
|
||||||
|
with patch(
|
||||||
|
"chromy.handlers.count_collection.count_collection",
|
||||||
|
return_value=7,
|
||||||
|
) as mocked:
|
||||||
|
result = _invoke(["c", "notes"])
|
||||||
|
|
||||||
|
mocked.assert_called_once_with("notes")
|
||||||
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
self.assertEqual(
|
||||||
|
result.stdout,
|
||||||
|
"The 'notes' collection contains 7 records.\n",
|
||||||
|
)
|
||||||
|
|
||||||
|
with self.subTest(alias="i"):
|
||||||
|
with patch(
|
||||||
|
"chromy.handlers.import_data.ingest_file",
|
||||||
|
return_value=3,
|
||||||
|
) as mocked:
|
||||||
|
result = _invoke(["i", "notes", "romeo_and_juliet.txt"])
|
||||||
|
|
||||||
|
mocked.assert_called_once_with(
|
||||||
|
"notes",
|
||||||
|
self._fixture_path("romeo_and_juliet.txt"),
|
||||||
|
)
|
||||||
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
self.assertEqual(
|
||||||
|
result.stdout,
|
||||||
|
"Added 3 records from 'romeo_and_juliet.txt' to collection 'notes'.\n"
|
||||||
|
"Imported 1 file(s) successfully; 0 failed.\n",
|
||||||
|
)
|
||||||
|
|
||||||
|
with self.subTest(alias="q"):
|
||||||
|
query_result = {"ids": [["1"]], "documents": [["hello"]]}
|
||||||
|
with (
|
||||||
|
patch(
|
||||||
|
"chromy.handlers.query.run_query",
|
||||||
|
return_value=query_result,
|
||||||
|
) as mocked,
|
||||||
|
patch(
|
||||||
|
"chromy.handlers.query.format_query_result",
|
||||||
|
return_value=["1"],
|
||||||
|
),
|
||||||
|
):
|
||||||
|
result = _invoke(["q", "notes", "Where is Romeo?"])
|
||||||
|
|
||||||
|
mocked.assert_called_once_with("notes", "Where is Romeo?")
|
||||||
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
self.assertEqual(result.stdout, "1\n")
|
||||||
|
|
||||||
|
with self.subTest(alias="del"):
|
||||||
|
with patch(
|
||||||
|
"chromy.handlers.delete_collection.delete_data",
|
||||||
|
return_value=2,
|
||||||
|
) as mocked:
|
||||||
|
result = _invoke(
|
||||||
|
["del", "notes", "--where", "file_name=play.txt"],
|
||||||
|
)
|
||||||
|
|
||||||
|
mocked.assert_called_once_with("notes", {"file_name": "play.txt"})
|
||||||
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
self.assertEqual(
|
||||||
|
result.stdout,
|
||||||
|
"Deleted 2 record(s) from collection 'notes' where "
|
||||||
|
"file_name=play.txt.\n",
|
||||||
|
)
|
||||||
|
|
||||||
def test_import_data(self) -> None:
|
def test_import_data(self) -> None:
|
||||||
with patch(
|
with patch(
|
||||||
"chromy.handlers.import_data.ingest_file",
|
"chromy.handlers.import_data.ingest_file",
|
||||||
@@ -252,6 +352,14 @@ class CliTests(unittest.TestCase):
|
|||||||
self.assertNotEqual(result.exit_code, 0)
|
self.assertNotEqual(result.exit_code, 0)
|
||||||
self.assertIn("Missing option", result.output)
|
self.assertIn("Missing option", result.output)
|
||||||
|
|
||||||
|
def test_cli_without_arguments_prints_error_and_help(self) -> None:
|
||||||
|
result = _invoke([])
|
||||||
|
|
||||||
|
self.assertEqual(result.exit_code, 1)
|
||||||
|
self.assertIn("Error: Missing command.", result.stdout)
|
||||||
|
self.assertIn("Usage:", result.stdout)
|
||||||
|
self.assertIn("Commands", result.stdout)
|
||||||
|
|
||||||
def test_cli_help_documents_chroma_folder_env_var(self) -> None:
|
def test_cli_help_documents_chroma_folder_env_var(self) -> None:
|
||||||
result = _invoke(["--help"])
|
result = _invoke(["--help"])
|
||||||
|
|
||||||
@@ -260,6 +368,19 @@ class CliTests(unittest.TestCase):
|
|||||||
self.assertIn("parent directory", result.stdout)
|
self.assertIn("parent directory", result.stdout)
|
||||||
self.assertIn("<CHROMA_FOLDER>/chroma", result.stdout)
|
self.assertIn("<CHROMA_FOLDER>/chroma", result.stdout)
|
||||||
|
|
||||||
|
def test_cli_help_documents_command_aliases(self) -> None:
|
||||||
|
result = _invoke(["--help"])
|
||||||
|
|
||||||
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
self.assertIn("Command aliases", result.stdout)
|
||||||
|
self.assertIn("list-collections: lc", result.stdout)
|
||||||
|
self.assertIn("create-collection: cc", result.stdout)
|
||||||
|
self.assertIn("delete-collection: dc", result.stdout)
|
||||||
|
self.assertIn("count: c", result.stdout)
|
||||||
|
self.assertIn("import: i", result.stdout)
|
||||||
|
self.assertIn("query: q", result.stdout)
|
||||||
|
self.assertIn("delete: del", result.stdout)
|
||||||
|
|
||||||
def test_cli_surfaces_chroma_path_errors(self) -> None:
|
def test_cli_surfaces_chroma_path_errors(self) -> None:
|
||||||
with patch(
|
with patch(
|
||||||
"chromy.handlers.list_collections.list_collections",
|
"chromy.handlers.list_collections.list_collections",
|
||||||
|
|||||||
Reference in New Issue
Block a user