chore(types): add more discriminator metadata (#825)

Author: RobertCraigie

src/anthropic/lib/streaming/_beta_types.py

@@ -1,5 +1,5 @@
 from typing import Union
-from typing_extensions import Literal
+from typing_extensions import Literal, Annotated
 
 from ..._models import BaseModel
 from ...types.beta import (
@@ -12,6 +12,7 @@
     BetaRawContentBlockDeltaEvent,
     BetaRawContentBlockStartEvent,
 )
+from ..._utils._transform import PropertyInfo
 
 
 class BetaTextEvent(BaseModel):
@@ -53,13 +54,16 @@ class BetaContentBlockStopEvent(BetaRawContentBlockStopEvent):
     content_block: BetaContentBlock
 
 
-BetaMessageStreamEvent = Union[
-    BetaTextEvent,
-    BetaInputJsonEvent,
-    BetaRawMessageStartEvent,
-    BetaRawMessageDeltaEvent,
-    BetaMessageStopEvent,
-    BetaRawContentBlockStartEvent,
-    BetaRawContentBlockDeltaEvent,
-    BetaContentBlockStopEvent,
+BetaMessageStreamEvent = Annotated[
+    Union[
+        BetaTextEvent,
+        BetaInputJsonEvent,
+        BetaRawMessageStartEvent,
+        BetaRawMessageDeltaEvent,
+        BetaMessageStopEvent,
+        BetaRawContentBlockStartEvent,
+        BetaRawContentBlockDeltaEvent,
+        BetaContentBlockStopEvent,
+    ],
+    PropertyInfo(discriminator="type"),
 ]

src/anthropic/lib/streaming/_types.py

@@ -1,5 +1,5 @@
 from typing import Union
-from typing_extensions import Literal
+from typing_extensions import Literal, Annotated
 
 from ...types import (
     Message,
@@ -12,6 +12,7 @@
     RawContentBlockStopEvent,
 )
 from ..._models import BaseModel
+from ..._utils._transform import PropertyInfo
 
 
 class TextEvent(BaseModel):
@@ -53,13 +54,16 @@ class ContentBlockStopEvent(RawContentBlockStopEvent):
     content_block: ContentBlock
 
 
-MessageStreamEvent = Union[
-    TextEvent,
-    InputJsonEvent,
-    RawMessageStartEvent,
-    RawMessageDeltaEvent,
-    MessageStopEvent,
-    RawContentBlockStartEvent,
-    RawContentBlockDeltaEvent,
-    ContentBlockStopEvent,
+MessageStreamEvent = Annotated[
+    Union[
+        TextEvent,
+        InputJsonEvent,
+        RawMessageStartEvent,
+        RawMessageDeltaEvent,
+        MessageStopEvent,
+        RawContentBlockStartEvent,
+        RawContentBlockDeltaEvent,
+        ContentBlockStopEvent,
+    ],
+    PropertyInfo(discriminator="type"),
 ]

src/anthropic/resources/beta/messages/messages.py

@@ -983,7 +983,6 @@ def stream(
             stream_cls=Stream[BetaRawMessageStreamEvent],
         )
         return BetaMessageStreamManager(make_request)
-    
 
     def count_tokens(
         self,

src/anthropic/types/tool_param.py

@@ -16,6 +16,7 @@ class InputSchemaTyped(TypedDict, total=False):
 
     properties: Optional[object]
 
+
 set_pydantic_config(InputSchemaTyped, {"extra": "allow"})
 
 InputSchema: TypeAlias = Union[InputSchemaTyped, Dict[str, object]]