"""
Agent creation with routing between local KB and web search.
"""
from langchain.agents import initialize_agent, Tool
from langchain_core.prompts import ChatPromptTemplate
from langchain_ollama import OllamaLLM
from tools import search_local_kb, web_search

# Define tools
local_tool = Tool(name="search_local_kb", func=search_local_kb, description="Search local knowledge base")
web_tool = Tool(name="web_search", func=web_search, description="Web search via Tavily")

llm = OllamaLLM(model="llama3")

prompt = ChatPromptTemplate.from_messages([
    ("system", "You are an assistant that can answer questions using either local knowledge or the web. Use the appropriate tool and indicate source in your response."),
    ("human", "{input}"),
])

agent_executor = initialize_agent([local_tool, web_tool], llm, agent="zero-shot-react-description", verbose=True)

# Example usage:
if __name__ == "__main__":
    while True:
        q = input("Query> ")
        if q.lower() in {"exit", "quit"}:
            break
        res = agent_executor.run(q)
        print(res)