feat: tools search_local_kb (chromadb) + web_search (tavily)
This commit is contained in:
@@ -1,30 +1,57 @@
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
"""Инструменты агента: локальный поиск (ChromaDB) и веб-поиск (Tavily)."""
|
||||
from langchain.tools import tool
|
||||
from langchain_ollama import OllamaEmbeddings
|
||||
from langchain_chroma import Chroma
|
||||
from langchain_tavily import TavilySearchResults
|
||||
from langchain_tavily import TavilySearch
|
||||
|
||||
from vectorstore import create_vectorstore
|
||||
_vectorstore: Chroma | None = None
|
||||
|
||||
|
||||
def init_tools(vectorstore: Chroma) -> None:
|
||||
"""Привязать vectorstore к инструментам (вызывается из main.py)."""
|
||||
global _vectorstore
|
||||
_vectorstore = vectorstore
|
||||
|
||||
# Global vector store instance
|
||||
VECTORSTORE = create_vectorstore()
|
||||
|
||||
@tool
|
||||
def search_local_kb(query: str, top_k: int = 3) -> str:
|
||||
"""Semantic search in the local ChromaDB knowledge base."""
|
||||
retriever = VECTORSTORE.as_retriever(search_kwargs={"k": top_k})
|
||||
docs = retriever.get_relevant_documents(query)
|
||||
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 "No relevant local knowledge found."
|
||||
return "\n\n".join(doc.page_content for doc in 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:
|
||||
"""Web search using Tavily."""
|
||||
tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY"))
|
||||
results = tavily.run(query)
|
||||
if not results:
|
||||
return "No web results found."
|
||||
return "\n\n".join(f"{r['title']}\n{r['content']}" for r in results)
|
||||
"""Поиск актуальной информации в интернете через 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
|
||||
|
||||
Reference in New Issue
Block a user