Add ai.ModelLike

examples/apps/durable_agent_temporal/backend/app/worker.py

@@ -109,6 +109,8 @@ class DurableAgent(ai.Agent):
     async def loop(self, context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]:
         tools_data = [tool.model_dump(mode="json") for tool in context.tools]
 
+        # A durable step serializes the model: it must be an ai.Model.
+        assert isinstance(context.model, ai.Model)
         while context.keep_running():
             result = await temporalio.workflow.execute_activity(
                 llm_activity,

examples/apps/durable_agent_workflows/backend/app/worker.py

@@ -92,6 +92,8 @@ class DurableAgent(ai.Agent):
     async def loop(self, context: ai.Context) -> AsyncGenerator[ai.events.AgentEvent]:
         tools_data = [tool.model_dump(mode="json") for tool in context.tools]
 
+        # A durable step serializes the model: it must be an ai.Model.
+        assert isinstance(context.model, ai.Model)
         while context.keep_running():
             result = await llm_step(
                 context.model.model_dump(mode="json"),

src/ai/__init__.py

@@ -74,6 +74,7 @@
     InferenceRequestParams,
     MinPSamplerParams,
     Model,
+    ModelLike,
     ModelProviderDefault,
     OutputParams,
     Provider,
@@ -138,6 +139,7 @@
     "InstallationError",
     "MinPSamplerParams",
     "Model",
+    "ModelLike",
     "ModelProviderDefault",
     "OutputParams",
     "Provider",

src/ai/agents/agent.py

@@ -1039,10 +1039,29 @@ async def _iterate(self) -> AsyncGenerator[events_.ToolCallResult]:
             yield res
 
 
+def _validate_model(value: Any) -> models.ModelLike:
+    # Another model-like is kept as it is; a dict from JSON is an ai.Model.
+    if isinstance(value, models.ModelLike) and not isinstance(
+        value, models.Model
+    ):
+        return value
+    return models.Model.model_validate(value)
+
+
+def _serialize_model(value: Any, info: pydantic.SerializationInfo) -> Any:
+    if isinstance(value, pydantic.BaseModel):
+        return value.model_dump(mode=info.mode)
+    return {"id": value.id, "provider_name": value.provider_name}
+
+
 class Context(pydantic.BaseModel):
     """Everything that goes into the LLM."""
 
-    model: models.Model
+    model: Annotated[
+        models.ModelLike,
+        pydantic.PlainValidator(_validate_model),
+        pydantic.PlainSerializer(_serialize_model),
+    ]
     messages: list[types.messages.Message]
     tools: list[Tool]
     output_type: type[pydantic.BaseModel] | None = pydantic.Field(
@@ -1493,7 +1512,7 @@ async def loop(
     @overload
     def run(
         self,
-        model: models.Model,
+        model: models.ModelLike,
         messages: list[types.messages.Message],
         *,
         params: models.InferenceRequestParams | None = None,
@@ -1502,7 +1521,7 @@ def run(
     @overload
     def run[T: pydantic.BaseModel](
         self,
-        model: models.Model,
+        model: models.ModelLike,
         messages: list[types.messages.Message],
         *,
         output_type: type[T],
@@ -1511,7 +1530,7 @@ def run[T: pydantic.BaseModel](
     ) -> AbstractAsyncContextManager[AgentStream[T]]: ...
     def run(
         self,
-        model: models.Model,
+        model: models.ModelLike,
         messages: list[types.messages.Message],
         *,
         output_type: type[pydantic.BaseModel] | None = None,
@@ -1556,7 +1575,7 @@ def run(
     @contextlib.asynccontextmanager
     async def _run(
         self,
-        model: models.Model,
+        model: models.ModelLike,
         messages: list[types.messages.Message],
         *,
         output_type: type[pydantic.BaseModel] | None,
@@ -1610,7 +1629,7 @@ async def _stream() -> AsyncGenerator[events_.AgentEvent]:
                     agent=type(self).__name__,
                     model=model.id,
                     messages=list(messages),
-                    provider=model.provider.name,
+                    provider=model.provider_name,
                     tool_names=[t.name for t in context.tools],
                     output_type=(
                         output_type.__name__

src/ai/models/__init__.py

@@ -39,7 +39,7 @@
     probe,
     stream,
 )
-from .core.model import Model, get_model
+from .core.model import Model, ModelLike, get_model
 from .core.params import (
     DEFAULT,
     GLOBAL,
@@ -85,6 +85,7 @@
     "InferenceRequestParams",
     "MinPSamplerParams",
     "Model",
+    "ModelLike",
     "ModelProviderDefault",
     "OutputParams",
     "Provider",

src/ai/models/core/api.py

@@ -22,11 +22,11 @@
 
 from ... import errors, types
 from ... import experimental_telemetry as telemetry
+from . import model as model_
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator, AsyncIterator, Sequence
 
-    from . import model as model_
     from . import params as params_
 
 # Stream output type.  Defaults to ``str``: when the stream was opened
@@ -37,7 +37,7 @@
 
 @dataclasses.dataclass(frozen=True)
 class _StreamRequest:
-    model: model_.Model
+    model: model_.ModelLike
     messages: list[types.messages.Message]
     tools: Sequence[types.tools.Tool] | None = None
     output_type: type[pydantic.BaseModel] | None = None
@@ -66,18 +66,26 @@ async def _do_stream(
         self,
         request: _StreamRequest,
     ) -> AsyncGenerator[types.events.Event]:
-        async for ev in request.model.provider.stream(
-            request.model,
-            request.messages,
-            tools=request.tools,
-            output_type=request.output_type,
-            params=request.params,
-        ):
-            yield ev
+        # Closed with the outer generator, not left to the garbage
+        # collector: a model-like with state (a worker) cleans up in it.
+        async with contextlib.aclosing(
+            request.model.stream(
+                request.messages,
+                tools=request.tools,
+                output_type=request.output_type,
+                params=request.params,
+            )
+        ) as inner:
+            async for ev in inner:
+                yield ev
 
     async def _do_generate(
         self, request: _StreamRequest
     ) -> types.messages.Message:
+        # Only a provider-backed model has a native generate; anything else
+        # that is model-like is drained from its stream by the caller.
+        if not isinstance(request.model, model_.Model):
+            raise NotImplementedError
         return await request.model.provider.generate(
             request.model,
             request.messages,
@@ -307,7 +315,7 @@ class StreamContext(Protocol):
     """
 
     @property
-    def model(self) -> model_.Model: ...
+    def model(self) -> model_.ModelLike: ...
     @property
     def messages(self) -> list[types.messages.Message]: ...
     @property
@@ -333,23 +341,23 @@ def stream[T: pydantic.BaseModel](
 ) -> AbstractAsyncContextManager[Stream[T]]: ...
 @overload
 def stream(
-    model: model_.Model,
+    model: model_.ModelLike,
     messages: list[types.messages.Message],
     *,
     tools: Sequence[types.tools.Tool] | None = None,
     params: params_.InferenceRequestParams | None = None,
 ) -> AbstractAsyncContextManager[Stream[str]]: ...
 @overload
 def stream[T: pydantic.BaseModel](
-    model: model_.Model,
+    model: model_.ModelLike,
     messages: list[types.messages.Message],
     *,
     tools: Sequence[types.tools.Tool] | None = None,
     output_type: type[T],
     params: params_.InferenceRequestParams | None = None,
 ) -> AbstractAsyncContextManager[Stream[T]]: ...
 def stream(
-    model: model_.Model | None = None,
+    model: model_.ModelLike | None = None,
     messages: list[types.messages.Message] | None = None,
     *,
     context: StreamContext | None = None,
@@ -401,7 +409,7 @@ def stream(
 @contextlib.asynccontextmanager
 async def _stream(
     *,
-    model: model_.Model,
+    model: model_.ModelLike,
     messages: list[types.messages.Message],
     tools: Sequence[types.tools.Tool] | None,
     output_type: type[pydantic.BaseModel] | None,
@@ -412,7 +420,7 @@ async def _stream(
         model=model.id,
         messages=list(messages),
         params=params,
-        provider=model.provider.name,
+        provider=model.provider_name,
         tool_names=[t.name for t in tools] if tools is not None else None,
         output_type=output_type.__name__ if output_type is not None else None,
     )
@@ -467,15 +475,15 @@ async def experimental_generate(
 ) -> types.messages.Message: ...
 @overload
 async def experimental_generate(
-    model: model_.Model,
+    model: model_.ModelLike,
     messages: list[types.messages.Message],
     *,
     tools: Sequence[types.tools.Tool] | None = None,
     output_type: type[pydantic.BaseModel] | None = None,
     params: params_.InferenceRequestParams | None = None,
 ) -> types.messages.Message: ...
 async def experimental_generate(
-    model: model_.Model | None = None,
+    model: model_.ModelLike | None = None,
     messages: list[types.messages.Message] | None = None,
     *,
     context: StreamContext | None = None,
@@ -525,7 +533,7 @@ async def experimental_generate(
         model=model.id,
         messages=list(messages),
         params=params,
-        provider=model.provider.name,
+        provider=model.provider_name,
         tool_names=[t.name for t in tools] if tools is not None else None,
         output_type=output_type.__name__ if output_type is not None else None,
     )

src/ai/models/core/model.py

@@ -1,13 +1,25 @@
 """Model metadata types."""
 
+from __future__ import annotations
+
 import os
 import weakref
-from typing import TYPE_CHECKING, Any, Self, cast
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Protocol,
+    Self,
+    cast,
+    runtime_checkable,
+)
 
 import pydantic
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import AsyncGenerator, Callable, Sequence
+
+    from ... import types
+    from . import params as params_
 
 from ... import _modelsdev
 from ...errors import ConfigurationError
@@ -16,6 +28,36 @@
 _DEFAULT_MODEL_ENV = "AI_SDK_DEFAULT_MODEL"
 
 
+@runtime_checkable
+class ModelLike(Protocol):
+    """Anything that answers a conversation the way a model does.
+
+    `ai.stream`, `ai.experimental_generate` and `ai.Agent.run` take a
+    `ModelLike`. `ai.Model` is one; anything with the same `stream` is
+    another, a coding-agent CLI for example.
+
+    `stream` receives the WHOLE conversation on every call and yields the
+    events of ONE assistant message, ending with `StreamEnd`.
+    """
+
+    @property
+    def id(self) -> str: ...
+
+    @property
+    def provider_name(self) -> str:
+        """Who answers: a provider name, or a worker kind. For telemetry."""
+        ...
+
+    def stream(
+        self,
+        messages: list[types.messages.Message],
+        *,
+        tools: Sequence[types.tools.Tool] | None = None,
+        output_type: type[pydantic.BaseModel] | None = None,
+        params: params_.InferenceRequestParams | None = None,
+    ) -> AsyncGenerator[types.events.Event]: ...
+
+
 def _exclude_if_none(v: Any) -> bool:
     # pydantic-core's SchemaSerializer doesn't GC-traverse ``exclude_if``
     # callables, so a closure stored there directly forms an uncollectable
@@ -50,6 +92,27 @@ def __hash__(self) -> int:
     def with_protocol(self, protocol: base.ProviderProtocol[Any]) -> Self:
         return self.model_copy(update={"protocol": protocol})
 
+    @property
+    def provider_name(self) -> str:
+        return self.provider.name
+
+    def stream(
+        self,
+        messages: list[types.messages.Message],
+        *,
+        tools: Sequence[types.tools.Tool] | None = None,
+        output_type: type[pydantic.BaseModel] | None = None,
+        params: params_.InferenceRequestParams | None = None,
+    ) -> AsyncGenerator[types.events.Event]:
+        """Return the raw provider stream. Use `ai.stream(model, ...)`."""
+        return self.provider.stream(
+            self,
+            messages,
+            tools=tools,
+            output_type=output_type,
+            params=params,
+        )
+
 
 def get_model(
     model_id: str | None = None,

tests/models/core/test_api.py

@@ -67,6 +67,7 @@ def test_context_json_roundtrip_includes_model() -> None:
 
     restored = ai.Context.model_validate_json(context.model_dump_json())
 
+    assert isinstance(restored.model, ai.Model)
     assert restored.model.id == "anthropic/claude-sonnet-4.6"
     assert isinstance(restored.model.provider, ai.providers.GatewayProvider)