06_agent_pydantic.py

import json
import os
import subprocess
from typing import Any, Dict, List

from pydantic_ai import Agent
from pydantic_ai.agent import AgentRunResult
from pydantic_ai.models.openai import OpenAIModel
from pydantic_ai.providers.openai import OpenAIProvider

# Set up Ollama LLM
OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
model = os.environ.get("MODEL", "qwen3:4b")
model = OpenAIModel(
    model_name=model, provider=OpenAIProvider(base_url=f"{OLLAMA_URL}/v1")
)


def list_transactions(email: str) -> list[dict]:
    """List user transaction by email"""
    try:
        result = subprocess.run(
            ["dundie", "list", "--email", email, "--asjson"],
            capture_output=True,
            text=True,
            check=True,
        )
        return json.loads(result.stdout.strip())
    except subprocess.CalledProcessError as e:
        print(f"Error listing transactions: {e}")
        return []


def sum_transactions(transactions: List[Dict[str, Any]]) -> float:
    """
    Sum the amounts of a list of transactions.
    """
    return sum(t["value"] for t in transactions)


agent = Agent(
    model,
    tools=[list_transactions, sum_transactions],
)


def invoke_agent(prompt: str) -> AgentRunResult:
    return agent.run_sync(prompt)


def main():
    while True:
        prompt = input("Enter your prompt: ")
        if not prompt:
            print("Prompt cannot be empty.")
            continue
        if prompt.strip() in ["exit", "quit", "q"]:
            print("Exiting...")
            break
        result = invoke_agent(prompt)
        print("-" * 50)
        print(result)


if __name__ == "__main__":
    main()