Skip to content

Commit 7ac1df0

Browse files
fix: Ensure stop_reason is captured by yielding on RawMessageDeltaEvent
The previous implementation had a timing issue where stop_reason arrived after the last chunk was yielded. This fix ensures that when we receive RawMessageDeltaEvent (which contains the stop_reason), we yield an additional chunk with the updated metadata. This addresses the test failure identified by @AstraBert where: - Content chunks were yielded before stop_reason arrived - stop_reason was None in the last chunk - RawMessageDeltaEvent came after ContentBlockStopEvent Now the flow is: 1. ContentBlockDeltaEvent -> yield chunks with content 2. ContentBlockStopEvent -> yield final content chunk 3. RawMessageDeltaEvent -> yield metadata update with stop_reason 4. Tests now pass with stop_reason correctly captured
1 parent 88555d1 commit 7ac1df0

File tree

1 file changed

+30
-0
lines changed
  • llama-index-integrations/llms/llama-index-llms-anthropic/llama_index/llms/anthropic

1 file changed

+30
-0
lines changed

llama-index-integrations/llms/llama-index-llms-anthropic/llama_index/llms/anthropic/base.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -619,6 +619,21 @@ def gen() -> Generator[AnthropicChatResponse, None, None]:
619619
}
620620
if hasattr(r, "delta") and hasattr(r.delta, "stop_reason"):
621621
stop_reason = r.delta.stop_reason
622+
623+
# Yield a final chunk with updated metadata including stop_reason
624+
yield AnthropicChatResponse(
625+
message=ChatMessage(
626+
role=role,
627+
blocks=content,
628+
additional_kwargs={
629+
"usage": usage_metadata if usage_metadata else None,
630+
"stop_reason": stop_reason,
631+
},
632+
),
633+
citations=cur_citations,
634+
delta="",
635+
raw=dict(r),
636+
)
622637
elif isinstance(r, RawMessageStopEvent):
623638
# Final event - no additional data to capture
624639
pass
@@ -851,6 +866,21 @@ async def gen() -> ChatResponseAsyncGen:
851866
}
852867
if hasattr(r, "delta") and hasattr(r.delta, "stop_reason"):
853868
stop_reason = r.delta.stop_reason
869+
870+
# Yield a final chunk with updated metadata including stop_reason
871+
yield AnthropicChatResponse(
872+
message=ChatMessage(
873+
role=role,
874+
blocks=content,
875+
additional_kwargs={
876+
"usage": usage_metadata if usage_metadata else None,
877+
"stop_reason": stop_reason,
878+
},
879+
),
880+
citations=cur_citations,
881+
delta="",
882+
raw=dict(r),
883+
)
854884
elif isinstance(r, RawMessageStopEvent):
855885
# Final event - no additional data to capture
856886
pass

0 commit comments

Comments
 (0)