feat: solution for 'Агент с RAG-памятью'

src/vector_store.py

@@ -1,29 +1,42 @@
-from langchain_qdrant import Qdrant
-from langchain.schema import Document
-import config
+"""
+Vector store implementation using Qdrant.
+"""
 
-class QdrantVectorStore:
+from qdrant_client import QdrantClient
+from qdrant_client.http import models as qdrant_models
+from langchain.vectorstores import Qdrant
+from config import QDRANT_HOST, QDRANT_PORT, QDRANT_COLLECTION_NAME
+from embeddings import get_ollama_embeddings
+
+def get_qdrant_client() -> QdrantClient:
     """
-    Wrapper around langchain_qdrant.Qdrant to provide a simple interface
-    for adding documents and retrieving a retriever.
+    Creates a Qdrant client connected to the local Qdrant instance.
     """
-    def __init__(self, embeddings, collection_name: str = None):
-        self.collection_name = collection_name or config.QDRANT_COLLECTION
-        self.qdrant = Qdrant(
-            url=f"http://{config.QDRANT_HOST}:{config.QDRANT_PORT}",
-            api_key=config.QDRANT_API_KEY,
-            collection_name=self.collection_name,
-            embeddings=embeddings,
+    return QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
+
+def ensure_collection(client: QdrantClient, collection_name: str, vector_size: int = 768) -> None:
+    """
+    Ensures that the specified collection exists in Qdrant.
+    If it does not exist, it will be created with the given vector size.
+    """
+    if not client.has_collection(collection_name):
+        client.recreate_collection(
+            collection_name=collection_name,
+            vectors_config=qdrant_models.VectorParams(
+                size=vector_size,
+                distance="Cosine"
+            )
         )
 
-    def add_documents(self, documents: list[Document]):
-        """
-        Add a list of langchain Document objects to the Qdrant collection.
-        """
-        self.qdrant.add_documents(documents)
-
-    def get_retriever(self):
-        """
-        Return a retriever that can be used with LangChain chains.
-        """
-        return self.qdrant.as_retriever()
+def get_vector_store() -> Qdrant:
+    """
+    Returns a Qdrant vector store instance ready for use with LangChain.
+    """
+    client = get_qdrant_client()
+    ensure_collection(client, QDRANT_COLLECTION_NAME)
+    embeddings = get_ollama_embeddings()
+    return Qdrant(
+        client=client,
+        collection_name=QDRANT_COLLECTION_NAME,
+        embeddings=embeddings
+    )