sync from accepted repo: knowledge_base.py
This commit is contained in:
+31
-20
@@ -13,7 +13,7 @@ from qdrant_client import QdrantClient
|
|||||||
from qdrant_client.http.models import Distance, VectorParams
|
from qdrant_client.http.models import Distance, VectorParams
|
||||||
|
|
||||||
COLLECTION_NAME = "knowledge_base"
|
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")
|
OLLAMA_BASE_URL = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
|
||||||
EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
|
EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
|
||||||
|
|
||||||
@@ -29,9 +29,15 @@ class KnowledgeBase:
|
|||||||
self.collection_name = collection_name
|
self.collection_name = collection_name
|
||||||
self.qdrant_path = Path(qdrant_path)
|
self.qdrant_path = Path(qdrant_path)
|
||||||
self.qdrant_path.mkdir(parents=True, exist_ok=True)
|
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.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._ensure_collection()
|
||||||
self.vector_store = QdrantVectorStore(
|
self.vector_store = QdrantVectorStore(
|
||||||
client=self.client,
|
client=self.client,
|
||||||
@@ -48,32 +54,37 @@ class KnowledgeBase:
|
|||||||
vectors_config=VectorParams(size=len(sample), distance=Distance.COSINE),
|
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:
|
def add_document(self, content: str, title: str) -> int:
|
||||||
chunks = self._chunk_document(content, title)
|
"""Добавляет документ (с чанкингом) в базу. Возвращает число чанков."""
|
||||||
self.vector_store.add_documents(chunks)
|
chunks = self.splitter.create_documents(
|
||||||
return len(chunks)
|
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]]:
|
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)
|
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:
|
for doc, score in hits:
|
||||||
result.append(
|
results.append(
|
||||||
{
|
{
|
||||||
"title": doc.metadata.get("title", "без названия"),
|
"title": doc.metadata.get("title", "без названия"),
|
||||||
"content": doc.page_content[:400],
|
"content": doc.page_content,
|
||||||
"score": round(float(score), 4),
|
"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:
|
def build_knowledge_base(
|
||||||
"""Фабрика для автопроверки."""
|
qdrant_path: str | Path = DEFAULT_QDRANT_PATH,
|
||||||
return KnowledgeBase()
|
) -> KnowledgeBase:
|
||||||
|
"""Сборка RAG-хранилища — точка входа для автопроверки."""
|
||||||
|
return KnowledgeBase(qdrant_path=qdrant_path)
|
||||||
|
|||||||
Reference in New Issue
Block a user