agentic calculator 😂

ai_calc.py

from pydantic import BaseModel
from pydantic_ai import Agent
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.ollama import OllamaProvider

ollama_model = OpenAIChatModel(
    model_name='llama3.2:latest',
    provider=OllamaProvider(base_url='http://localhost:11434/v1'),
)

class SumModel(BaseModel):
    a: int
    b: int

class SumResultModel(BaseModel):
    result: int

def compute_sum(input=BaseModel):
    a = input.a
    b = input.b
    agent = Agent(ollama_model, output_type=SumResultModel)
    result = agent.run_sync(f"Sum of {a} + {b}")
    return result.output.result

def test_sum():
    a, b = 3, 4
    assert compute_sum(SumModel(a=a, b=b)) == a + b

if __name__ == "__main__":
    import pytest
    pytest.main([__file__, "--disable-warnings"])