Добавить qdrant_client.py

qdrant_client.py

@@ -0,0 +1,175 @@
+# qdrant_client.py
+"""
+Модуль инициализации клиента Qdrant и простых операций с ним.
+Используется в агенте для RAG‑памяти.
+
+Поддержка:
+- Создание/получение коллекции
+- Добавление документов (с чанками) с эмбеддингами Ollama
+- Семантический поиск по запросу
+
+Требуемые зависимости:
+    pip install langchain langchain-qdrant langchain-ollama qdrant-client
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+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:
+    """
+    Обёртка над qdrant_client и langchain_qdrant для упрощения работы с коллекцией.
+    """
+
+    def __init__(self, host: str = QDRANT_HOST, port: int = QDRANT_PORT,
+                 collection_name: str = COLLECTION_NAME) -> None:
+        self.host = host
+        self.port = port
+        self.collection_name = collection_name
+
+        # Инициализируем низкоуровневый клиент
+        self.client = _QdrantClient(host=self.host, port=self.port)
+
+        # Создаём коллекцию при необходимости
+        if not self._collection_exists():
+            self._create_collection()
+
+        # Обёртка LangChain для упрощённого добавления/поиска
+        self.vector_store = QdrantVectorStore(
+            client=self.client,
+            collection_name=self.collection_name,
+            embedding=OllamaEmbeddings(model=EMBEDDING_MODEL),
+        )
+
+    # ----------------------------------------------------------------------- #
+    # Внутренние вспомогательные методы
+    # ----------------------------------------------------------------------- #
+
+    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_documents(self, documents: Iterable[str], titles: Iterable[str]) -> None:
+        """
+        Добавляет документы в коллекцию после разбиения на чанки.
+
+        :param documents: Итерируемый список текстов.
+        :param titles: Заголовки документов (один к одному с documents).
+        """
+        splitter = RecursiveCharacterTextSplitter(**TEXT_SPLITTER_PARAMS)
+        all_chunks: List[Tuple[str, str]] = []
+
+        for doc_text, title in zip(documents, titles):
+            chunks = splitter.split_text(doc_text)
+            # Для каждого чанка сохраняем заголовок как метаданные
+            all_chunks.extend([(chunk, title) for chunk in chunks])
+
+        texts, metas = zip(*all_chunks)
+
+        self.vector_store.add_texts(
+            texts=texts,
+            metadatas=[{"title": meta} for meta in metas]
+        )
+
+    def search(self, query: str, limit: int = 5) -> List[dict]:
+        """
+        Семантический поиск по запросу.
+
+        :param query: Текст запроса.
+        :param limit: Количество результатов.
+        :return: Список словарей с полями 'text', 'score' и метаданными.
+        """
+        results = self.vector_store.similarity_search_with_score(query, k=limit)
+        return [
+            {
+                "text": r[0].page_content,
+                "score": r[1],
+                "metadata": r[0].metadata
+            }
+            for r in results
+        ]
+
+    # ----------------------------------------------------------------------- #
+    # Утилиты для загрузки файлов из директории
+    # ----------------------------------------------------------------------- #
+
+    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)
+
+
+# --------------------------------------------------------------------------- #
+# Пример использования (можно удалить при импорте в другие модули)
+# --------------------------------------------------------------------------- #
+
+if __name__ == "__main__":
+    client = QdrantClient()
+    # Загрузка из папки ./data
+    client.load_from_directory("./data")
+
+    # Поиск по запросу
+    query_text = "Какой смысл жизни?"
+    results = client.search(query_text, limit=3)
+    for i, res in enumerate(results, 1):
+        print(f"\nРезультат {i}:")
+        print(f"Score: {res['score']:.4f}")
+        print(f"Title: {res['metadata'].get('title')}")
+        print(f"Text snippet: {res['text'][:200]}...")