anthropic 0.72.0__py3-none-any.whl → 0.73.0__py3-none-any.whl

anthropic/lib/streaming/_beta_messages.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+import builtins
 from types import TracebackType
-from typing import TYPE_CHECKING, Any, Type, Callable, cast
+from typing import TYPE_CHECKING, Any, Type, Generic, Callable, cast
 from typing_extensions import Self, Iterator, Awaitable, AsyncIterator, assert_never
 
 import httpx
@@ -11,23 +12,27 @@ from anthropic.types.beta.beta_tool_use_block import BetaToolUseBlock
 from anthropic.types.beta.beta_mcp_tool_use_block import BetaMCPToolUseBlock
 from anthropic.types.beta.beta_server_tool_use_block import BetaServerToolUseBlock
 
+from ..._types import NOT_GIVEN, NotGiven
 from ..._utils import consume_sync_iterator, consume_async_iterator
 from ..._models import build, construct_type, construct_type_unchecked
 from ._beta_types import (
-    BetaTextEvent,
     BetaCitationEvent,
     BetaThinkingEvent,
     BetaInputJsonEvent,
     BetaSignatureEvent,
-    BetaMessageStopEvent,
-    BetaMessageStreamEvent,
-    BetaContentBlockStopEvent,
+    ParsedBetaTextEvent,
+    ParsedBetaMessageStopEvent,
+    ParsedBetaMessageStreamEvent,
+    ParsedBetaContentBlockStopEvent,
 )
 from ..._streaming import Stream, AsyncStream
-from ...types.beta import BetaMessage, BetaContentBlock, BetaRawMessageStreamEvent
+from ...types.beta import BetaRawMessageStreamEvent
+from ..._utils._utils import is_given
+from .._parse._response import ResponseFormatT, parse_text
+from ...types.beta.parsed_beta_message import ParsedBetaMessage, ParsedBetaContentBlock
 
 
-class BetaMessageStream:
+class BetaMessageStream(Generic[ResponseFormatT]):
     text_stream: Iterator[str]
     """Iterator over just the text deltas in the stream.
 
@@ -38,11 +43,16 @@ class BetaMessageStream:
     ```
     """
 
-    def __init__(self, raw_stream: Stream[BetaRawMessageStreamEvent]) -> None:
+    def __init__(
+        self,
+        raw_stream: Stream[BetaRawMessageStreamEvent],
+        output_format: ResponseFormatT | NotGiven,
+    ) -> None:
         self._raw_stream = raw_stream
         self.text_stream = self.__stream_text__()
         self._iterator = self.__stream__()
-        self.__final_message_snapshot: BetaMessage | None = None
+        self.__final_message_snapshot: ParsedBetaMessage[ResponseFormatT] | None = None
+        self.__output_format = output_format
 
     @property
     def response(self) -> httpx.Response:
@@ -52,10 +62,10 @@ class BetaMessageStream:
     def request_id(self) -> str | None:
         return self.response.headers.get("request-id")  # type: ignore[no-any-return]
 
-    def __next__(self) -> BetaMessageStreamEvent:
+    def __next__(self) -> ParsedBetaMessageStreamEvent[ResponseFormatT]:
         return self._iterator.__next__()
 
-    def __iter__(self) -> Iterator[BetaMessageStreamEvent]:
+    def __iter__(self) -> Iterator[ParsedBetaMessageStreamEvent[ResponseFormatT]]:
         for item in self._iterator:
             yield item
 
@@ -78,7 +88,7 @@ class BetaMessageStream:
         """
         self._raw_stream.close()
 
-    def get_final_message(self) -> BetaMessage:
+    def get_final_message(self) -> ParsedBetaMessage[ResponseFormatT]:
         """Waits until the stream has been read to completion and returns
         the accumulated `Message` object.
         """
@@ -113,16 +123,17 @@ class BetaMessageStream:
 
     # properties
     @property
-    def current_message_snapshot(self) -> BetaMessage:
+    def current_message_snapshot(self) -> ParsedBetaMessage[ResponseFormatT]:
         assert self.__final_message_snapshot is not None
         return self.__final_message_snapshot
 
-    def __stream__(self) -> Iterator[BetaMessageStreamEvent]:
+    def __stream__(self) -> Iterator[ParsedBetaMessageStreamEvent[ResponseFormatT]]:
         for sse_event in self._raw_stream:
             self.__final_message_snapshot = accumulate_event(
                 event=sse_event,
                 current_snapshot=self.__final_message_snapshot,
                 request_headers=self.response.request.headers,
+                output_format=self.__output_format,
             )
 
             events_to_fire = build_events(event=sse_event, message_snapshot=self.current_message_snapshot)
@@ -135,7 +146,7 @@ class BetaMessageStream:
                 yield chunk.delta.text
 
 
-class BetaMessageStreamManager:
+class BetaMessageStreamManager(Generic[ResponseFormatT]):
     """Wrapper over MessageStream that is returned by `.stream()`.
 
     ```py
@@ -148,13 +159,16 @@ class BetaMessageStreamManager:
     def __init__(
         self,
         api_request: Callable[[], Stream[BetaRawMessageStreamEvent]],
+        *,
+        output_format: ResponseFormatT | NotGiven,
     ) -> None:
-        self.__stream: BetaMessageStream | None = None
+        self.__stream: BetaMessageStream[ResponseFormatT] | None = None
         self.__api_request = api_request
+        self.__output_format = output_format
 
-    def __enter__(self) -> BetaMessageStream:
+    def __enter__(self) -> BetaMessageStream[ResponseFormatT]:
         raw_stream = self.__api_request()
-        self.__stream = BetaMessageStream(raw_stream)
+        self.__stream = BetaMessageStream(raw_stream, output_format=self.__output_format)
         return self.__stream
 
     def __exit__(
@@ -167,7 +181,7 @@ class BetaMessageStreamManager:
             self.__stream.close()
 
 
-class BetaAsyncMessageStream:
+class BetaAsyncMessageStream(Generic[ResponseFormatT]):
     text_stream: AsyncIterator[str]
     """Async iterator over just the text deltas in the stream.
 
@@ -178,11 +192,16 @@ class BetaAsyncMessageStream:
     ```
     """
 
-    def __init__(self, raw_stream: AsyncStream[BetaRawMessageStreamEvent]) -> None:
+    def __init__(
+        self,
+        raw_stream: AsyncStream[BetaRawMessageStreamEvent],
+        output_format: ResponseFormatT | NotGiven,
+    ) -> None:
         self._raw_stream = raw_stream
         self.text_stream = self.__stream_text__()
         self._iterator = self.__stream__()
-        self.__final_message_snapshot: BetaMessage | None = None
+        self.__final_message_snapshot: ParsedBetaMessage[ResponseFormatT] | None = None
+        self.__output_format = output_format
 
     @property
     def response(self) -> httpx.Response:
@@ -192,10 +211,10 @@ class BetaAsyncMessageStream:
     def request_id(self) -> str | None:
         return self.response.headers.get("request-id")  # type: ignore[no-any-return]
 
-    async def __anext__(self) -> BetaMessageStreamEvent:
+    async def __anext__(self) -> ParsedBetaMessageStreamEvent[ResponseFormatT]:
         return await self._iterator.__anext__()
 
-    async def __aiter__(self) -> AsyncIterator[BetaMessageStreamEvent]:
+    async def __aiter__(self) -> AsyncIterator[ParsedBetaMessageStreamEvent[ResponseFormatT]]:
         async for item in self._iterator:
             yield item
 
@@ -218,7 +237,7 @@ class BetaAsyncMessageStream:
         """
         await self._raw_stream.close()
 
-    async def get_final_message(self) -> BetaMessage:
+    async def get_final_message(self) -> ParsedBetaMessage[ResponseFormatT]:
         """Waits until the stream has been read to completion and returns
         the accumulated `Message` object.
         """
@@ -253,16 +272,17 @@ class BetaAsyncMessageStream:
 
     # properties
     @property
-    def current_message_snapshot(self) -> BetaMessage:
+    def current_message_snapshot(self) -> ParsedBetaMessage[ResponseFormatT]:
         assert self.__final_message_snapshot is not None
         return self.__final_message_snapshot
 
-    async def __stream__(self) -> AsyncIterator[BetaMessageStreamEvent]:
+    async def __stream__(self) -> AsyncIterator[ParsedBetaMessageStreamEvent[ResponseFormatT]]:
         async for sse_event in self._raw_stream:
             self.__final_message_snapshot = accumulate_event(
                 event=sse_event,
                 current_snapshot=self.__final_message_snapshot,
                 request_headers=self.response.request.headers,
+                output_format=self.__output_format,
             )
 
             events_to_fire = build_events(event=sse_event, message_snapshot=self.current_message_snapshot)
@@ -275,7 +295,7 @@ class BetaAsyncMessageStream:
                 yield chunk.delta.text
 
 
-class BetaAsyncMessageStreamManager:
+class BetaAsyncMessageStreamManager(Generic[ResponseFormatT]):
     """Wrapper over BetaAsyncMessageStream that is returned by `.stream()`
     so that an async context manager can be used without `await`ing the
     original client call.
@@ -290,13 +310,16 @@ class BetaAsyncMessageStreamManager:
     def __init__(
         self,
         api_request: Awaitable[AsyncStream[BetaRawMessageStreamEvent]],
+        *,
+        output_format: ResponseFormatT | NotGiven = NOT_GIVEN,
     ) -> None:
-        self.__stream: BetaAsyncMessageStream | None = None
+        self.__stream: BetaAsyncMessageStream[ResponseFormatT] | None = None
         self.__api_request = api_request
+        self.__output_format = output_format
 
-    async def __aenter__(self) -> BetaAsyncMessageStream:
+    async def __aenter__(self) -> BetaAsyncMessageStream[ResponseFormatT]:
         raw_stream = await self.__api_request
-        self.__stream = BetaAsyncMessageStream(raw_stream)
+        self.__stream = BetaAsyncMessageStream(raw_stream, output_format=self.__output_format)
         return self.__stream
 
     async def __aexit__(
@@ -312,16 +335,18 @@ class BetaAsyncMessageStreamManager:
 def build_events(
     *,
     event: BetaRawMessageStreamEvent,
-    message_snapshot: BetaMessage,
-) -> list[BetaMessageStreamEvent]:
-    events_to_fire: list[BetaMessageStreamEvent] = []
+    message_snapshot: ParsedBetaMessage[ResponseFormatT],
+) -> list[ParsedBetaMessageStreamEvent[ResponseFormatT]]:
+    events_to_fire: list[ParsedBetaMessageStreamEvent[ResponseFormatT]] = []
 
     if event.type == "message_start":
         events_to_fire.append(event)
     elif event.type == "message_delta":
         events_to_fire.append(event)
     elif event.type == "message_stop":
-        events_to_fire.append(build(BetaMessageStopEvent, type="message_stop", message=message_snapshot))
+        events_to_fire.append(
+            build(ParsedBetaMessageStopEvent[ResponseFormatT], type="message_stop", message=message_snapshot)
+        )
     elif event.type == "content_block_start":
         events_to_fire.append(event)
     elif event.type == "content_block_delta":
@@ -332,7 +357,7 @@ def build_events(
             if content_block.type == "text":
                 events_to_fire.append(
                     build(
-                        BetaTextEvent,
+                        ParsedBetaTextEvent,
                         type="text",
                         text=event.delta.text,
                         snapshot=content_block.text,
@@ -385,9 +410,14 @@ def build_events(
     elif event.type == "content_block_stop":
         content_block = message_snapshot.content[event.index]
 
-        events_to_fire.append(
-            build(BetaContentBlockStopEvent, type="content_block_stop", index=event.index, content_block=content_block),
+        event_to_fire = build(
+            ParsedBetaContentBlockStopEvent,
+            type="content_block_stop",
+            index=event.index,
+            content_block=content_block,
         )
+
+        events_to_fire.append(event_to_fire)
     else:
         # we only want exhaustive checking for linters, not at runtime
         if TYPE_CHECKING:  # type: ignore[unreachable]
@@ -408,9 +438,10 @@ TRACKS_TOOL_INPUT = (
 def accumulate_event(
     *,
     event: BetaRawMessageStreamEvent,
-    current_snapshot: BetaMessage | None,
+    current_snapshot: ParsedBetaMessage[ResponseFormatT] | None,
     request_headers: httpx.Headers,
-) -> BetaMessage:
+    output_format: ResponseFormatT | NotGiven = NOT_GIVEN,
+) -> ParsedBetaMessage[ResponseFormatT]:
     if not isinstance(cast(Any, event), BaseModel):
         event = cast(  # pyright: ignore[reportUnnecessaryCast]
             BetaRawMessageStreamEvent,
@@ -420,11 +451,15 @@ def accumulate_event(
             ),
         )
         if not isinstance(cast(Any, event), BaseModel):
-            raise TypeError(f"Unexpected event runtime type, after deserialising twice - {event} - {type(event)}")
+            raise TypeError(
+                f"Unexpected event runtime type, after deserialising twice - {event} - {builtins.type(event)}"
+            )
 
     if current_snapshot is None:
         if event.type == "message_start":
-            return BetaMessage.construct(**cast(Any, event.message.to_dict()))
+            return cast(
+                ParsedBetaMessage[ResponseFormatT], ParsedBetaMessage.construct(**cast(Any, event.message.to_dict()))
+            )
 
         raise RuntimeError(f'Unexpected event order, got {event.type} before "message_start"')
 
@@ -432,8 +467,8 @@ def accumulate_event(
         # TODO: check index
         current_snapshot.content.append(
             cast(
-                BetaContentBlock,
-                construct_type(type_=BetaContentBlock, value=event.content_block.model_dump()),
+                Any,  # Pydantic does not support generic unions at runtime
+                construct_type(type_=ParsedBetaContentBlock, value=event.content_block.model_dump()),
             ),
         )
     elif event.type == "content_block_delta":
@@ -481,6 +516,10 @@ def accumulate_event(
             # we only want exhaustive checking for linters, not at runtime
             if TYPE_CHECKING:  # type: ignore[unreachable]
                 assert_never(event.delta)
+    elif event.type == "content_block_stop":
+        content_block = current_snapshot.content[event.index]
+        if content_block.type == "text" and is_given(output_format):
+            content_block.parsed_output = parse_text(content_block.text, output_format)
     elif event.type == "message_delta":
         current_snapshot.container = event.delta.container
         current_snapshot.stop_reason = event.delta.stop_reason

anthropic/lib/streaming/_beta_types.py

@@ -1,10 +1,10 @@
-from typing import Union
+from typing import TYPE_CHECKING, Any, Dict, Union, Generic, cast
 from typing_extensions import List, Literal, Annotated
 
-from ..._models import BaseModel
+import jiter
+
+from ..._models import BaseModel, GenericModel
 from ...types.beta import (
-    BetaMessage,
-    BetaContentBlock,
     BetaRawMessageStopEvent,
     BetaRawMessageDeltaEvent,
     BetaRawMessageStartEvent,
@@ -12,11 +12,13 @@ from ...types.beta import (
     BetaRawContentBlockDeltaEvent,
     BetaRawContentBlockStartEvent,
 )
+from .._parse._response import ResponseFormatT
 from ..._utils._transform import PropertyInfo
+from ...types.beta.parsed_beta_message import ParsedBetaMessage, ParsedBetaContentBlock
 from ...types.beta.beta_citations_delta import Citation
 
 
-class BetaTextEvent(BaseModel):
+class ParsedBetaTextEvent(BaseModel):
     type: Literal["text"]
 
     text: str
@@ -25,6 +27,9 @@ class BetaTextEvent(BaseModel):
     snapshot: str
     """The entire accumulated text"""
 
+    def parsed_snapshot(self) -> Dict[str, Any]:
+        return cast(Dict[str, Any], jiter.from_json(self.snapshot.encode("utf-8"), partial_mode="trailing-strings"))
+
 
 class BetaCitationEvent(BaseModel):
     type: Literal["citation"]
@@ -70,31 +75,34 @@ class BetaInputJsonEvent(BaseModel):
     """
 
 
-class BetaMessageStopEvent(BetaRawMessageStopEvent):
+class ParsedBetaMessageStopEvent(BetaRawMessageStopEvent, GenericModel, Generic[ResponseFormatT]):
     type: Literal["message_stop"]
 
-    message: BetaMessage
+    message: ParsedBetaMessage[ResponseFormatT]
 
 
-class BetaContentBlockStopEvent(BetaRawContentBlockStopEvent):
+class ParsedBetaContentBlockStopEvent(BetaRawContentBlockStopEvent, GenericModel, Generic[ResponseFormatT]):
     type: Literal["content_block_stop"]
 
-    content_block: BetaContentBlock
+    if TYPE_CHECKING:
+        content_block: ParsedBetaContentBlock[ResponseFormatT]
+    else:
+        content_block: ParsedBetaContentBlock
 
 
-BetaMessageStreamEvent = Annotated[
+ParsedBetaMessageStreamEvent = Annotated[
     Union[
-        BetaTextEvent,
+        ParsedBetaTextEvent,
         BetaCitationEvent,
         BetaThinkingEvent,
         BetaSignatureEvent,
         BetaInputJsonEvent,
         BetaRawMessageStartEvent,
         BetaRawMessageDeltaEvent,
-        BetaMessageStopEvent,
+        ParsedBetaMessageStopEvent[ResponseFormatT],
         BetaRawContentBlockStartEvent,
         BetaRawContentBlockDeltaEvent,
-        BetaContentBlockStopEvent,
+        ParsedBetaContentBlockStopEvent[ResponseFormatT],
     ],
     PropertyInfo(discriminator="type"),
 ]