Update src/agent.py

src/agent.py

@@ -1,32 +1,35 @@
-"""Agent that uses the RAG tools.
+"""Agent construction for the RAG system.
 
-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``.
+This module builds a LangChain agent that uses a local knowledge base.
+The agent is created using :func:`langchain.agents.create_agent`.
 """
 
-from __future__ import annotations
-
 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
 
-# Instantiate the LLM.
-llm = ChatOllama(model="llama3", temperature=0.7)
-
-# Create the agent.
-agent = create_agent(
-    model=llm,
-    tools=[search_knowledge_base, add_to_knowledge_base],
-    verbose=True,
+# LLM configuration
+LLM_MODEL = "llama3"
+SYSTEM_PROMPT = (
+    "You are an assistant that uses a local knowledge base. "
+    "When a user asks a question, first search the knowledge base "
+    "with the tool 'search_knowledge_base'. If the information is not "
+    "sufficient, ask clarifying questions. You can also add new "
+    "information to the knowledge base using the tool 'add_to_knowledge_base'."
 )
 
-# Helper function to run a user query.
+# Create the LLM with system prompt
+llm = ChatOllama(model=LLM_MODEL, system=SYSTEM_PROMPT)
+
+# Build the agent
+_tools = [search_knowledge_base, add_to_knowledge_base]
+agent = create_agent(llm=llm, tools=_tools, verbose=True)
+
+# Helper function to run a query through the agent
 
 def run_query(user_input: str) -> str:
-    """Invoke the agent with a user message and return the response.
+    """Run a user query through the agent and return the response.
 
     Parameters
     ----------
@@ -36,18 +39,14 @@ def run_query(user_input: str) -> str:
     Returns
     -------
     str
-        The agent's reply.
+        The agent's response.
     """
-    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", "")
+    # The agent expects a dict with an "input" key
+    result = agent.run({"input": user_input})
+    # The result may be a string or a dict; convert to string
+    if isinstance(result, dict):
+        return str(result)
+    return result
 
-# End of src/agent.py
+# Expose the agent for external use
+__all__ = ["agent", "run_query"]