fix: tools.py — Экзамен: RAG-агент с ChromaDB и веб-поиском
This commit is contained in:
@@ -0,0 +1,28 @@
|
|||||||
|
import os
|
||||||
|
from langchain.tools import tool
|
||||||
|
from langchain_tavily import TavilySearchResults
|
||||||
|
from langchain_chroma import Chroma
|
||||||
|
|
||||||
|
@tool
|
||||||
|
def search_local_kb(query: str, top_k: int = 3) -> str:
|
||||||
|
"""Search the local knowledge base for relevant information."""
|
||||||
|
vectorstore = Chroma(
|
||||||
|
collection_name="knowledge",
|
||||||
|
persist_directory="./chroma_db",
|
||||||
|
)
|
||||||
|
docs = vectorstore.similarity_search(query, k=top_k)
|
||||||
|
if not docs:
|
||||||
|
return "No results found in local knowledge base."
|
||||||
|
return "\n".join(f\"{i+1}. {doc.page_content[:200]}...\" for i, doc in enumerate(docs))
|
||||||
|
|
||||||
|
@tool
|
||||||
|
def web_search(query: str) -> str:
|
||||||
|
"""Search the web for up-to-date information."""
|
||||||
|
tavily = TavilySearchResults(
|
||||||
|
api_key=os.getenv("TAVILY_API_KEY"),
|
||||||
|
max_results=3,
|
||||||
|
)
|
||||||
|
results = tavily.run(query)
|
||||||
|
if not results:
|
||||||
|
return "No results found on the web."
|
||||||
|
return "\n".join(f\"{i+1}. {res['title']}\\n{res['url']}\\n{res['content'][:200]}...\" for i, res in enumerate(results))
|
||||||
Reference in New Issue
Block a user