Update agent.py

agent.py

@@ -1,76 +1,37 @@
-"""RAG agent implementation.
+"""Agent construction for the RAG system.
 
-This module exposes two factory functions:
+The agent uses the modern LangChain 1.x interfaces.
-
+It is built from a ChatOllama LLM and the tools defined in ``rag_tools``.
-* ``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_ollama import ChatOllama
+from langchain.agents import AgentExecutor
 
-from langchain.agents import AgentExecutor, create_openai_tools_agent
+from .rag_tools import search_knowledge_base, add_to_knowledge_base
-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]
 
+# LLM configuration
+LLM_MODEL = "llama3"
 SYSTEM_PROMPT = (
-    "You are an assistant that has access to a knowledge base.  Use the "
+    "You are an assistant that uses a local knowledge base. "
-    "provided tools to search and add information.  If you need to "
+    "When a user asks a question, first search the knowledge base "
-    "retrieve information, call the search_knowledge_base tool.  If you "
+    "with the tool 'search_knowledge_base'. If the information is not "
-    "need to store new data, call add_to_knowledge_base.  Do not "
+    "sufficient, ask clarifying questions. You can also add new "
-    "make up facts."
+    "information to the knowledge base using the tool 'add_to_knowledge_base'."
 )
+# Pass the system prompt directly to the LLM
+llm = ChatOllama(model=LLM_MODEL, system=SYSTEM_PROMPT)
 
-
+# Build the agent executor
-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.
+    """Return an AgentExecutor configured with the LLM and tools."""
+    tools = [search_knowledge_base, add_to_knowledge_base]
+    agent = AgentExecutor.from_llm_and_tools(
+        llm=llm,
+        tools=tools,
+        verbose=True,
+    )
+    return agent
 
-    The function name is kept for backward compatibility with older
+# Alias for compatibility with tests that expect `create_agent`
-    examples that expected ``create_agent_executor``.
+create_agent = 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"])