diff --git a/tools.py b/tools.py index 09a7f16..34dfb04 100644 --- a/tools.py +++ b/tools.py @@ -1,57 +1,41 @@ -""" -tools.py — RAG-инструменты для агента. - -Два инструмента через декоратор @tool: - • search_knowledge_base — семантический поиск в базе знаний - • add_to_knowledge_base — добавление документа в базу -""" - -from langchain_core.tools import tool -from vector_store import add_documents, search +from langchain.tools import tool +from vector_store import search_documents, add_documents @tool def search_knowledge_base(query: str, max_results: int = 5) -> str: - """ - Выполняет семантический поиск в базе знаний. + """Semantic search in the knowledge base. Use this tool to find relevant information. Args: - query: поисковый запрос на естественном языке. - max_results: максимальное количество возвращаемых результатов (по умолчанию 5). + query: The search query string. + max_results: Maximum number of results to return (default 5). Returns: - Строка с найденными фрагментами и их оценками релевантности. + Formatted string with search results and relevance scores. """ - results = search(query, max_results=max_results) - + results = search_documents(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) + return "No results found in the knowledge base." + output_lines = [f"Found {len(results)} result(s):\n"] + for i, r in enumerate(results, 1): + title = r["metadata"].get("title", "Unknown") + score = r["score"] + content = r["content"] + output_lines.append(f"[{i}] Title: {title} | Score: {score}") + output_lines.append(f" {content}\n") + return "\n".join(output_lines) @tool -def add_to_knowledge_base(content: str, title: str = "") -> str: - """ - Добавляет новый документ в базу знаний. +def add_to_knowledge_base(content: str, title: str) -> str: + """Add a document to the knowledge base. Use this tool to store new information. Args: - content: текстовое содержимое документа. - title: название / источник документа (необязательно). + content: The full text content of the document to add. + title: A descriptive title for the document. Returns: - Сообщение об успехе с количеством созданных чанков. + Confirmation message with the number of chunks stored. """ - n_chunks = add_documents(content, title=title) - return ( - f"Документ «{title or 'без названия'}» успешно добавлен в базу знаний. " - f"Создано чанков: {n_chunks}." - ) \ No newline at end of file + num_chunks = add_documents(content=content, title=title) + return f"Successfully added document '{title}' to the knowledge base ({num_chunks} chunk(s) stored)." \ No newline at end of file