Move yield of metrics chunk after generation chunk

[ihmaws]

Move yield of metrics chunk after generation chunk

• when using mistral and streaming is enabled,the final chunk includes a stop_reason. There is nothing to say this final chunk doesn't also include some generated text. The existing implementation would result in that final chunk never getting sent back
• this update moves the yield of the metrics chunk after the generation chunk
• also included a change to include invocation metrics for cohere models

Closes #215

[SanjayTiwaryMLAI]

@3coins, Need urgent help here as there is issue with the production deployment GenAI Chatbot for Government e market place. (https://github.com/ihmaws)has done deep dive on the issue and suggested few resolution. Can you look and please accept the PR.

This is first India Public sector generative AI chatbot, customer is asking for resolution on urgent basis

[3coins]

@ihmaws
Thanks for submitting this change. One minor comment about removing the block for cohere, looks good otherwise.

[3coins]

This ends up adding an additional chunk to the response with an empty text. For example:

# Last text chunk with content
content='.' additional_kwargs={} response_metadata={'stop_reason': None} id='run-98ddf143-b333-47de-b8e5-b0dd55832cf6'

# empty content chunk processed for case when metrics are present
content='' additional_kwargs={} response_metadata={'stop_reason': 'stop', 'amazon-bedrock-invocationMetrics': {'inputTokenCount': 33, 'outputTokenCount': 10, 'invocationLatency': 230, 'firstByteLatency': 139}} id='run-98ddf143-b333-47de-b8e5-b0dd55832cf6'

# This one is processed with the _get_invocation_metrics_chunk
content='' additional_kwargs={} response_metadata={} id='run-98ddf143-b333-47de-b8e5-b0dd55832cf6' usage_metadata={'input_tokens': 33, 'output_tokens': 10, 'total_tokens': 43}

Ideally, we should be able to extract both the content and metrics in one go inside the _stream_response_to_generation_chunk method, not sure why we ended up with 2 different statements. I am not sure if Bedrock or any of the models guarantee the metrics as a last chunk; if that ends up not the case, this logic will break because of the return statements. However, this change seems harmless, maybe we can tackle the other issue in a separate PR.

[3coins]

@ihmaws
One of integration tests is failing, can you verify if this passes for you, otherwise, update it.

======================================================= FAILURES ========================================================
__________________________________________________ test_bedrock_invoke __________________________________________________

chat = ChatBedrock(client=<botocore.client.BedrockRuntime object at 0x121fccd90>, region_name='us-west-2', model_id='anthropic.claude-3-sonnet-20240229-v1:0', model_kwargs={'temperature': 0})

    @pytest.mark.scheduled
    def test_bedrock_invoke(chat: ChatBedrock) -> None:
        """Test invoke tokens from ChatBedrock."""
        result = chat.invoke("I'm Pickle Rick", config=dict(tags=["foo"]))
        assert isinstance(result.content, str)
        assert "usage" in result.additional_kwargs
>       assert result.additional_kwargs["usage"]["prompt_tokens"] == 12
E       assert 13 == 12

tests/integration_tests/chat_models/test_bedrock.py:240: AssertionError

Use this to verify:

poetry run pytest tests/integration_tests/chat_models/test_bedrock.py -k "test_bedrock_invoke"

[ihmaws]

@ihmaws One of integration tests is failing, can you verify if this passes for you, otherwise, update it.

Interesting, I'm not seeing that test failure. Let me try and have a look. Can you also run the test one more time and confirm you've got the latest code?

[3coins]

Interesting, I'm not seeing that test failure. Let me try and have a look. Can you also run the test one more time and confirm you've got the latest code?

Let me try pulling in latest.