Pydantic AI Compatibility

pip install "pydantic-ai>=1.100"
export VLMRUN_API_KEY="<VLMRUN_API_KEY>"

import os
import uuid

from pydantic_ai import Agent, AudioUrl, DocumentUrl, ImageUrl, RunContext, VideoUrl
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.profiles.openai import OpenAIModelProfile
from pydantic_ai.providers.openai import OpenAIProvider

model = OpenAIChatModel(
    "vlmrun-orion-2:auto",
    provider=OpenAIProvider(
        base_url="https://api.vlm.run/v1/openai",
        api_key=os.environ["VLMRUN_API_KEY"],
    ),
    profile=OpenAIModelProfile(
        openai_supports_strict_tool_definition=False,
        openai_unsupported_model_settings=("max_completion_tokens",),
        openai_chat_supports_document_input=True,
        openai_chat_supports_file_urls=True,
        openai_chat_audio_input_encoding="uri",
    ),
)

# Stock OpenAIChatModel rejects VideoUrl; map it to VLM's video_url part.
async def _map_video_url(item: VideoUrl):
    return {"type": "video_url", "video_url": {"url": item.url}}

model._map_video_url_item = _map_video_url

agent = Agent(
    model,
    system_prompt=(
        "Use the server execute_code tool for attached documents, images, "
        "audio, or video. Call lookup_purchase_order only when asked to "
        "reconcile an invoice. Never invent field values."
    ),
)

@agent.tool
def lookup_purchase_order(ctx: RunContext[None], invoice_number: str) -> dict:
    """Return the purchase order your ERP has on file for an invoice number."""
    return {"invoice_number": invoice_number, "po_number": "PO-4417", "total": 4560.00}

def run_kwargs() -> dict:
    return {
        "model_settings": {
            "extra_body": {
                "session_id": str(uuid.uuid4()),
                "tool_execution": "hybrid",
            }
        }
    }

INVOICE = (
    "https://storage.googleapis.com/vlm-data-public-prod/hub/examples/"
    "document.invoice/wordpress-pdf-invoice-plugin-sample.pdf"
)

result = agent.run_sync(
    ["Extract vendor, invoice number, date, and total.", DocumentUrl(INVOICE)],
    **run_kwargs(),
)
print(result.output)

DONUTS = (
    "https://storage.googleapis.com/vlm-data-public-prod/hub/examples/"
    "image.object-detection/donuts.png"
)

result = agent.run_sync(
    ["Count the donuts and briefly describe the image.", ImageUrl(DONUTS)],
    **run_kwargs(),
)
print(result.output)

AUDIO = (
    "https://storage.googleapis.com/vlm-data-public-prod/hub/examples/"
    "audio.transcription-summary/two_minute_rules.mp3"
)

result = agent.run_sync(
    ["Summarize this audio in 2 sentences.", AudioUrl(AUDIO)],
    **run_kwargs(),
)
print(result.output)

import asyncio
from pydantic_ai.usage import UsageLimits

VIDEO = (
    "https://storage.googleapis.com/vlm-data-public-prod/hub/examples/"
    "video.transcription/bakery.mp4"
)

async def run_streamed(prompt) -> str:
    async with agent.iter(
        prompt,
        usage_limits=UsageLimits(request_limit=20),
        **run_kwargs(),
    ) as run:
        async for node in run:
            if Agent.is_model_request_node(node):
                async with node.stream(run.ctx) as stream:
                    async for _ in stream:
                        pass
        return run.result.output

print(asyncio.run(run_streamed(["Describe this video in 2 sentences.", VideoUrl(VIDEO)])))

result = agent.run_sync(
    [
        "Extract the invoice number and total with execute_code. "
        "Then call lookup_purchase_order once and say if the totals match.",
        DocumentUrl(INVOICE),
    ],
    **run_kwargs(),
)
print(result.output)