Merge pull request #55 from vercel-labs/fix-models

Rework models api

Author: anbuzin

examples/coding-agent/1_raw_stream.py

@@ -19,11 +19,13 @@
 
 async def main() -> None:
     try:
-        async for event in ai.models.stream(model, messages):
-            if isinstance(event, ai.TextDelta):
-                print(event.chunk, end="", flush=True)
+        async with ai.stream(model, messages) as s:
+            async for event in s:
+                if isinstance(event, ai.TextDelta):
+                    print(event.chunk, end="", flush=True)
         print()
     finally:
+        # Explicit clients need explicit cleanup.
         await client.aclose()
 
 

examples/samples/explicit_client.py

@@ -1,8 +1,9 @@
 """Inline image generation — LLM that outputs images alongside text.
 
 Models like Gemini 3 Pro Image can generate images as part of their
-language model response. The images arrive as FileParts on the final
-MessageEnd message.
+language model response. The images arrive as ``FileEvent`` events
+during the stream and end up as ``FilePart``s on the aggregated
+``Stream.message``.
 """
 
 import asyncio
@@ -22,20 +23,18 @@
 
 
 async def main() -> None:
-    last_msg: ai.Message | None = None
-
-    # Stream — text deltas arrive as events, images arrive on MessageEnd
-    async for event in ai.stream(model, messages):
-        if isinstance(event, ai.TextDelta):
-            print(event.chunk, end="", flush=True)
-        elif isinstance(event, ai.MessageEnd):
-            last_msg = event.message
+    # Stream — text deltas arrive as TextDelta events, generated images
+    # arrive as FileEvent events and accumulate on s.message.
+    async with ai.stream(model, messages) as s:
+        async for event in s:
+            if isinstance(event, ai.TextDelta):
+                print(event.chunk, end="", flush=True)
 
     print()
 
-    # Check for images in the final message
-    if last_msg and last_msg.images:
-        for i, img in enumerate(last_msg.images):
+    # Check for images in the aggregated message.
+    if s.message.images:
+        for i, img in enumerate(s.message.images):
             filename = f"inline_{i}.png"
             data = (
                 img.data if isinstance(img.data, bytes) else base64.b64decode(img.data)

examples/samples/inline_image.py

@@ -20,9 +20,10 @@
 
 
 async def main() -> None:
-    async for event in ai.stream(model, messages):
-        if isinstance(event, ai.TextDelta):
-            print(event.chunk, end="", flush=True)
+    async with ai.stream(model, messages) as s:
+        async for event in s:
+            if isinstance(event, ai.TextDelta):
+                print(event.chunk, end="", flush=True)
     print()
 
 

examples/samples/multimodal_input.py

@@ -13,9 +13,10 @@
 
 
 async def main() -> None:
-    async for event in ai.stream(model, messages):
-        if isinstance(event, ai.TextDelta):
-            print(event.chunk, end="", flush=True)
+    async with ai.stream(model, messages) as s:
+        async for event in s:
+            if isinstance(event, ai.TextDelta):
+                print(event.chunk, end="", flush=True)
     print()
 
 

examples/samples/stream.py

@@ -20,15 +20,18 @@ class Recipe(pydantic.BaseModel):
 
 
 async def main() -> None:
-    # Stream with structured output — watch JSON arrive, get validated at the end
-    async for event in ai.stream(model, messages, output_type=Recipe):
-        if isinstance(event, ai.TextDelta):
-            print(event.chunk, end="", flush=True)
-        elif isinstance(event, ai.MessageEnd) and event.message.output:
-            recipe: Recipe = event.message.output
-            print(f"\n\nParsed recipe: {recipe.name}")
-            print(f"  Ingredients: {', '.join(recipe.ingredients)}")
-            print(f"  Prep time: {recipe.prep_time_minutes} min")
+    # Stream with structured output — watch JSON arrive, get validated at the end.
+    async with ai.stream(model, messages, output_type=Recipe) as s:
+        async for event in s:
+            if isinstance(event, ai.TextDelta):
+                print(event.chunk, end="", flush=True)
+
+    # After iteration, s.output is the validated pydantic model.
+    recipe: Recipe | None = s.output
+    if recipe is not None:
+        print(f"\n\nParsed recipe: {recipe.name}")
+        print(f"  Ingredients: {', '.join(recipe.ingredients)}")
+        print(f"  Prep time: {recipe.prep_time_minutes} min")
 
 
 if __name__ == "__main__":

examples/samples/structured_output.py

@@ -24,15 +24,17 @@
 
 
 async def main() -> None:
-    # Stream with tools — the model may emit tool calls
-    async for event in ai.stream(model, messages, tools=[get_weather]):
-        if isinstance(event, ai.TextDelta):
-            print(event.chunk, end="", flush=True)
-        elif isinstance(event, ai.MessageEnd):
-            for tc in event.message.tool_calls:
-                print(f"\nTool call: {tc.tool_name}({tc.tool_args})")
+    # Stream with tools — the model may emit tool calls.
+    async with ai.stream(model, messages, tools=[get_weather]) as s:
+        async for event in s:
+            if isinstance(event, ai.TextDelta):
+                print(event.chunk, end="", flush=True)
     print()
 
+    # After iteration, s.tool_calls collects every tool call from the response.
+    for tc in s.tool_calls:
+        print(f"Tool call: {tc.tool_name}({tc.tool_args})")
+
 
 if __name__ == "__main__":
     asyncio.run(main())

examples/samples/tools_schema.py

@@ -17,10 +17,15 @@
 from .middleware import AgentRunContext, Middleware
 from .models import (
     Client,
+    Executor,
+    GenerateExecutor,
+    GenerateRequest,
     ImageParams,
     Model,
     Provider,
-    StreamResult,
+    Stream,
+    StreamExecutor,
+    StreamRequest,
     VideoParams,
     ai_gateway,
     anthropic,
@@ -32,22 +37,20 @@
 
 # Re-export core types
 from .types import (
-    End,
     Event,
+    FileEvent,
     FilePart,
     HookPart,
     HookResolution,
     HookSuspention,
     Message,
-    MessageEnd,
-    MessageStart,
     Part,
     ReasoningDelta,
     ReasoningEnd,
     ReasoningPart,
     ReasoningStart,
-    Start,
-    StreamResultLike,
+    StreamEnd,
+    StreamStart,
     StructuredOutputPart,
     TextDelta,
     TextEnd,
@@ -74,12 +77,10 @@
 
 __all__ = [
     # Types (from types/)
-    "Start",
-    "End",
     "Event",
     "Message",
-    "MessageStart",
-    "MessageEnd",
+    "StreamStart",
+    "StreamEnd",
     "Part",
     "TextPart",
     "TextStart",
@@ -94,6 +95,7 @@
     "ReasoningStart",
     "ReasoningDelta",
     "ReasoningEnd",
+    "FileEvent",
     "FilePart",
     "HookPart",
     "HookSuspention",
@@ -116,8 +118,12 @@
     "ImageParams",
     "VideoParams",
     "Client",
-    "StreamResult",
-    "StreamResultLike",
+    "Stream",
+    "StreamRequest",
+    "GenerateRequest",
+    "Executor",
+    "StreamExecutor",
+    "GenerateExecutor",
     "check_connection",
     "stream",
     "generate",

src/ai/__init__.py

@@ -13,6 +13,7 @@
 from .. import middleware as middleware_
 from .. import models, types
 from ..types import builders
+from . import events as events_
 from . import runtime
 
 
@@ -199,21 +200,23 @@ def resolve(self, tool_parts: list[types.ToolCallPart]) -> list[ToolCall]:
         ]
 
 
-StreamItem = types.Event | types.Message
+StreamItem = events_.AgentEvent | types.Message
 
 
 class LoopFn(Protocol):
     def __call__(self, context: Context) -> AsyncGenerator[StreamItem]: ...
 
 
-async def _message_events(message: types.Message) -> AsyncGenerator[types.Event]:
-    yield types.MessageStart(message=message)
-    yield types.MessageEnd(message=message)
+async def _message_events(
+    message: types.Message,
+) -> AsyncGenerator[events_.AgentEvent]:
+    yield events_.MessageStart(message=message)
+    yield events_.MessageEnd(message=message)
 
 
 async def _coerce_events(
     source: AsyncIterable[StreamItem],
-) -> AsyncGenerator[types.Event]:
+) -> AsyncGenerator[events_.AgentEvent]:
     async for item in source:
         if isinstance(item, types.Message):
             async for event in _message_events(item):
@@ -222,15 +225,23 @@ async def _coerce_events(
             yield item
 
 
-async def _default_loop(context: Context) -> AsyncGenerator[types.Event]:
+async def _default_loop(context: Context) -> AsyncGenerator[events_.AgentEvent]:
     while True:
         stream = models.stream(
             context.model,
             context.messages,
             tools=context.tools,
         )
-        async for event in stream:
-            yield event
+        async for stream_event in stream:
+            yield stream_event
+
+        # Bridge: emit MessageStart/MessageEnd around the assistant message
+        # the model stream just produced, so _collect_messages and downstream
+        # consumers (AI-SDK outbound, label stamping) see the same boundary
+        # events they did under the previous adapter contract.
+        if stream.message is not None and stream.message.parts:
+            async for boundary in _message_events(stream.message):
+                yield boundary
 
         tool_calls = context.resolve(stream.tool_calls)
         if not tool_calls:
@@ -244,14 +255,14 @@ async def _default_loop(context: Context) -> AsyncGenerator[types.Event]:
         # Left un-stamped: the tool result is the input of the *next* turn,
         # so the next stream() call will stamp it with that turn's id.
         tool_msg = builders.tool_message(*(t.result() for t in tasks))
-        async for event in _message_events(tool_msg):
-            yield event
+        async for boundary in _message_events(tool_msg):
+            yield boundary
 
 
 async def _collect_messages(
     source: AsyncIterable[StreamItem],
     messages: list[types.Message],
-) -> AsyncGenerator[types.Event]:
+) -> AsyncGenerator[events_.AgentEvent]:
     """Intercept yielded events and collect MessageEnd messages into *messages*.
 
     This runs on the **producer** side (same coroutine as the loop function),
@@ -260,7 +271,7 @@ async def _collect_messages(
     happened on the consumer side of the runtime queue.
     """
     async for event in _coerce_events(source):
-        if isinstance(event, types.MessageEnd):
+        if isinstance(event, events_.MessageEnd):
             message = event.message
             for i, existing in enumerate(messages):
                 if existing.id == message.id:
@@ -292,7 +303,7 @@ async def yield_from(source: AsyncIterable[StreamItem]) -> str:
     last: types.Message | None = None
     async for item in _coerce_events(source):
         await rt.put_event(item)
-        if isinstance(item, types.MessageEnd):
+        if isinstance(item, events_.MessageEnd):
             last = item.message
     return last.text if last else ""
 
@@ -325,7 +336,7 @@ async def run(
         *,
         label: str | None = None,
         middleware: list[middleware_.Middleware] | None = None,
-    ) -> AsyncGenerator[types.Event]:
+    ) -> AsyncGenerator[events_.AgentEvent]:
         """Run the agent loop, yielding events to the consumer.
 
         Args:
@@ -349,7 +360,7 @@ async def run(
 
         async def _real(
             call: middleware_.AgentRunContext,
-        ) -> AsyncGenerator[types.Event]:
+        ) -> AsyncGenerator[events_.AgentEvent]:
             context = Context(
                 model=call.model,
                 messages=list(call.messages),
@@ -359,8 +370,8 @@ async def _real(
             async for event in runtime.run(source):
                 if call.label is not None:
                     event_message: types.Message | None = None
-                    if isinstance(event, types.MessageEnd) or (
-                        isinstance(event, types.MessageStart)
+                    if isinstance(event, events_.MessageEnd) or (
+                        isinstance(event, events_.MessageStart)
                         and event.message is not None
                     ):
                         event_message = event.message