phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

return traces in bedrock guardrails when enabled

Browse source
This commit is contained in:
Ishaan Jaff 2024-08-16 11:35:43 -07:00
parent 89ba7b3e11
commit 9851fa7b1b
1 changed files with 8 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/utils.py
View file

@ -10226,6 +10226,14 @@ class CustomStreamWrapper:
return return
elif self.received_finish_reason is not None: elif self.received_finish_reason is not None:
if self.sent_last_chunk is True: if self.sent_last_chunk is True:
# Bedrock returns the guardrail trace in the last chunk - we want to return this here
if (
self.custom_llm_provider == "bedrock"
and "trace" in model_response
):
return model_response
# Default - return StopIteration
raise StopIteration raise StopIteration
# flush any remaining holding chunk # flush any remaining holding chunk
if len(self.holding_chunk) > 0: if len(self.holding_chunk) > 0: