sync from accepted repo: knowledge_base.py

knowledge_base.py

@@ -13,7 +13,7 @@ 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")
+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")
 
@@ -29,9 +29,15 @@ class KnowledgeBase:
         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.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.splitter = RecursiveCharacterTextSplitter(
+            chunk_size=500,
+            chunk_overlap=50,
+        )
         self._ensure_collection()
         self.vector_store = QdrantVectorStore(
             client=self.client,
@@ -48,32 +54,37 @@ class KnowledgeBase:
             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)
+        """Добавляет документ (с чанкингом) в базу. Возвращает число чанков."""
+        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)
-        result: list[dict[str, Any]] = []
+        results: list[dict[str, Any]] = []
         for doc, score in hits:
-            result.append(
+            results.append(
                 {
                     "title": doc.metadata.get("title", "без названия"),
-                    "content": doc.page_content[:400],
+                    "content": doc.page_content,
                     "score": round(float(score), 4),
                 }
             )
-        return result
+        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() -> KnowledgeBase:
-    """Фабрика для автопроверки."""
-    return KnowledgeBase()
+def build_knowledge_base(
+    qdrant_path: str | Path = DEFAULT_QDRANT_PATH,
+) -> KnowledgeBase:
+    """Сборка RAG-хранилища — точка входа для автопроверки."""
+    return KnowledgeBase(qdrant_path=qdrant_path)