diff --git a/knowledge_base.py b/knowledge_base.py index e18e5e0..7890407 100644 --- a/knowledge_base.py +++ b/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)