mirror of
https://github.com/hwchase17/langchain.git
synced 2025-09-06 21:43:44 +00:00
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:
@@ -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"
|
||||
|
Reference in New Issue
Block a user