fix: align with accepted structure (rag_tools.py)
This commit is contained in:
+29
-13
@@ -1,29 +1,41 @@
|
||||
"""RAG-инструменты для агента."""
|
||||
"""RAG-инструменты агента."""
|
||||
from __future__ import annotations
|
||||
|
||||
from langchain.tools import tool
|
||||
|
||||
from vector_store import add_document_to_store, search_store
|
||||
try:
|
||||
from knowledge_base import KnowledgeBase, build_knowledge_base
|
||||
except ModuleNotFoundError:
|
||||
from .knowledge_base import KnowledgeBase, build_knowledge_base
|
||||
|
||||
_kb: KnowledgeBase | None = None
|
||||
|
||||
|
||||
def get_knowledge_base() -> KnowledgeBase:
|
||||
"""Ленивая инициализация — не подключаем Ollama при импорте модуля."""
|
||||
global _kb
|
||||
if _kb is None:
|
||||
_kb = build_knowledge_base()
|
||||
return _kb
|
||||
|
||||
|
||||
@tool
|
||||
def search_knowledge_base(query: str, max_results: int = 4) -> str:
|
||||
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
||||
"""Семантический поиск в базе знаний.
|
||||
|
||||
Args:
|
||||
query: поисковый запрос
|
||||
max_results: максимум фрагментов в ответе
|
||||
"""
|
||||
hits = search_store(query, max_results=max_results)
|
||||
if not hits:
|
||||
return "В базе знаний ничего не найдено."
|
||||
lines = []
|
||||
for doc, score in hits:
|
||||
title = doc.metadata.get("title", "без названия")
|
||||
results = get_knowledge_base().search(query, max_results=max_results)
|
||||
if not results:
|
||||
return "По запросу ничего не найдено."
|
||||
lines = [f"Найдено: {len(results)} фрагмент(ов)\n"]
|
||||
for idx, item in enumerate(results, start=1):
|
||||
lines.append(
|
||||
f"[релевантность={score:.3f}] {title}: {doc.page_content[:400]}"
|
||||
f"{idx}. [{item['title']}] (score={item['score']})\n{item['content']}\n"
|
||||
)
|
||||
return "\n\n".join(lines)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@tool
|
||||
@@ -34,5 +46,9 @@ def add_to_knowledge_base(content: str, title: str) -> str:
|
||||
content: текст документа
|
||||
title: заголовок / имя документа
|
||||
"""
|
||||
chunks = add_document_to_store(content, title)
|
||||
return f"Документ '{title}' добавлен ({chunks} чанков)."
|
||||
if not content.strip():
|
||||
return "Ошибка: пустой content."
|
||||
if not title.strip():
|
||||
return "Ошибка: пустой title."
|
||||
chunks = get_knowledge_base().add_document(content.strip(), title.strip())
|
||||
return f"Документ «{title}» добавлен ({chunks} чанк(ов))."
|
||||
|
||||
Reference in New Issue
Block a user