52 lines
1.8 KiB
Python
52 lines
1.8 KiB
Python
"""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
|
|
|
|
|
|
@tool
|
|
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
|
"""Семантический поиск в базе знаний.
|
|
|
|
Args:
|
|
query: Поисковый запрос на естественном языке.
|
|
max_results: Максимум результатов (по умолчанию 5).
|
|
"""
|
|
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"{idx}. [{item['title']}] (score={item['score']})\n{item['content']}\n"
|
|
)
|
|
return "\n".join(lines)
|
|
|
|
|
|
@tool
|
|
def add_to_knowledge_base(content: str, title: str) -> str:
|
|
"""Добавление документа в базу знаний.
|
|
|
|
Args:
|
|
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} чанк(ов))."
|