diff --git a/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/base.py b/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/base.py index 493b3c2033995d259b7f562f8393fbb5216f3ed5..706772fec648a3a1f88328dcdbc55ee2d28d3a4c 100644 --- a/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/base.py +++ b/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/base.py @@ -454,7 +454,7 @@ class BedrockConverse(FunctionCallingLLM): all_kwargs = self._get_all_kwargs(**kwargs) # invoke LLM in AWS Bedrock Converse with retry - response = await converse_with_retry_async( + response_gen = await converse_with_retry_async( session=self._asession, config=self._config, messages=converse_messages, @@ -467,7 +467,7 @@ class BedrockConverse(FunctionCallingLLM): async def gen() -> ChatResponseAsyncGen: content = {} role = MessageRole.ASSISTANT - async for chunk in response["stream"]: + async for chunk in response_gen: if content_block_delta := chunk.get("contentBlockDelta"): content_delta = content_block_delta["delta"] content = join_two_dicts(content, content_delta) @@ -489,7 +489,7 @@ class BedrockConverse(FunctionCallingLLM): }, ), delta=content_delta.get("text", ""), - raw=response, + raw=chunk, ) elif content_block_start := chunk.get("contentBlockStart"): tool_use = content_block_start["toolUse"] @@ -511,7 +511,7 @@ class BedrockConverse(FunctionCallingLLM): "status": status, }, ), - raw=response, + raw=chunk, ) return gen() diff --git a/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/utils.py b/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/utils.py index 433753e6f2f55676d0b263a98ed5ef98c85b76f9..31be9c1d1795339dd88d4a4c87004cafa6b273aa 100644 --- a/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/utils.py +++ b/llama-index-integrations/llms/llama-index-llms-bedrock-converse/llama_index/llms/bedrock_converse/utils.py @@ -357,15 +357,27 @@ async def converse_with_retry_async( converse_kwargs, {k: v for k, v in kwargs.items() if k != "tools"} ) + ## NOTE: Returning the generator directly from converse_stream doesn't work + # So, we have to use two separate functions for streaming and non-streaming + # This differs from the synchronous version, and is a bit of a hack + # Further investigation is needed + @retry_decorator async def _conversion_with_retry(**kwargs: Any) -> Any: - # the async boto3 client needs to be defined inside this async with, otherwise it will raise an error async with session.client("bedrock-runtime", config=config) as client: - if stream: - return await client.converse_stream(**kwargs) return await client.converse(**kwargs) - return await _conversion_with_retry(**converse_kwargs) + @retry_decorator + async def _conversion_stream_with_retry(**kwargs: Any) -> Any: + async with session.client("bedrock-runtime", config=config) as client: + response = await client.converse_stream(**kwargs) + async for event in response["stream"]: + yield event + + if stream: + return _conversion_stream_with_retry(**converse_kwargs) + else: + return await _conversion_with_retry(**converse_kwargs) def join_two_dicts(dict1: Dict[str, Any], dict2: Dict[str, Any]) -> Dict[str, Any]: diff --git a/llama-index-integrations/llms/llama-index-llms-bedrock-converse/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-bedrock-converse/pyproject.toml index 81b63353e3fa97d044c6bde9a49d7b79fd5ae309..00d62b04e4fac087373bdeebff92f67a2847cf3a 100644 --- a/llama-index-integrations/llms/llama-index-llms-bedrock-converse/pyproject.toml +++ b/llama-index-integrations/llms/llama-index-llms-bedrock-converse/pyproject.toml @@ -27,7 +27,7 @@ exclude = ["**/BUILD"] license = "MIT" name = "llama-index-llms-bedrock-converse" readme = "README.md" -version = "0.3.8" +version = "0.3.9" [tool.poetry.dependencies] python = ">=3.8.1,<4.0"