From e299fef931f82b32cd871717676d25de75767775 Mon Sep 17 00:00:00 2001 From: Andrey Buzin Date: Thu, 5 Nov 2026 19:02:42 -0800 Subject: [PATCH] Add ai.ModelLike --- .../backend/app/worker.py | 2 + .../backend/app/worker.py | 2 + src/ai/__init__.py | 2 + src/ai/agents/agent.py | 31 ++++++-- src/ai/models/__init__.py | 3 +- src/ai/models/core/api.py | 46 ++++++----- src/ai/models/core/model.py | 67 +++++++++++++++- tests/models/core/test_api.py | 1 + tests/models/core/test_model.py | 77 +++++++++++++++++++ 9 files changed, 203 insertions(+), 28 deletions(-) create mode 100644 tests/models/core/test_model.py diff --git a/examples/apps/durable_agent_temporal/backend/app/worker.py b/examples/apps/durable_agent_temporal/backend/app/worker.py index 58aac72f..fe10f33b 100644 --- a/examples/apps/durable_agent_temporal/backend/app/worker.py +++ b/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, diff --git a/examples/apps/durable_agent_workflows/backend/app/worker.py b/examples/apps/durable_agent_workflows/backend/app/worker.py index 0b88d193..23eed150 100644 --- a/examples/apps/durable_agent_workflows/backend/app/worker.py +++ b/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"), diff --git a/src/ai/__init__.py b/src/ai/__init__.py index b10faf3e..8fada860 100644 --- a/src/ai/__init__.py +++ b/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", diff --git a/src/ai/agents/agent.py b/src/ai/agents/agent.py index ec53583c..df063fd1 100644 --- a/src/ai/agents/agent.py +++ b/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__ diff --git a/src/ai/models/__init__.py b/src/ai/models/__init__.py index 1b3b99d8..9c664b6b 100644 --- a/src/ai/models/__init__.py +++ b/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", diff --git a/src/ai/models/core/api.py b/src/ai/models/core/api.py index 4e6f10bf..64d295ba 100644 --- a/src/ai/models/core/api.py +++ b/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,7 +341,7 @@ 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, @@ -341,7 +349,7 @@ def stream( ) -> 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, @@ -349,7 +357,7 @@ def stream[T: pydantic.BaseModel]( 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,7 +475,7 @@ 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, @@ -475,7 +483,7 @@ async def experimental_generate( 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, ) diff --git a/src/ai/models/core/model.py b/src/ai/models/core/model.py index e944767d..a49dd8e6 100644 --- a/src/ai/models/core/model.py +++ b/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, diff --git a/tests/models/core/test_api.py b/tests/models/core/test_api.py index 67925265..7eccd056 100644 --- a/tests/models/core/test_api.py +++ b/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) diff --git a/tests/models/core/test_model.py b/tests/models/core/test_model.py new file mode 100644 index 00000000..8281de75 --- /dev/null +++ b/tests/models/core/test_model.py @@ -0,0 +1,77 @@ +"""`ai.ModelLike`: anything with a model's `stream` goes where a model goes.""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator, Sequence +from typing import Any + +import pydantic + +import ai +from ai.types import events +from ai.types import messages as m +from ai.types import tools as tools_ + + +class Echo: + """Not an `ai.Model`: it answers with the last user message's text.""" + + def __init__(self) -> None: + self.calls: list[list[m.Message]] = [] + + @property + def id(self) -> str: + return "echo" + + @property + def provider_name(self) -> str: + return "echo" + + async def stream( + self, + messages: list[m.Message], + *, + tools: Sequence[tools_.Tool] | None = None, + output_type: type[pydantic.BaseModel] | None = None, + params: Any = None, + ) -> AsyncGenerator[events.Event]: + self.calls.append(list(messages)) + text = next(x.text for x in reversed(messages) if x.role == "user") + yield events.StreamStart() + yield events.TextStart(block_id="t") + yield events.TextDelta(chunk=text, block_id="t") + yield events.TextEnd(block_id="t") + yield events.StreamEnd(finish_reason="stop") + + +def test_a_model_and_a_stand_in_are_model_like() -> None: + assert isinstance(Echo(), ai.ModelLike) + assert isinstance(ai.get_model("openai/gpt-5"), ai.ModelLike) + + +async def test_stream_takes_a_model_like() -> None: + async with ai.stream(Echo(), [ai.user_message("hello")]) as s: + async for _ in s: + pass + assert s.text == "hello" + + +async def test_generate_takes_a_model_like() -> None: + message = await ai.experimental_generate(Echo(), [ai.user_message("hi")]) + assert message.text == "hi" + + +async def test_an_agent_runs_on_a_model_like() -> None: + model = Echo() + async with ai.Agent().run(model, [ai.user_message("ping")]) as run: + async for _ in run: + pass + assert [x.role for x in run.messages] == ["user", "assistant"] + assert run.messages[-1].text == "ping" + assert len(model.calls) == 1 + + +def test_a_context_keeps_a_model_like_as_it_is() -> None: + model = Echo() + context = ai.Context(model=model, messages=[], tools=[]) + assert context.model is model