diff --git a/agent.py b/agent.py index 73ba556..93f1780 100644 --- a/agent.py +++ b/agent.py @@ -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"])