27 lines
1.1 KiB
Python
27 lines
1.1 KiB
Python
"""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
|