39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
"""RAG-инструменты для агента."""
|
|
from __future__ import annotations
|
|
|
|
from langchain.tools import tool
|
|
|
|
from vector_store import add_document_to_store, search_store
|
|
|
|
|
|
@tool
|
|
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
|
"""Семантический поиск в базе знаний.
|
|
|
|
Args:
|
|
query: поисковый запрос
|
|
max_results: максимум фрагментов в ответе
|
|
"""
|
|
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"[релевантность={score:.3f}] {title}: {doc.page_content[:400]}"
|
|
)
|
|
return "\n\n".join(lines)
|
|
|
|
|
|
@tool
|
|
def add_to_knowledge_base(content: str, title: str) -> str:
|
|
"""Добавляет документ в векторную базу знаний.
|
|
|
|
Args:
|
|
content: текст документа
|
|
title: заголовок / имя документа
|
|
"""
|
|
chunks = add_document_to_store(content, title)
|
|
return f"Документ '{title}' добавлен ({chunks} чанков)."
|