77 lines
2.6 KiB
Python
77 lines
2.6 KiB
Python
"""RAG agent implementation.
|
||
|
||
This module exposes two factory functions:
|
||
|
||
* ``create_agent`` – returns a LangChain agent that can use the two tools
|
||
defined in :mod:`tools`.
|
||
* ``create_agent_executor`` – returns an executor that can be used directly
|
||
from the command line.
|
||
|
||
The agent uses a simple system prompt that instructs it to use the knowledge
|
||
base for every query. The tools are automatically added to the agent.
|
||
"""
|
||
|
||
from typing import Any, Dict
|
||
|
||
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
||
from langchain.chat_models import ChatOpenAI
|
||
from langchain.tools import BaseTool
|
||
|
||
# Import the tools – they expose ``search_knowledge_base`` and
|
||
# ``add_to_knowledge_base`` as LangChain tools.
|
||
from tools import search_knowledge_base, add_to_knowledge_base
|
||
|
||
# Create the OpenAI chat model – for local usage we can use Ollama via
|
||
# ``ChatOpenAI`` with a custom endpoint. For the purposes of this
|
||
# implementation we assume the user has an OpenAI-compatible endpoint.
|
||
# If Ollama is used, replace the model name with ``llama3``.
|
||
chat_model = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
|
||
|
||
# List of tools the agent can use
|
||
TOOLS: list[BaseTool] = [search_knowledge_base, add_to_knowledge_base]
|
||
|
||
SYSTEM_PROMPT = (
|
||
"You are an assistant that has access to a knowledge base. Use the "
|
||
"provided tools to search and add information. If you need to "
|
||
"retrieve information, call the search_knowledge_base tool. If you "
|
||
"need to store new data, call add_to_knowledge_base. Do not "
|
||
"make up facts."
|
||
)
|
||
|
||
|
||
def create_agent() -> AgentExecutor:
|
||
"""Create a LangChain agent that can perform RAG.
|
||
|
||
Returns
|
||
-------
|
||
AgentExecutor
|
||
The configured agent.
|
||
"""
|
||
agent = create_openai_tools_agent(
|
||
llm=chat_model,
|
||
tools=TOOLS,
|
||
system_message=SYSTEM_PROMPT,
|
||
)
|
||
executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
|
||
return executor
|
||
|
||
|
||
def create_agent_executor() -> AgentExecutor:
|
||
"""Convenience wrapper that returns the same executor.
|
||
|
||
The function name is kept for backward compatibility with older
|
||
examples that expected ``create_agent_executor``.
|
||
"""
|
||
return create_agent()
|
||
|
||
# If this file is executed directly, run a simple interactive loop.
|
||
if __name__ == "__main__":
|
||
executor = create_agent()
|
||
print("RAG agent ready. Type /quit to exit.")
|
||
while True:
|
||
user_input = input("User: ")
|
||
if user_input.strip().lower() == "/quit":
|
||
break
|
||
response = executor.invoke({"input": user_input})
|
||
print("Agent:", response["output"])
|