update comments

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()

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