diff --git a/rag_tools.py b/rag_tools.py index 4a03c84..ae7a062 100644 --- a/rag_tools.py +++ b/rag_tools.py @@ -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} чанк(ов))."