"""
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 [])
    ]