79 lines
3.0 KiB
Python
79 lines
3.0 KiB
Python
"""Векторное хранилище знаний: 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)
|