1 files changed, 15 insertions, 4 deletions
diff --git a/main.py b/main.py
index 5a70788..eebc89d 100644
--- a/main.py
+++ b/main.py
@@ -1,7 +1,8 @@
+from typing import Any
 from langchain_ollama import ChatOllama
 from langchain.agents import create_agent
-from langchain_community.tools.file_management import FileSearchTool
-
+from langchain_community.tools.file_management import ListDirectoryTool
+from langchain.agents.middleware import after_model, AgentState
 
 GEMMA3 = "gemma3"
 LLAMA3 = "llama3.1:8b"
@@ -14,12 +15,22 @@ model = ChatOllama(
     num_predict=256,
 )
 
+
+@after_model
+def log_response(state: AgentState, runtime) -> dict[str, Any] | None:
+    print(f"Model returned: {state['messages'][-1].content}")
+    return None
+
+
 agent = create_agent(
     model,
     tools=[
-        FileSearchTool(),
+        ListDirectoryTool(),
     ],
     system_prompt="You are a helpful assistent.",
+    middleware=[
+        log_response,
+    ],
 )
 
 result = agent.invoke(
@@ -27,7 +38,7 @@ result = agent.invoke(
         "messages": [
             {
                 "role": "user",
-                "content": "What files match *.py in .?",
+                "content": "What files are in .?",
             }
         ]
     }