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
+The agent is built with ``create_agent`` from ``langchain.agents`` and is
-encourages the assistant to search the knowledge base before answering.
+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
+# Create the agent.
-TOOLS: List = [search_knowledge_base, add_to_knowledge_base]
+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.
+# Helper function to run a user query.
-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}"),
-])
 
-# Create the tool‑calling agent
+def run_query(user_input: str) -> str:
-agent = create_tool_calling_agent(llm, TOOLS, prompt)
+    """Invoke the agent with a user message and return the response.
-
-# 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.
 
     Parameters
     ----------
-    messages: List[dict]
+    user_input: str
-        Each message is a dict with keys ``role`` ("user" or "assistant") and ``content``.
+        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
+    result = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
-    chat_history = [(msg["role"], msg["content"]) for msg in messages]
+    # The result is a dict with a list of messages. The last message is the assistant's reply.
-    input_text = messages[-1]["content"] if messages else ""
+    messages = result.get("messages", [])
-    result = agent_executor.invoke({"input": input_text, "chat_history": chat_history})
+    if not messages:
-    return result["output"]
+        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