add knowledge_base.py

2026-05-26 07:24:27 +00:00
parent cb830f66e0
commit 08442b3bc0
1 changed files with 78 additions and 0 deletions
@@ -0,0 +1,78 @@
+"""Векторное хранилище знаний: 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 = "./qdrant_storage"
+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.embeddings = OllamaEmbeddings(model=EMBED_MODEL)
+        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)