Skip to content

Commit

Permalink
fixes due to conflicts
Browse files Browse the repository at this point in the history
  • Loading branch information
davidsbatista committed Aug 12, 2024
1 parent 6fbce0a commit f3694dd
Showing 1 changed file with 3 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -198,18 +198,18 @@ def run(
body=json.dumps(body), modelId=self.model, accept="application/json", contentType="application/json"
)
response_body = json.loads(response.get("body").read().decode("utf-8"))
responses = self.model_adapter.get_responses(response_body=response_body)
replies = self.model_adapter.get_responses(response_body=response_body)
except ClientError as exception:
msg = f"Could not inference Amazon Bedrock model {self.model} due: {exception}"
raise AmazonBedrockInferenceError(msg) from exception

# rename the meta key to be inline with OpenAI meta output keys
for response in responses:
for response in replies:
if response.meta is not None and "usage" in response.meta:
response.meta["usage"]["prompt_tokens"] = response.meta["usage"].pop("input_tokens")
response.meta["usage"]["completion_tokens"] = response.meta["usage"].pop("output_tokens")

return responses
return {"replies": replies}

@classmethod
def get_model_adapter(cls, model: str) -> Optional[Type[BedrockModelChatAdapter]]:
Expand Down

0 comments on commit f3694dd

Please sign in to comment.