08_mcp_sse_pydantic.py

import asyncio
import os
from typing import Any, Dict, List

from pydantic_ai import Agent
from pydantic_ai.agent import AgentRunResult
from pydantic_ai.mcp import MCPServerHTTP
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")
)

MCP_URL = os.environ.get("MCP_URL", "http://localhost:8000/sse")
dundie_mcp = MCPServerHTTP(MCP_URL)


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=[sum_transactions], mcp_servers=[dundie_mcp])


async def invoke_agent(prompt: str) -> AgentRunResult:
    return await agent.run(prompt)


async def main():
    async with agent.run_mcp_servers():
        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 = await invoke_agent(prompt)
            print(result)


if __name__ == "__main__":
    asyncio.run(main())