57 lines
2.1 KiB
Python
57 lines
2.1 KiB
Python
"""
|
|
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}."
|
|
) |