Repository: jerryjliu/gpt_index
File: llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/base.py (Python)

Code before:
                        ),
                        raw=chunk,
                    )
                elif message_stop := chunk.get("messageStop"):
                    # Handle messageStop event - this contains the stop reason
                    # We don't yield here, just track the event
                    pass
                elif metadata := chunk.get("metadata"):
                    # Handle metadata event - this contains the final token usage
                    if usage := metadata.get("usage"):
                        final_usage = usage
                        # Yield a final response with correct token usage
                        yield ChatResponse(
                            message=ChatMessage(
                                role=role,
                                content=content.get("text", ""),
                                additional_kwargs={
                                    "tool_calls": tool_calls,
                                    "tool_call_id": [
                                        tc.get("toolUseId", "") for tc in tool_calls
                                    ],
                                    "status": [],
                                },
                            ),
                            delta="",
                            raw={"usage": final_usage},
                            additional_kwargs=self._get_response_token_counts({"usage": final_usage}),
                        )

        return gen()

    @llm_completion_callback()
    def stream_complete(
        self, prompt: str, formatted: bool = False, **kwargs: Any
    ) -> CompletionResponseGen:
        stream_complete_fn = stream_chat_to_completion_decorator(self.stream_chat)
        return stream_complete_fn(prompt, **kwargs)

    @llm_chat_callback()
    async def achat(
        self, messages: Sequence[ChatMessage], **kwargs: Any
    ) -> ChatResponse:
        # convert Llama Index messages to AWS Bedrock Converse messages
        converse_messages, system_prompt = messages_to_converse_messages(messages)
        all_kwargs = self._get_all_kwargs(**kwargs)

        # invoke LLM in AWS Bedrock Converse with retry
        response = await converse_with_retry_async(
            session=self._asession,
            config=self._config,
            messages=converse_messages,

Review comment (suggestion):
I am pretty sure that `raw` here should contain the raw response from BedrockConverse and not the usage metadata(?)
I would add that only to additional_kwargs, as you do below

Code after:
                        raw=chunk,
                    )
                elif message_stop := chunk.get("messageStop"):
                    # Handle messageStop event - this contains the stop reason
                    # We don't yield here, just track the event
                    pass
                elif metadata := chunk.get("metadata"):
                    # Handle metadata event - this contains the final token usage
                    if usage := metadata.get("usage"):
                        # Yield a final response with correct token usage
                        yield ChatResponse(
                            message=ChatMessage(
                                role=role,
                                content=content.get("text", ""),
                                additional_kwargs={
                                    "tool_calls": tool_calls,
                                    "tool_call_id": [
                                        tc.get("toolUseId", "") for tc in tool_calls
                                    ],
                                    "status": [],
                                },
                            ),
                            delta="",
                            raw=chunk,
                            additional_kwargs=self._get_response_token_counts(metadata),
                        )

        return gen()

    @llm_completion_callback()
    def stream_complete(
        self, prompt: str, formatted: bool = False, **kwargs: Any
    ) -> CompletionResponseGen:
        stream_complete_fn = stream_chat_to_completion_decorator(self.stream_chat)
        return stream_complete_fn(prompt, **kwargs)

    @llm_chat_callback()
    async def achat(
        self, messages: Sequence[ChatMessage], **kwargs: Any
    ) -> ChatResponse:
        # convert Llama Index messages to AWS Bedrock Converse messages
        converse_messages, system_prompt = messages_to_converse_messages(messages)
        all_kwargs = self._get_all_kwargs(**kwargs)

        # invoke LLM in AWS Bedrock Converse with retry
        response = await converse_with_retry_async(
            session=self._asession,
            config=self._config,
            messages=converse_messages,
            system_prompt=system_prompt,
            max_retries=self.max_retries,