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

{ai_sdk_stream_python-0.2.0a2 → 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.0a2
+Version: 0.2.0a3
 Summary: A Python package for AI SDK streaming utilities
 Author: Shloimy Wiesel
 Requires-Dist: fastapi>=0.128.8

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

@@ -1,6 +1,6 @@
 [project]
 name = "ai-sdk-stream-python"
-version = "0.2.0-a.2"
+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.0a2 → 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.0a2 → 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.0a2 → ai_sdk_stream_python-0.2.0a3}/src/ai_sdk_stream_python/context.py

@@ -60,9 +60,19 @@ from typing import Any, ClassVar, Generic, TypeVar
 
 from pydantic import BaseModel
 
-from .collect import SourceRecord, StreamRecord, ToolCallRecord
+from .collect import (
+    DataPartRecord,
+    FileRecord,
+    SourceRecord,
+    StreamRecord,
+    ToolCallRecord,
+)
 from .events import (
+    AbortEvent,
     BaseEvent,
+    DataEvent,
+    ErrorEvent,
+    FileEvent,
     FinishEvent,
     FinishStepEvent,
     ReasoningDeltaEvent,
@@ -75,6 +85,7 @@ from .events import (
     TextEndEvent,
     TextStartEvent,
     ToolInputAvailableEvent,
+    ToolInputDeltaEvent,
     ToolInputStartEvent,
     ToolOutputAvailableEvent,
     ToolOutputErrorEvent,
@@ -307,6 +318,63 @@ class StreamContext(Generic[_InfoT]):
             )
         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:
@@ -329,6 +397,51 @@ class StreamContext(Generic[_InfoT]):
             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,
@@ -377,16 +490,33 @@ class StreamContext(Generic[_InfoT]):
         )
         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.0a2 → 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",
 ]