diff --git a/chroma_functions.py b/chroma_functions.py index 0807cf0..df9ec93 100644 --- a/chroma_functions.py +++ b/chroma_functions.py @@ -1,8 +1,9 @@ -import chromadb -from chromadb.errors import NotFoundError from typing import List from uuid import uuid4 +import chromadb +from chromadb.errors import NotFoundError + from embed import EmbeddingRecord @@ -55,3 +56,7 @@ def add_data(collection: str, data: List[EmbeddingRecord]) -> None: documents=[record["text"] for record in data], embeddings=[record["embedding"] for record in data], ) + + +def query_data(collection_name: str, texts: list[str]): + raise NotImplementedError() diff --git a/cli_parser.py b/cli_parser.py index 30d0177..f59b193 100644 --- a/cli_parser.py +++ b/cli_parser.py @@ -5,12 +5,14 @@ def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Inspect local Chroma collections.") subparsers = parser.add_subparsers(dest="command") + # List existing collections subparsers.add_parser( "list-collections", aliases=["lc"], help="List all collections stored in the local Chroma database.", ) + # Create a new collection create_parser = subparsers.add_parser( "create-collection", aliases=["cc"], @@ -18,6 +20,7 @@ def build_parser() -> argparse.ArgumentParser: ) create_parser.add_argument("name", help="Name of the collection to create.") + # Delete a collection delete_parser = subparsers.add_parser( "delete-collection", aliases=["dc"], @@ -25,6 +28,7 @@ def build_parser() -> argparse.ArgumentParser: ) delete_parser.add_argument("name", help="Name of the collection to delete.") + # Count documents in a collection count_parser = subparsers.add_parser( "count", aliases=["co"], @@ -32,12 +36,24 @@ def build_parser() -> argparse.ArgumentParser: ) count_parser.add_argument("name", help="Name of the collection to count.") + # Add documents to a collection add_parser = subparsers.add_parser( "add-data", aliases=["ad"], help="Chunk, embed, and add a file to a collection in the local Chroma database.", ) add_parser.add_argument("collection", help="Name of the target collection.") - add_parser.add_argument("file", help="Path to the file to chunk and add to the collection.") + add_parser.add_argument( + "file", help="Path to the file to chunk and add to the collection." + ) + + # Query doc + query_parser = subparsers.add_parser( + "query", + aliases=["q"], + help="Query a collection with given text/s.", + ) + query_parser.add_argument("collection", help="Name of the target collection.") + query_parser.add_argument("texts", help="The text/s to query.") return parser