Merge pull request #2148 from pipecat-ai/filipi/aws_bedrock
Refactoring AWSBedrockLLMService to work async
This commit is contained in:
@@ -25,6 +25,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
||||
|
||||
### 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
|
||||
longer than the idle timeout duration. This allows you to use the
|
||||
`UserIdleProcessor` in conjunction with function calls that take a while to
|
||||
|
||||
@@ -42,7 +42,7 @@ Website = "https://pipecat.ai"
|
||||
[project.optional-dependencies]
|
||||
anthropic = [ "anthropic~=0.49.0" ]
|
||||
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" ]
|
||||
azure = [ "azure-cognitiveservices-speech~=1.42.0"]
|
||||
cartesia = [ "cartesia~=2.0.3", "websockets~=13.1" ]
|
||||
|
||||
@@ -55,7 +55,7 @@ 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
|
||||
except ModuleNotFoundError as e:
|
||||
@@ -749,13 +749,17 @@ 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._settings = {
|
||||
@@ -903,70 +907,74 @@ 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)
|
||||
async with self._aws_session.client(
|
||||
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
|
||||
tool_use_block = None
|
||||
json_accumulator = ""
|
||||
# Process the streaming response
|
||||
tool_use_block = None
|
||||
json_accumulator = ""
|
||||
|
||||
function_calls = []
|
||||
for event in response["stream"]:
|
||||
self.reset_watchdog()
|
||||
function_calls = []
|
||||
|
||||
# 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"]
|
||||
)
|
||||
async for event in response["stream"]:
|
||||
self.reset_watchdog()
|
||||
|
||||
# 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 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 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 {}
|
||||
# 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 = ""
|
||||
|
||||
# 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,
|
||||
# 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}")
|
||||
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)
|
||||
# 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:
|
||||
|
||||
@@ -30,7 +30,7 @@ from pipecat.transcriptions.language import Language
|
||||
from pipecat.utils.tracing.service_decorators import traced_tts
|
||||
|
||||
try:
|
||||
import boto3
|
||||
import aioboto3
|
||||
from botocore.exceptions import BotoCoreError, ClientError
|
||||
except ModuleNotFoundError as e:
|
||||
logger.error(f"Exception: {e}")
|
||||
@@ -177,13 +177,25 @@ class AWSPollyTTSService(TTSService):
|
||||
|
||||
params = params or AWSPollyTTSService.InputParams()
|
||||
|
||||
self._polly_client = boto3.client(
|
||||
"polly",
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=api_key,
|
||||
aws_session_token=aws_session_token,
|
||||
region_name=region,
|
||||
)
|
||||
# Get credentials from environment variables if not provided
|
||||
self._aws_params = {
|
||||
"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_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 = {
|
||||
"engine": params.engine,
|
||||
"language": self.language_to_service_language(params.language)
|
||||
@@ -199,24 +211,6 @@ class AWSPollyTTSService(TTSService):
|
||||
|
||||
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:
|
||||
"""Check if this service can generate processing metrics.
|
||||
|
||||
@@ -279,14 +273,6 @@ class AWSPollyTTSService(TTSService):
|
||||
Yields:
|
||||
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}]")
|
||||
|
||||
try:
|
||||
@@ -309,30 +295,32 @@ class AWSPollyTTSService(TTSService):
|
||||
# Filter out None values
|
||||
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:
|
||||
logger.error(f"{self} No audio data returned")
|
||||
yield None
|
||||
return
|
||||
audio_data = await self._resampler.resample(audio_data, 16000, self.sample_rate)
|
||||
|
||||
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]
|
||||
if len(chunk) > 0:
|
||||
await self.stop_ttfb_metrics()
|
||||
frame = TTSAudioRawFrame(chunk, self.sample_rate, 1)
|
||||
yield frame
|
||||
|
||||
yield TTSStoppedFrame()
|
||||
for i in range(0, len(audio_data), CHUNK_SIZE):
|
||||
chunk = audio_data[i : i + CHUNK_SIZE]
|
||||
if len(chunk) > 0:
|
||||
await self.stop_ttfb_metrics()
|
||||
frame = TTSAudioRawFrame(chunk, self.sample_rate, 1)
|
||||
yield frame
|
||||
|
||||
yield TTSStoppedFrame()
|
||||
except (BotoCoreError, ClientError) as error:
|
||||
logger.exception(f"{self} error generating TTS: {error}")
|
||||
error_message = f"AWS Polly TTS error: {str(error)}"
|
||||
|
||||
Reference in New Issue
Block a user