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

src/vector_store.py

@@ -0,0 +1,97 @@
+"""
+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 .config import (
+    QDRANT_HOST,
+    QDRANT_PORT,
+    QDRANT_API_KEY,
+    QDRANT_COLLECTION,
+)
+
+
+class QdrantVectorStore(VectorStore):
+    """
+    A wrapper around langchain_qdrant.Qdrant that implements the
+    VectorStore interface expected by LangChain chains.
+    """
+
+    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,
+            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]:
+        """
+        Perform a similarity search against the Qdrant store.
+        """
+        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
+        ]
+
+    # 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)