AI Python Custom Provider Agent Skill

from collections.abc import AsyncGenerator, Sequence
from typing import Any, Literal

import pydantic
import ai


class MyProtocol(ai.ProviderProtocol[Any]):
    protocol_class_id: Literal["my_protocol"] = "my_protocol"

    def stream(
        self,
        client: Any,
        model: ai.Model,
        messages: list[ai.messages.Message],
        *,
        tools: Sequence[ai.tools.Tool] | None = None,
        output_type: type[pydantic.BaseModel] | None = None,
        params: ai.InferenceRequestParams | None = None,
        provider: str,
    ) -> AsyncGenerator[ai.events.Event]:
        return self._stream(client, model, messages, tools=tools)

    async def _stream(
        self,
        client: Any,
        model: ai.Model,
        messages: list[ai.messages.Message],
        *,
        tools: Sequence[ai.tools.Tool] | None,
    ) -> AsyncGenerator[ai.events.Event]:
        yield ai.events.StreamStart()
        yield ai.events.TextStart(block_id="text")
        yield ai.events.TextDelta(block_id="text", chunk="Hello")
        yield ai.events.TextEnd(block_id="text")
        yield ai.events.StreamEnd()


class MyProvider(ai.Provider[Any]):
    provider_class_id: Literal["my_provider"] = "my_provider"
    name: str = "my"
    default_base_url: str = "https://example.invalid"

    def __init__(self, *, client: Any) -> None:
        super().__init__()
        self._set_client(client)

    def default_protocol(self) -> ai.ProviderProtocol[Any]:
        return MyProtocol()

    async def list_models(self) -> list[str]:
        return ["my-model"]

    async def probe(self, model: ai.Model) -> None:
        return None


model = ai.Model(id="my-model", provider=MyProvider(client=client))

yield ai.events.ToolStart(tool_call_id=tcid, tool_name=name)
yield ai.events.ToolDelta(tool_call_id=tcid, chunk=args_json)
yield ai.events.ToolEnd(
    tool_call_id=tcid,
    tool_call=ai.messages.DUMMY_TOOL_CALL,
)