diff --git a/tools.py b/tools.py index 0c54f66..cb8160f 100644 --- a/tools.py +++ b/tools.py @@ -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