add query

chroma_functions.py

@@ -2,6 +2,7 @@ from typing import List
 from uuid import uuid4
 
 import chromadb
+from chromadb.api.types import QueryResult
 from chromadb.errors import NotFoundError
 
 from embed import EmbeddingRecord
@@ -58,5 +59,21 @@ def add_data(collection: str, data: List[EmbeddingRecord]) -> None:
     )
 
 
-def query_data(collection_name: str, texts: list[str]):
+def query_data(collection_name: str, texts: list[str]) -> QueryResult:
-    raise NotImplementedError()
+    if not texts:
+        return {
+            "ids": [],
+            "documents": [],
+            "metadatas": [],
+            "distances": [],
+            "embeddings": [],
+        }
+
+    client = chromadb.PersistentClient()
+
+    try:
+        collection = client.get_collection(name=collection_name)
+    except NotFoundError:
+        raise
+
+    return collection.query(query_texts=texts)

main.py

@@ -9,6 +9,7 @@ from chroma_functions import (
     create_collection,
     delete_collection,
     list_collections,
+    query_data,
 )
 from chunk_functions import chunk_file
 from cli_parser import build_parser
@@ -80,6 +81,17 @@ def main() -> int:
 
         return 0
 
+    if args.command in {"query", "q"}:
+        try:
+            result = query_data(args.collection, [args.texts])
+        except NotFoundError:
+            print(f"Collection '{args.collection}' does not exist.")
+            return 1
+
+        print(result)
+
+        return 0
+
     print("Nothing to do. Use -h to see available commands.")
 
     return 0