from smolagents import LiteLLMModel, ToolCallingAgent


def main() -> None:
    model = LiteLLMModel(
        model_id="ollama_chat/qwen2:7b",
        api_base="http://127.0.0.1:11434",
        temperature=0.3,
    )

    # ToolCallingAgent is more reliable than CodeAgent for general chat-like prompts
    # with local Ollama models, because it avoids strict code-block parsing.
    agent = ToolCallingAgent(
        tools=[],
        model=model,
        max_steps=3,
    )

    prompt = "Explain how AI agents work in 3 bullet points."

    try:
        answer = agent.run(prompt)
        print(answer)
    except Exception:
        # Some Ollama models intermittently skip tool-call JSON.
        # Fallback to direct model generation while still using smolagents.
        response = model.generate(
            messages=[
                {"role": "system", "content": "You are a concise and helpful assistant."},
                {"role": "user", "content": prompt},
            ]
        )
        print(response.content)


if __name__ == "__main__":
    main()