feat: add stream support for Bedrock #191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Merged

miri-bar merged 14 commits into main from bedrock_streaming

Aug 6, 2024

Contributor

miri-bar commented Jul 28, 2024

No description provided.

miri-bar added 2 commits

July 28, 2024 14:31


          feat: add stream support for bedrock

b0dc5a8


          chore: fix error

6b06c8b

github-actions bot added documentation aws size:xl labels


          chore: fix stream examples

f63df64

github-actions bot added size:xl and removed size:xl labels


          Merge branch 'main' into bedrock_streaming

b43131a

github-actions bot added size:xl and removed size:xl labels


          fix: add fix for last chunk on bedrock, add tests

f778eea

github-actions bot added size:xl and removed size:xl labels


          test: fix unittests

b01df9e

github-actions bot added size:xl and removed size:xl labels


          test: remove comments

661c5f5

github-actions bot added size:xl and removed size:xl labels


          test: fix integration test

ae2c38e

github-actions bot added size:xl and removed size:xl labels

miri-bar requested a review from a team

July 29, 2024 15:30


          Merge branch 'main' into bedrock_streaming

0c7a7d0

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/ai21_bedrock_client.py Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/studio/resources/studio_resource.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/http_client/async_http_client.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/errors.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/ai21_bedrock_client.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/models/ai21_base_model.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

tests/integration_tests/clients/bedrock/test_chat_completions.py Outdated Show resolved Hide resolved

miri-bar and others added 2 commits

August 4, 2024 16:42


          Merge branch 'main' into bedrock_streaming

85b9137


          chore: cr comments

miri-bar requested a review from Josephasafg

August 4, 2024 14:04

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/stream/stream.py Show resolved Hide resolved

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated Show resolved Hide resolved

miri-bar added 2 commits

August 5, 2024 18:17


          refactor: cr comments

8616f4e


          refactor: async stream example

4f48869

Josephasafg reviewed

View reviewed changes

ai21/clients/bedrock/_stream_decoder.py Outdated

    
                      chunk_dict = {**chunk_dict, **bedrock_metrics_dict}

                      return json.dumps(chunk_dict)

                  def _process_chunks(self, event_stream_buffer, chunk):

Contributor

Josephasafg Aug 5, 2024

Missing types

Josephasafg reviewed

View reviewed changes

ai21/errors.py Outdated

    
              class StreamingDecodeError(AI21Error):

                  def __init__(self, chunk: str, error_message: Optional[str] = None):

                      message = f"Failed to decode chunk: {chunk} in stream. Please check the stream format"

                      message = f"Failed to decode chunk: {chunk} in stream. Please check the stream format. "

Contributor

Josephasafg Aug 5, 2024

Its weird to add a space at the end of the string "just in case"
You can do something like -

if error_message:
    message = f"{message} Error: {error_message}"

Josephasafg reviewed

View reviewed changes

ai21/stream/stream_commons.py Outdated

    
                      pass

              class _SSEDecoder(SSEDecoderBase):

Contributor

Josephasafg Aug 5, 2024

If this class is protected then I assume the rest should be too no?

Josephasafg approved these changes

View reviewed changes

github-actions bot added the lgtm label


          refactor: cr comments

fad2c49

miri-bar merged commit 5392b52 into main

miri-bar deleted the bedrock_streaming branch

August 6, 2024 05:59

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

aws documentation lgtm size:xl