fix: align with accepted structure (knowledge_base.py)
This commit is contained in:
@@ -0,0 +1,79 @@
|
||||
"""Векторное хранилище знаний: 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_data")
|
||||
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=100)
|
||||
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 _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)
|
||||
|
||||
def search(self, query: str, max_results: int = 5) -> list[dict[str, Any]]:
|
||||
hits = self.vector_store.similarity_search_with_score(query, k=max_results)
|
||||
result: list[dict[str, Any]] = []
|
||||
for doc, score in hits:
|
||||
result.append(
|
||||
{
|
||||
"title": doc.metadata.get("title", "без названия"),
|
||||
"content": doc.page_content[:400],
|
||||
"score": round(float(score), 4),
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def build_knowledge_base() -> KnowledgeBase:
|
||||
"""Фабрика для автопроверки."""
|
||||
return KnowledgeBase()
|
||||
Reference in New Issue
Block a user