Update agent.py

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"])