llama-index-llms-bedrock-converse 0.10.6__py3-none-any.whl → 0.11.0__py3-none-any.whl

llama_index/llms/bedrock_converse/base.py

@@ -4,6 +4,7 @@ from typing import (
     Callable,
     Dict,
     List,
+    Literal,
     Optional,
     Sequence,
     Tuple,
@@ -23,6 +24,7 @@ from llama_index.core.base.llms.types import (
     MessageRole,
     TextBlock,
     ThinkingBlock,
+    ToolCallBlock,
 )
 from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.callbacks import CallbackManager
@@ -157,6 +159,13 @@ class BedrockConverse(FunctionCallingLLM):
     guardrail_version: Optional[str] = Field(
         description="The version number for the guardrail. The value can also be DRAFT"
     )
+    guardrail_stream_processing_mode: Optional[Literal["sync", "async"]] = Field(
+        description=(
+            "The stream processing mode to use when leveraging a guardrail in a streaming request (ConverseStream). "
+            "If set, the specified mode will be included in the request's guardrail configuration object, altering the streaming response behavior. "
+            "If a value is not provided, no mode will be explicitly included in the request's guardrail configuration object, and thus Amazon Bedrock's default, Synchronous Mode, will be used."
+        )
+    )
     application_inference_profile_arn: Optional[str] = Field(
         description="The ARN of an application inference profile to invoke in place of the model. If provided, make sure the model argument refers to the same one underlying the application inference profile."
     )
@@ -211,6 +220,7 @@ class BedrockConverse(FunctionCallingLLM):
         output_parser: Optional[BaseOutputParser] = None,
         guardrail_identifier: Optional[str] = None,
         guardrail_version: Optional[str] = None,
+        guardrail_stream_processing_mode: Optional[Literal["sync", "async"]] = None,
         application_inference_profile_arn: Optional[str] = None,
         trace: Optional[str] = None,
         thinking: Optional[ThinkingDict] = None,
@@ -263,6 +273,7 @@ class BedrockConverse(FunctionCallingLLM):
             botocore_config=botocore_config,
             guardrail_identifier=guardrail_identifier,
             guardrail_version=guardrail_version,
+            guardrail_stream_processing_mode=guardrail_stream_processing_mode,
             application_inference_profile_arn=application_inference_profile_arn,
             trace=trace,
             thinking=thinking,
@@ -355,7 +366,7 @@ class BedrockConverse(FunctionCallingLLM):
     def _get_content_and_tool_calls(
         self, response: Optional[Dict[str, Any]] = None, content: Dict[str, Any] = None
     ) -> Tuple[
-        List[Union[TextBlock, ThinkingBlock]], Dict[str, Any], List[str], List[str]
+        List[Union[TextBlock, ThinkingBlock, ToolCallBlock]], List[str], List[str]
     ]:
         assert response is not None or content is not None, (
             f"Either response or content must be provided. Got response: {response}, content: {content}"
@@ -363,10 +374,9 @@ class BedrockConverse(FunctionCallingLLM):
         assert response is None or content is None, (
             f"Only one of response or content should be provided. Got response: {response}, content: {content}"
         )
-        tool_calls = []
         tool_call_ids = []
         status = []
-        blocks = []
+        blocks: List[TextBlock | ThinkingBlock | ToolCallBlock] = []
         if content is not None:
             content_list = [content]
         else:
@@ -391,7 +401,13 @@ class BedrockConverse(FunctionCallingLLM):
                     tool_usage["toolUseId"] = content_block["toolUseId"]
                 if "name" not in tool_usage:
                     tool_usage["name"] = content_block["name"]
-                tool_calls.append(tool_usage)
+                blocks.append(
+                    ToolCallBlock(
+                        tool_name=tool_usage.get("name", ""),
+                        tool_call_id=tool_usage.get("toolUseId"),
+                        tool_kwargs=tool_usage.get("input", {}),
+                    )
+                )
             if tool_result := content_block.get("toolResult", None):
                 for tool_result_content in tool_result["content"]:
                     if text := tool_result_content.get("text", None):
@@ -399,7 +415,7 @@ class BedrockConverse(FunctionCallingLLM):
                 tool_call_ids.append(tool_result_content.get("toolUseId", ""))
                 status.append(tool_result.get("status", ""))
 
-        return blocks, tool_calls, tool_call_ids, status
+        return blocks, tool_call_ids, status
 
     @llm_chat_callback()
     def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
@@ -426,16 +442,13 @@ class BedrockConverse(FunctionCallingLLM):
             **all_kwargs,
         )
 
-        blocks, tool_calls, tool_call_ids, status = self._get_content_and_tool_calls(
-            response
-        )
+        blocks, tool_call_ids, status = self._get_content_and_tool_calls(response)
 
         return ChatResponse(
             message=ChatMessage(
                 role=MessageRole.ASSISTANT,
                 blocks=blocks,
                 additional_kwargs={
-                    "tool_calls": tool_calls,
                     "tool_call_id": tool_call_ids,
                     "status": status,
                 },
@@ -474,6 +487,7 @@ class BedrockConverse(FunctionCallingLLM):
             stream=True,
             guardrail_identifier=self.guardrail_identifier,
             guardrail_version=self.guardrail_version,
+            guardrail_stream_processing_mode=self.guardrail_stream_processing_mode,
             trace=self.trace,
             **all_kwargs,
         )
@@ -529,7 +543,7 @@ class BedrockConverse(FunctionCallingLLM):
                                     current_tool_call, tool_use_delta
                                 )
 
-                    blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                    blocks: List[Union[TextBlock, ThinkingBlock, ToolCallBlock]] = [
                         TextBlock(text=content.get("text", ""))
                     ]
                     if thinking != "":
@@ -542,13 +556,21 @@ class BedrockConverse(FunctionCallingLLM):
                                 },
                             ),
                         )
+                    if tool_calls:
+                        for tool_call in tool_calls:
+                            blocks.append(
+                                ToolCallBlock(
+                                    tool_kwargs=tool_call.get("input", {}),
+                                    tool_name=tool_call.get("name", ""),
+                                    tool_call_id=tool_call.get("toolUseId"),
+                                )
+                            )
 
                     yield ChatResponse(
                         message=ChatMessage(
                             role=role,
                             blocks=blocks,
                             additional_kwargs={
-                                "tool_calls": tool_calls,
                                 "tool_call_id": [
                                     tc.get("toolUseId", "") for tc in tool_calls
                                 ],
@@ -568,7 +590,7 @@ class BedrockConverse(FunctionCallingLLM):
                         # Add to our list of tool calls
                         tool_calls.append(current_tool_call)
 
-                    blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                    blocks: List[Union[TextBlock, ThinkingBlock, ToolCallBlock]] = [
                         TextBlock(text=content.get("text", ""))
                     ]
                     if thinking != "":
@@ -582,12 +604,21 @@ class BedrockConverse(FunctionCallingLLM):
                             ),
                         )
 
+                    if tool_calls:
+                        for tool_call in tool_calls:
+                            blocks.append(
+                                ToolCallBlock(
+                                    tool_kwargs=tool_call.get("input", {}),
+                                    tool_name=tool_call.get("name", ""),
+                                    tool_call_id=tool_call.get("toolUseId"),
+                                )
+                            )
+
                     yield ChatResponse(
                         message=ChatMessage(
                             role=role,
                             blocks=blocks,
                             additional_kwargs={
-                                "tool_calls": tool_calls,
                                 "tool_call_id": [
                                     tc.get("toolUseId", "") for tc in tool_calls
                                 ],
@@ -604,7 +635,7 @@ class BedrockConverse(FunctionCallingLLM):
                     # Handle metadata event - this contains the final token usage
                     if usage := metadata.get("usage"):
                         # Yield a final response with correct token usage
-                        blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                        blocks: List[Union[TextBlock, ThinkingBlock, ToolCallBlock]] = [
                             TextBlock(text=content.get("text", ""))
                         ]
                         if thinking != "":
@@ -617,13 +648,21 @@ class BedrockConverse(FunctionCallingLLM):
                                     },
                                 ),
                             )
+                        if tool_calls:
+                            for tool_call in tool_calls:
+                                blocks.append(
+                                    ToolCallBlock(
+                                        tool_kwargs=tool_call.get("input", {}),
+                                        tool_name=tool_call.get("name", ""),
+                                        tool_call_id=tool_call.get("toolUseId"),
+                                    )
+                                )
 
                         yield ChatResponse(
                             message=ChatMessage(
                                 role=role,
                                 blocks=blocks,
                                 additional_kwargs={
-                                    "tool_calls": tool_calls,
                                     "tool_call_id": [
                                         tc.get("toolUseId", "") for tc in tool_calls
                                     ],
@@ -668,21 +707,19 @@ class BedrockConverse(FunctionCallingLLM):
             stream=False,
             guardrail_identifier=self.guardrail_identifier,
             guardrail_version=self.guardrail_version,
+            guardrail_stream_processing_mode=self.guardrail_stream_processing_mode,
             trace=self.trace,
             boto_client_kwargs=self._boto_client_kwargs,
             **all_kwargs,
         )
 
-        blocks, tool_calls, tool_call_ids, status = self._get_content_and_tool_calls(
-            response
-        )
+        blocks, tool_call_ids, status = self._get_content_and_tool_calls(response)
 
         return ChatResponse(
             message=ChatMessage(
                 role=MessageRole.ASSISTANT,
                 blocks=blocks,
                 additional_kwargs={
-                    "tool_calls": tool_calls,
                     "tool_call_id": tool_call_ids,
                     "status": status,
                 },
@@ -777,7 +814,7 @@ class BedrockConverse(FunctionCallingLLM):
                                 current_tool_call = join_two_dicts(
                                     current_tool_call, tool_use_delta
                                 )
-                    blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                    blocks: List[Union[TextBlock, ThinkingBlock, ToolCallBlock]] = [
                         TextBlock(text=content.get("text", ""))
                     ]
                     if thinking != "":
@@ -791,12 +828,21 @@ class BedrockConverse(FunctionCallingLLM):
                             ),
                         )
 
+                    if tool_calls:
+                        for tool_call in tool_calls:
+                            blocks.append(
+                                ToolCallBlock(
+                                    tool_kwargs=tool_call.get("input", {}),
+                                    tool_name=tool_call.get("name", ""),
+                                    tool_call_id=tool_call.get("toolUseId"),
+                                )
+                            )
+
                     yield ChatResponse(
                         message=ChatMessage(
                             role=role,
                             blocks=blocks,
                             additional_kwargs={
-                                "tool_calls": tool_calls,
                                 "tool_call_id": [
                                     tc.get("toolUseId", "") for tc in tool_calls
                                 ],
@@ -816,7 +862,7 @@ class BedrockConverse(FunctionCallingLLM):
                         # Add to our list of tool calls
                         tool_calls.append(current_tool_call)
 
-                    blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                    blocks: List[Union[TextBlock, ThinkingBlock, ToolCallBlock]] = [
                         TextBlock(text=content.get("text", ""))
                     ]
                     if thinking != "":
@@ -830,12 +876,21 @@ class BedrockConverse(FunctionCallingLLM):
                             ),
                         )
 
+                    if tool_calls:
+                        for tool_call in tool_calls:
+                            blocks.append(
+                                ToolCallBlock(
+                                    tool_kwargs=tool_call.get("input", {}),
+                                    tool_name=tool_call.get("name", ""),
+                                    tool_call_id=tool_call.get("toolUseId"),
+                                )
+                            )
+
                     yield ChatResponse(
                         message=ChatMessage(
                             role=role,
                             blocks=blocks,
                             additional_kwargs={
-                                "tool_calls": tool_calls,
                                 "tool_call_id": [
                                     tc.get("toolUseId", "") for tc in tool_calls
                                 ],
@@ -852,7 +907,7 @@ class BedrockConverse(FunctionCallingLLM):
                     # Handle metadata event - this contains the final token usage
                     if usage := metadata.get("usage"):
                         # Yield a final response with correct token usage
-                        blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                        blocks: List[Union[TextBlock, ThinkingBlock, ToolCallBlock]] = [
                             TextBlock(text=content.get("text", ""))
                         ]
                         if thinking != "":
@@ -866,12 +921,21 @@ class BedrockConverse(FunctionCallingLLM):
                                 ),
                             )
 
+                        if tool_calls:
+                            for tool_call in tool_calls:
+                                blocks.append(
+                                    ToolCallBlock(
+                                        tool_kwargs=tool_call.get("input", {}),
+                                        tool_name=tool_call.get("name", ""),
+                                        tool_call_id=tool_call.get("toolUseId"),
+                                    )
+                                )
+
                         yield ChatResponse(
                             message=ChatMessage(
                                 role=role,
                                 blocks=blocks,
                                 additional_kwargs={
-                                    "tool_calls": tool_calls,
                                     "tool_call_id": [
                                         tc.get("toolUseId", "") for tc in tool_calls
                                     ],
@@ -948,7 +1012,11 @@ class BedrockConverse(FunctionCallingLLM):
         **kwargs: Any,
     ) -> List[ToolSelection]:
         """Predict and call the tool."""
-        tool_calls = response.message.additional_kwargs.get("tool_calls", [])
+        tool_calls = [
+            block
+            for block in response.message.blocks
+            if isinstance(block, ToolCallBlock)
+        ]
 
         if len(tool_calls) < 1:
             if error_on_no_tool_call:
@@ -960,26 +1028,23 @@ class BedrockConverse(FunctionCallingLLM):
 
         tool_selections = []
         for tool_call in tool_calls:
-            if "toolUseId" not in tool_call or "name" not in tool_call:
-                raise ValueError("Invalid tool call.")
-
             # handle empty inputs
             argument_dict = {}
-            if "input" in tool_call and isinstance(tool_call["input"], str):
+            if isinstance(tool_call.tool_kwargs, str):
                 # TODO parse_partial_json is not perfect
                 try:
-                    argument_dict = parse_partial_json(tool_call["input"])
+                    argument_dict = parse_partial_json(tool_call.tool_kwargs)
                 except ValueError:
                     argument_dict = {}
-            elif "input" in tool_call and isinstance(tool_call["input"], dict):
-                argument_dict = tool_call["input"]
+            elif isinstance(tool_call.tool_kwargs, dict):
+                argument_dict = tool_call.tool_kwargs
             else:
                 continue
 
             tool_selections.append(
                 ToolSelection(
-                    tool_id=tool_call["toolUseId"],
-                    tool_name=tool_call["name"],
+                    tool_id=tool_call.tool_call_id or "",
+                    tool_name=tool_call.tool_name,
                     tool_kwargs=argument_dict,
                 )
             )

llama_index/llms/bedrock_converse/utils.py

@@ -32,6 +32,7 @@ from llama_index.core.base.llms.types import (
     DocumentBlock,
     CachePoint,
     ThinkingBlock,
+    ToolCallBlock,
 )
 
 
@@ -303,6 +304,23 @@ def _content_block_to_bedrock_format(
     elif isinstance(block, AudioBlock):
         logger.warning("Audio blocks are not supported in Bedrock Converse API.")
         return None
+    elif isinstance(block, ToolCallBlock):
+        if isinstance(block.tool_kwargs, str):
+            try:
+                tool_input = json.loads(block.tool_kwargs or "{}")
+            except json.JSONDecodeError:
+                tool_input = {}
+        else:
+            tool_input = block.tool_kwargs
+
+        return {
+            "toolUse": {
+                "input": tool_input,
+                "toolUseId": block.tool_call_id or "",
+                "name": block.tool_name,
+            }
+        }
+
     else:
         logger.warning(f"Unsupported block type: {type(block)}")
         return None
@@ -345,7 +363,9 @@ def messages_to_converse_messages(
     converse_messages = []
     system_prompt = []
     current_system_prompt = ""
+
     for message in messages:
+        unique_tool_calls = []
         if message.role == MessageRole.SYSTEM:
             # we iterate over blocks, if content was used, the blocks are added anyway
             for block in message.blocks:
@@ -402,6 +422,13 @@ def messages_to_converse_messages(
                 )
                 if bedrock_format_block:
                     content.append(bedrock_format_block)
+                    if "toolUse" in bedrock_format_block:
+                        unique_tool_calls.append(
+                            (
+                                bedrock_format_block["toolUse"]["toolUseId"],
+                                bedrock_format_block["toolUse"]["name"],
+                            )
+                        )
 
             if content:
                 converse_messages.append(
@@ -411,6 +438,7 @@ def messages_to_converse_messages(
                     }
                 )
 
+        # keep this code here for compatibility with older chat histories
         # convert tool calls to the AWS Bedrock Converse format
         # NOTE tool calls might show up within any message,
         # e.g. within assistant message or in consecutive tool calls,
@@ -418,25 +446,28 @@ def messages_to_converse_messages(
         tool_calls = message.additional_kwargs.get("tool_calls", [])
         content = []
         for tool_call in tool_calls:
-            assert "toolUseId" in tool_call, f"`toolUseId` not found in {tool_call}"
-            assert "input" in tool_call, f"`input` not found in {tool_call}"
-            assert "name" in tool_call, f"`name` not found in {tool_call}"
-            tool_input = tool_call["input"] if tool_call["input"] else {}
-            if isinstance(tool_input, str):
-                try:
-                    tool_input = json.loads(tool_input or "{}")
-                except json.JSONDecodeError:
-                    tool_input = {}
-
-            content.append(
-                {
-                    "toolUse": {
-                        "input": tool_input,
-                        "toolUseId": tool_call["toolUseId"],
-                        "name": tool_call["name"],
-                    }
-                }
-            )
+            try:
+                assert "toolUseId" in tool_call
+                assert "input" in tool_call
+                assert "name" in tool_call
+                if (tool_call["toolUseId"], tool_call["name"]) not in unique_tool_calls:
+                    tool_input = tool_call["input"] if tool_call["input"] else {}
+                    if isinstance(tool_input, str):
+                        try:
+                            tool_input = json.loads(tool_input or "{}")
+                        except json.JSONDecodeError:
+                            tool_input = {}
+                    content.append(
+                        {
+                            "toolUse": {
+                                "input": tool_input,
+                                "toolUseId": tool_call["toolUseId"],
+                                "name": tool_call["name"],
+                            }
+                        }
+                    )
+            except AssertionError:
+                continue
         if len(content) > 0:
             converse_messages.append(
                 {
@@ -499,9 +530,15 @@ def tools_to_converse_tools(
 
 
 def force_single_tool_call(response: ChatResponse) -> None:
-    tool_calls = response.message.additional_kwargs.get("tool_calls", [])
+    tool_calls = [
+        block for block in response.message.blocks if isinstance(block, ToolCallBlock)
+    ]
     if len(tool_calls) > 1:
-        response.message.additional_kwargs["tool_calls"] = [tool_calls[0]]
+        response.message.blocks = [
+            block
+            for block in response.message.blocks
+            if not isinstance(block, ToolCallBlock)
+        ] + [tool_calls[0]]
 
 
 def _create_retry_decorator(client: Any, max_retries: int) -> Callable[[Any], Any]:
@@ -563,6 +600,7 @@ def converse_with_retry(
     stream: bool = False,
     guardrail_identifier: Optional[str] = None,
     guardrail_version: Optional[str] = None,
+    guardrail_stream_processing_mode: Optional[Literal["sync", "async"]] = None,
     trace: Optional[str] = None,
     **kwargs: Any,
 ) -> Any:
@@ -603,6 +641,10 @@ def converse_with_retry(
         converse_kwargs["guardrailConfig"]["guardrailVersion"] = guardrail_version
         if trace:
             converse_kwargs["guardrailConfig"]["trace"] = trace
+        if guardrail_stream_processing_mode and stream:
+            converse_kwargs["guardrailConfig"]["streamProcessingMode"] = (
+                guardrail_stream_processing_mode
+            )
 
     converse_kwargs = join_two_dicts(
         converse_kwargs,
@@ -644,6 +686,7 @@ async def converse_with_retry_async(
     stream: bool = False,
     guardrail_identifier: Optional[str] = None,
     guardrail_version: Optional[str] = None,
+    guardrail_stream_processing_mode: Optional[Literal["sync", "async"]] = None,
     trace: Optional[str] = None,
     boto_client_kwargs: Optional[Dict[str, Any]] = None,
     **kwargs: Any,
@@ -690,6 +733,10 @@ async def converse_with_retry_async(
         converse_kwargs["guardrailConfig"]["guardrailVersion"] = guardrail_version
         if trace:
             converse_kwargs["guardrailConfig"]["trace"] = trace
+        if guardrail_stream_processing_mode and stream:
+            converse_kwargs["guardrailConfig"]["streamProcessingMode"] = (
+                guardrail_stream_processing_mode
+            )
     converse_kwargs = join_two_dicts(
         converse_kwargs,
         {

{llama_index_llms_bedrock_converse-0.10.6.dist-info → llama_index_llms_bedrock_converse-0.11.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llama-index-llms-bedrock-converse
-Version: 0.10.6
+Version: 0.11.0
 Summary: llama-index llms bedrock converse integration
 Author-email: Your Name <you@example.com>
 License-Expression: MIT
@@ -8,7 +8,7 @@ License-File: LICENSE
 Requires-Python: <4.0,>=3.9
 Requires-Dist: aioboto3<16,>=15.0.0
 Requires-Dist: boto3<2,>=1.38.27
-Requires-Dist: llama-index-core<0.15,>=0.14.3
+Requires-Dist: llama-index-core<0.15,>=0.14.5
 Description-Content-Type: text/markdown
 
 # LlamaIndex Llms Integration: Bedrock Converse

llama_index_llms_bedrock_converse-0.11.0.dist-info/RECORD

@@ -0,0 +1,7 @@
+llama_index/llms/bedrock_converse/__init__.py,sha256=xE3ZHLXqFr7TTTgQlYH9bLLPRZAV3dJyiz_iUFXBfak,98
+llama_index/llms/bedrock_converse/base.py,sha256=tMex6AHod6YofxaLX9RsucVHBjNU-a_KP9OAgFvRHXo,45059
+llama_index/llms/bedrock_converse/utils.py,sha256=XgAWPSGrEft8bh9NMzhmSIjrKveXXjzIkBOJmeHf9xc,29747
+llama_index_llms_bedrock_converse-0.11.0.dist-info/METADATA,sha256=YgCAv1cu4yxHtvQZm6M2w27-tK_8RWy8wGFWRz5ZqB0,7834
+llama_index_llms_bedrock_converse-0.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+llama_index_llms_bedrock_converse-0.11.0.dist-info/licenses/LICENSE,sha256=JPQLUZD9rKvCTdu192Nk0V5PAwklIg6jANii3UmTyMs,1065
+llama_index_llms_bedrock_converse-0.11.0.dist-info/RECORD,,

llama_index_llms_bedrock_converse-0.10.6.dist-info/RECORD

@@ -1,7 +0,0 @@
-llama_index/llms/bedrock_converse/__init__.py,sha256=xE3ZHLXqFr7TTTgQlYH9bLLPRZAV3dJyiz_iUFXBfak,98
-llama_index/llms/bedrock_converse/base.py,sha256=MISQTYCwaUIinyOGZhOdPYQYQAFrRSLAkiTEE-8f21A,41478
-llama_index/llms/bedrock_converse/utils.py,sha256=-SDwcbgcX_W7p6SEt6-3HH_cAsKxG_4e2wHhi5kFKJU,27893
-llama_index_llms_bedrock_converse-0.10.6.dist-info/METADATA,sha256=0y3xka1zYDazsks_i17_00wj89QwvAu2tBB-s4k2sXE,7834
-llama_index_llms_bedrock_converse-0.10.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-llama_index_llms_bedrock_converse-0.10.6.dist-info/licenses/LICENSE,sha256=JPQLUZD9rKvCTdu192Nk0V5PAwklIg6jANii3UmTyMs,1065
-llama_index_llms_bedrock_converse-0.10.6.dist-info/RECORD,,