diff --git a/knowledge_base.py b/knowledge_base.py deleted file mode 100644 index 7890407..0000000 --- a/knowledge_base.py +++ /dev/null @@ -1,90 +0,0 @@ -"""Векторное хранилище знаний: 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)