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))