This commit is contained in:
+5
-6
@@ -5,7 +5,7 @@ from typing import Annotated, Callable
|
|||||||
import typer
|
import typer
|
||||||
from chromadb.errors import InternalError, NotFoundError
|
from chromadb.errors import InternalError, NotFoundError
|
||||||
|
|
||||||
from chromy.handlers.add_data import handle_add_data
|
from chromy.handlers.import_data import handle_import
|
||||||
from chromy.handlers.count_collection import handle_count_collection
|
from chromy.handlers.count_collection import handle_count_collection
|
||||||
from chromy.handlers.create_collection import handle_create_collection
|
from chromy.handlers.create_collection import handle_create_collection
|
||||||
from chromy.handlers.delete_collection import (
|
from chromy.handlers.delete_collection import (
|
||||||
@@ -88,22 +88,21 @@ def count(
|
|||||||
|
|
||||||
|
|
||||||
@app.command(
|
@app.command(
|
||||||
"add-data",
|
"import",
|
||||||
help="Chunk, embed, and add a file to a collection in the local Chroma database.",
|
help="Chunk, embed, and add a file to a collection in the local Chroma database.",
|
||||||
)
|
)
|
||||||
def add_data(
|
def import_data(
|
||||||
collection: Annotated[
|
collection: Annotated[
|
||||||
str,
|
str,
|
||||||
typer.Argument(help="Name of the target collection."),
|
typer.Argument(help="Name of the target collection."),
|
||||||
],
|
],
|
||||||
file: Annotated[
|
file: Annotated[
|
||||||
str,
|
str,
|
||||||
typer.Argument(
|
typer.Argument(help="Path to the file to chunk and add to the collection."),
|
||||||
help="Path to the file to chunk and add to the collection."),
|
|
||||||
],
|
],
|
||||||
) -> None:
|
) -> None:
|
||||||
try:
|
try:
|
||||||
_run(lambda: handle_add_data(collection, file))
|
_run(lambda: handle_import(collection, file))
|
||||||
except NotFoundError:
|
except NotFoundError:
|
||||||
_fail(f"Collection '{collection}' does not exist.")
|
_fail(f"Collection '{collection}' does not exist.")
|
||||||
except FileNotFoundError:
|
except FileNotFoundError:
|
||||||
|
|||||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
|||||||
from chromy.utilities import ingest_file
|
from chromy.utilities import ingest_file
|
||||||
|
|
||||||
|
|
||||||
def handle_add_data(collection: str, file: str) -> int:
|
def handle_import(collection: str, file: str) -> int:
|
||||||
records_added = ingest_file(collection, file)
|
records_added = ingest_file(collection, file)
|
||||||
print(f"Added {records_added} records to collection '{collection}'.")
|
print(f"Added {records_added} records to collection '{collection}'.")
|
||||||
return 0
|
return 0
|
||||||
+3
-3
@@ -53,12 +53,12 @@ class CliTests(unittest.TestCase):
|
|||||||
self.assertEqual(result.exit_code, 0)
|
self.assertEqual(result.exit_code, 0)
|
||||||
self.assertEqual(result.stdout, "7\n")
|
self.assertEqual(result.stdout, "7\n")
|
||||||
|
|
||||||
def test_add_data(self) -> None:
|
def test_import_data(self) -> None:
|
||||||
with patch(
|
with patch(
|
||||||
"chromy.handlers.add_data.ingest_file",
|
"chromy.handlers.import_data.ingest_file",
|
||||||
return_value=3,
|
return_value=3,
|
||||||
) as ingest_file:
|
) as ingest_file:
|
||||||
result = _invoke(["add-data", "notes", "romeo_and_juliet.txt"])
|
result = _invoke(["import", "notes", "romeo_and_juliet.txt"])
|
||||||
|
|
||||||
ingest_file.assert_called_once_with("notes", "romeo_and_juliet.txt")
|
ingest_file.assert_called_once_with("notes", "romeo_and_juliet.txt")
|
||||||
self.assertEqual(result.exit_code, 0)
|
self.assertEqual(result.exit_code, 0)
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ from contextlib import redirect_stdout
|
|||||||
from typing import TypeVar
|
from typing import TypeVar
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
from chromy.handlers.add_data import handle_add_data
|
from chromy.handlers.import_data import handle_import
|
||||||
from chromy.handlers.count_collection import handle_count_collection
|
from chromy.handlers.count_collection import handle_count_collection
|
||||||
from chromy.handlers.create_collection import handle_create_collection
|
from chromy.handlers.create_collection import handle_create_collection
|
||||||
from chromy.handlers.delete_collection import (
|
from chromy.handlers.delete_collection import (
|
||||||
@@ -83,13 +83,13 @@ class HandlerTests(unittest.TestCase):
|
|||||||
self.assertEqual(exit_code, 0)
|
self.assertEqual(exit_code, 0)
|
||||||
self.assertEqual(output, "7\n")
|
self.assertEqual(output, "7\n")
|
||||||
|
|
||||||
def test_add_data_uses_typed_input(self) -> None:
|
def test_import_data_uses_typed_input(self) -> None:
|
||||||
with patch(
|
with patch(
|
||||||
"chromy.handlers.add_data.ingest_file",
|
"chromy.handlers.import_data.ingest_file",
|
||||||
return_value=3,
|
return_value=3,
|
||||||
) as ingest_file:
|
) as ingest_file:
|
||||||
exit_code, output = _capture_output(
|
exit_code, output = _capture_output(
|
||||||
handle_add_data,
|
handle_import,
|
||||||
"notes",
|
"notes",
|
||||||
"romeo_and_juliet.txt",
|
"romeo_and_juliet.txt",
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user