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

Merge pull request #5243 from BerriAI/litellm_add_bedrock_traces_in_response

Browse source 
[Feat] Add bedrock Guardrail `traces ` in response when trace=enabled
This commit is contained in:
Ishaan Jaff 2024-08-16 14:49:20 -07:00 • committed by GitHub
commit cff01b2de3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 90 additions and 23 deletions
Download patch file Download diff file Expand all files Collapse all files

19
litellm/utils.py
View file

@ -9581,7 +9581,8 @@ class CustomStreamWrapper:
):
if self.received_finish_reason is not None:
raise StopIteration
if "provider_specific_fields" not in chunk:
raise StopIteration
anthropic_response_obj: GChunk = chunk
completion_obj["content"] = anthropic_response_obj["text"]
if anthropic_response_obj["is_finished"]:
@ -9604,6 +9605,14 @@ class CustomStreamWrapper:
):
completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]]
if (
"provider_specific_fields" in anthropic_response_obj
and anthropic_response_obj["provider_specific_fields"] is not None
):
for key, value in anthropic_response_obj[
"provider_specific_fields"
].items():
setattr(model_response, key, value)
response_obj = anthropic_response_obj
elif (
self.custom_llm_provider
@ -10219,6 +10228,14 @@ class CustomStreamWrapper:
return
elif self.received_finish_reason is not None:
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
# flush any remaining holding chunk
if len(self.holding_chunk) > 0: