community[minor]: Bedrock async methods (#12477)

Description: Added support for asynchronous streaming in the Bedrock
class and corresponding tests.

Primarily:
  async def aprepare_output_stream
    async def _aprepare_input_and_invoke_stream
    async def _astream
    async def _acall

I've ensured that the code adheres to the project's linting and
formatting standards by running make format, make lint, and make test.

Issue: #12054, #11589

Dependencies: None

Tag maintainer: @baskaryan 

Twitter handle: @dominic_lovric

---------

Co-authored-by: Piyush Jain <piyushjain@duck.com>
This commit is contained in:
DL
2024-01-22 23:44:49 +01:00
committed by GitHub
parent d6275e47f2
commit b9e7f6f38a
2 changed files with 220 additions and 16 deletions

View File

@@ -1,6 +1,14 @@
import json
from typing import AsyncGenerator, Dict
from unittest.mock import MagicMock, patch
import pytest
from langchain_community.llms.bedrock import ALTERNATION_ERROR, _human_assistant_format
from langchain_community.llms.bedrock import (
ALTERNATION_ERROR,
Bedrock,
_human_assistant_format,
)
TEST_CASES = {
"""Hey""": """
@@ -250,3 +258,51 @@ def test__human_assistant_format() -> None:
else:
output = _human_assistant_format(input_text)
assert output == expected_output
# Sample mock streaming response data
MOCK_STREAMING_RESPONSE = [
{"chunk": {"bytes": b'{"text": "nice"}'}},
{"chunk": {"bytes": b'{"text": " to meet"}'}},
{"chunk": {"bytes": b'{"text": " you"}'}},
]
async def async_gen_mock_streaming_response() -> AsyncGenerator[Dict, None]:
for item in MOCK_STREAMING_RESPONSE:
yield item
@pytest.mark.asyncio
async def test_bedrock_async_streaming_call() -> None:
# Mock boto3 import
mock_boto3 = MagicMock()
mock_boto3.Session.return_value.client.return_value = (
MagicMock()
) # Mocking the client method of the Session object
with patch.dict(
"sys.modules", {"boto3": mock_boto3}
): # Mocking boto3 at the top level using patch.dict
# Mock the `Bedrock` class's method that invokes the model
mock_invoke_method = MagicMock(return_value=async_gen_mock_streaming_response())
with patch.object(
Bedrock, "_aprepare_input_and_invoke_stream", mock_invoke_method
):
# Instantiate the Bedrock LLM
llm = Bedrock(
client=None,
model_id="anthropic.claude-v2",
streaming=True,
)
# Call the _astream method
chunks = [
json.loads(chunk["chunk"]["bytes"])["text"] # type: ignore
async for chunk in llm._astream("Hey, how are you?")
]
# Assertions
assert len(chunks) == 3
assert chunks[0] == "nice"
assert chunks[1] == " to meet"
assert chunks[2] == " you"