ai-sdk-stream-python 0.2.0a1__tar.gz → 0.2.0a3__tar.gz

{ai_sdk_stream_python-0.2.0a1 → ai_sdk_stream_python-0.2.0a3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: ai-sdk-stream-python
-Version: 0.2.0a1
+Version: 0.2.0a3
 Summary: A Python package for AI SDK streaming utilities
 Author: Shloimy Wiesel
 Requires-Dist: fastapi>=0.128.8
@@ -63,6 +63,7 @@ async def my_work(ctx):
 | **Typed events** | All 16 v6 protocol events as Pydantic models |
 | **Lifecycle auto-management** | `start`, `start-step`, `text-start` etc. are emitted automatically |
 | **Shared state** | `ctx.store.get/set()` — dot-path key-value store shared across modules |
+| **Custom information** | `ctx.info` — typed, read-only Pydantic model for request-scoped metadata |
 | **Pass as parameter** | `ctx` flows through your services like a logger or DB session |
 | **Stream collection** | `collect=True` records all emitted content into `ctx.record` for DB persistence |
 | **Low-level escape hatch** | `ctx.write(event)` / `ctx.write_event_to_stream(ev)` for raw control |
@@ -123,6 +124,33 @@ plan = await ctx.store.get("user.plan", default="free")  # with default
 
 The store uses an `asyncio.Lock` — safe to use across concurrent coroutines.
 
+### Custom information (`ctx.info`)
+
+Pass a Pydantic model at construction time to carry static, read-only request-scoped data (e.g. `user_id`, `rate_limit`, `tenant_id`) through every service layer without threading extra arguments:
+
+```python
+from pydantic import BaseModel
+from ai_sdk_stream_python import StreamContext
+
+class RequestInfo(BaseModel):
+    user_id: str
+    rate_limit: int
+
+# Typed constructor — IDE infers ctx.info as RequestInfo
+ctx: StreamContext[RequestInfo] = StreamContext(
+    custom_information=RequestInfo(user_id="u_42", rate_limit=100)
+)
+
+# In any service layer that receives ctx:
+if ctx.info is not None:
+    print(ctx.info.user_id)    # "u_42"
+    print(ctx.info.rate_limit) # 100
+```
+
+- `ctx.info` is **read-only** (no setter). For mutable runtime state use `ctx.store`.
+- Defaults to `None` when `custom_information` is not passed.
+- `StreamContext` is generic — annotate as `StreamContext[YourModel]` for full IDE support.
+
 ### Writing to the stream
 
 ```python
@@ -226,6 +254,7 @@ ctx.message_id           # the message ID in the start event
 ctx.current_text_id      # ID of open text part, or None
 ctx.current_reasoning_id # ID of open reasoning part, or None
 ctx.is_finished          # True after finish()/abort()
+ctx.info                 # custom_information model, or None
 ctx.response_headers     # dict with x-vercel-ai-ui-message-stream: v1
 ```
 
@@ -367,9 +396,9 @@ npm run dev
 uv run pytest tests/ -v
 ```
 
-51 tests covering: basic lifecycle, reasoning ↔ text transitions, tool calls,
+56 tests covering: basic lifecycle, reasoning ↔ text transitions, tool calls,
 multi-step flows, source events, edge cases (double finish, abort, write after finish),
-StateStore integration, and stream collection (`collect=True`).
+StateStore integration, stream collection (`collect=True`), and custom information (`ctx.info`).
 
 ---
 

{ai_sdk_stream_python-0.2.0a1 → ai_sdk_stream_python-0.2.0a3}/README.md

@@ -53,6 +53,7 @@ async def my_work(ctx):
 | **Typed events** | All 16 v6 protocol events as Pydantic models |
 | **Lifecycle auto-management** | `start`, `start-step`, `text-start` etc. are emitted automatically |
 | **Shared state** | `ctx.store.get/set()` — dot-path key-value store shared across modules |
+| **Custom information** | `ctx.info` — typed, read-only Pydantic model for request-scoped metadata |
 | **Pass as parameter** | `ctx` flows through your services like a logger or DB session |
 | **Stream collection** | `collect=True` records all emitted content into `ctx.record` for DB persistence |
 | **Low-level escape hatch** | `ctx.write(event)` / `ctx.write_event_to_stream(ev)` for raw control |
@@ -113,6 +114,33 @@ plan = await ctx.store.get("user.plan", default="free")  # with default
 
 The store uses an `asyncio.Lock` — safe to use across concurrent coroutines.
 
+### Custom information (`ctx.info`)
+
+Pass a Pydantic model at construction time to carry static, read-only request-scoped data (e.g. `user_id`, `rate_limit`, `tenant_id`) through every service layer without threading extra arguments:
+
+```python
+from pydantic import BaseModel
+from ai_sdk_stream_python import StreamContext
+
+class RequestInfo(BaseModel):
+    user_id: str
+    rate_limit: int
+
+# Typed constructor — IDE infers ctx.info as RequestInfo
+ctx: StreamContext[RequestInfo] = StreamContext(
+    custom_information=RequestInfo(user_id="u_42", rate_limit=100)
+)
+
+# In any service layer that receives ctx:
+if ctx.info is not None:
+    print(ctx.info.user_id)    # "u_42"
+    print(ctx.info.rate_limit) # 100
+```
+
+- `ctx.info` is **read-only** (no setter). For mutable runtime state use `ctx.store`.
+- Defaults to `None` when `custom_information` is not passed.
+- `StreamContext` is generic — annotate as `StreamContext[YourModel]` for full IDE support.
+
 ### Writing to the stream
 
 ```python
@@ -216,6 +244,7 @@ ctx.message_id           # the message ID in the start event
 ctx.current_text_id      # ID of open text part, or None
 ctx.current_reasoning_id # ID of open reasoning part, or None
 ctx.is_finished          # True after finish()/abort()
+ctx.info                 # custom_information model, or None
 ctx.response_headers     # dict with x-vercel-ai-ui-message-stream: v1
 ```
 
@@ -357,9 +386,9 @@ npm run dev
 uv run pytest tests/ -v
 ```
 
-51 tests covering: basic lifecycle, reasoning ↔ text transitions, tool calls,
+56 tests covering: basic lifecycle, reasoning ↔ text transitions, tool calls,
 multi-step flows, source events, edge cases (double finish, abort, write after finish),
-StateStore integration, and stream collection (`collect=True`).
+StateStore integration, stream collection (`collect=True`), and custom information (`ctx.info`).
 
 ---
 

{ai_sdk_stream_python-0.2.0a1 → ai_sdk_stream_python-0.2.0a3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ai-sdk-stream-python"
-version = "0.2.0-a.1"
+version = "0.2.0-a.3"
 description = "A Python package for AI SDK streaming utilities"
 readme = "README.md"
 authors = [

{ai_sdk_stream_python-0.2.0a1 → ai_sdk_stream_python-0.2.0a3}/src/ai_sdk_stream_python/__init__.py

@@ -42,12 +42,18 @@ Quickstart::
         )
 """
 
+from .collect import DataPartRecord as DataPartRecord
+from .collect import FileRecord as FileRecord
 from .collect import SourceRecord as SourceRecord
 from .collect import StreamRecord as StreamRecord
 from .collect import ToolCallRecord as ToolCallRecord
 from .context import StreamContext, ToolCallHandle
 from .events import (
+    AbortEvent,
     BaseEvent,
+    DataEvent,
+    ErrorEvent,
+    FileEvent,
     FinishEvent,
     FinishStepEvent,
     ReasoningDeltaEvent,
@@ -77,6 +83,8 @@ __all__ = [
     "StreamRecord",
     "ToolCallRecord",
     "SourceRecord",
+    "FileRecord",
+    "DataPartRecord",
     # Base / union
     "BaseEvent",
     "UIMessageStreamEvent",
@@ -101,4 +109,11 @@ __all__ = [
     "ToolOutputErrorEvent",
     # Sources
     "SourceUrlEvent",
+    # Files
+    "FileEvent",
+    # Custom data parts
+    "DataEvent",
+    # Error / Abort
+    "ErrorEvent",
+    "AbortEvent",
 ]

{ai_sdk_stream_python-0.2.0a1 → ai_sdk_stream_python-0.2.0a3}/src/ai_sdk_stream_python/collect.py

@@ -32,6 +32,23 @@ class SourceRecord:
     title: str | None = None
 
 
+@dataclass
+class FileRecord:
+    """A file/image emitted via ``write_file``."""
+
+    url: str
+    media_type: str
+
+
+@dataclass
+class DataPartRecord:
+    """A non-transient custom data part emitted via ``write_data``."""
+
+    name: str
+    data: dict[str, Any]
+    id: str | None = None
+
+
 @dataclass
 class StreamRecord:
     """
@@ -66,6 +83,8 @@ class StreamRecord:
     reasoning: str = ""
     tool_calls: list[ToolCallRecord] = field(default_factory=list)
     sources: list[SourceRecord] = field(default_factory=list)
+    files: list[FileRecord] = field(default_factory=list)
+    data_parts: list[DataPartRecord] = field(default_factory=list)
     finish_reason: str | None = None
     step_count: int = 0
 
@@ -93,9 +112,30 @@ class StreamRecord:
                 }
                 for s in self.sources
             ],
+            "files": [
+                {
+                    "url": f.url,
+                    "media_type": f.media_type,
+                }
+                for f in self.files
+            ],
+            "data_parts": [
+                {
+                    "name": dp.name,
+                    "data": dp.data,
+                    "id": dp.id,
+                }
+                for dp in self.data_parts
+            ],
             "finish_reason": self.finish_reason,
             "step_count": self.step_count,
         }
 
 
-__all__ = ["SourceRecord", "StreamRecord", "ToolCallRecord"]
+__all__ = [
+    "DataPartRecord",
+    "FileRecord",
+    "SourceRecord",
+    "StreamRecord",
+    "ToolCallRecord",
+]

{ai_sdk_stream_python-0.2.0a1 → ai_sdk_stream_python-0.2.0a3}/src/ai_sdk_stream_python/context.py

@@ -56,11 +56,23 @@ import asyncio
 import uuid
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Generic, TypeVar
 
-from .collect import SourceRecord, StreamRecord, ToolCallRecord
+from pydantic import BaseModel
+
+from .collect import (
+    DataPartRecord,
+    FileRecord,
+    SourceRecord,
+    StreamRecord,
+    ToolCallRecord,
+)
 from .events import (
+    AbortEvent,
     BaseEvent,
+    DataEvent,
+    ErrorEvent,
+    FileEvent,
     FinishEvent,
     FinishStepEvent,
     ReasoningDeltaEvent,
@@ -73,12 +85,15 @@ from .events import (
     TextEndEvent,
     TextStartEvent,
     ToolInputAvailableEvent,
+    ToolInputDeltaEvent,
     ToolInputStartEvent,
     ToolOutputAvailableEvent,
     ToolOutputErrorEvent,
 )
 from .state import StateStore
 
+_InfoT = TypeVar("_InfoT", bound=BaseModel)
+
 
 @dataclass
 class ToolCallHandle:
@@ -88,7 +103,7 @@ class ToolCallHandle:
     toolName: str
 
 
-class StreamContext:
+class StreamContext(Generic[_InfoT]):
     """
     A stateful context for producing a Vercel AI SDK v6 UIMessageStream.
 
@@ -100,6 +115,11 @@ class StreamContext:
     ----------
     store : StateStore
         Async key-value store shared across the background task and any caller.
+    info : _InfoT | None
+        Static, read-only metadata supplied at construction time (e.g. user_id,
+        rate_limit).  Pass any Pydantic ``BaseModel`` instance as
+        ``custom_information=`` — it is available unchanged for the entire
+        lifetime of the context.  Defaults to ``None`` when not provided.
     """
 
     response_headers: ClassVar[dict[str, str]] = {
@@ -114,10 +134,12 @@ class StreamContext:
         message_id: str | None = None,
         *,
         collect: bool = False,
+        custom_information: _InfoT | None = None,
     ) -> None:
         self._message_id: str = message_id or str(uuid.uuid4())
         self._queue: asyncio.Queue[BaseEvent | None] = asyncio.Queue()
         self.store: StateStore = StateStore()
+        self._info: _InfoT | None = custom_information
 
         # Lifecycle state
         self._started: bool = False
@@ -151,6 +173,18 @@ class StreamContext:
     def is_finished(self) -> bool:
         return self._finished
 
+    @property
+    def info(self) -> _InfoT | None:
+        """
+        Static, read-only metadata supplied at construction time.
+
+        Returns the ``custom_information`` value passed to ``__init__``, or
+        ``None`` if none was provided.  Useful for carrying request-scoped
+        data (e.g. ``user_id``, ``rate_limit``) through service layers without
+        threading extra function arguments.
+        """
+        return self._info
+
     @property
     def record(self) -> StreamRecord | None:
         """
@@ -284,6 +318,63 @@ class StreamContext:
             )
         return ToolCallHandle(toolCallId=tcid, toolName=tool_name)
 
+    async def start_tool_input(
+        self,
+        tool_name: str,
+        *,
+        tool_call_id: str | None = None,
+    ) -> ToolCallHandle:
+        """
+        Emit ``tool-input-start`` and return a handle for streaming deltas.
+
+        Use this when tool arguments arrive incrementally (e.g. from an LLM
+        stream).  Follow with :meth:`stream_tool_input_delta` calls and
+        finish with :meth:`finish_tool_input`.
+        """
+        await self._ensure_step_open()
+        await self._ensure_text_closed()
+        await self._ensure_reasoning_closed()
+        tcid = tool_call_id or str(uuid.uuid4())
+        self._queue.put_nowait(ToolInputStartEvent(toolCallId=tcid, toolName=tool_name))
+        if self._record is not None:
+            self._record.tool_calls.append(
+                ToolCallRecord(tool_call_id=tcid, tool_name=tool_name, input={})
+            )
+        return ToolCallHandle(toolCallId=tcid, toolName=tool_name)
+
+    async def stream_tool_input_delta(
+        self, tool_call_id: str, input_text_delta: str
+    ) -> None:
+        """Emit a ``tool-input-delta`` for an in-progress tool call."""
+        self.write_event_to_stream(
+            ToolInputDeltaEvent(
+                toolCallId=tool_call_id, inputTextDelta=input_text_delta
+            )
+        )
+
+    async def finish_tool_input(
+        self,
+        tool_call_id: str,
+        tool_name: str,
+        input: dict[str, Any],
+    ) -> None:
+        """
+        Emit ``tool-input-available`` to close a streaming tool call.
+
+        Updates the collected :class:`~collect.ToolCallRecord` with the
+        final *input* if collection is enabled.
+        """
+        if self._record is not None:
+            for tc in self._record.tool_calls:
+                if tc.tool_call_id == tool_call_id:
+                    tc.input = input
+                    break
+        self.write_event_to_stream(
+            ToolInputAvailableEvent(
+                toolCallId=tool_call_id, toolName=tool_name, input=input
+            )
+        )
+
     async def complete_tool_call(self, tool_call_id: str, output: Any) -> None:
         """Emit ``tool-output-available`` with the tool result."""
         if self._record is not None:
@@ -306,6 +397,51 @@ class StreamContext:
             ToolOutputErrorEvent(toolCallId=tool_call_id, error=error)
         )
 
+    async def write_data(
+        self,
+        name: str,
+        data: dict[str, Any],
+        *,
+        id: str | None = None,
+        transient: bool = False,
+    ) -> None:
+        """
+        Emit a custom data part (``data-{name}`` type).
+
+        *name* is validated: it must be non-empty and contain only
+        alphanumeric characters, hyphens, or underscores.
+        Non-transient parts are collected in ``ctx.record.data_parts``.
+        Transient parts are only available through the ``onData`` callback
+        on the frontend; they are not stored in message history.
+        """
+        if not name or not all(c.isalnum() or c in "-_" for c in name):
+            raise ValueError(
+                f"Invalid data part name {name!r}. "
+                "Use only alphanumeric characters, hyphens, or underscores."
+            )
+        await self._ensure_started()
+        event = DataEvent(
+            type=f"data-{name}",
+            data=data,
+            id=id,
+            transient=transient or None,
+        )
+        if self._record is not None and not transient:
+            self._record.data_parts.append(DataPartRecord(name=name, data=data, id=id))
+        self.write_event_to_stream(event)
+
+    async def write_file(self, url: str, media_type: str) -> None:
+        """
+        Emit a ``file`` event (image, PDF, or other file content).
+
+        Auto-emits ``start`` and ``start-step`` if not yet open.
+        On the frontend this produces a ``FileUIPart`` in ``message.parts``.
+        """
+        await self._ensure_step_open()
+        if self._record is not None:
+            self._record.files.append(FileRecord(url=url, media_type=media_type))
+        self.write_event_to_stream(FileEvent(url=url, mediaType=media_type))
+
     async def write_source(
         self,
         source_id: str,
@@ -354,16 +490,33 @@ class StreamContext:
         )
         self._queue.put_nowait(None)  # sentinel → stream() yields [DONE]
 
-    async def abort(self) -> None:
+    async def abort(self, reason: str | None = None) -> None:
         """
-        Terminate the stream immediately without a proper finish event.
+        Emit an ``abort`` event and terminate the stream.
 
         Use in error-handling paths where the normal ``finish()`` flow
-        cannot be reached.
+        cannot be reached.  The optional *reason* string is forwarded to
+        the frontend so ``useChat`` can surface why the stream stopped.
+        Safe to call multiple times; subsequent calls are no-ops.
+        """
+        if self._finished:
+            return
+        self._finished = True
+        self._queue.put_nowait(AbortEvent(reason=reason))
+        self._queue.put_nowait(None)
+
+    async def error(self, error_text: str) -> None:
+        """
+        Emit an ``error`` event and terminate the stream.
+
+        The AI SDK v6 ``useChat`` hook surfaces this via its ``error`` object.
+        Safe to call multiple times; subsequent calls are no-ops.
         """
         if self._finished:
             return
+        await self._ensure_started()
         self._finished = True
+        self._queue.put_nowait(ErrorEvent(errorText=error_text))
         self._queue.put_nowait(None)
 
     # ── SSE async generator ───────────────────────────────────────────────────

{ai_sdk_stream_python-0.2.0a1 → ai_sdk_stream_python-0.2.0a3}/src/ai_sdk_stream_python/events.py

@@ -129,6 +129,47 @@ class SourceUrlEvent(BaseEvent):
     title: str | None = None
 
 
+# ── Files ──────────────────────────────────────────────────────────────────────
+
+
+class FileEvent(BaseEvent):
+    type: Literal["file"] = "file"
+    url: str
+    mediaType: str
+
+
+# ── Error ──────────────────────────────────────────────────────────────────────
+
+
+class ErrorEvent(BaseEvent):
+    type: Literal["error"] = "error"
+    errorText: str
+
+
+# ── Custom data parts ─────────────────────────────────────────────────────────
+
+
+class DataEvent(BaseEvent):
+    """Custom data part with a dynamic ``data-{name}`` type.
+
+    Not included in the ``UIMessageStreamEvent`` discriminated union because
+    the ``type`` field is dynamic.  Use ``ctx.write_data()`` to emit these.
+    """
+
+    type: str  # "data-{name}", validated by ctx.write_data()
+    data: dict[str, Any]
+    id: str | None = None
+    transient: bool | None = None
+
+
+# ── Abort ──────────────────────────────────────────────────────────────────────
+
+
+class AbortEvent(BaseEvent):
+    type: Literal["abort"] = "abort"
+    reason: str | None = None
+
+
 # ── Discriminated union ────────────────────────────────────────────────────────
 
 UIMessageStreamEvent = Annotated[
@@ -146,6 +187,9 @@ UIMessageStreamEvent = Annotated[
     | ToolOutputAvailableEvent
     | ToolOutputErrorEvent
     | SourceUrlEvent
+    | FileEvent
+    | ErrorEvent
+    | AbortEvent
     | FinishStepEvent
     | FinishEvent,
     Field(discriminator="type"),
@@ -170,4 +214,8 @@ __all__ = [
     "ToolOutputAvailableEvent",
     "ToolOutputErrorEvent",
     "SourceUrlEvent",
+    "FileEvent",
+    "DataEvent",
+    "ErrorEvent",
+    "AbortEvent",
 ]