From 8f7d7648b701d5397812df2d747ad508a85723bf Mon Sep 17 00:00:00 2001 From: lonpatovaadelina Date: Tue, 26 May 2026 09:21:07 +0000 Subject: [PATCH] =?UTF-8?q?=D0=90=D0=B3=D0=B5=D0=BD=D1=82=20=D1=81=20RAG?= =?UTF-8?q?=E2=80=91=D0=BF=D0=B0=D0=BC=D1=8F=D1=82=D1=8C=D1=8E:=20vector?= =?UTF-8?q?=5Fstore.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../vector_store.py | 196 ++++++++++++++++++ 1 file changed, 196 insertions(+) create mode 100644 solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/vector_store.py diff --git a/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/vector_store.py b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/vector_store.py new file mode 100644 index 0000000..3fba5b1 --- /dev/null +++ b/solutions/6a02e23da6fe2e4ac16acf65_Агент_с_RAG_памятью/vector_store.py @@ -0,0 +1,196 @@ +<|channel|>final code<|message|>```python +# vector_store.py +""" +Модуль работы с Qdrant, Ollama и чанкингом. +Содержит: +- инициализацию клиента и коллекции; +- добавление документов (с разбиением на чанки); +- поиск по семантической схожести. + +Используем: +- langchain_qdrant.QdrantVectorStore +- langchain_ollama.OllamaEmbeddings +- RecursiveCharacterTextSplitter +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Iterable, List, Tuple + +from langchain_core.documents import Document +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_qdrant import QdrantVectorStore +from langchain_ollama import OllamaEmbeddings +from qdrant_client import 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_base") + +EMBEDDING_MODEL: str = os.getenv("OLLAMA_EMBEDDINGS", "nomic-embed-text") +DIMENSIONS: int = 512 # размер эмбеддингов nomic-embed-text + +# --------------------------------------------------------------------------- # +# Класс VectorStore +# --------------------------------------------------------------------------- # + +class VectorStore: + """ + Обёртка над QdrantVectorStore с Ollama‑эмбеддингами и чанкингом. + """ + + def __init__( + self, + host: str = QDRANT_HOST, + port: int = QDRANT_PORT, + collection_name: str = COLLECTION_NAME, + embedding_model: str = EMBEDDING_MODEL, + dimensions: int = DIMENSIONS, + ) -> None: + """ + Создаёт клиент Qdrant и коллекцию (если её нет). + """ + self.client = QdrantClient(host=host, port=port) + self.collection_name = collection_name + + # Проверяем наличие коллекции + if not self._collection_exists(): + self._create_collection(dimensions) + + # Создаём объект VectorStore + self.store = QdrantVectorStore( + client=self.client, + collection_name=self.collection_name, + embedding=OllamaEmbeddings(model=embedding_model), + ) + + # --------------------------------------------------------------------- # + # Внутренние методы для работы с коллекцией + # --------------------------------------------------------------------- # + + def _collection_exists(self) -> bool: + """Проверяем, существует ли коллекция.""" + return self.collection_name in self.client.get_collections().collections + + def _create_collection(self, dimensions: int) -> None: + """Создаём новую коллекцию с заданными параметрами.""" + self.client.recreate_collection( + collection_name=self.collection_name, + vectors_config=VectorParams(size=dimensions, distance=Distance.COSINE), + ) + + # --------------------------------------------------------------------- # + # Чанкинг + # --------------------------------------------------------------------- # + + @staticmethod + def _split_text(text: str) -> List[Document]: + """ + Разбивает текст на чанки с метаданными. + Возвращает список Document. + """ + splitter = RecursiveCharacterTextSplitter( + chunk_size=1000, + chunk_overlap=200, + separators=["\n\n", "\n", " ", ""], + ) + texts = splitter.split_text(text) + return [Document(page_content=t) for t in texts] + + # --------------------------------------------------------------------- # + # Добавление документов + # --------------------------------------------------------------------- # + + def add_document(self, content: str, title: str | None = None) -> List[str]: + """ + Добавляет документ в базу: + * разбивает на чанки; + * сохраняет каждый чанк с метаданными. + Возвращает список ID добавленных чанков. + """ + docs = self._split_text(content) + # Добавляем заголовок как metadata, если он передан + if title: + for doc in docs: + doc.metadata["title"] = title + + ids = self.store.add_documents(docs) + return ids + + def add_documents_from_path(self, path: str | Path) -> List[str]: + """ + Загружает все текстовые файлы из директории (рекурсивно) + и добавляет их в базу. + Возвращает список всех ID. + """ + p = Path(path) + if not p.is_dir(): + raise ValueError(f"Путь {path} не является директорией") + + all_ids: List[str] = [] + for file in p.rglob("*"): + if file.suffix.lower() in {".txt", ".md"}: + text = file.read_text(encoding="utf-8") + ids = self.add_document(text, title=file.stem) + all_ids.extend(ids) + return all_ids + + # --------------------------------------------------------------------- # + # Поиск + # --------------------------------------------------------------------- # + + def search( + self, + query: str, + k: int = 5, + filter_metadata: dict | None = None, + ) -> List[Tuple[str, float]]: + """ + Семантический поиск по базе. + Возвращает список (document_id, score) упорядоченных по убыванию релевантности. + """ + results = self.store.similarity_search_with_score( + query=query, + k=k, + filter=filter_metadata, + ) + return [(doc.metadata.get("id", ""), score) for doc, score in results] + + # --------------------------------------------------------------------- # + # Удаление (необязательно) + # --------------------------------------------------------------------- # + + def delete_document(self, doc_id: str) -> None: + """Удаляет документ по ID.""" + self.client.delete( + collection_name=self.collection_name, + points_selector={"ids": [doc_id]}, + ) + +# --------------------------------------------------------------------------- # +# Пример использования (можно закомментировать при импорте) +# --------------------------------------------------------------------------- # + +if __name__ == "__main__": + store = VectorStore() + # Добавляем примерный документ + doc_text = """ + LangChain – это библиотека для создания LLM‑агентов. + Она позволяет легко интегрировать внешние сервисы, такие как базы данных, + поисковые движки и другие инструменты. В этом примере мы используем Qdrant + в качестве хранилища эмбеддингов и Ollama для генерации эмбеддингов. + """ + store.add_document(doc_text, title="LangChain Overview") + + # Поиск по запросу + results = store.search("что такое LangChain", k=3) + print("Результаты поиска:") + for doc_id, score in results: + print(f"ID: {doc_id} | Score: {score:.4f}") \ No newline at end of file