51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
from langchain_ollama import ChatOllama
|
|
from langchain.agents import create_react_agent, AgentExecutor
|
|
from langchain_core.prompts import PromptTemplate
|
|
from tools import search_knowledge_base, add_to_knowledge_base
|
|
|
|
LLM_MODEL = "llama3"
|
|
|
|
SYSTEM_PROMPT = """You are a helpful AI assistant with access to a knowledge base.
|
|
Always use the knowledge base tools to search for relevant information before answering questions.
|
|
When you receive new information that should be remembered, add it to the knowledge base.
|
|
|
|
You have access to the following tools:
|
|
|
|
{tools}
|
|
|
|
Use the following format:
|
|
|
|
Question: the input question you must answer
|
|
Thought: you should always think about what to do
|
|
Action: the action to take, should be one of [{tool_names}]
|
|
Action Input: the input to the action
|
|
Observation: the result of the action
|
|
... (this Thought/Action/Action Input/Observation can repeat N times)
|
|
Thought: I now know the final answer
|
|
Final Answer: the final answer to the original input question
|
|
|
|
Begin!
|
|
|
|
Question: {input}
|
|
Thought:{agent_scratchpad}"""
|
|
|
|
|
|
def create_rag_agent() -> AgentExecutor:
|
|
llm = ChatOllama(model=LLM_MODEL, temperature=0)
|
|
tools = [search_knowledge_base, add_to_knowledge_base]
|
|
prompt = PromptTemplate.from_template(SYSTEM_PROMPT)
|
|
agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
|
|
agent_executor = AgentExecutor(
|
|
agent=agent,
|
|
tools=tools,
|
|
verbose=True,
|
|
handle_parsing_errors=True,
|
|
max_iterations=10,
|
|
)
|
|
return agent_executor
|
|
|
|
|
|
def run_agent(query: str) -> str:
|
|
agent = create_rag_agent()
|
|
result = agent.invoke({"input": query})
|
|
return result.get("output", "") |