41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
"""
|
|
Agent tools for local KB search and web search via Tavily.
|
|
"""
|
|
|
|
from typing import List, Dict
|
|
|
|
from langchain.agents import tool
|
|
from langchain_tavily import TavilySearch
|
|
from langchain.schema.document import Document
|
|
|
|
# Local KB search tool
|
|
@tool("search_local_kb")
|
|
def search_local_kb(query: str, top_k: int = 3) -> List[Dict]:
|
|
"""Semantic search in the local ChromaDB collection.
|
|
|
|
Returns a list of dicts with keys ``text`` and ``source``.
|
|
"""
|
|
from vectorstore import create_vectorstore
|
|
|
|
# Assume the collection is already created and persisted
|
|
store = create_vectorstore()
|
|
results = store.query(query_texts=[query], n_results=top_k)
|
|
return [
|
|
{"text": doc.page_content, "source": doc.metadata.get("source", "unknown")}
|
|
for doc in results
|
|
]
|
|
|
|
# Web search tool using Tavily.
|
|
@tool("web_search")
|
|
def web_search(query: str) -> List[Dict]:
|
|
"""Search the web with Tavily and return a list of result snippets.
|
|
|
|
Requires environment variable TAVILY_API_KEY.
|
|
"""
|
|
tav = TavilySearch()
|
|
results = tav.search(query, max_results=5)
|
|
return [
|
|
{"text": r["title"] + ": " + r["snippet"] if isinstance(r, dict) else str(r)}
|
|
for r in (results if isinstance(results, list) else [])
|
|
]
|