Merge pull request #2148 from pipecat-ai/filipi/aws_bedrock

Refactoring AWSBedrockLLMService to work async
This commit is contained in:
Filipi da Silva Fuchter
2025-07-08 12:14:28 -03:00
committed by GitHub
4 changed files with 115 additions and 116 deletions

View File

@@ -25,6 +25,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Changed ### Changed
- Refactored `AWSBedrockLLMService` and `AWSPollyTTSService` to work
asynchronously using `aioboto3` instead of the `boto3` library.
- The `UserIdleProcessor` now handles the scenario where function calls take - The `UserIdleProcessor` now handles the scenario where function calls take
longer than the idle timeout duration. This allows you to use the longer than the idle timeout duration. This allows you to use the
`UserIdleProcessor` in conjunction with function calls that take a while to `UserIdleProcessor` in conjunction with function calls that take a while to

View File

@@ -42,7 +42,7 @@ Website = "https://pipecat.ai"
[project.optional-dependencies] [project.optional-dependencies]
anthropic = [ "anthropic~=0.49.0" ] anthropic = [ "anthropic~=0.49.0" ]
assemblyai = [ "websockets~=13.1" ] assemblyai = [ "websockets~=13.1" ]
aws = [ "boto3~=1.37.16", "websockets~=13.1" ] aws = [ "aioboto3~=15.0.0", "websockets~=13.1" ]
aws-nova-sonic = [ "aws_sdk_bedrock_runtime~=0.0.2" ] aws-nova-sonic = [ "aws_sdk_bedrock_runtime~=0.0.2" ]
azure = [ "azure-cognitiveservices-speech~=1.42.0"] azure = [ "azure-cognitiveservices-speech~=1.42.0"]
cartesia = [ "cartesia~=2.0.3", "websockets~=13.1" ] cartesia = [ "cartesia~=2.0.3", "websockets~=13.1" ]

View File

@@ -55,7 +55,7 @@ from pipecat.services.llm_service import LLMService
from pipecat.utils.tracing.service_decorators import traced_llm from pipecat.utils.tracing.service_decorators import traced_llm
try: try:
import boto3 import aioboto3
import httpx import httpx
from botocore.config import Config from botocore.config import Config
except ModuleNotFoundError as e: except ModuleNotFoundError as e:
@@ -749,13 +749,17 @@ class AWSBedrockLLMService(LLMService):
read_timeout=300, # 5 minutes read_timeout=300, # 5 minutes
retries={"max_attempts": 3}, retries={"max_attempts": 3},
) )
session = boto3.Session(
aws_access_key_id=aws_access_key, self._aws_session = aioboto3.Session()
aws_secret_access_key=aws_secret_key,
aws_session_token=aws_session_token, # Store AWS session parameters for creating client in async context
region_name=aws_region, self._aws_params = {
) "aws_access_key_id": aws_access_key,
self._client = session.client(service_name="bedrock-runtime", config=client_config) "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.set_model_name(model)
self._settings = { self._settings = {
@@ -903,70 +907,74 @@ class AWSBedrockLLMService(LLMService):
logger.debug(f"Calling AWS Bedrock model with: {request_params}") logger.debug(f"Calling AWS Bedrock model with: {request_params}")
# Call AWS Bedrock with streaming async with self._aws_session.client(
response = self._client.converse_stream(**request_params) service_name="bedrock-runtime", **self._aws_params
) as client:
# Call AWS Bedrock with streaming
response = await client.converse_stream(**request_params)
await self.stop_ttfb_metrics() await self.stop_ttfb_metrics()
# Process the streaming response # Process the streaming response
tool_use_block = None tool_use_block = None
json_accumulator = "" json_accumulator = ""
function_calls = [] function_calls = []
for event in response["stream"]:
self.reset_watchdog()
# Handle text content async for event in response["stream"]:
if "contentBlockDelta" in event: self.reset_watchdog()
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 # Handle text content
elif "contentBlockStart" in event: if "contentBlockDelta" in event:
content_block_start = event["contentBlockStart"]["start"] delta = event["contentBlockDelta"]["delta"]
if "toolUse" in content_block_start: if "text" in delta:
tool_use_block = { await self.push_frame(LLMTextFrame(delta["text"]))
"id": content_block_start["toolUse"].get("toolUseId", ""), completion_tokens_estimate += self._estimate_tokens(delta["text"])
"name": content_block_start["toolUse"].get("name", ""), elif "toolUse" in delta and "input" in delta["toolUse"]:
} # Handle partial JSON for tool use
json_accumulator = "" json_accumulator += delta["toolUse"]["input"]
completion_tokens_estimate += self._estimate_tokens(
delta["toolUse"]["input"]
)
# Handle message completion with tool use # Handle tool use start
elif "messageStop" in event and "stopReason" in event["messageStop"]: elif "contentBlockStart" in event:
if event["messageStop"]["stopReason"] == "tool_use" and tool_use_block: content_block_start = event["contentBlockStart"]["start"]
try: if "toolUse" in content_block_start:
arguments = json.loads(json_accumulator) if json_accumulator else {} tool_use_block = {
"id": content_block_start["toolUse"].get("toolUseId", ""),
"name": content_block_start["toolUse"].get("name", ""),
}
json_accumulator = ""
# Only call function if it's not the no_operation tool # Handle message completion with tool use
if not using_noop_tool: elif "messageStop" in event and "stopReason" in event["messageStop"]:
function_calls.append( if event["messageStop"]["stopReason"] == "tool_use" and tool_use_block:
FunctionCallFromLLM( try:
context=context, arguments = json.loads(json_accumulator) if json_accumulator else {}
tool_call_id=tool_use_block["id"],
function_name=tool_use_block["name"], # Only call function if it's not the no_operation tool
arguments=arguments, 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:
else: logger.debug("Ignoring no_operation tool call")
logger.debug("Ignoring no_operation tool call") except json.JSONDecodeError:
except json.JSONDecodeError: logger.error(f"Failed to parse tool arguments: {json_accumulator}")
logger.error(f"Failed to parse tool arguments: {json_accumulator}")
# Handle usage metrics if available # Handle usage metrics if available
if "metadata" in event and "usage" in event["metadata"]: if "metadata" in event and "usage" in event["metadata"]:
usage = event["metadata"]["usage"] usage = event["metadata"]["usage"]
prompt_tokens += usage.get("inputTokens", 0) prompt_tokens += usage.get("inputTokens", 0)
completion_tokens += usage.get("outputTokens", 0) completion_tokens += usage.get("outputTokens", 0)
cache_read_input_tokens += usage.get("cacheReadInputTokens", 0) cache_read_input_tokens += usage.get("cacheReadInputTokens", 0)
cache_creation_input_tokens += usage.get("cacheWriteInputTokens", 0) cache_creation_input_tokens += usage.get("cacheWriteInputTokens", 0)
await self.run_function_calls(function_calls) await self.run_function_calls(function_calls)
except asyncio.CancelledError: except asyncio.CancelledError:

View File

@@ -30,7 +30,7 @@ from pipecat.transcriptions.language import Language
from pipecat.utils.tracing.service_decorators import traced_tts from pipecat.utils.tracing.service_decorators import traced_tts
try: try:
import boto3 import aioboto3
from botocore.exceptions import BotoCoreError, ClientError from botocore.exceptions import BotoCoreError, ClientError
except ModuleNotFoundError as e: except ModuleNotFoundError as e:
logger.error(f"Exception: {e}") logger.error(f"Exception: {e}")
@@ -177,13 +177,25 @@ class AWSPollyTTSService(TTSService):
params = params or AWSPollyTTSService.InputParams() params = params or AWSPollyTTSService.InputParams()
self._polly_client = boto3.client( # Get credentials from environment variables if not provided
"polly", self._aws_params = {
aws_access_key_id=aws_access_key_id, "aws_access_key_id": aws_access_key_id or os.getenv("AWS_ACCESS_KEY_ID"),
aws_secret_access_key=api_key, "aws_secret_access_key": api_key or os.getenv("AWS_SECRET_ACCESS_KEY"),
aws_session_token=aws_session_token, "aws_session_token": aws_session_token or os.getenv("AWS_SESSION_TOKEN"),
region_name=region, "region_name": region or os.getenv("AWS_REGION", "us-east-1"),
) }
# Validate that we have the required credentials
if (
not self._aws_params["aws_access_key_id"]
or not self._aws_params["aws_secret_access_key"]
):
raise ValueError(
"AWS credentials not found. Please provide them either through constructor parameters "
"or set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables."
)
self._aws_session = aioboto3.Session()
self._settings = { self._settings = {
"engine": params.engine, "engine": params.engine,
"language": self.language_to_service_language(params.language) "language": self.language_to_service_language(params.language)
@@ -199,24 +211,6 @@ class AWSPollyTTSService(TTSService):
self.set_voice(voice_id) self.set_voice(voice_id)
# Get credentials from environment variables if not provided
self._credentials = {
"aws_access_key_id": aws_access_key_id or os.getenv("AWS_ACCESS_KEY_ID"),
"aws_secret_access_key": api_key or os.getenv("AWS_SECRET_ACCESS_KEY"),
"aws_session_token": aws_session_token or os.getenv("AWS_SESSION_TOKEN"),
"region": region or os.getenv("AWS_REGION", "us-east-1"),
}
# Validate that we have the required credentials
if (
not self._credentials["aws_access_key_id"]
or not self._credentials["aws_secret_access_key"]
):
raise ValueError(
"AWS credentials not found. Please provide them either through constructor parameters "
"or set AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables."
)
def can_generate_metrics(self) -> bool: def can_generate_metrics(self) -> bool:
"""Check if this service can generate processing metrics. """Check if this service can generate processing metrics.
@@ -279,14 +273,6 @@ class AWSPollyTTSService(TTSService):
Yields: Yields:
Frame: Audio frames containing the synthesized speech. Frame: Audio frames containing the synthesized speech.
""" """
def read_audio_data(**args):
response = self._polly_client.synthesize_speech(**args)
if "AudioStream" in response:
audio_data = response["AudioStream"].read()
return audio_data
return None
logger.debug(f"{self}: Generating TTS [{text}]") logger.debug(f"{self}: Generating TTS [{text}]")
try: try:
@@ -309,30 +295,32 @@ class AWSPollyTTSService(TTSService):
# Filter out None values # Filter out None values
filtered_params = {k: v for k, v in params.items() if v is not None} filtered_params = {k: v for k, v in params.items() if v is not None}
audio_data = await asyncio.to_thread(read_audio_data, **filtered_params) async with self._aws_session.client("polly", **self._aws_params) as polly:
response = await polly.synthesize_speech(**filtered_params)
if "AudioStream" in response:
# Get the streaming body and read it
stream = response["AudioStream"]
audio_data = await stream.read()
else:
logger.error(f"{self} No audio stream in response")
audio_data = None
if not audio_data: audio_data = await self._resampler.resample(audio_data, 16000, self.sample_rate)
logger.error(f"{self} No audio data returned")
yield None
return
audio_data = await self._resampler.resample(audio_data, 16000, self.sample_rate) await self.start_tts_usage_metrics(text)
await self.start_tts_usage_metrics(text) yield TTSStartedFrame()
yield TTSStartedFrame() CHUNK_SIZE = self.chunk_size
CHUNK_SIZE = self.chunk_size for i in range(0, len(audio_data), CHUNK_SIZE):
chunk = audio_data[i : i + CHUNK_SIZE]
for i in range(0, len(audio_data), CHUNK_SIZE): if len(chunk) > 0:
chunk = audio_data[i : i + CHUNK_SIZE] await self.stop_ttfb_metrics()
if len(chunk) > 0: frame = TTSAudioRawFrame(chunk, self.sample_rate, 1)
await self.stop_ttfb_metrics() yield frame
frame = TTSAudioRawFrame(chunk, self.sample_rate, 1)
yield frame
yield TTSStoppedFrame()
yield TTSStoppedFrame()
except (BotoCoreError, ClientError) as error: except (BotoCoreError, ClientError) as error:
logger.exception(f"{self} error generating TTS: {error}") logger.exception(f"{self} error generating TTS: {error}")
error_message = f"AWS Polly TTS error: {str(error)}" error_message = f"AWS Polly TTS error: {str(error)}"