llama-index-llms-bedrock-converse 0.9.5__py3-none-any.whl → 0.10.1__py3-none-any.whl

llama_index/llms/bedrock_converse/base.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import (
     Any,
     Callable,
@@ -20,6 +21,8 @@ from llama_index.core.base.llms.types import (
     CompletionResponseGen,
     LLMMetadata,
     MessageRole,
+    TextBlock,
+    ThinkingBlock,
 )
 from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.callbacks import CallbackManager
@@ -46,6 +49,8 @@ from llama_index.llms.bedrock_converse.utils import (
     join_two_dicts,
     messages_to_converse_messages,
     tools_to_converse_tools,
+    is_reasoning,
+    ThinkingDict,
 )
 
 if TYPE_CHECKING:
@@ -158,6 +163,10 @@ class BedrockConverse(FunctionCallingLLM):
     trace: Optional[str] = Field(
         description="Specifies whether to enable or disable the Bedrock trace. If enabled, you can see the full Bedrock trace."
     )
+    thinking: Optional[ThinkingDict] = Field(
+        description="Specifies the thinking configuration of a reasoning model. Only applicable to Anthropic and DeepSeek models",
+        default=None,
+    )
     additional_kwargs: Dict[str, Any] = Field(
         default_factory=dict,
         description="Additional kwargs for the bedrock invokeModel request.",
@@ -200,6 +209,7 @@ class BedrockConverse(FunctionCallingLLM):
         guardrail_version: Optional[str] = None,
         application_inference_profile_arn: Optional[str] = None,
         trace: Optional[str] = None,
+        thinking: Optional[ThinkingDict] = None,
     ) -> None:
         additional_kwargs = additional_kwargs or {}
         callback_manager = callback_manager or CallbackManager([])
@@ -213,6 +223,13 @@ class BedrockConverse(FunctionCallingLLM):
             "botocore_session": botocore_session,
         }
 
+        if not is_reasoning(model) and thinking is not None:
+            thinking = None
+            warnings.warn(
+                "You set thinking parameters for a non-reasoning models, they will be ignored",
+                UserWarning,
+            )
+
         super().__init__(
             temperature=temperature,
             max_tokens=max_tokens,
@@ -243,6 +260,7 @@ class BedrockConverse(FunctionCallingLLM):
             guardrail_version=guardrail_version,
             application_inference_profile_arn=application_inference_profile_arn,
             trace=trace,
+            thinking=thinking,
         )
 
         self._config = None
@@ -330,7 +348,9 @@ class BedrockConverse(FunctionCallingLLM):
 
     def _get_content_and_tool_calls(
         self, response: Optional[Dict[str, Any]] = None, content: Dict[str, Any] = None
-    ) -> Tuple[str, Dict[str, Any], List[str], List[str]]:
+    ) -> Tuple[
+        List[Union[TextBlock, ThinkingBlock]], Dict[str, Any], 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}"
         )
@@ -340,14 +360,26 @@ class BedrockConverse(FunctionCallingLLM):
         tool_calls = []
         tool_call_ids = []
         status = []
-        text_content = ""
+        blocks = []
         if content is not None:
             content_list = [content]
         else:
             content_list = response["output"]["message"]["content"]
+
         for content_block in content_list:
             if text := content_block.get("text", None):
-                text_content += text
+                blocks.append(TextBlock(text=text))
+            if thinking := content_block.get("reasoningContent", None):
+                blocks.append(
+                    ThinkingBlock(
+                        content=thinking.get("reasoningText", {}).get("text", None),
+                        additional_information={
+                            "signature": thinking.get("reasoningText", {}).get(
+                                "signature", None
+                            )
+                        },
+                    )
+                )
             if tool_usage := content_block.get("toolUse", None):
                 if "toolUseId" not in tool_usage:
                     tool_usage["toolUseId"] = content_block["toolUseId"]
@@ -361,7 +393,7 @@ class BedrockConverse(FunctionCallingLLM):
                 tool_call_ids.append(tool_result_content.get("toolUseId", ""))
                 status.append(tool_result.get("status", ""))
 
-        return text_content, tool_calls, tool_call_ids, status
+        return blocks, tool_calls, tool_call_ids, status
 
     @llm_chat_callback()
     def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
@@ -370,6 +402,8 @@ class BedrockConverse(FunctionCallingLLM):
             messages, self.model
         )
         all_kwargs = self._get_all_kwargs(**kwargs)
+        if self.thinking is not None:
+            all_kwargs["thinking"] = self.thinking
 
         # invoke LLM in AWS Bedrock Converse with retry
         response = converse_with_retry(
@@ -386,14 +420,14 @@ class BedrockConverse(FunctionCallingLLM):
             **all_kwargs,
         )
 
-        content, tool_calls, tool_call_ids, status = self._get_content_and_tool_calls(
+        blocks, tool_calls, tool_call_ids, status = self._get_content_and_tool_calls(
             response
         )
 
         return ChatResponse(
             message=ChatMessage(
                 role=MessageRole.ASSISTANT,
-                content=content,
+                blocks=blocks,
                 additional_kwargs={
                     "tool_calls": tool_calls,
                     "tool_call_id": tool_call_ids,
@@ -420,6 +454,8 @@ class BedrockConverse(FunctionCallingLLM):
             messages, self.model
         )
         all_kwargs = self._get_all_kwargs(**kwargs)
+        if self.thinking is not None:
+            all_kwargs["thinking"] = self.thinking
 
         # invoke LLM in AWS Bedrock Converse with retry
         response = converse_with_retry(
@@ -441,12 +477,22 @@ class BedrockConverse(FunctionCallingLLM):
             tool_calls = []  # Track tool calls separately
             current_tool_call = None  # Track the current tool call being built
             role = MessageRole.ASSISTANT
+            thinking = ""
+            thinking_signature = ""
 
             for chunk in response["stream"]:
                 if content_block_delta := chunk.get("contentBlockDelta"):
                     content_delta = content_block_delta["delta"]
                     content = join_two_dicts(content, content_delta)
 
+                    if "reasoningContent" in content_delta:
+                        thinking += content_delta.get("reasoningContent", {}).get(
+                            "text", ""
+                        )
+                        thinking_signature += content_delta.get(
+                            "reasoningContent", {}
+                        ).get("signature", "")
+
                     # If this delta contains tool call info, update current tool call
                     if "toolUse" in content_delta:
                         tool_use_delta = content_delta["toolUse"]
@@ -477,10 +523,24 @@ class BedrockConverse(FunctionCallingLLM):
                                     current_tool_call, tool_use_delta
                                 )
 
+                    blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                        TextBlock(text=content.get("text", ""))
+                    ]
+                    if thinking != "":
+                        blocks.insert(
+                            0,
+                            ThinkingBlock(
+                                content=thinking,
+                                additional_information={
+                                    "signature": thinking_signature
+                                },
+                            ),
+                        )
+
                     yield ChatResponse(
                         message=ChatMessage(
                             role=role,
-                            content=content.get("text", ""),
+                            blocks=blocks,
                             additional_kwargs={
                                 "tool_calls": tool_calls,
                                 "tool_call_id": [
@@ -502,10 +562,24 @@ class BedrockConverse(FunctionCallingLLM):
                         # Add to our list of tool calls
                         tool_calls.append(current_tool_call)
 
+                    blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                        TextBlock(text=content.get("text", ""))
+                    ]
+                    if thinking != "":
+                        blocks.insert(
+                            0,
+                            ThinkingBlock(
+                                content=thinking,
+                                additional_information={
+                                    "signature": thinking_signature
+                                },
+                            ),
+                        )
+
                     yield ChatResponse(
                         message=ChatMessage(
                             role=role,
-                            content=content.get("text", ""),
+                            blocks=blocks,
                             additional_kwargs={
                                 "tool_calls": tool_calls,
                                 "tool_call_id": [
@@ -524,10 +598,24 @@ 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]] = [
+                            TextBlock(text=content.get("text", ""))
+                        ]
+                        if thinking != "":
+                            blocks.insert(
+                                0,
+                                ThinkingBlock(
+                                    content=thinking,
+                                    additional_information={
+                                        "signature": thinking_signature
+                                    },
+                                ),
+                            )
+
                         yield ChatResponse(
                             message=ChatMessage(
                                 role=role,
-                                content=content.get("text", ""),
+                                blocks=blocks,
                                 additional_kwargs={
                                     "tool_calls": tool_calls,
                                     "tool_call_id": [
@@ -559,6 +647,8 @@ class BedrockConverse(FunctionCallingLLM):
             messages, self.model
         )
         all_kwargs = self._get_all_kwargs(**kwargs)
+        if self.thinking is not None:
+            all_kwargs["thinking"] = self.thinking
 
         # invoke LLM in AWS Bedrock Converse with retry
         response = await converse_with_retry_async(
@@ -577,14 +667,14 @@ class BedrockConverse(FunctionCallingLLM):
             **all_kwargs,
         )
 
-        content, tool_calls, tool_call_ids, status = self._get_content_and_tool_calls(
+        blocks, tool_calls, tool_call_ids, status = self._get_content_and_tool_calls(
             response
         )
 
         return ChatResponse(
             message=ChatMessage(
                 role=MessageRole.ASSISTANT,
-                content=content,
+                blocks=blocks,
                 additional_kwargs={
                     "tool_calls": tool_calls,
                     "tool_call_id": tool_call_ids,
@@ -611,6 +701,8 @@ class BedrockConverse(FunctionCallingLLM):
             messages, self.model
         )
         all_kwargs = self._get_all_kwargs(**kwargs)
+        if self.thinking is not None:
+            all_kwargs["thinking"] = self.thinking
 
         # invoke LLM in AWS Bedrock Converse with retry
         response_gen = await converse_with_retry_async(
@@ -634,12 +726,22 @@ class BedrockConverse(FunctionCallingLLM):
             tool_calls = []  # Track tool calls separately
             current_tool_call = None  # Track the current tool call being built
             role = MessageRole.ASSISTANT
+            thinking = ""
+            thinking_signature = ""
 
             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)
 
+                    if "reasoningContent" in content_delta:
+                        thinking += content_delta.get("reasoningContent", {}).get(
+                            "text", ""
+                        )
+                        thinking_signature += content_delta.get(
+                            "reasoningContent", {}
+                        ).get("signature", "")
+
                     # If this delta contains tool call info, update current tool call
                     if "toolUse" in content_delta:
                         tool_use_delta = content_delta["toolUse"]
@@ -669,11 +771,24 @@ class BedrockConverse(FunctionCallingLLM):
                                 current_tool_call = join_two_dicts(
                                     current_tool_call, tool_use_delta
                                 )
+                    blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                        TextBlock(text=content.get("text", ""))
+                    ]
+                    if thinking != "":
+                        blocks.insert(
+                            0,
+                            ThinkingBlock(
+                                content=thinking,
+                                additional_information={
+                                    "signature": thinking_signature
+                                },
+                            ),
+                        )
 
                     yield ChatResponse(
                         message=ChatMessage(
                             role=role,
-                            content=content.get("text", ""),
+                            blocks=blocks,
                             additional_kwargs={
                                 "tool_calls": tool_calls,
                                 "tool_call_id": [
@@ -695,10 +810,24 @@ class BedrockConverse(FunctionCallingLLM):
                         # Add to our list of tool calls
                         tool_calls.append(current_tool_call)
 
+                    blocks: List[Union[TextBlock, ThinkingBlock]] = [
+                        TextBlock(text=content.get("text", ""))
+                    ]
+                    if thinking != "":
+                        blocks.insert(
+                            0,
+                            ThinkingBlock(
+                                content=thinking,
+                                additional_information={
+                                    "signature": thinking_signature
+                                },
+                            ),
+                        )
+
                     yield ChatResponse(
                         message=ChatMessage(
                             role=role,
-                            content=content.get("text", ""),
+                            blocks=blocks,
                             additional_kwargs={
                                 "tool_calls": tool_calls,
                                 "tool_call_id": [
@@ -717,10 +846,24 @@ 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]] = [
+                            TextBlock(text=content.get("text", ""))
+                        ]
+                        if thinking != "":
+                            blocks.insert(
+                                0,
+                                ThinkingBlock(
+                                    content=thinking,
+                                    additional_information={
+                                        "signature": thinking_signature
+                                    },
+                                ),
+                            )
+
                         yield ChatResponse(
                             message=ChatMessage(
                                 role=role,
-                                content=content.get("text", ""),
+                                blocks=blocks,
                                 additional_kwargs={
                                     "tool_calls": tool_calls,
                                     "tool_call_id": [

llama_index/llms/bedrock_converse/utils.py

@@ -1,7 +1,18 @@
 import base64
 import json
 import logging
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Literal,
+    Union,
+)
+from typing_extensions import TypedDict
 from tenacity import (
     before_sleep_log,
     retry,
@@ -20,6 +31,7 @@ from llama_index.core.base.llms.types import (
     AudioBlock,
     DocumentBlock,
     CachePoint,
+    ThinkingBlock,
 )
 
 
@@ -151,6 +163,19 @@ BEDROCK_PROMPT_CACHING_SUPPORTED_MODELS = (
     "amazon.nova-micro-v1:0",
 )
 
+BEDROCK_REASONING_MODELS = (
+    "anthropic.claude-3-7-sonnet-20250219-v1:0",
+    "anthropic.claude-opus-4-20250514-v1:0",
+    "anthropic.claude-sonnet-4-20250514-v1:0",
+    "anthropic.claude-sonnet-4-5-20250929-v1:0",
+    "deepseek.r1-v1:0",
+)
+
+
+def is_reasoning(model_name: str) -> bool:
+    model_name = get_model_name(model_name)
+    return model_name in BEDROCK_REASONING_MODELS
+
 
 def get_model_name(model_name: str) -> str:
     """Extract base model name from region-prefixed model identifier."""
@@ -220,6 +245,22 @@ def _content_block_to_bedrock_format(
         return {
             "text": block.text,
         }
+    elif isinstance(block, ThinkingBlock):
+        if block.content:
+            thinking_data = {
+                "reasoningContent": {"reasoningText": {"text": block.content}}
+            }
+            if (
+                "signature" in block.additional_information
+                and block.additional_information["signature"]
+            ):
+                thinking_data["reasoningContent"]["reasoningText"]["signature"] = (
+                    block.additional_information["signature"]
+                )
+
+            return thinking_data
+        else:
+            return None
     elif isinstance(block, DocumentBlock):
         if not block.data:
             file_buffer = block.resolve_document()
@@ -518,6 +559,10 @@ def converse_with_retry(
             "temperature": temperature,
         },
     }
+    if "thinking" in kwargs:
+        converse_kwargs["additionalModelRequestFields"] = {
+            "thinking": kwargs["thinking"]
+        }
     if system_prompt:
         if isinstance(system_prompt, str):
             # if the system prompt is a simple text (for retro compatibility)
@@ -547,7 +592,14 @@ def converse_with_retry(
         {
             k: v
             for k, v in kwargs.items()
-            if k not in ["tools", "guardrail_identifier", "guardrail_version", "trace"]
+            if k
+            not in [
+                "tools",
+                "guardrail_identifier",
+                "guardrail_version",
+                "trace",
+                "thinking",
+            ]
         },
     )
 
@@ -589,6 +641,10 @@ async def converse_with_retry_async(
             "temperature": temperature,
         },
     }
+    if "thinking" in kwargs:
+        converse_kwargs["additionalModelRequestFields"] = {
+            "thinking": kwargs["thinking"]
+        }
 
     if system_prompt:
         if isinstance(system_prompt, str):
@@ -622,7 +678,14 @@ async def converse_with_retry_async(
         {
             k: v
             for k, v in kwargs.items()
-            if k not in ["tools", "guardrail_identifier", "guardrail_version", "trace"]
+            if k
+            not in [
+                "tools",
+                "guardrail_identifier",
+                "guardrail_version",
+                "trace",
+                "thinking",
+            ]
         },
     )
     _boto_client_kwargs = {}
@@ -688,3 +751,8 @@ def join_two_dicts(dict1: Dict[str, Any], dict2: Dict[str, Any]) -> Dict[str, An
             else:
                 new_dict[key] += value
     return new_dict
+
+
+class ThinkingDict(TypedDict):
+    type: Literal["enabled"]
+    budget_tokens: int

{llama_index_llms_bedrock_converse-0.9.5.dist-info → llama_index_llms_bedrock_converse-0.10.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llama-index-llms-bedrock-converse
-Version: 0.9.5
+Version: 0.10.1
 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.13.0
+Requires-Dist: llama-index-core<0.15,>=0.14.3
 Description-Content-Type: text/markdown
 
 # LlamaIndex Llms Integration: Bedrock Converse

llama_index_llms_bedrock_converse-0.10.1.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=tXQMmgyYZmrudRsQ7uofvQyIKNUTQtTpcdny1EoMCB0,41080
+llama_index/llms/bedrock_converse/utils.py,sha256=O0z1eJVjX_ZdghESiHfpx1KxmS3PQJIjSAUJtUsnH4c,27248
+llama_index_llms_bedrock_converse-0.10.1.dist-info/METADATA,sha256=w6TjnF4jlxMVcXBLEJzSjEUfCWPZ2xkjS84VbF6hbA0,7834
+llama_index_llms_bedrock_converse-0.10.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+llama_index_llms_bedrock_converse-0.10.1.dist-info/licenses/LICENSE,sha256=JPQLUZD9rKvCTdu192Nk0V5PAwklIg6jANii3UmTyMs,1065
+llama_index_llms_bedrock_converse-0.10.1.dist-info/RECORD,,

llama_index_llms_bedrock_converse-0.9.5.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=SHGRmAYcMUMAjmnHWFLYKPzvsAAjCF_A1Mvc7s9I7IM,35233
-llama_index/llms/bedrock_converse/utils.py,sha256=Ly-s3mROVreinvYmRcAJU7MksSHqeTEa1tnY3na17wg,25565
-llama_index_llms_bedrock_converse-0.9.5.dist-info/METADATA,sha256=ALn1SYaHR7aYuALhDShDhbWBDHOzw9RAHQI1iS0xhLM,7833
-llama_index_llms_bedrock_converse-0.9.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-llama_index_llms_bedrock_converse-0.9.5.dist-info/licenses/LICENSE,sha256=JPQLUZD9rKvCTdu192Nk0V5PAwklIg6jANii3UmTyMs,1065
-llama_index_llms_bedrock_converse-0.9.5.dist-info/RECORD,,