Rework models and agents

.github/workflows/ci.yml

@@ -21,5 +21,6 @@ jobs:
       - run: uv run ruff format --check src tests
       - run: uv run ruff check src tests
       - run: uv run mypy src tests
+      - run: uv run pyright src tests
 
       - run: uv run pytest

CLAUDE.md

@@ -1,4 +1,35 @@
+# ai
+
+## development guidelines
+
 1. use `uv` to manage the project; `uv add` and `uv remove` to manage dependencies, `uv run` to run
 2. after making changes run lint and typecheck: `uv run ruff check --fix src tests` and `uv run mypy src tests`
 3. import by module (except `typing`) to improve readability via namespacing
 4. treat `stream_step` and `stream_loop` as user code. they are convenience functions that could be reimplemented by the user, they *must* stay clean.
+
+## design principles
+
+### 1. maximize composability
+
+provide simple lego bricks that the user can build their feature with. each block should do one thing and be reasonably decoupled from the rest.
+expose correct primitives to make it easy to modify behavior without rewriting it from scratch.
+
+- *example*: `agents` module provides `@ai.stream`, `@ai.tool` and `@ai.hook` that can be combined into an arbitrarily complex agent graph using plain python.
+- *can the user rewrite this feature in plain python using the existing primitives?*
+
+### 2. minimize dsl-ness and frameworkiness
+
+express features in a way that doesn't require the user to read documentation and learn the framework. glue things together using python.
+handle complexity inside the framework instead of delegating it to users.
+
+- *example*: `Runtime` does the heavy lifting so that multi-agent graphs can be expressed using python `asyncio`.
+- *does this require the user to learn a framework-specific concept that has a direct python equivalent?*
+
+### 3. keep data model simple
+
+ensure state is easy to serialize and deserialize, modify, and compose at any level of granularity.
+move normalization and translation complexity inside the framework and keep the public data model minimal.
+
+- *example*: public data model consists of a single unified `Message` type. the framework does not expose events and other intermediate steps unless the user is writing a custom adapter.
+
+

examples/fastapi-vite/backend/agent.py

@@ -16,10 +16,11 @@ async def talk_to_mothership(question: str) -> str:
     return f"Mothership says: {question} -> Soon."
 
 
-def get_llm() -> ai.LanguageModel:
-    """Create the LLM instance."""
-    return ai.ai_gateway.GatewayModel(model="anthropic/claude-opus-4.6")
-
+MODEL = ai.Model(
+    id="anthropic/claude-opus-4.6",
+    adapter="ai-gateway-v3",
+    provider="ai-gateway",
+)
 
 TOOLS: list[ai.Tool[..., Any]] = [talk_to_mothership]
 
@@ -43,10 +44,17 @@ async def _execute_with_approval(
         tc.set_error("Tool call was denied by the user.")
 
 
+chat_agent = ai.agent(
+    model=MODEL,
+    system="",
+    tools=TOOLS,
+)
+
+
+@chat_agent.loop
 async def graph(
-    llm: ai.LanguageModel,
+    agent: ai.Agent,
     messages: list[ai.Message],
-    tools: list[ai.Tool[..., Any]],
 ) -> ai.StreamResult:
     """Agent graph with human-in-the-loop tool approval.
 
@@ -58,7 +66,7 @@ async def graph(
     local_messages = list(messages)
 
     while True:
-        result = await ai.stream_step(llm, local_messages, tools)
+        result = await ai.stream_step(agent.model, local_messages, agent.tools)
 
         if not result.tool_calls:
             return result

examples/fastapi-vite/backend/main.py

@@ -51,20 +51,12 @@ async def chat(request: ChatRequest) -> fastapi.responses.StreamingResponse:
     session_id = request.session_id or "default"
     checkpoint_key = f"checkpoint:{session_id}"
 
-    llm = agent.get_llm()
-
     checkpoint = None
     saved = await file_storage.get(checkpoint_key)
     if saved:
         checkpoint = ai.Checkpoint.model_validate(saved)
 
-    result = ai.run(
-        agent.graph,
-        llm,
-        messages,
-        agent.TOOLS,
-        checkpoint=checkpoint,
-    )
+    result = agent.chat_agent.run(messages, checkpoint=checkpoint)
 
     async def stream_response() -> AsyncGenerator[str]:
         async for chunk in ai.ai_sdk_ui.to_sse_stream(result):

examples/models/buffer.py

@@ -0,0 +1,32 @@
+"""Buffered response — drain the stream, get the final message."""
+
+import asyncio
+
+from vercel_ai_sdk import models as m
+from vercel_ai_sdk.types import messages as messages_
+
+model = m.Model(
+    id="anthropic/claude-sonnet-4",
+    adapter="ai-gateway-v3",
+    provider="ai-gateway",
+)
+
+messages = [
+    messages_.Message(
+        role="user",
+        parts=[messages_.TextPart(text="What is 2 + 2?")],
+    ),
+]
+
+
+async def main() -> None:
+    result = await m.buffer(m.stream(model, messages))
+    print(result.text)
+    if result.usage:
+        print(
+            f"tokens: {result.usage.input_tokens} in, {result.usage.output_tokens} out"
+        )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/models/direct_adapter.py

@@ -0,0 +1,42 @@
+"""Direct adapter call — bypass the registry, call the adapter function directly."""
+
+import asyncio
+import os
+
+from vercel_ai_sdk import models as m
+from vercel_ai_sdk.models import ai_gateway as ai_gateway_v3
+from vercel_ai_sdk.types import messages as messages_
+
+model = m.Model(
+    id="anthropic/claude-sonnet-4",
+    adapter="ai-gateway-v3",
+    provider="ai-gateway",
+)
+
+client = m.Client(
+    base_url="https://ai-gateway.vercel.sh/v3/ai",
+    api_key=os.environ["AI_GATEWAY_API_KEY"],
+)
+
+messages = [
+    messages_.Message(
+        role="user",
+        parts=[messages_.TextPart(text="Say hello in three languages.")],
+    ),
+]
+
+
+async def main() -> None:
+    # Call the adapter function directly — no registry lookup, no auto-client.
+    # This is the lowest level of the API.
+    try:
+        async for msg in ai_gateway_v3.stream(client, model, messages):
+            if msg.text_delta:
+                print(msg.text_delta, end="", flush=True)
+        print()
+    finally:
+        await client.aclose()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/models/explicit_client.py

@@ -0,0 +1,41 @@
+"""Explicit client — bring your own auth and base URL."""
+
+import asyncio
+import os
+
+from vercel_ai_sdk import models as m
+from vercel_ai_sdk.types import messages as messages_
+
+model = m.Model(
+    id="anthropic/claude-sonnet-4",
+    adapter="ai-gateway-v3",
+    provider="ai-gateway",
+)
+
+# Explicit client — useful for custom auth, proxies, or self-hosted gateways.
+client = m.Client(
+    base_url="https://ai-gateway.vercel.sh/v3/ai",
+    api_key=os.environ["AI_GATEWAY_API_KEY"],
+    headers={"X-Custom-Header": "example"},
+)
+
+messages = [
+    messages_.Message(
+        role="user",
+        parts=[messages_.TextPart(text="Hello!")],
+    ),
+]
+
+
+async def main() -> None:
+    try:
+        async for msg in m.stream(model, messages, client=client):
+            if msg.text_delta:
+                print(msg.text_delta, end="", flush=True)
+        print()
+    finally:
+        await client.aclose()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/models/image_generation.py

@@ -0,0 +1,52 @@
+"""Image generation — dedicated image model via generate()."""
+
+import asyncio
+import base64
+import pathlib
+
+from vercel_ai_sdk import models as m
+from vercel_ai_sdk.types import messages as messages_
+
+model = m.Model(
+    id="google/imagen-4.0-generate-001",
+    adapter="ai-gateway-v3",
+    provider="ai-gateway",
+    capabilities=("image",),
+)
+
+messages = [
+    messages_.Message(
+        role="user",
+        parts=[
+            messages_.TextPart(
+                text=(
+                    "Anime girl with twin tails and cat ears, wearing a "
+                    "sailor school uniform, striking a victory pose in front "
+                    "of a futuristic Tokyo skyline at night, neon lights "
+                    "reflecting in her eyes, digital art style"
+                )
+            ),
+        ],
+    ),
+]
+
+
+async def main() -> None:
+    result = await m.generate(model, messages, m.ImageParams(n=2, aspect_ratio="16:9"))
+
+    print(f"Generated {len(result.images)} image(s)")
+    for i, img in enumerate(result.images):
+        filename = f"generated_{i}.png"
+        data = img.data if isinstance(img.data, bytes) else base64.b64decode(img.data)
+        pathlib.Path(filename).write_bytes(data)
+        print(f"  {filename}: {img.media_type}, {len(data)} bytes")
+
+    if result.usage:
+        print(
+            f"Usage: {result.usage.input_tokens} input, "
+            f"{result.usage.output_tokens} output tokens"
+        )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/models/inline_image.py

@@ -0,0 +1,74 @@
+"""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 in the streamed
+Message.
+"""
+
+import asyncio
+import base64
+import pathlib
+
+from vercel_ai_sdk import models as m
+from vercel_ai_sdk.types import messages as messages_
+
+# This is a language model that can also output images inline.
+model = m.Model(
+    id="google/gemini-3-pro-image",
+    adapter="ai-gateway-v3",
+    provider="ai-gateway",
+    capabilities=("text", "image"),
+)
+
+messages = [
+    messages_.Message(
+        role="system",
+        parts=[
+            messages_.TextPart(
+                text=(
+                    "You are an anime art assistant. When asked to draw or create "
+                    "an image, generate it in a soft pastel anime style."
+                )
+            ),
+        ],
+    ),
+    messages_.Message(
+        role="user",
+        parts=[
+            messages_.TextPart(
+                text=(
+                    "Draw an anime girl with long silver hair and violet eyes, "
+                    "sitting in a field of cherry blossoms at sunset."
+                )
+            ),
+        ],
+    ),
+]
+
+
+async def main() -> None:
+    last_msg: messages_.Message | None = None
+
+    # Stream — text deltas arrive as usual, images arrive as FileParts
+    async for msg in m.stream(model, messages):
+        if msg.text_delta:
+            print(msg.text_delta, end="", flush=True)
+        last_msg = msg
+
+    print()
+
+    # Check for images in the final message
+    if last_msg and last_msg.images:
+        for i, img in enumerate(last_msg.images):
+            filename = f"inline_{i}.png"
+            data = (
+                img.data if isinstance(img.data, bytes) else base64.b64decode(img.data)
+            )
+            pathlib.Path(filename).write_bytes(data)
+            print(f"Saved {filename} ({img.media_type}, {len(data)} bytes)")
+    else:
+        print("No images were generated in this response.")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/models/multimodal_input.py

@@ -0,0 +1,38 @@
+"""Multimodal input — send a local image to the model and ask about it."""
+
+import asyncio
+import pathlib
+
+from vercel_ai_sdk import models as m
+from vercel_ai_sdk.types import messages as messages_
+
+model = m.Model(
+    id="anthropic/claude-sonnet-4",
+    adapter="ai-gateway-v3",
+    provider="ai-gateway",
+)
+
+# Load a local image file (replace with your own path)
+image_path = pathlib.Path("sample_image.jpg")
+image_data = image_path.read_bytes()
+
+messages = [
+    messages_.Message(
+        role="user",
+        parts=[
+            messages_.TextPart(text="Describe this image in detail."),
+            messages_.FilePart(data=image_data, media_type="image/jpeg"),
+        ],
+    ),
+]
+
+
+async def main() -> None:
+    async for msg in m.stream(model, messages):
+        if msg.text_delta:
+            print(msg.text_delta, end="", flush=True)
+    print()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())