Update src/agent.py

src/agent.py

@@ -1,54 +1,53 @@
-"""Agent implementation using LangChain tool‑calling.
+"""Agent that uses the RAG tools.
 
-The agent is built with the modern LangChain interface. It uses a system prompt that
-encourages the assistant to search the knowledge base before answering.
+The agent is built with ``create_agent`` from ``langchain.agents`` and is
+configured to use the local ``ChatOllama`` model (``llama3``).  It has two
+tools: ``search_knowledge_base`` and ``add_to_knowledge_base``.
 """
 
-from typing import List
+from __future__ import annotations
 
-from langchain_core.prompts import ChatPromptTemplate
-from langchain.agents import AgentExecutor, create_tool_calling_agent
 from langchain_ollama import ChatOllama
+from langchain.agents import create_agent
 
+# Import the tools from the package.
 from .tools import search_knowledge_base, add_to_knowledge_base
 
-# Create the LLM that will power the agent. We use Ollama's local llama3 model.
+# Instantiate the LLM.
 llm = ChatOllama(model="llama3", temperature=0.7)
 
-# Register tools
-TOOLS: List = [search_knowledge_base, add_to_knowledge_base]
+# Create the agent.
+agent = create_agent(
+    model=llm,
+    tools=[search_knowledge_base, add_to_knowledge_base],
+    verbose=True,
+)
 
-# Prompt template that tells the agent to use the tools when needed.
-prompt = ChatPromptTemplate.from_messages([
-    ("system", "You are a helpful assistant. Use the provided tools to answer the user."),
-    ("placeholder", "{chat_history}"),
-    ("human", "{input}"),
-    ("placeholder", "{agent_scratchpad}"),
-])
+# Helper function to run a user query.
 
-# Create the tool‑calling agent
-agent = create_tool_calling_agent(llm, TOOLS, prompt)
-
-# Wrap it in an executor that will handle the conversation loop
-agent_executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
-
-# Expose a simple function that can be called from a CLI or other entry point
-
-def run_agent(messages: List[dict]):
-    """Run the agent on a list of chat messages.
+def run_query(user_input: str) -> str:
+    """Invoke the agent with a user message and return the response.
 
     Parameters
     ----------
-    messages: List[dict]
-        Each message is a dict with keys ``role`` ("user" or "assistant") and ``content``.
+    user_input: str
+        The raw user message.
 
     Returns
     -------
     str
-        The assistant's reply.
+        The agent's reply.
     """
-    # Convert the list of messages into the format expected by the executor
-    chat_history = [(msg["role"], msg["content"]) for msg in messages]
-    input_text = messages[-1]["content"] if messages else ""
-    result = agent_executor.invoke({"input": input_text, "chat_history": chat_history})
-    return result["output"]
+    result = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
+    # The result is a dict with a list of messages. The last message is the assistant's reply.
+    messages = result.get("messages", [])
+    if not messages:
+        return "No response."
+    # Find the last assistant message.
+    for msg in reversed(messages):
+        if msg.get("role") == "assistant":
+            return msg.get("content", "")
+    # Fallback to the first message.
+    return messages[-1].get("content", "")
+
+# End of src/agent.py