Pydantic AI Compatibility

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

import os
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,
    ),
)

from pydantic_ai import DocumentUrl, VideoUrl

def use_vlmrun_media_parts(model: OpenAIChatModel) -> None:
    """Send Pydantic AI media inputs as VLM Run content parts."""

    async def _document(item: DocumentUrl):
        return {"type": "file_url", "file_url": {"url": item.url, "detail": "auto"}}

    async def _video(item: VideoUrl):
        return {"type": "video_url", "video_url": {"url": item.url, "detail": "auto"}}

    model._map_document_url_item = _document
    model._map_video_url_item = _video

use_vlmrun_media_parts(model)

from pydantic_ai import Agent, RunContext

agent = Agent(
    model,
    system_prompt=(
        "You are a document and video analyst. Use the server execute_code tool "
        "for anything involving the attached document, image, or video, and report "
        "only values it returns. Call lookup_purchase_order exactly once only when "
        "the user explicitly asks you to reconcile an invoice against the ERP. If "
        "its result is already in the conversation, answer from that result and do "
        "not call it again. 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."""
    # Replace with a real query against your system of record.
    return {"invoice_number": invoice_number, "po_number": "PO-4417", "total": 4560.00}

import uuid

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

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

result = agent.run_sync(
    [
        "Extract the vendor, invoice number, invoice date, and total from this "
        "invoice. Use execute_code with the attached document ref.",
        DocumentUrl(INVOICE_PDF),
    ],
    **hybrid_run_kwargs(),
)

print(result.output)

import asyncio

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

output = asyncio.run(
    run_streamed(
        agent,
        [
            "Summarize this video segment by segment. For each segment, give the start "
            "and end timestamp and one sentence describing what happens.",
            VideoUrl(BAKERY_VIDEO),
        ],
        hybrid_run_kwargs(),
    )
)

result = agent.run_sync(
    [
        "Extract the invoice number and total from this invoice with execute_code. "
        "Then call lookup_purchase_order exactly once with that invoice number. "
        "Compare the two totals and say whether they match. Do not call the tool again.",
        DocumentUrl(INVOICE_PDF),
    ],
    **hybrid_run_kwargs(),
)

from pydantic_ai import ImageUrl

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

result = agent.run_sync(
    [
        "Detect every donut in this image, draw bounding boxes, and return the "
        "count along with the annotated image.",
        ImageUrl(DONUTS_IMAGE),
    ],
    **hybrid_run_kwargs(),
)

import asyncio
from pydantic_ai.usage import UsageLimits

async def run_streamed(agent: Agent, prompt, run_kwargs: dict) -> 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 request_stream:
                    async for _ in request_stream:
                        pass
        return run.result.output

output = asyncio.run(
    run_streamed(
        agent,
        ["Summarize this invoice.", DocumentUrl(INVOICE_PDF)],
        hybrid_run_kwargs(),
    )
)