From 05ae8d3ffa11e07f8868e1e81c387ce4419dddfa Mon Sep 17 00:00:00 2001 From: Adithya Suresh Date: Fri, 4 Apr 2025 05:36:09 +0000 Subject: [PATCH] Removed OpenAI based context formatting --- src/pipecat/services/aws/llm.py | 580 ++++++++++++++++---------------- 1 file changed, 286 insertions(+), 294 deletions(-) diff --git a/src/pipecat/services/aws/llm.py b/src/pipecat/services/aws/llm.py index 2f762e9bd..cb21eccaa 100644 --- a/src/pipecat/services/aws/llm.py +++ b/src/pipecat/services/aws/llm.py @@ -46,16 +46,6 @@ from pipecat.processors.aggregators.openai_llm_context import ( from pipecat.processors.frame_processor import FrameDirection from pipecat.services.ai_services import LLMService -try: - from anthropic import NOT_GIVEN, 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]`. " - + "Also, set `ANTHROPIC_API_KEY` environment variable." - ) - raise Exception(f"Missing module: {e}") - @dataclass class BedrockContextAggregatorPair: @@ -69,288 +59,6 @@ class BedrockContextAggregatorPair: return self._assistant -class BedrockLLMService(LLMService): - """This class implements inference with AWS Bedrock models including Amazon Nova and Anthropic Claude. - - Requires AWS credentials to be configured in the environment or through boto3 configuration. - """ - class InputParams(BaseModel): - max_tokens: Optional[int] = Field(default_factory=lambda: 4096, ge=1) - temperature: Optional[float] = Field(default_factory=lambda: 0.7, ge=0.0, le=1.0) - top_p: Optional[float] = Field(default_factory=lambda: 0.999, ge=0.0, le=1.0) - stop_sequences: Optional[List[str]] = Field(default_factory=lambda: []) - latency: Optional[str] = Field(default_factory=lambda: "standard") - additional_model_request_fields: Optional[Dict[str, Any]] = Field(default_factory=dict) - - def __init__( - self, - *, - aws_access_key: str, - aws_secret_key: str, - aws_session_token: Optional[str] = None, - aws_region: str = "us-east-1", - model: str, - params: InputParams = InputParams(), - client_config: Optional[Config] = None, - **kwargs, - ): - super().__init__(**kwargs) - - # Initialize the Bedrock client - if not client_config: - client_config = Config( - connect_timeout=300, # 5 minutes - 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.set_model_name(model) - self._settings = { - "max_tokens": params.max_tokens, - "temperature": params.temperature, - "top_p": params.top_p, - "latency": params.latency, - "additional_model_request_fields": params.additional_model_request_fields if isinstance(params.additional_model_request_fields, dict) else {}, - } - - # Determine model provider from model ID - self.model_provider = self._get_model_provider(model) - logger.info(f"Using AWS Bedrock model: {model} from provider: {self.model_provider}") - - def _get_model_provider(self, model: str) -> str: - """Determine the model provider from the model ID""" - if "anthropic." in model: - return "anthropic" - elif "amazon." in model: - return "amazon" - else: - raise ValueError(f"Unsupported model: {model}. Only Anthropic Claude and Amazon Nova model families are supported.") - - def can_generate_metrics(self) -> bool: - return True - - def create_context_aggregator( - self, - context: OpenAILLMContext, - *, - user_kwargs: Mapping[str, Any] = {}, - assistant_kwargs: Mapping[str, Any] = {}, - ) -> BedrockContextAggregatorPair: - """Create an instance of BedrockContextAggregatorPair from an - OpenAILLMContext. Constructor keyword arguments for both the user and - assistant aggregators can be provided. - - Args: - context (OpenAILLMContext): The LLM context. - user_kwargs (Mapping[str, Any], optional): Additional keyword - arguments for the user context aggregator constructor. Defaults - to an empty mapping. - assistant_kwargs (Mapping[str, Any], optional): Additional keyword - arguments for the assistant context aggregator - constructor. Defaults to an empty mapping. - - Returns: - BedrockContextAggregatorPair: A pair of context aggregators, one - for the user and one for the assistant, encapsulated in an - BedrockContextAggregatorPair. - """ - context.set_llm_adapter(self.get_llm_adapter()) - - if isinstance(context, OpenAILLMContext): - context = BedrockLLMContext.from_openai_context(context) - - user = BedrockUserContextAggregator(context, **user_kwargs) - assistant = BedrockAssistantContextAggregator(context, **assistant_kwargs) - return BedrockContextAggregatorPair(_user=user, _assistant=assistant) - - async def _process_context(self, context: "BedrockLLMContext"): - # Usage tracking - prompt_tokens = 0 - completion_tokens = 0 - completion_tokens_estimate = 0 - use_completion_tokens_estimate = False - - try: - await self.push_frame(LLMFullResponseStartFrame()) - await self.start_processing_metrics() - - # logger.debug( - # f"{self}: Generating chat with Bedrock model {self.model_name} | [{context.get_messages_for_logging()}]" - # ) - - await self.start_ttfb_metrics() - - # Set up inference config - inference_config = { - "maxTokens": self._settings["max_tokens"], - "temperature": self._settings["temperature"], - "topP": self._settings["top_p"], - } - - # Prepare request parameters - request_params = { - "modelId": self.model_name, - "messages": context.messages, - "inferenceConfig": inference_config, - "additionalModelRequestFields": self._settings["additional_model_request_fields"] - } - - # Add system message - request_params["system"] = [{"text": context.system}] - - # Add tools if present - if context.tools: - tool_config = { - "tools": context.tools - } - - # Add tool_choice if specified - if context.tool_choice: - if context.tool_choice == "auto": - tool_config["toolChoice"] = {"auto": {}} - elif context.tool_choice == "none": - # Skip adding toolChoice for "none" - pass - elif isinstance(context.tool_choice, dict) and "function" in context.tool_choice: - tool_config["toolChoice"] = { - "tool": { - "name": context.tool_choice["function"]["name"] - } - } - - request_params["toolConfig"] = tool_config - - # Add performance config if latency is specified - if self._settings["latency"] in ["standard", "optimized"]: - request_params["performanceConfig"] = { - "latency": self._settings["latency"] - } - - logger.debug(f"Calling Bedrock model with: {request_params}") - - # Call 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 = "" - - 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 {} - await self.call_function( - context=context, - tool_call_id=tool_use_block["id"], - function_name=tool_use_block["name"], - arguments=arguments, - ) - 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) - - except asyncio.CancelledError: - # If we're interrupted, we won't get a complete usage report. So set our flag to use the - # token estimate. The reraise the exception so all the processors running in this task - # also get cancelled. - use_completion_tokens_estimate = True - raise - except httpx.TimeoutException: - await self._call_event_handler("on_completion_timeout") - except Exception as e: - logger.exception(f"{self} exception: {e}") - finally: - await self.stop_processing_metrics() - await self.push_frame(LLMFullResponseEndFrame()) - comp_tokens = ( - completion_tokens - if not use_completion_tokens_estimate - else completion_tokens_estimate - ) - await self._report_usage_metrics( - prompt_tokens=prompt_tokens, - completion_tokens=comp_tokens, - ) - - async def process_frame(self, frame: Frame, direction: FrameDirection): - await super().process_frame(frame, direction) - - context = None - if isinstance(frame, OpenAILLMContextFrame): - context = BedrockLLMContext.upgrade_to_bedrock(frame.context) - elif isinstance(frame, LLMMessagesFrame): - context = BedrockLLMContext.from_messages(frame.messages) - elif isinstance(frame, VisionImageRawFrame): - # This is only useful in very simple pipelines because it creates - # a new context. Generally we want a context manager to catch - # UserImageRawFrames coming through the pipeline and add them - # to the context. - context = BedrockLLMContext.from_image_frame(frame) - elif isinstance(frame, LLMUpdateSettingsFrame): - await self._update_settings(frame.settings) - else: - await self.push_frame(frame, direction) - - if context: - await self._process_context(context) - - def _estimate_tokens(self, text: str) -> int: - return int(len(re.split(r"[^\w]+", text)) * 1.3) - - async def _report_usage_metrics( - self, - prompt_tokens: int, - completion_tokens: int, - ): - if prompt_tokens or completion_tokens: - tokens = LLMTokenUsage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ) - await self.start_llm_usage_metrics(tokens) - - class BedrockLLMContext(OpenAILLMContext): def __init__( self, @@ -358,7 +66,7 @@ class BedrockLLMContext(OpenAILLMContext): tools: Optional[List[dict]] = None, tool_choice: Optional[dict] = None, *, - system: Union[str, NotGiven] = NOT_GIVEN, + system: Optional[str] = None, ): super().__init__(messages=messages, tools=tools, tool_choice=tool_choice) self.system = system @@ -375,6 +83,7 @@ class BedrockLLMContext(OpenAILLMContext): @classmethod def from_openai_context(cls, openai_context: OpenAILLMContext): + logger.debug("from_openai_context called") self = cls( messages=openai_context.messages, tools=openai_context.tools, @@ -621,6 +330,7 @@ class BedrockLLMContext(OpenAILLMContext): merging consecutive messages with the same role, and ensuring proper content formatting. """ # Handle system message if present at the beginning + logger.debug(f"_restructure_from_bedrock_messages: {self.messages}") if self.messages and self.messages[0]["role"] == "system": if len(self.messages) == 1: self.messages[0]["role"] = "user" @@ -653,6 +363,7 @@ class BedrockLLMContext(OpenAILLMContext): self.messages.extend(merged_messages) def _restructure_from_openai_messages(self): + logger.debug(f"_restructure_from_openai_messages: {self.messages}") # first, map across self._messages calling self.from_standard_message(m) to modify messages in place try: self._messages[:] = [self.from_standard_message(m) for m in self._messages] @@ -794,4 +505,285 @@ class BedrockAssistantContextAggregator(LLMAssistantContextAggregator): image=frame.image, text=frame.request.context, ) - \ No newline at end of file + + +class BedrockLLMService(LLMService): + """This class implements inference with AWS Bedrock models including Amazon Nova and Anthropic Claude. + + Requires AWS credentials to be configured in the environment or through boto3 configuration. + """ + class InputParams(BaseModel): + max_tokens: Optional[int] = Field(default_factory=lambda: 4096, ge=1) + temperature: Optional[float] = Field(default_factory=lambda: 0.7, ge=0.0, le=1.0) + top_p: Optional[float] = Field(default_factory=lambda: 0.999, ge=0.0, le=1.0) + stop_sequences: Optional[List[str]] = Field(default_factory=lambda: []) + latency: Optional[str] = Field(default_factory=lambda: "standard") + additional_model_request_fields: Optional[Dict[str, Any]] = Field(default_factory=dict) + + def __init__( + self, + *, + aws_access_key: Optional[str] = None, + aws_secret_key: Optional[str] = None, + aws_session_token: Optional[str] = None, + aws_region: str = "us-east-1", + model: str, + params: InputParams = InputParams(), + client_config: Optional[Config] = None, + **kwargs, + ): + super().__init__(**kwargs) + + # Initialize the Bedrock client + if not client_config: + client_config = Config( + connect_timeout=300, # 5 minutes + 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.set_model_name(model) + self._settings = { + "max_tokens": params.max_tokens, + "temperature": params.temperature, + "top_p": params.top_p, + "latency": params.latency, + "additional_model_request_fields": params.additional_model_request_fields if isinstance(params.additional_model_request_fields, dict) else {}, + } + + # Determine model provider from model ID + self.model_provider = self._get_model_provider(model) + logger.info(f"Using AWS Bedrock model: {model} from provider: {self.model_provider}") + + def _get_model_provider(self, model: str) -> str: + """Determine the model provider from the model ID""" + if "anthropic." in model: + return "anthropic" + elif "amazon." in model: + return "amazon" + else: + raise ValueError(f"Unsupported model: {model}. Only Anthropic Claude and Amazon Nova model families are supported.") + + def can_generate_metrics(self) -> bool: + return True + + def create_context_aggregator( + self, + context: BedrockLLMContext, + *, + user_kwargs: Mapping[str, Any] = {}, + assistant_kwargs: Mapping[str, Any] = {}, + ) -> BedrockContextAggregatorPair: + """Create an instance of BedrockContextAggregatorPair from an + OpenAILLMContext. Constructor keyword arguments for both the user and + assistant aggregators can be provided. + + Args: + context (OpenAILLMContext): The LLM context. + user_kwargs (Mapping[str, Any], optional): Additional keyword + arguments for the user context aggregator constructor. Defaults + to an empty mapping. + assistant_kwargs (Mapping[str, Any], optional): Additional keyword + arguments for the assistant context aggregator + constructor. Defaults to an empty mapping. + + Returns: + BedrockContextAggregatorPair: A pair of context aggregators, one + for the user and one for the assistant, encapsulated in an + BedrockContextAggregatorPair. + """ + context.set_llm_adapter(self.get_llm_adapter()) + + if isinstance(context, OpenAILLMContext) and not isinstance(context, BedrockLLMContext): + context = BedrockLLMContext.from_openai_context(context) + + user = BedrockUserContextAggregator(context, **user_kwargs) + assistant = BedrockAssistantContextAggregator(context, **assistant_kwargs) + return BedrockContextAggregatorPair(_user=user, _assistant=assistant) + + async def _process_context(self, context: "BedrockLLMContext"): + # Usage tracking + prompt_tokens = 0 + completion_tokens = 0 + completion_tokens_estimate = 0 + use_completion_tokens_estimate = False + + try: + await self.push_frame(LLMFullResponseStartFrame()) + await self.start_processing_metrics() + + # logger.debug( + # f"{self}: Generating chat with Bedrock model {self.model_name} | [{context.get_messages_for_logging()}]" + # ) + + await self.start_ttfb_metrics() + + # Set up inference config + inference_config = { + "maxTokens": self._settings["max_tokens"], + "temperature": self._settings["temperature"], + "topP": self._settings["top_p"], + } + + # Prepare request parameters + request_params = { + "modelId": self.model_name, + "messages": context.messages, + "inferenceConfig": inference_config, + "additionalModelRequestFields": self._settings["additional_model_request_fields"] + } + + # Add system message + request_params["system"] = [{"text": context.system}] + + # Add tools if present + if context.tools: + tool_config = { + "tools": context.tools + } + + # Add tool_choice if specified + if context.tool_choice: + if context.tool_choice == "auto": + tool_config["toolChoice"] = {"auto": {}} + elif context.tool_choice == "none": + # Skip adding toolChoice for "none" + pass + elif isinstance(context.tool_choice, dict) and "function" in context.tool_choice: + tool_config["toolChoice"] = { + "tool": { + "name": context.tool_choice["function"]["name"] + } + } + + request_params["toolConfig"] = tool_config + + # Add performance config if latency is specified + if self._settings["latency"] in ["standard", "optimized"]: + request_params["performanceConfig"] = { + "latency": self._settings["latency"] + } + + logger.debug(f"Calling Bedrock model with: {request_params}") + + # Call 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 = "" + + 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 {} + await self.call_function( + context=context, + tool_call_id=tool_use_block["id"], + function_name=tool_use_block["name"], + arguments=arguments, + ) + 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) + + except asyncio.CancelledError: + # If we're interrupted, we won't get a complete usage report. So set our flag to use the + # token estimate. The reraise the exception so all the processors running in this task + # also get cancelled. + use_completion_tokens_estimate = True + raise + except httpx.TimeoutException: + await self._call_event_handler("on_completion_timeout") + except Exception as e: + logger.exception(f"{self} exception: {e}") + finally: + await self.stop_processing_metrics() + await self.push_frame(LLMFullResponseEndFrame()) + comp_tokens = ( + completion_tokens + if not use_completion_tokens_estimate + else completion_tokens_estimate + ) + await self._report_usage_metrics( + prompt_tokens=prompt_tokens, + completion_tokens=comp_tokens, + ) + + async def process_frame(self, frame: Frame, direction: FrameDirection): + await super().process_frame(frame, direction) + + context = None + if isinstance(frame, OpenAILLMContextFrame): + context = BedrockLLMContext.upgrade_to_bedrock(frame.context) + elif isinstance(frame, LLMMessagesFrame): + context = BedrockLLMContext.from_messages(frame.messages) + elif isinstance(frame, VisionImageRawFrame): + # This is only useful in very simple pipelines because it creates + # a new context. Generally we want a context manager to catch + # UserImageRawFrames coming through the pipeline and add them + # to the context. + context = BedrockLLMContext.from_image_frame(frame) + elif isinstance(frame, LLMUpdateSettingsFrame): + await self._update_settings(frame.settings) + else: + await self.push_frame(frame, direction) + + if context: + await self._process_context(context) + + def _estimate_tokens(self, text: str) -> int: + return int(len(re.split(r"[^\w]+", text)) * 1.3) + + async def _report_usage_metrics( + self, + prompt_tokens: int, + completion_tokens: int, + ): + if prompt_tokens or completion_tokens: + tokens = LLMTokenUsage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ) + await self.start_llm_usage_metrics(tokens)