diff --git a/knowledge_base.py b/knowledge_base.py new file mode 100644 index 0000000..e18e5e0 --- /dev/null +++ b/knowledge_base.py @@ -0,0 +1,79 @@ +"""Векторное хранилище знаний: 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_data") +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=100) + 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 _chunk_document(self, content: str, title: str) -> list[Document]: + docs = self.splitter.create_documents(texts=[content], metadatas=[{"title": title}]) + for idx, doc in enumerate(docs): + doc.metadata["chunk_index"] = idx + doc.metadata["source"] = title + return docs + + def add_document(self, content: str, title: str) -> int: + chunks = self._chunk_document(content, title) + self.vector_store.add_documents(chunks) + return len(chunks) + + def search(self, query: str, max_results: int = 5) -> list[dict[str, Any]]: + hits = self.vector_store.similarity_search_with_score(query, k=max_results) + result: list[dict[str, Any]] = [] + for doc, score in hits: + result.append( + { + "title": doc.metadata.get("title", "без названия"), + "content": doc.page_content[:400], + "score": round(float(score), 4), + } + ) + return result + + +def build_knowledge_base() -> KnowledgeBase: + """Фабрика для автопроверки.""" + return KnowledgeBase()