"""RAG agent with ChromaDB local search and Tavily web search."""
from langchain_ollama import ChatOllama
from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent
from tools import search_local_kb, web_search

llm = ChatOllama(model="llama3", temperature=0.0)

SYSTEM_PROMPT = (
    "You are a helpful AI assistant with access to two information sources:\n"
    "1. Local knowledge base (ChromaDB) - use search_local_kb for locally stored documents.\n"
    "2. Web search (Tavily) - use web_search for current or general information.\n\n"
    "Always choose the most appropriate source and indicate in your answer 
    which source you used: [Source: Local KB] or [Source: Web]."
)

agent = create_react_agent(
    model=llm,
    tools=[search_local_kb, web_search],
    state_modifier=SYSTEM_PROMPT,
)

def run_agent(user_input: str) -> str:
    """Send user_input to the ReAct agent and return its final response."""
    result = agent.invoke({"messages": [HumanMessage(content=user_input)]})
    return result["messages"][-1].content