pass trace through for bedrock guardrails

This commit is contained in:
Ishaan Jaff 2024-08-16 09:10:56 -07:00
parent 026c7194f8
commit 89ba7b3e11
3 changed files with 82 additions and 23 deletions

View file

@ -9579,7 +9579,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"]:
@ -9602,6 +9603,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