add file

qdrant_client.py

@@ -1,36 +1,22 @@
-from qdrant_client import QdrantClient
-from qdrant_client.http import models
+from chromadb import Client as ChromaClient
 from langchain_ollama import OllamaEmbeddings
 
-class QdrantStore:
-    def __init__(self, url="http://localhost:6333", collection_name="rag_collection"):
-        self.client = QdrantClient(url=url)
+class ChromaStore:
+    def __init__(self, collection_name="rag_collection"):
+        self.client = ChromaClient()
         self.collection_name = collection_name
         self._ensure_collection()
 
     def _ensure_collection(self):
-        if self.collection_name not in self.client.get_collections().collections:
-            self.client.recreate_collection(
-                collection_name=self.collection_name,
-                vectors_config=models.VectorParams(size=384, distance=models.Distance.COSINE),
-            )
+        if self.collection_name not in [c.name for c in self.client.get_collections()]:
+            self.client.create_collection(name=self.collection_name, metadata={})
 
     def add_documents(self, documents, titles):
         embeddings = OllamaEmbeddings(model="nomic-embed-text")
         vectors = embeddings.embed_documents(documents)
         payload = [{"title": t} for t in titles]
-        self.client.upsert(
-            collection_name=self.collection_name,
-            points=models.Batch(points=[models.PointStruct(id=i, vector=v, payload=p) for i, (v, p) in enumerate(zip(vectors, payload))])
-        )
+        self.client.upsert(collection_name=self.collection_name, documents=documents, ids=[str(i) for i in range(len(documents))], metadatas=payload)
 
     def search(self, query, limit=5):
-        embeddings = OllamaEmbeddings(model="nomic-embed-text")
-        query_vector = embeddings.embed_query(query)
-        results = self.client.search(
-            collection_name=self.collection_name,
-            query_vector=query_vector,
-            limit=limit,
-            with_payload=True,
-        )
-        return [r.payload for r in results]
+        results = self.client.query(collection_name=self.collection_name, query_text=query, n_results=limit, include_metadata=True)
+        return [r['metadata'] for r in results]