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

src/vector_store.py

@@ -1,97 +1,29 @@
-"""
-Vector store implementation using Qdrant via langchain-qdrant.
-Provides a simple interface for adding documents and performing
-similarity search. Embeddings are generated using OpenAIEmbeddings
-by default, but can be overridden by passing a custom embedding
-function.
-"""
-
-from __future__ import annotations
-
-from typing import Iterable, List, Optional
-
-from langchain.embeddings import OpenAIEmbeddings
 from langchain_qdrant import Qdrant
-from langchain.vectorstores import VectorStore
-from langchain_core.documents import Document
+from langchain.schema import Document
+import config
 
-from .config import (
-    QDRANT_HOST,
-    QDRANT_PORT,
-    QDRANT_API_KEY,
-    QDRANT_COLLECTION,
-)
-
-
-class QdrantVectorStore(VectorStore):
+class QdrantVectorStore:
     """
-    A wrapper around langchain_qdrant.Qdrant that implements the
-    VectorStore interface expected by LangChain chains.
+    Wrapper around langchain_qdrant.Qdrant to provide a simple interface
+    for adding documents and retrieving a retriever.
     """
-
-    def __init__(
-        self,
-        embeddings: Optional[OpenAIEmbeddings] = None,
-        collection_name: str = QDRANT_COLLECTION,
-    ):
-        self.embeddings = embeddings or OpenAIEmbeddings()
-        self.collection_name = collection_name
-
-        # Initialize Qdrant client
-        self.client = Qdrant(
-            host=QDRANT_HOST,
-            port=QDRANT_PORT,
-            api_key=QDRANT_API_KEY,
+    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,
-        )
-
-    def add_documents(self, documents: Iterable[Document]) -> None:
-        """
-        Add a collection of documents to the Qdrant store.
-        """
-        texts = [doc.page_content for doc in documents]
-        metadatas = [doc.metadata for doc in documents]
-        ids = [doc.id for doc in documents if doc.id is not None]
-
-        # Embed the documents
-        embeddings = self.embeddings.embed_documents(texts)
-
-        # Upsert into Qdrant
-        self.client.upsert(
             embeddings=embeddings,
-            documents=texts,
-            metadatas=metadatas,
-            ids=ids,
         )
 
-    def similarity_search(
-        self,
-        query: str,
-        k: int = 5,
-        filter: Optional[dict] = None,
-    ) -> List[Document]:
+    def add_documents(self, documents: list[Document]):
         """
-        Perform a similarity search against the Qdrant store.
+        Add a list of langchain Document objects to the Qdrant collection.
         """
-        query_embedding = self.embeddings.embed_query(query)
-        results = self.client.search(
-            query_embedding=query_embedding,
-            limit=k,
-            filter=filter,
-        )
-        # Convert results to Document objects
-        return [
-            Document(
-                page_content=result["payload"]["text"],
-                metadata=result["payload"],
-                id=result["id"],
-            )
-            for result in results
-        ]
+        self.qdrant.add_documents(documents)
 
-    # The following methods are required by the VectorStore interface
-    def embed_query(self, query: str) -> List[float]:
-        return self.embeddings.embed_query(query)
-
-    def embed_documents(self, documents: List[str]) -> List[List[float]]:
-        return self.embeddings.embed_documents(documents)
+    def get_retriever(self):
+        """
+        Return a retriever that can be used with LangChain chains.
+        """
+        return self.qdrant.as_retriever()