from typing import Any
from langchain_ollama import ChatOllama
from langchain.agents import create_agent
from langchain_community.tools.file_management import ListDirectoryTool
from langchain.agents.middleware import after_model, AgentState

GEMMA3 = "gemma3"
LLAMA3 = "llama3.1:8b"


model = ChatOllama(
    model=LLAMA3,
    validate_model_on_init=True,
    temperature=0.8,
    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=[
        ListDirectoryTool(),
    ],
    system_prompt="You are a helpful assistent.",
    middleware=[
        log_response,
    ],
)

result = agent.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What files are in .?",
            }
        ]
    }
)

for message in result["messages"]:
    print(message.content)