revert 1748a39f01

revert Обновить vector_store.py

vector_store.py

@@ -1,26 +1,42 @@
-from langchain_community.vectorstores import Chroma
+from langchain_qdrant import QdrantVectorStore
 from langchain_ollama import OllamaEmbeddings
-from langchain.schema import Document
+from langchain_core.documents import Document
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+import uuid
 
-COLLECTION_NAME = "rag_collection"
+COLLECTION_NAME = "knowledge_base"
 EMBEDDING_MODEL = "nomic-embed-text"
-CHROMA_PERSIST_DIR = "./chroma_db"
+VECTOR_SIZE = 768
 
 
-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 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),
         )
 
-    def add_documents(self, docs: list[Document]) -> None:
-        self.client.add_documents(docs)
 
-    def search(self, query: str, limit: int = 5) -> list[Document]:
-        return self.client.similarity_search(query, k=limit)
+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 add_documents(content: str, title: str) -> int:
@@ -37,20 +53,21 @@ def add_documents(content: str, title: str) -> int:
         )
         for i, chunk in enumerate(chunks)
     ]
-    store = ChromaStore()
+    store = get_vector_store()
     store.add_documents(docs)
     return len(docs)
 
 
 def search_documents(query: str, max_results: int = 5) -> list[dict]:
-    store = ChromaStore()
-    results = store.search(query, limit=max_results)
+    store = get_vector_store()
+    results = store.similarity_search_with_relevance_scores(query, k=max_results)
     output = []
-    for doc in results:
+    for doc, score in results:
         output.append(
             {
                 "content": doc.page_content,
                 "metadata": doc.metadata,
+                "score": round(score, 4),
             }
         )
     return output