diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/qdrant_client.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/qdrant_client.py index 7790ed0..6c93f17 100644 --- a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/qdrant_client.py +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/qdrant_client.py @@ -1,154 +1,175 @@ # qdrant_client.py """ -Модуль инициализации клиента Qdrant, добавления документов и поиска по эмбеддингам Ollama. -Использует: -- langchain_qdrant.QdrantVectorStore для работы с коллекцией Qdrant -- langchain_ollama.OllamaEmbeddings для генерации векторных представлений -- RecursiveCharacterTextSplitter для разбиения больших текстов на чанки +Модуль инициализации клиента Qdrant и простых операций с ним. +Используется в агенте для RAG‑памяти. -Функциональность: -1. Инициализация клиента и создание/получение коллекции. -2. Добавление документа (с заголовком) в базу знаний с автоматическим чанкингом. -3. Поиск по семантическому запросу с ограничением количества результатов. +Поддержка: +- Создание/получение коллекции +- Добавление документов (с чанками) с эмбеддингами Ollama +- Семантический поиск по запросу -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__( - 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: перекрытие между чанками - """ + 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 - # Параметры подключения к Qdrant - host = host or os.getenv("QDRANT_HOST", "localhost") - port = port or int(os.getenv("QDRANT_PORT", 6333)) + # Инициализируем низкоуровневый клиент + self.client = _QdrantClient(host=self.host, port=self.port) - # Создаём объект эмбеддинговой модели Ollama - self.embeddings = OllamaEmbeddings(model=embedding_model) + # Создаём коллекцию при необходимости + if not self._collection_exists(): + self._create_collection() - # Инициализируем QdrantVectorStore - self.store = QdrantVectorStore( - url=f"http://{host}:{port}", + # Обёртка LangChain для упрощённого добавления/поиска + self.vector_store = QdrantVectorStore( + 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: - """ - Добавляет документ в базу знаний. Текст разбивается на чанки, - каждому присваиваются метаданные (title и номер чанка). + # ----------------------------------------------------------------------- # + # Публичные методы + # ----------------------------------------------------------------------- # - :param content: полный текст документа - :param title: заголовок/имя документа + def add_documents(self, documents: Iterable[str], titles: Iterable[str]) -> None: """ - # Разбиваем контент на чанки - chunks = self.splitter.split_text(content) + Добавляет документы в коллекцию после разбиения на чанки. - documents = [] - for idx, chunk in enumerate(chunks): - doc = { - "page_content": chunk, - "metadata": {"title": title, "chunk_index": idx}, + :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 } - documents.append(doc) + for r in results + ] - # Добавляем векторные представления чанков - 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]]: + def load_from_directory(self, directory: str | Path) -> None: """ - Семантический поиск по базе знаний. + Загружает все текстовые файлы из указанной папки в коллекцию. - :param query: поисковый запрос - :param max_results: максимальное количество результатов - :param filter_metadata: словарь метаданных для фильтрации (например, {"title": "MyDoc"}) - :return: список найденных документов с полями: - - page_content - - metadata - - score + :param directory: Путь к каталогу. """ - # Выполняем поиск векторного пространства - results = self.store.similarity_search_with_score( - query=query, - k=max_results, - filter=filter_metadata, - ) + dir_path = Path(directory) + if not dir_path.is_dir(): + raise ValueError(f"Папка {directory} не существует") - # Преобразуем к удобному формату - formatted = [] - for doc, score in results: - formatted.append( - { - "content": doc.page_content, - "metadata": doc.metadata, - "score": score, - } - ) - return formatted + 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") - 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) + # Поиск по запросу + query_text = "Какой смысл жизни?" + results = client.search(query_text, limit=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("Content snippet:", res["content"][:200], "...") \ No newline at end of file + print(f"Text snippet: {res['text'][:200]}...") \ No newline at end of file