Обновить vector_store.py

vector_store.py

@@ -1,42 +1,26 @@
-from langchain_qdrant import QdrantVectorStore
+from langchain_community.vectorstores import Chroma
 from langchain_ollama import OllamaEmbeddings
-from langchain_core.documents import Document
-from qdrant_client import QdrantClient
-from qdrant_client.models import Distance, VectorParams
+from langchain.schema import Document
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-import uuid
 
-COLLECTION_NAME = "knowledge_base"
+COLLECTION_NAME = "rag_collection"
 EMBEDDING_MODEL = "nomic-embed-text"
-VECTOR_SIZE = 768
+CHROMA_PERSIST_DIR = "./chroma_db"
 
 
-def get_embeddings():
-    return OllamaEmbeddings(model=EMBEDDING_MODEL)
-
-
-def get_qdrant_client():
-    return QdrantClient(host="localhost", port=6333)
-
-
-def init_collection(client: QdrantClient):
-    collections = [c.name for c in client.get_collections().collections]
-    if COLLECTION_NAME not in collections:
-        client.create_collection(
-            collection_name=COLLECTION_NAME,
-            vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
+class ChromaStore:
+    def __init__(self, collection_name: str = COLLECTION_NAME):
+        self.client = Chroma(
+            embedding_function=OllamaEmbeddings(model=EMBEDDING_MODEL),
+            collection_name=collection_name,
+            persist_directory=CHROMA_PERSIST_DIR,
         )
 
+    def add_documents(self, docs: list[Document]) -> None:
+        self.client.add_documents(docs)
 
-def get_vector_store() -> QdrantVectorStore:
-    client = get_qdrant_client()
-    init_collection(client)
-    embeddings = get_embeddings()
-    return QdrantVectorStore(
-        client=client,
-        collection_name=COLLECTION_NAME,
-        embedding=embeddings,
-    )
+    def search(self, query: str, limit: int = 5) -> list[Document]:
+        return self.client.similarity_search(query, k=limit)
 
 
 def add_documents(content: str, title: str) -> int:
@@ -53,21 +37,20 @@ def add_documents(content: str, title: str) -> int:
         )
         for i, chunk in enumerate(chunks)
     ]
-    store = get_vector_store()
+    store = ChromaStore()
     store.add_documents(docs)
     return len(docs)
 
 
 def search_documents(query: str, max_results: int = 5) -> list[dict]:
-    store = get_vector_store()
-    results = store.similarity_search_with_relevance_scores(query, k=max_results)
+    store = ChromaStore()
+    results = store.search(query, limit=max_results)
     output = []
-    for doc, score in results:
+    for doc in results:
         output.append(
             {
                 "content": doc.page_content,
                 "metadata": doc.metadata,
-                "score": round(score, 4),
             }
         )
     return output