from langchain_ollama import ChatOllama
from langchain.agents import create_tool_calling_agent
from langchain.agents import AgentExecutor
from tools import search_local_kb, web_search


def create_agent():
    """Create an agent with tools for local and web search."""
    llm = ChatOllama(model="llama3", temperature=0)
    
    tools = [search_local_kb, web_search]
    
    system_prompt = """You are a helpful AI assistant with access to two search tools:

1. search_local_kb - Search in the local knowledge base (ChromaDB) for information from stored documents
2. web_search - Search the web via Tavily for current information and recent news

Choose the appropriate tool based on the user's question:
- For questions about local documents, notes, or stored knowledge: use search_local_kb
- For questions about current events, recent news, or facts that may not be in local documents: use web_search

Always indicate the source of your answer in the response: 'Source: chromadb' or 'Source: tavily'
"""
    
    agent = create_tool_calling_agent(llm, tools, system_prompt)
    agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    
    return agent_executor