add agent.py

agent.py

@@ -0,0 +1,26 @@
+"""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