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

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

@@ -1,154 +1,175 @@
 # qdrant_client.py
 """
-Модуль инициализации клиента Qdrant, добавления документов и поиска по эмбеддингам Ollama.
+Модуль инициализации клиента Qdrant и простых операций с ним.
-Использует:
+Используется в агенте для RAG‑памяти.
-- langchain_qdrant.QdrantVectorStore для работы с коллекцией Qdrant
-- langchain_ollama.OllamaEmbeddings для генерации векторных представлений
-- RecursiveCharacterTextSplitter для разбиения больших текстов на чанки
 
-Функциональность:
+Поддержка:
-1. Инициализация клиента и создание/получение коллекции.
+- Создание/получение коллекции
-2. Добавление документа (с заголовком) в базу знаний с автоматическим чанкингом.
+- Добавление документов (с чанками) с эмбеддингами Ollama
-3. Поиск по семантическому запросу с ограничением количества результатов.
+- Семантический поиск по запросу
 
-Author: ChatGPT
+Требуемые зависимости:
+    pip install langchain langchain-qdrant langchain-ollama qdrant-client
 """
 
 from __future__ import annotations
 
 import os
 from pathlib import Path
-from typing import List, Dict, Any
+from typing import Iterable, List, Tuple
 
 from langchain_ollama import OllamaEmbeddings
 from langchain_qdrant import QdrantVectorStore
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from qdrant_client import QdrantClient as _QdrantClient
+from qdrant_client.http.models import Distance, VectorParams
 
+# --------------------------------------------------------------------------- #
+# Конфигурация (можно переопределить через переменные окружения)
+# --------------------------------------------------------------------------- #
+
+QDRANT_HOST: str = os.getenv("QDRANT_HOST", "localhost")
+QDRANT_PORT: int = int(os.getenv("QDRANT_PORT", "6333"))
+COLLECTION_NAME: str = os.getenv("QDRANT_COLLECTION", "rag_knowledge")
+
+EMBEDDING_MODEL: str = os.getenv("OLLAMA_EMBEDDING_MODEL", "nomic-embed-text")
+TEXT_SPLITTER_PARAMS = {
+    "chunk_size": 1000,
+    "chunk_overlap": 200,
+}
+
+# --------------------------------------------------------------------------- #
+# Класс клиента Qdrant
+# --------------------------------------------------------------------------- #
 
 class QdrantClient:
     """
-    Класс-обёртка над QdrantVectorStore.
+    Обёртка над qdrant_client и langchain_qdrant для упрощения работы с коллекцией.
     """
 
-    def __init__(
+    def __init__(self, host: str = QDRANT_HOST, port: int = QDRANT_PORT,
-        self,
+                 collection_name: str = COLLECTION_NAME) -> None:
-        collection_name: str = "knowledge_base",
+        self.host = host
-        host: str | None = None,
+        self.port = port
-        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")
+        self.client = _QdrantClient(host=self.host, port=self.port)
-        port = port or int(os.getenv("QDRANT_PORT", 6333))
 
-        # Создаём объект эмбеддинговой модели Ollama
+        # Создаём коллекцию при необходимости
-        self.embeddings = OllamaEmbeddings(model=embedding_model)
+        if not self._collection_exists():
+            self._create_collection()
 
-        # Инициализируем QdrantVectorStore
+        # Обёртка LangChain для упрощённого добавления/поиска
-        self.store = QdrantVectorStore(
+        self.vector_store = QdrantVectorStore(
-            url=f"http://{host}:{port}",
+            client=self.client,
             collection_name=self.collection_name,
-            embedding_function=self.embeddings.embed_query,  # используем embed_query для совместимости
+            embedding=OllamaEmbeddings(model=EMBEDDING_MODEL),
         )
 
-        # Создаём splitter для разбиения текста на чанки
+    # ----------------------------------------------------------------------- #
-        self.splitter = RecursiveCharacterTextSplitter(
+    # Внутренние вспомогательные методы
-            chunk_size=chunk_size,
+    # ----------------------------------------------------------------------- #
-            chunk_overlap=chunk_overlap,
+
+    def _collection_exists(self) -> bool:
+        """Проверяем наличие коллекции."""
+        try:
+            self.client.get_collection(name=self.collection_name)
+            return True
+        except Exception:
+            return False
+
+    def _create_collection(self) -> None:
+        """Создаём коллекцию с параметрами по умолчанию."""
+        vector_params = VectorParams(size=384, distance=Distance.COSINE)
+        self.client.create_collection(
+            collection_name=self.collection_name,
+            vectors_config=vector_params
         )
 
-    def add_document(self, content: str, title: str) -> None:
+    # ----------------------------------------------------------------------- #
+    # Публичные методы
+    # ----------------------------------------------------------------------- #
+
+    def add_documents(self, documents: Iterable[str], titles: Iterable[str]) -> None:
         """
-        Добавляет документ в базу знаний. Текст разбивается на чанки,
+        Добавляет документы в коллекцию после разбиения на чанки.
-        каждому присваиваются метаданные (title и номер чанка).
 
-        :param content: полный текст документа
+        :param documents: Итерируемый список текстов.
-        :param title: заголовок/имя документа
+        :param titles: Заголовки документов (один к одному с documents).
         """
-        # Разбиваем контент на чанки
+        splitter = RecursiveCharacterTextSplitter(**TEXT_SPLITTER_PARAMS)
-        chunks = self.splitter.split_text(content)
+        all_chunks: List[Tuple[str, str]] = []
 
-        documents = []
+        for doc_text, title in zip(documents, titles):
-        for idx, chunk in enumerate(chunks):
+            chunks = splitter.split_text(doc_text)
-            doc = {
+            # Для каждого чанка сохраняем заголовок как метаданные
-                "page_content": chunk,
+            all_chunks.extend([(chunk, title) for chunk in chunks])
-                "metadata": {"title": title, "chunk_index": idx},
-            }
-            documents.append(doc)
 
-        # Добавляем векторные представления чанков
+        texts, metas = zip(*all_chunks)
-        self.store.add_documents(documents)
 
-    def search(
+        self.vector_store.add_texts(
-        self,
+            texts=texts,
-        query: str,
+            metadatas=[{"title": meta} for meta in metas]
-        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,
         )
 
-        # Преобразуем к удобному формату
+    def search(self, query: str, limit: int = 5) -> List[dict]:
-        formatted = []
+        """
-        for doc, score in results:
+        Семантический поиск по запросу.
-            formatted.append(
+
+        :param query: Текст запроса.
+        :param limit: Количество результатов.
+        :return: Список словарей с полями 'text', 'score' и метаданными.
+        """
+        results = self.vector_store.similarity_search_with_score(query, k=limit)
+        return [
             {
-                    "content": doc.page_content,
+                "text": r[0].page_content,
-                    "metadata": doc.metadata,
+                "score": r[1],
-                    "score": score,
+                "metadata": r[0].metadata
             }
-            )
+            for r in results
-        return formatted
+        ]
+
+    # ----------------------------------------------------------------------- #
+    # Утилиты для загрузки файлов из директории
+    # ----------------------------------------------------------------------- #
+
+    def load_from_directory(self, directory: str | Path) -> None:
+        """
+        Загружает все текстовые файлы из указанной папки в коллекцию.
+
+        :param directory: Путь к каталогу.
+        """
+        dir_path = Path(directory)
+        if not dir_path.is_dir():
+            raise ValueError(f"Папка {directory} не существует")
+
+        docs = []
+        titles = []
+
+        for file in dir_path.glob("**/*.txt"):
+            text = file.read_text(encoding="utf-8")
+            docs.append(text)
+            titles.append(file.stem)
+
+        self.add_documents(docs, titles)
 
 
 # --------------------------------------------------------------------------- #
-# Пример использования (не входит в публичный API модуля)
+# Пример использования (можно удалить при импорте в другие модули)
 # --------------------------------------------------------------------------- #
 
 if __name__ == "__main__":
-    # Инициализируем клиент
     client = QdrantClient()
+    # Загрузка из папки ./data
+    client.load_from_directory("./data")
 
-    # Добавляем пример документа из файла
+    # Поиск по запросу
-    sample_path = Path("sample.txt")
+    query_text = "Какой смысл жизни?"
-    if sample_path.exists():
+    results = client.search(query_text, limit=3)
-        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"\nРезультат {i}:")
+        print(f"Score: {res['score']:.4f}")
         print(f"Title: {res['metadata'].get('title')}")
-        print(f"Chunk index: {res['metadata'].get('chunk_index')}")
+        print(f"Text snippet: {res['text'][:200]}...")
-        print("Content snippet:", res["content"][:200], "...")