Update agent.py
This commit is contained in:
@@ -1,77 +1,76 @@
|
||||
"""
|
||||
Main entry point for the RAG agent.
|
||||
"""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.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict
|
||||
|
||||
from langchain_ollama import ChatOllama
|
||||
from langchain.agents import create_agent
|
||||
from langchain.agents.agent_toolkits import BaseToolkit
|
||||
from langchain.agents.agent_types import AgentType
|
||||
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
||||
from langchain.chat_models import ChatOpenAI
|
||||
from langchain.tools import BaseTool
|
||||
|
||||
from rag_tools import search_knowledge_base, add_to_knowledge_base
|
||||
from qdrant_store import load_directory
|
||||
# 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
|
||||
|
||||
# Load environment variables if any
|
||||
from dotenv import load_dotenv
|
||||
load_dotenv()
|
||||
# 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)
|
||||
|
||||
# Configuration
|
||||
LLM_MODEL = "llama3"
|
||||
KNOWLEDGE_DIR = os.getenv("KNOWLEDGE_DIR", "./knowledge")
|
||||
# List of tools the agent can use
|
||||
TOOLS: list[BaseTool] = [search_knowledge_base, add_to_knowledge_base]
|
||||
|
||||
# Ensure knowledge directory exists and load documents
|
||||
Path(KNOWLEDGE_DIR).mkdir(parents=True, exist_ok=True)
|
||||
load_directory(KNOWLEDGE_DIR)
|
||||
|
||||
# Define tools
|
||||
class SearchTool(BaseTool):
|
||||
name = "search_knowledge_base"
|
||||
description = "Perform semantic search in the knowledge base."
|
||||
func = search_knowledge_base
|
||||
|
||||
class AddTool(BaseTool):
|
||||
name = "add_to_knowledge_base"
|
||||
description = "Add a new document to the knowledge base."
|
||||
func = add_to_knowledge_base
|
||||
|
||||
# Simple toolkit
|
||||
class RAGToolkit(BaseToolkit):
|
||||
def get_tools(self):
|
||||
return [SearchTool(), AddTool()]
|
||||
|
||||
def get_base_prompt(self):
|
||||
return None
|
||||
|
||||
# Create LLM
|
||||
llm = ChatOllama(model=LLM_MODEL)
|
||||
|
||||
# System prompt instructing the agent to use the knowledge base
|
||||
SYSTEM_PROMPT = """
|
||||
You are an assistant that uses a knowledge base. When answering user queries, first search the knowledge base with the search_knowledge_base tool. If the information is not sufficient, ask the user for clarification. You can also add new documents to the knowledge base using add_to_knowledge_base.
|
||||
"""
|
||||
|
||||
# Create agent
|
||||
agent = create_agent(
|
||||
llm=llm,
|
||||
toolkit=RAGToolkit(),
|
||||
system_prompt=SYSTEM_PROMPT,
|
||||
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
||||
verbose=True,
|
||||
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."
|
||||
)
|
||||
|
||||
async def main():
|
||||
print("RAG Agent ready. Type your query (or 'quit' to exit).")
|
||||
while True:
|
||||
user_input = input("\n> ")
|
||||
if user_input.lower() in {"quit", "exit", "q"}:
|
||||
print("Goodbye!")
|
||||
break
|
||||
response = await agent.ainvoke(user_input)
|
||||
print("\nAssistant:", response)
|
||||
|
||||
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__":
|
||||
asyncio.run(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"])
|
||||
|
||||
Reference in New Issue
Block a user