43 lines
1.4 KiB
Python
43 lines
1.4 KiB
Python
import os
|
|
from typing import Any
|
|
|
|
from langchain.tools import tool
|
|
from langchain_tavily import TavilySearch
|
|
|
|
from vectorstore import create_vectorstore
|
|
|
|
|
|
def _format_local_docs(results: list[Any]) -> str:
|
|
if not results:
|
|
return "Источник: chromadb\nНет релевантных локальных документов."
|
|
|
|
parts = ["Источник: chromadb"]
|
|
for index, doc in enumerate(results, start=1):
|
|
source = doc.metadata.get("source", "неизвестно")
|
|
parts.append(f"{index}. ({source}) {doc.page_content}")
|
|
return "\n".join(parts)
|
|
|
|
|
|
@tool
|
|
def search_local_kb(query: str, top_k: int = 3) -> str:
|
|
"""Search the local ChromaDB knowledge base for internal course notes and local documents."""
|
|
collection = create_vectorstore()
|
|
retriever = collection.as_retriever(search_kwargs={"k": top_k})
|
|
results = retriever.invoke(query)
|
|
return _format_local_docs(results)
|
|
|
|
|
|
@tool
|
|
def web_search(query: str) -> str:
|
|
"""Search the web with Tavily for fresh facts, current news, or anything not covered by local documents."""
|
|
api_key = os.getenv("TAVILY_API_KEY")
|
|
if not api_key:
|
|
return "Источник: tavily\nTAVILY_API_KEY не задан, поиск в интернете недоступен."
|
|
|
|
search_tool = TavilySearch(
|
|
max_results=5,
|
|
topic="general",
|
|
)
|
|
result = search_tool.invoke({"query": query})
|
|
return f"Источник: tavily\n{result}"
|