Агент с RAG‑памятью: qdrant_client.py

solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/qdrant_client.py

@@ -0,0 +1,154 @@
+# qdrant_client.py
+"""
+Модуль инициализации клиента Qdrant, добавления документов и поиска по эмбеддингам Ollama.
+Использует:
+- langchain_qdrant.QdrantVectorStore для работы с коллекцией Qdrant
+- langchain_ollama.OllamaEmbeddings для генерации векторных представлений
+- RecursiveCharacterTextSplitter для разбиения больших текстов на чанки
+
+Функциональность:
+1. Инициализация клиента и создание/получение коллекции.
+2. Добавление документа (с заголовком) в базу знаний с автоматическим чанкингом.
+3. Поиск по семантическому запросу с ограничением количества результатов.
+
+Author: ChatGPT
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import List, Dict, Any
+
+from langchain_ollama import OllamaEmbeddings
+from langchain_qdrant import QdrantVectorStore
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+
+class QdrantClient:
+    """
+    Класс-обёртка над QdrantVectorStore.
+    """
+
+    def __init__(
+        self,
+        collection_name: str = "knowledge_base",
+        host: str | None = None,
+        port: int | None = None,
+        embedding_model: str = "nomic-embed-text",
+        chunk_size: int = 1000,
+        chunk_overlap: int = 200,
+    ) -> None:
+        """
+        Инициализация клиента Qdrant и подготовка коллекции.
+
+        :param collection_name: имя коллекции в Qdrant
+        :param host: хост Qdrant (по умолчанию localhost)
+        :param port: порт Qdrant (по умолчанию 6333)
+        :param embedding_model: название модели Ollama для эмбеддингов
+        :param chunk_size: максимальная длина чанка в символах
+        :param chunk_overlap: перекрытие между чанками
+        """
+        self.collection_name = collection_name
+
+        # Параметры подключения к Qdrant
+        host = host or os.getenv("QDRANT_HOST", "localhost")
+        port = port or int(os.getenv("QDRANT_PORT", 6333))
+
+        # Создаём объект эмбеддинговой модели Ollama
+        self.embeddings = OllamaEmbeddings(model=embedding_model)
+
+        # Инициализируем QdrantVectorStore
+        self.store = QdrantVectorStore(
+            url=f"http://{host}:{port}",
+            collection_name=self.collection_name,
+            embedding_function=self.embeddings.embed_query,  # используем embed_query для совместимости
+        )
+
+        # Создаём splitter для разбиения текста на чанки
+        self.splitter = RecursiveCharacterTextSplitter(
+            chunk_size=chunk_size,
+            chunk_overlap=chunk_overlap,
+        )
+
+    def add_document(self, content: str, title: str) -> None:
+        """
+        Добавляет документ в базу знаний. Текст разбивается на чанки,
+        каждому присваиваются метаданные (title и номер чанка).
+
+        :param content: полный текст документа
+        :param title: заголовок/имя документа
+        """
+        # Разбиваем контент на чанки
+        chunks = self.splitter.split_text(content)
+
+        documents = []
+        for idx, chunk in enumerate(chunks):
+            doc = {
+                "page_content": chunk,
+                "metadata": {"title": title, "chunk_index": idx},
+            }
+            documents.append(doc)
+
+        # Добавляем векторные представления чанков
+        self.store.add_documents(documents)
+
+    def search(
+        self,
+        query: str,
+        max_results: int = 5,
+        filter_metadata: Dict[str, Any] | None = None,
+    ) -> List[Dict[str, Any]]:
+        """
+        Семантический поиск по базе знаний.
+
+        :param query: поисковый запрос
+        :param max_results: максимальное количество результатов
+        :param filter_metadata: словарь метаданных для фильтрации (например, {"title": "MyDoc"})
+        :return: список найденных документов с полями:
+                 - page_content
+                 - metadata
+                 - score
+        """
+        # Выполняем поиск векторного пространства
+        results = self.store.similarity_search_with_score(
+            query=query,
+            k=max_results,
+            filter=filter_metadata,
+        )
+
+        # Преобразуем к удобному формату
+        formatted = []
+        for doc, score in results:
+            formatted.append(
+                {
+                    "content": doc.page_content,
+                    "metadata": doc.metadata,
+                    "score": score,
+                }
+            )
+        return formatted
+
+
+# --------------------------------------------------------------------------- #
+# Пример использования (не входит в публичный API модуля)
+# --------------------------------------------------------------------------- #
+
+if __name__ == "__main__":
+    # Инициализируем клиент
+    client = QdrantClient()
+
+    # Добавляем пример документа из файла
+    sample_path = Path("sample.txt")
+    if sample_path.exists():
+        text = sample_path.read_text(encoding="utf-8")
+        client.add_document(content=text, title=sample_path.stem)
+
+    # Выполняем поиск
+    query_str = "What is the main idea of the document?"
+    results = client.search(query=query_str, max_results=3)
+    for i, res in enumerate(results, 1):
+        print(f"\nResult {i} (score={res['score']:.4f})")
+        print(f"Title: {res['metadata'].get('title')}")
+        print(f"Chunk index: {res['metadata'].get('chunk_index')}")
+        print("Content snippet:", res["content"][:200], "...")