Merge pull request #2442 from pipecat-ai/mb/retry-completion

retry_on_timeout: Anthropic, AWS Bedrock
This commit is contained in:
Mark Backman
2025-08-19 11:23:43 -07:00
committed by GitHub
3 changed files with 69 additions and 5 deletions

View File

@@ -11,8 +11,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
- Added `MistralLLMService`, using Mistral's chat completion API.
- For `OpenAILLMService` and its subclasses, added the ability to retry
executing a chat completion after a timeout period. The new args are
- Added the ability to retry executing a chat completion after a timeout period
for `OpenAILLMService` and its subclasses, `AnthropicLLMService`, and
`AWSBedrockLLMService`. The LLM services accept new args:
`retry_timeout_secs` and `retry_on_timeout`. This feature is disabled by
default.

View File

@@ -57,7 +57,7 @@ from pipecat.utils.asyncio.watchdog_async_iterator import WatchdogAsyncIterator
from pipecat.utils.tracing.service_decorators import traced_llm
try:
from anthropic import NOT_GIVEN, AsyncAnthropic, NotGiven
from anthropic import NOT_GIVEN, APITimeoutError, AsyncAnthropic, NotGiven
except ModuleNotFoundError as e:
logger.error(f"Exception: {e}")
logger.error("In order to use Anthropic, you need to `pip install pipecat-ai[anthropic]`.")
@@ -133,6 +133,8 @@ class AnthropicLLMService(LLMService):
model: str = "claude-sonnet-4-20250514",
params: Optional[InputParams] = None,
client=None,
retry_timeout_secs: Optional[float] = 5.0,
retry_on_timeout: Optional[bool] = False,
**kwargs,
):
"""Initialize the Anthropic LLM service.
@@ -142,6 +144,8 @@ class AnthropicLLMService(LLMService):
model: Model name to use. Defaults to "claude-sonnet-4-20250514".
params: Optional model parameters for inference.
client: Optional custom Anthropic client instance.
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)
@@ -150,6 +154,8 @@ class AnthropicLLMService(LLMService):
api_key=api_key
) # if the client is provided, use it and remove it, otherwise create a new one
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,
"enable_prompt_caching_beta": params.enable_prompt_caching_beta or False,
@@ -167,6 +173,31 @@ class AnthropicLLMService(LLMService):
"""
return True
async def _create_message_stream(self, api_call, params):
"""Create message stream with optional timeout and retry.
Args:
api_call: The Anthropic API method to call.
params: Parameters for the API call.
Returns:
Async stream of message events.
"""
if self._retry_on_timeout:
try:
response = await asyncio.wait_for(
api_call(**params), timeout=self._retry_timeout_secs
)
return response
except (APITimeoutError, asyncio.TimeoutError):
# Retry, this time without a timeout so we get a response
logger.debug(f"{self}: Retrying message creation due to timeout")
response = await api_call(**params)
return response
else:
response = await api_call(**params)
return response
@property
def enable_prompt_caching_beta(self) -> bool:
"""Check if prompt caching beta feature is enabled.
@@ -250,7 +281,7 @@ class AnthropicLLMService(LLMService):
params.update(self._settings["extra"])
response = await api_call(**params)
response = await self._create_message_stream(api_call, params)
await self.stop_ttfb_metrics()

View File

@@ -58,6 +58,7 @@ try:
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(
@@ -724,6 +725,8 @@ 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.
@@ -736,6 +739,8 @@ class AWSBedrockLLMService(LLMService):
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)
@@ -762,6 +767,8 @@ class AWSBedrockLLMService(LLMService):
}
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,
@@ -782,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(
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,
@@ -911,7 +943,7 @@ class AWSBedrockLLMService(LLMService):
service_name="bedrock-runtime", **self._aws_params
) as client:
# Call AWS Bedrock with streaming
response = await client.converse_stream(**request_params)
response = await self._create_converse_stream(client, request_params)
await self.stop_ttfb_metrics()