58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
"""Инструменты агента: локальный поиск (ChromaDB) и веб-поиск (Tavily)."""
|
|
from langchain.tools import tool
|
|
from langchain_chroma import Chroma
|
|
from langchain_tavily import TavilySearch
|
|
|
|
_vectorstore: Chroma | None = None
|
|
|
|
|
|
def init_tools(vectorstore: Chroma) -> None:
|
|
"""Привязать vectorstore к инструментам (вызывается из main.py)."""
|
|
global _vectorstore
|
|
_vectorstore = vectorstore
|
|
|
|
|
|
@tool
|
|
def search_local_kb(query: str, top_k: int = 5) -> str:
|
|
"""Семантический поиск по локальной базе знаний (ChromaDB).
|
|
Используй для вопросов по загруженным документам и конспектам.
|
|
"""
|
|
if _vectorstore is None:
|
|
return "База знаний не инициализирована."
|
|
docs = _vectorstore.similarity_search(query, k=top_k)
|
|
if not docs:
|
|
return "Документов по запросу не найдено в базе знаний."
|
|
parts = [
|
|
f"[{i + 1}] ({d.metadata.get('source', 'unknown')})
|
|
{d.page_content}"
|
|
for i, d in enumerate(docs)
|
|
]
|
|
return "Источник: chromadb
|
|
|
|
" + "
|
|
|
|
".join(parts)
|
|
|
|
|
|
@tool
|
|
def web_search(query: str) -> str:
|
|
"""Поиск актуальной информации в интернете через Tavily.
|
|
Используй для новостей, актуальных событий, фактов из сети.
|
|
"""
|
|
tavily = TavilySearch(max_results=5)
|
|
results = tavily.invoke({"query": query})
|
|
if isinstance(results, list):
|
|
parts = [
|
|
f"[{i + 1}] {r.get('title', '')}
|
|
{r.get('content', r.get('snippet', ''))}"
|
|
for i, r in enumerate(results)
|
|
]
|
|
content = "
|
|
|
|
".join(parts)
|
|
else:
|
|
content = str(results)
|
|
return "Источник: tavily
|
|
|
|
" + content
|