"""Векторное хранилище знаний: Qdrant + Ollama embeddings + чанкинг.""" from __future__ import annotations import os from pathlib import Path from typing import Any from langchain_core.documents import Document from langchain_ollama import OllamaEmbeddings from langchain_qdrant import QdrantVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter from qdrant_client import QdrantClient from qdrant_client.http.models import Distance, VectorParams COLLECTION_NAME = "knowledge_base" DEFAULT_QDRANT_PATH = os.getenv("QDRANT_PATH", "./qdrant_storage") OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434") EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text") class KnowledgeBase: """Локальная RAG-база на Qdrant с эмбеддингами Ollama.""" def __init__( self, qdrant_path: str | Path = DEFAULT_QDRANT_PATH, collection_name: str = COLLECTION_NAME, ) -> None: self.collection_name = collection_name self.qdrant_path = Path(qdrant_path) self.qdrant_path.mkdir(parents=True, exist_ok=True) self.embeddings = OllamaEmbeddings( model=EMBED_MODEL, base_url=OLLAMA_BASE_URL, ) self.client = QdrantClient(path=str(self.qdrant_path)) self.splitter = RecursiveCharacterTextSplitter( chunk_size=500, chunk_overlap=50, ) self._ensure_collection() self.vector_store = QdrantVectorStore( client=self.client, collection_name=self.collection_name, embedding=self.embeddings, ) def _ensure_collection(self) -> None: if self.client.collection_exists(self.collection_name): return sample = self.embeddings.embed_query("init") self.client.create_collection( collection_name=self.collection_name, vectors_config=VectorParams(size=len(sample), distance=Distance.COSINE), ) def add_document(self, content: str, title: str) -> int: """Добавляет документ (с чанкингом) в базу. Возвращает число чанков.""" chunks = self.splitter.create_documents( texts=[content], metadatas=[{"title": title}], ) ids = self.vector_store.add_documents(chunks) return len(ids) def search(self, query: str, max_results: int = 5) -> list[dict[str, Any]]: """Семантический поиск с оценкой релевантности (score).""" hits = self.vector_store.similarity_search_with_score(query, k=max_results) results: list[dict[str, Any]] = [] for doc, score in hits: results.append( { "title": doc.metadata.get("title", "без названия"), "content": doc.page_content, "score": round(float(score), 4), } ) return results def add_file(self, file_path: Path) -> int: text = file_path.read_text(encoding="utf-8") title = file_path.stem return self.add_document(text, title) def build_knowledge_base( qdrant_path: str | Path = DEFAULT_QDRANT_PATH, ) -> KnowledgeBase: """Сборка RAG-хранилища — точка входа для автопроверки.""" return KnowledgeBase(qdrant_path=qdrant_path)