Add src/agent.py

src/agent.py

@@ -0,0 +1,54 @@
+"""Agent implementation using LangChain tool‑calling.
+
+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.
+"""
+
+from typing import List
+
+from langchain_core.prompts import ChatPromptTemplate
+from langchain.agents import AgentExecutor, create_tool_calling_agent
+from langchain_ollama import ChatOllama
+
+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.
+llm = ChatOllama(model="llama3", temperature=0.7)
+
+# Register tools
+TOOLS: List = [search_knowledge_base, add_to_knowledge_base]
+
+# 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}"),
+])
+
+# 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.
+
+    Parameters
+    ----------
+    messages: List[dict]
+        Each message is a dict with keys ``role`` ("user" or "assistant") and ``content``.
+
+    Returns
+    -------
+    str
+        The assistant'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"]