@@ -91,9 +91,12 @@ class _RecordedUsage:
9191
9292class _StreamSpanContext :
9393 def __init__ (
94- self , stream : "Union[Stream, MessageStream, AsyncStream, AsyncMessageStream]"
94+ self ,
95+ stream : "Union[Stream, MessageStream, AsyncStream, AsyncMessageStream]" ,
96+ guaranteed_streaming_state : bool = False ,
9597 ) -> None :
9698 self ._stream = stream
99+ self ._guaranteed_streaming_state = guaranteed_streaming_state
97100
98101 def __enter__ (self ) -> "_StreamSpanContext" :
99102 return self
@@ -108,7 +111,9 @@ def __exit__(
108111 if not hasattr (self ._stream , "_span" ):
109112 return
110113
111- if not hasattr (self ._stream , "_model" ):
114+ if not self ._guaranteed_streaming_state and not hasattr (
115+ self ._stream , "_model"
116+ ):
112117 self ._stream ._span .__exit__ (exc_type , exc_val , exc_tb )
113118 del self ._stream ._span
114119 return
@@ -484,7 +489,7 @@ def _wrap_synchronous_message_iterator(
484489 Sets information received while iterating the response stream on the AI Client Span.
485490 Responsible for closing the AI Client Span unless the span has already been closed in the close() patch.
486491 """
487- with _StreamSpanContext (stream ):
492+ with _StreamSpanContext (stream , guaranteed_streaming_state = True ):
488493 for event in iterator :
489494 # Message and content types are aliases for corresponding Raw* types, introduced in
490495 # https://github.com/anthropics/anthropic-sdk-python/commit/bc9d11cd2addec6976c46db10b7c89a8c276101a
@@ -514,7 +519,7 @@ async def _wrap_asynchronous_message_iterator(
514519 Sets information received while iterating the response stream on the AI Client Span.
515520 Responsible for closing the AI Client Span unless the span has already been closed in the close() patch.
516521 """
517- with _StreamSpanContext (stream ):
522+ with _StreamSpanContext (stream , guaranteed_streaming_state = True ):
518523 async for event in iterator :
519524 # Message and content types are aliases for corresponding Raw* types, introduced in
520525 # https://github.com/anthropics/anthropic-sdk-python/commit/bc9d11cd2addec6976c46db10b7c89a8c276101a
0 commit comments