diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..c1e310b --- /dev/null +++ b/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