dv-pipecat-ai 0.0.74.dev770__py3-none-any.whl → 0.0.82.dev776__py3-none-any.whl

pipecat/services/aws/llm.py

@@ -55,9 +55,10 @@ from pipecat.services.llm_service import LLMService
 from pipecat.utils.tracing.service_decorators import traced_llm
 
 try:
-    import boto3
+    import aioboto3
     import httpx
     from botocore.config import Config
+    from botocore.exceptions import ReadTimeoutError
 except ModuleNotFoundError as e:
     logger.error(f"Exception: {e}")
     logger.error(
@@ -104,12 +105,6 @@ class AWSBedrockLLMContext(OpenAILLMContext):
     Extends OpenAI LLM context to handle AWS Bedrock's specific message format
     and system message handling. Manages conversion between OpenAI and Bedrock
     message formats.
-
-    Args:
-        messages: List of conversation messages in OpenAI format.
-        tools: List of available function calling tools.
-        tool_choice: Tool selection strategy or specific tool choice.
-        system: System message content for AWS Bedrock.
     """
 
     def __init__(
@@ -120,6 +115,14 @@ class AWSBedrockLLMContext(OpenAILLMContext):
         *,
         system: Optional[str] = None,
     ):
+        """Initialize AWS Bedrock LLM context.
+
+        Args:
+            messages: List of conversation messages in OpenAI format.
+            tools: List of available function calling tools.
+            tool_choice: Tool selection strategy or specific tool choice.
+            system: System message content for AWS Bedrock.
+        """
         super().__init__(messages=messages, tools=tools, tool_choice=tool_choice)
         self.system = system
 
@@ -205,20 +208,37 @@ class AWSBedrockLLMContext(OpenAILLMContext):
         Handles text content and function calls for both user and assistant messages.
 
         Args:
-            obj: Message in AWS Bedrock format:
-                {
-                    "role": "user/assistant",
-                    "content": [{"text": str} | {"toolUse": {...}} | {"toolResult": {...}}]
-                }
+            obj: Message in AWS Bedrock format.
 
         Returns:
-            List of messages in standard format:
-            [
+            List of messages in standard format.
+
+        Examples:
+            AWS Bedrock format input::
+
                 {
-                    "role": "user/assistant/tool",
-                    "content": [{"type": "text", "text": str}]
+                    "role": "assistant",
+                    "content": [
+                        {"text": "Hello"},
+                        {"toolUse": {"toolUseId": "123", "name": "search", "input": {"q": "test"}}}
+                    ]
                 }
-            ]
+
+            Standard format output::
+
+                [
+                    {"role": "assistant", "content": [{"type": "text", "text": "Hello"}]},
+                    {
+                        "role": "assistant",
+                        "tool_calls": [
+                            {
+                                "type": "function",
+                                "id": "123",
+                                "function": {"name": "search", "arguments": '{"q": "test"}'}
+                            }
+                        ]
+                    }
+                ]
         """
         role = obj.get("role")
         content = obj.get("content")
@@ -292,23 +312,38 @@ class AWSBedrockLLMContext(OpenAILLMContext):
         Empty text content is converted to "(empty)".
 
         Args:
-            message: Message in standard format:
+            message: Message in standard format.
+
+        Returns:
+            Message in AWS Bedrock format.
+
+        Examples:
+            Standard format input::
+
                 {
-                    "role": "user/assistant/tool",
-                    "content": str | [{"type": "text", ...}],
-                    "tool_calls": [{"id": str, "function": {"name": str, "arguments": str}}]
+                    "role": "assistant",
+                    "tool_calls": [
+                        {
+                            "id": "123",
+                            "function": {"name": "search", "arguments": '{"q": "test"}'}
+                        }
+                    ]
                 }
 
-        Returns:
-            Message in AWS Bedrock format:
-            {
-                "role": "user/assistant",
-                "content": [
-                    {"text": str} |
-                    {"toolUse": {"toolUseId": str, "name": str, "input": dict}} |
-                    {"toolResult": {"toolUseId": str, "content": [...], "status": str}}
-                ]
-            }
+            AWS Bedrock format output::
+
+                {
+                    "role": "assistant",
+                    "content": [
+                        {
+                            "toolUse": {
+                                "toolUseId": "123",
+                                "name": "search",
+                                "input": {"q": "test"}
+                            }
+                        }
+                    ]
+                }
         """
         if message["role"] == "tool":
             # Try to parse the content as JSON if it looks like JSON
@@ -656,16 +691,6 @@ class AWSBedrockLLMService(LLMService):
     Provides inference capabilities for AWS Bedrock models including Amazon Nova
     and Anthropic Claude. Supports streaming responses, function calling, and
     vision capabilities.
-
-    Args:
-        model: The AWS Bedrock model identifier to use.
-        aws_access_key: AWS access key ID. If None, uses default credentials.
-        aws_secret_key: AWS secret access key. If None, uses default credentials.
-        aws_session_token: AWS session token for temporary credentials.
-        aws_region: AWS region for the Bedrock service.
-        params: Model parameters and configuration.
-        client_config: Custom boto3 client configuration.
-        **kwargs: Additional arguments passed to parent LLMService.
     """
 
     # Overriding the default adapter to use the Anthropic one.
@@ -700,8 +725,24 @@ class AWSBedrockLLMService(LLMService):
         aws_region: str = "us-east-1",
         params: Optional[InputParams] = None,
         client_config: Optional[Config] = None,
+        retry_timeout_secs: Optional[float] = 5.0,
+        retry_on_timeout: Optional[bool] = False,
         **kwargs,
     ):
+        """Initialize the AWS Bedrock LLM service.
+
+        Args:
+            model: The AWS Bedrock model identifier to use.
+            aws_access_key: AWS access key ID. If None, uses default credentials.
+            aws_secret_key: AWS secret access key. If None, uses default credentials.
+            aws_session_token: AWS session token for temporary credentials.
+            aws_region: AWS region for the Bedrock service.
+            params: Model parameters and configuration.
+            client_config: Custom boto3 client configuration.
+            retry_timeout_secs: Request timeout in seconds for retry logic.
+            retry_on_timeout: Whether to retry the request once if it times out.
+            **kwargs: Additional arguments passed to parent LLMService.
+        """
         super().__init__(**kwargs)
 
         params = params or AWSBedrockLLMService.InputParams()
@@ -713,15 +754,21 @@ class AWSBedrockLLMService(LLMService):
                 read_timeout=300,  # 5 minutes
                 retries={"max_attempts": 3},
             )
-        session = boto3.Session(
-            aws_access_key_id=aws_access_key,
-            aws_secret_access_key=aws_secret_key,
-            aws_session_token=aws_session_token,
-            region_name=aws_region,
-        )
-        self._client = session.client(service_name="bedrock-runtime", config=client_config)
+
+        self._aws_session = aioboto3.Session()
+
+        # Store AWS session parameters for creating client in async context
+        self._aws_params = {
+            "aws_access_key_id": aws_access_key,
+            "aws_secret_access_key": aws_secret_key,
+            "aws_session_token": aws_session_token,
+            "region_name": aws_region,
+            "config": client_config,
+        }
 
         self.set_model_name(model)
+        self._retry_timeout_secs = retry_timeout_secs
+        self._retry_on_timeout = retry_on_timeout
         self._settings = {
             "max_tokens": params.max_tokens,
             "temperature": params.temperature,
@@ -742,6 +789,31 @@ class AWSBedrockLLMService(LLMService):
         """
         return True
 
+    async def _create_converse_stream(self, client, request_params):
+        """Create converse stream with optional timeout and retry.
+
+        Args:
+            client: The AWS Bedrock client instance.
+            request_params: Parameters for the converse_stream call.
+
+        Returns:
+            Async stream of response events.
+        """
+        if self._retry_on_timeout:
+            try:
+                response = await asyncio.wait_for(
+                    await client.converse_stream(**request_params), timeout=self._retry_timeout_secs
+                )
+                return response
+            except (ReadTimeoutError, asyncio.TimeoutError) as e:
+                # Retry, this time without a timeout so we get a response
+                logger.debug(f"{self}: Retrying converse_stream due to timeout")
+                response = await client.converse_stream(**request_params)
+                return response
+        else:
+            response = await client.converse_stream(**request_params)
+            return response
+
     def create_context_aggregator(
         self,
         context: OpenAILLMContext,
@@ -867,70 +939,72 @@ class AWSBedrockLLMService(LLMService):
 
             logger.debug(f"Calling AWS Bedrock model with: {request_params}")
 
-            # Call AWS Bedrock with streaming
-            response = self._client.converse_stream(**request_params)
-
-            await self.stop_ttfb_metrics()
-
-            # Process the streaming response
-            tool_use_block = None
-            json_accumulator = ""
-
-            function_calls = []
-            for event in response["stream"]:
-                self.reset_watchdog()
-
-                # Handle text content
-                if "contentBlockDelta" in event:
-                    delta = event["contentBlockDelta"]["delta"]
-                    if "text" in delta:
-                        await self.push_frame(LLMTextFrame(delta["text"]))
-                        completion_tokens_estimate += self._estimate_tokens(delta["text"])
-                    elif "toolUse" in delta and "input" in delta["toolUse"]:
-                        # Handle partial JSON for tool use
-                        json_accumulator += delta["toolUse"]["input"]
-                        completion_tokens_estimate += self._estimate_tokens(
-                            delta["toolUse"]["input"]
-                        )
-
-                # Handle tool use start
-                elif "contentBlockStart" in event:
-                    content_block_start = event["contentBlockStart"]["start"]
-                    if "toolUse" in content_block_start:
-                        tool_use_block = {
-                            "id": content_block_start["toolUse"].get("toolUseId", ""),
-                            "name": content_block_start["toolUse"].get("name", ""),
-                        }
-                        json_accumulator = ""
-
-                # Handle message completion with tool use
-                elif "messageStop" in event and "stopReason" in event["messageStop"]:
-                    if event["messageStop"]["stopReason"] == "tool_use" and tool_use_block:
-                        try:
-                            arguments = json.loads(json_accumulator) if json_accumulator else {}
-
-                            # Only call function if it's not the no_operation tool
-                            if not using_noop_tool:
-                                function_calls.append(
-                                    FunctionCallFromLLM(
-                                        context=context,
-                                        tool_call_id=tool_use_block["id"],
-                                        function_name=tool_use_block["name"],
-                                        arguments=arguments,
+            async with self._aws_session.client(
+                service_name="bedrock-runtime", **self._aws_params
+            ) as client:
+                # Call AWS Bedrock with streaming
+                response = await self._create_converse_stream(client, request_params)
+
+                await self.stop_ttfb_metrics()
+
+                # Process the streaming response
+                tool_use_block = None
+                json_accumulator = ""
+
+                function_calls = []
+
+                async for event in response["stream"]:
+                    # Handle text content
+                    if "contentBlockDelta" in event:
+                        delta = event["contentBlockDelta"]["delta"]
+                        if "text" in delta:
+                            await self.push_frame(LLMTextFrame(delta["text"]))
+                            completion_tokens_estimate += self._estimate_tokens(delta["text"])
+                        elif "toolUse" in delta and "input" in delta["toolUse"]:
+                            # Handle partial JSON for tool use
+                            json_accumulator += delta["toolUse"]["input"]
+                            completion_tokens_estimate += self._estimate_tokens(
+                                delta["toolUse"]["input"]
+                            )
+
+                    # Handle tool use start
+                    elif "contentBlockStart" in event:
+                        content_block_start = event["contentBlockStart"]["start"]
+                        if "toolUse" in content_block_start:
+                            tool_use_block = {
+                                "id": content_block_start["toolUse"].get("toolUseId", ""),
+                                "name": content_block_start["toolUse"].get("name", ""),
+                            }
+                            json_accumulator = ""
+
+                    # Handle message completion with tool use
+                    elif "messageStop" in event and "stopReason" in event["messageStop"]:
+                        if event["messageStop"]["stopReason"] == "tool_use" and tool_use_block:
+                            try:
+                                arguments = json.loads(json_accumulator) if json_accumulator else {}
+
+                                # Only call function if it's not the no_operation tool
+                                if not using_noop_tool:
+                                    function_calls.append(
+                                        FunctionCallFromLLM(
+                                            context=context,
+                                            tool_call_id=tool_use_block["id"],
+                                            function_name=tool_use_block["name"],
+                                            arguments=arguments,
+                                        )
                                     )
-                                )
-                            else:
-                                logger.debug("Ignoring no_operation tool call")
-                        except json.JSONDecodeError:
-                            logger.error(f"Failed to parse tool arguments: {json_accumulator}")
-
-                # Handle usage metrics if available
-                if "metadata" in event and "usage" in event["metadata"]:
-                    usage = event["metadata"]["usage"]
-                    prompt_tokens += usage.get("inputTokens", 0)
-                    completion_tokens += usage.get("outputTokens", 0)
-                    cache_read_input_tokens += usage.get("cacheReadInputTokens", 0)
-                    cache_creation_input_tokens += usage.get("cacheWriteInputTokens", 0)
+                                else:
+                                    logger.debug("Ignoring no_operation tool call")
+                            except json.JSONDecodeError:
+                                logger.error(f"Failed to parse tool arguments: {json_accumulator}")
+
+                    # Handle usage metrics if available
+                    if "metadata" in event and "usage" in event["metadata"]:
+                        usage = event["metadata"]["usage"]
+                        prompt_tokens += usage.get("inputTokens", 0)
+                        completion_tokens += usage.get("outputTokens", 0)
+                        cache_read_input_tokens += usage.get("cacheReadInputTokens", 0)
+                        cache_creation_input_tokens += usage.get("cacheWriteInputTokens", 0)
 
             await self.run_function_calls(function_calls)
         except asyncio.CancelledError: