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

2026-05-26 09:21:07 +00:00
parent 64ae908897
commit 8f7d7648b7
1 changed files with 196 additions and 0 deletions
@@ -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}")