From 7d8a7769f5e5b4ce20108a9ab307498f52ea2f85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 28 May 2026 16:43:27 +0000 Subject: [PATCH] add agent.py --- agent.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 agent.py 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