diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..09a7f16 --- /dev/null +++ b/tools.py @@ -0,0 +1,57 @@ +""" +tools.py — RAG-инструменты для агента. + +Два инструмента через декоратор @tool: + • search_knowledge_base — семантический поиск в базе знаний + • add_to_knowledge_base — добавление документа в базу +""" + +from langchain_core.tools import tool +from vector_store import add_documents, search + + +@tool +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """ + Выполняет семантический поиск в базе знаний. + + Args: + query: поисковый запрос на естественном языке. + max_results: максимальное количество возвращаемых результатов (по умолчанию 5). + + Returns: + Строка с найденными фрагментами и их оценками релевантности. + """ + results = search(query, max_results=max_results) + + if not results: + return "В базе знаний ничего не найдено по данному запросу." + + lines = [] + for i, (doc, score) in enumerate(results, start=1): + title = doc.metadata.get("title", "—") + chunk = doc.metadata.get("chunk_index", 0) + lines.append( + f"[{i}] (score={score:.3f}) «{title}» chunk#{chunk}\n{doc.page_content}" + ) + + return "\n\n".join(lines) + + +@tool +def add_to_knowledge_base(content: str, title: str = "") -> str: + """ + Добавляет новый документ в базу знаний. + + Args: + content: текстовое содержимое документа. + title: название / источник документа (необязательно). + + Returns: + Сообщение об успехе с количеством созданных чанков. + """ + n_chunks = add_documents(content, title=title) + return ( + f"Документ «{title or 'без названия'}» успешно добавлен в базу знаний. " + f"Создано чанков: {n_chunks}." + ) \ No newline at end of file