diff --git a/rag_tools.py b/rag_tools.py index 995c18b..0988ae5 100644 --- a/rag_tools.py +++ b/rag_tools.py @@ -1,19 +1,9 @@ -"""RAG-инструменты агента.""" +"""RAG-инструменты для агента.""" from __future__ import annotations from langchain.tools import tool -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 +from vector_store import add_document_to_store, search_store @tool @@ -21,31 +11,28 @@ def search_knowledge_base(query: str, max_results: int = 5) -> str: """Семантический поиск в базе знаний. Args: - query: Поисковый запрос на естественном языке. - max_results: Максимум результатов (по умолчанию 5). + query: поисковый запрос + max_results: максимум фрагментов в ответе """ - 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): + hits = search_store(query, max_results=max_results) + if not hits: + return "В базе знаний ничего не найдено." + lines = [] + for doc, score in hits: + title = doc.metadata.get("title", "без названия") lines.append( - f"{idx}. [{item['title']}] (score={item['score']})\n{item['content']}\n" + f"[релевантность={score:.3f}] {title}: {doc.page_content[:400]}" ) - return "\n".join(lines) + return "\n\n".join(lines) @tool def add_to_knowledge_base(content: str, title: str) -> str: - """Добавление документа в базу знаний. + """Добавляет документ в векторную базу знаний. Args: - content: Текст документа. - title: Заголовок / имя документа. + content: текст документа + title: заголовок / имя документа """ - 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} чанк(ов))." + chunks = add_document_to_store(content, title) + return f"Документ '{title}' добавлен ({chunks} чанков)."