diff --git a/docs/my-website/docs/providers/bedrock_agentcore.md b/docs/my-website/docs/providers/bedrock_agentcore.md new file mode 100644 index 0000000000..43df7f8251 --- /dev/null +++ b/docs/my-website/docs/providers/bedrock_agentcore.md @@ -0,0 +1,246 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Bedrock AgentCore + +Call Bedrock AgentCore in the OpenAI Request/Response format. + +| Property | Details | +|----------|---------| +| Description | Amazon Bedrock AgentCore provides direct access to hosted agent runtimes for executing agentic workflows with foundation models. | +| Provider Route on LiteLLM | `bedrock/agentcore/{AGENT_RUNTIME_ARN}` | +| Provider Doc | [AWS Bedrock AgentCore ↗](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgentRuntime.html) | + +## Quick Start + +### Model Format to LiteLLM + +To call a bedrock agent runtime through LiteLLM, use the following model format. + +Here the `model=bedrock/agentcore/` tells LiteLLM to call the bedrock `InvokeAgentRuntime` API. + +```shell showLineNumbers title="Model Format to LiteLLM" +bedrock/agentcore/{AGENT_RUNTIME_ARN} +``` + +**Example:** +- `bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent-runtime` + +You can find the Agent Runtime ARN in your AWS Bedrock console under AgentCore. + +### LiteLLM Python SDK + +```python showLineNumbers title="Basic AgentCore Completion" +import litellm + +# Make a completion request to your AgentCore runtime +response = litellm.completion( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent-runtime", + messages=[ + { + "role": "user", + "content": "Explain machine learning in simple terms" + } + ], +) + +print(response.choices[0].message.content) +print(f"Usage: {response.usage}") +``` + +```python showLineNumbers title="Streaming AgentCore Responses" +import litellm + +# Stream responses from your AgentCore runtime +response = litellm.completion( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent-runtime", + messages=[ + { + "role": "user", + "content": "What are the key principles of software architecture?" + } + ], + stream=True, +) + +for chunk in response: + if chunk.choices[0].delta.content: + print(chunk.choices[0].delta.content, end="") +``` + +### LiteLLM Proxy + +#### 1. Configure your model in config.yaml + + + + +```yaml showLineNumbers title="LiteLLM Proxy Configuration" +model_list: + - model_name: agentcore-runtime-1 + litellm_params: + model: bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent-runtime + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_region_name: us-west-2 + + - model_name: agentcore-runtime-2 + litellm_params: + model: bedrock/agentcore/arn:aws:bedrock-agentcore:us-east-1:987654321098:runtime/production-runtime + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_region_name: us-east-1 +``` + + + + +#### 2. Start the LiteLLM Proxy + +```bash showLineNumbers title="Start LiteLLM Proxy" +litellm --config config.yaml +``` + +#### 3. Make requests to your AgentCore runtimes + + + + +```bash showLineNumbers title="Basic AgentCore Request" +curl http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -d '{ + "model": "agentcore-runtime-1", + "messages": [ + { + "role": "user", + "content": "Summarize the main benefits of cloud computing" + } + ] + }' +``` + +```bash showLineNumbers title="Streaming AgentCore Request" +curl http://localhost:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -d '{ + "model": "agentcore-runtime-2", + "messages": [ + { + "role": "user", + "content": "Explain the differences between SQL and NoSQL databases" + } + ], + "stream": true + }' +``` + + + + + +```python showLineNumbers title="Using OpenAI SDK with LiteLLM Proxy" +from openai import OpenAI + +# Initialize client with your LiteLLM proxy URL +client = OpenAI( + base_url="http://localhost:4000", + api_key="your-litellm-api-key" +) + +# Make a completion request to your AgentCore runtime +response = client.chat.completions.create( + model="agentcore-runtime-1", + messages=[ + { + "role": "user", + "content": "What are best practices for API design?" + } + ] +) + +print(response.choices[0].message.content) +``` + +```python showLineNumbers title="Streaming with OpenAI SDK" +from openai import OpenAI + +client = OpenAI( + base_url="http://localhost:4000", + api_key="your-litellm-api-key" +) + +# Stream AgentCore responses +stream = client.chat.completions.create( + model="agentcore-runtime-2", + messages=[ + { + "role": "user", + "content": "Describe the microservices architecture pattern" + } + ], + stream=True +) + +for chunk in stream: + if chunk.choices[0].delta.content is not None: + print(chunk.choices[0].delta.content, end="") +``` + + + + +## Provider-specific Parameters + +AgentCore supports additional parameters that can be passed to customize the runtime invocation. + + + + +```python showLineNumbers title="Using AgentCore-specific parameters" +from litellm import completion + +response = litellm.completion( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent-runtime", + messages=[ + { + "role": "user", + "content": "Analyze this data and provide insights", + } + ], + qualifier="production", # PROVIDER-SPECIFIC: Runtime qualifier/version + runtimeSessionId="session-abc-123", # PROVIDER-SPECIFIC: Custom session ID +) +``` + + + + +```yaml showLineNumbers title="LiteLLM Proxy Configuration with Parameters" +model_list: + - model_name: agentcore-runtime-prod + litellm_params: + model: bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/my-agent-runtime + aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + aws_region_name: us-west-2 + qualifier: production +``` + + + + +### Available Parameters + +| Parameter | Type | Description | +|-----------|------|-------------| +| `qualifier` | string | Optional runtime qualifier/version to invoke a specific version of the agent runtime | +| `runtimeSessionId` | string | Optional custom session ID (must be 33+ characters). If not provided, LiteLLM generates one automatically | + +## Further Reading + +- [AWS Bedrock AgentCore Documentation](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgentRuntime.html) +- [LiteLLM Authentication to Bedrock](https://docs.litellm.ai/docs/providers/bedrock#boto3---authentication) + diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index ccba0bdb7c..ab1e6df125 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -494,6 +494,7 @@ const sidebars = { "providers/bedrock_embedding", "providers/bedrock_image_gen", "providers/bedrock_rerank", + "providers/bedrock_agentcore", "providers/bedrock_agents", "providers/bedrock_batches", "providers/bedrock_vector_store", diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 4c85443754..72e270428a 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -901,7 +901,7 @@ class BaseAWSLLM: api_base: Optional[str], aws_bedrock_runtime_endpoint: Optional[str], aws_region_name: str, - endpoint_type: Optional[Literal["runtime", "agent"]] = "runtime", + endpoint_type: Optional[Literal["runtime", "agent", "agentcore"]] = "runtime", ) -> Tuple[str, str]: env_aws_bedrock_runtime_endpoint = get_secret("AWS_BEDROCK_RUNTIME_ENDPOINT") if api_base is not None: @@ -935,7 +935,7 @@ class BaseAWSLLM: return endpoint_url, proxy_endpoint_url def _select_default_endpoint_url( - self, endpoint_type: Optional[Literal["runtime", "agent"]], aws_region_name: str + self, endpoint_type: Optional[Literal["runtime", "agent", "agentcore"]], aws_region_name: str ) -> str: """ Select the default endpoint url based on the endpoint type @@ -944,6 +944,8 @@ class BaseAWSLLM: """ if endpoint_type == "agent": return f"https://bedrock-agent-runtime.{aws_region_name}.amazonaws.com" + elif endpoint_type == "agentcore": + return f"https://bedrock-agentcore.{aws_region_name}.amazonaws.com" else: return f"https://bedrock-runtime.{aws_region_name}.amazonaws.com" @@ -1091,7 +1093,7 @@ class BaseAWSLLM: def _sign_request( self, - service_name: Literal["bedrock", "sagemaker"], + service_name: Literal["bedrock", "sagemaker", "bedrock-agentcore"], headers: dict, optional_params: dict, request_data: dict, diff --git a/litellm/llms/bedrock/chat/agentcore/__init__.py b/litellm/llms/bedrock/chat/agentcore/__init__.py new file mode 100644 index 0000000000..a2f1387620 --- /dev/null +++ b/litellm/llms/bedrock/chat/agentcore/__init__.py @@ -0,0 +1,4 @@ +from .transformation import AmazonAgentCoreConfig + +__all__ = ["AmazonAgentCoreConfig"] + diff --git a/litellm/llms/bedrock/chat/agentcore/sse_iterator.py b/litellm/llms/bedrock/chat/agentcore/sse_iterator.py new file mode 100644 index 0000000000..8e0e698e61 --- /dev/null +++ b/litellm/llms/bedrock/chat/agentcore/sse_iterator.py @@ -0,0 +1,150 @@ +""" +SSE Stream Iterator for Bedrock AgentCore. + +Handles Server-Sent Events (SSE) streaming responses from AgentCore. +""" + +import json +from typing import TYPE_CHECKING + +import httpx + +from litellm._logging import verbose_logger +from litellm._uuid import uuid +from litellm.types.llms.bedrock_agentcore import AgentCoreUsage +from litellm.types.utils import Delta, ModelResponse, StreamingChoices, Usage + +if TYPE_CHECKING: + pass + + +class AgentCoreSSEStreamIterator: + """Iterator for AgentCore SSE streaming responses.""" + + def __init__(self, response: httpx.Response, model: str): + self.response = response + self.model = model + self.finished = False + self.line_iterator = self.response.iter_lines() + + def __iter__(self): + return self + + def __next__(self) -> ModelResponse: + """Parse SSE events and yield ModelResponse chunks.""" + try: + for line in self.line_iterator: + line = line.strip() + + if not line or not line.startswith('data:'): + continue + + # Extract JSON from SSE line + json_str = line[5:].strip() + if not json_str: + continue + + try: + data = json.loads(json_str) + + # Skip non-dict data + if not isinstance(data, dict): + continue + + # Process content delta events + if "event" in data and isinstance(data["event"], dict): + event_payload = data["event"] + content_block_delta = event_payload.get("contentBlockDelta") + + if content_block_delta: + delta = content_block_delta.get("delta", {}) + text = delta.get("text", "") + + if text: + # Yield chunk with text + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=self.model, + object="chat.completion.chunk", + ) + + chunk.choices = [ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content=text, role="assistant"), + ) + ] + + return chunk + + # Check for metadata/usage + metadata = event_payload.get("metadata") + if metadata and "usage" in metadata: + # This is the final chunk with usage + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=self.model, + object="chat.completion.chunk", + ) + + chunk.choices = [ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ] + + usage_data: AgentCoreUsage = metadata["usage"] # type: ignore + setattr(chunk, "usage", Usage( + prompt_tokens=usage_data.get("inputTokens", 0), + completion_tokens=usage_data.get("outputTokens", 0), + total_tokens=usage_data.get("totalTokens", 0), + )) + + self.finished = True + return chunk + + # Check for final message (alternative finish signal) + if "message" in data and isinstance(data["message"], dict): + if not self.finished: + chunk = ModelResponse( + id=f"chatcmpl-{uuid.uuid4()}", + created=0, + model=self.model, + object="chat.completion.chunk", + ) + + chunk.choices = [ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(), + ) + ] + + self.finished = True + return chunk + + except json.JSONDecodeError: + verbose_logger.debug(f"Skipping non-JSON SSE line: {line[:100]}") + continue + + # Stream ended naturally + raise StopIteration + + except StopIteration: + raise + except httpx.StreamConsumed: + # This is expected when the stream has been fully consumed + raise StopIteration + except httpx.StreamClosed: + # This is expected when the stream is closed + raise StopIteration + except Exception as e: + verbose_logger.error(f"Error in AgentCore SSE stream: {str(e)}") + raise StopIteration + diff --git a/litellm/llms/bedrock/chat/agentcore/transformation.py b/litellm/llms/bedrock/chat/agentcore/transformation.py new file mode 100644 index 0000000000..1bfd2809a1 --- /dev/null +++ b/litellm/llms/bedrock/chat/agentcore/transformation.py @@ -0,0 +1,588 @@ +""" +Transformation for Bedrock AgentCore + +https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgentRuntime.html +""" + +import json +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union +from urllib.parse import quote + +import httpx + +from litellm._logging import verbose_logger +from litellm._uuid import uuid +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + convert_content_list_to_str, +) +from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException +from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +from litellm.llms.bedrock.chat.agentcore.sse_iterator import AgentCoreSSEStreamIterator +from litellm.llms.bedrock.common_utils import BedrockError +from litellm.types.llms.bedrock_agentcore import ( + AgentCoreMessage, + AgentCoreParsedResponse, + AgentCoreUsage, +) +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import Choices, Message, ModelResponse, Usage + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.utils import CustomStreamWrapper + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + HTTPHandler = Any + AsyncHTTPHandler = Any + CustomStreamWrapper = Any + + +class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM): + def __init__(self, **kwargs): + BaseConfig.__init__(self, **kwargs) + BaseAWSLLM.__init__(self, **kwargs) + + def get_supported_openai_params(self, model: str) -> List[str]: + """ + Bedrock AgentCore has 0 OpenAI compatible params + """ + return [] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + """ + Map OpenAI params to AgentCore params + """ + return optional_params + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete url for the request + """ + ### SET RUNTIME ENDPOINT ### + aws_bedrock_runtime_endpoint = optional_params.get( + "aws_bedrock_runtime_endpoint", None + ) + + # Extract ARN from model string + agent_runtime_arn = self._get_agent_runtime_arn(model) + + # Parse ARN to get region + region = self._extract_region_from_arn(agent_runtime_arn) + + # Build the base endpoint URL for AgentCore + # Note: We don't use get_runtime_endpoint as AgentCore has its own endpoint structure + if aws_bedrock_runtime_endpoint: + base_url = aws_bedrock_runtime_endpoint + else: + base_url = f"https://bedrock-agentcore.{region}.amazonaws.com" + + # Based on boto3 client.invoke_agent_runtime, the path is: + # /runtimes/{URL-ENCODED-ARN}/invocations?qualifier= + encoded_arn = quote(agent_runtime_arn, safe='') + endpoint_url = f"{base_url}/runtimes/{encoded_arn}/invocations" + + # Add qualifier as query parameter if provided + if "qualifier" in optional_params: + endpoint_url = f"{endpoint_url}?qualifier={optional_params['qualifier']}" + + return endpoint_url + + def sign_request( + self, + headers: dict, + optional_params: dict, + request_data: dict, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + stream: Optional[bool] = None, + fake_stream: Optional[bool] = None, + ) -> Tuple[dict, Optional[bytes]]: + return self._sign_request( + service_name="bedrock-agentcore", + headers=headers, + optional_params=optional_params, + request_data=request_data, + api_base=api_base, + model=model, + stream=stream, + fake_stream=fake_stream, + api_key=api_key, + ) + + def _get_agent_runtime_arn(self, model: str) -> str: + """ + Extract ARN from model string + model = "agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC" + returns: "arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC" + """ + parts = model.split("/", 1) + if len(parts) != 2 or parts[0] != "agentcore": + raise ValueError( + "Invalid model format. Expected format: 'model=bedrock/agentcore/arn:aws:bedrock-agentcore:region:account:runtime/runtime_id'" + ) + return parts[1] + + def _extract_region_from_arn(self, arn: str) -> str: + """ + Extract region from ARN + arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC + returns: us-west-2 + """ + parts = arn.split(":") + if len(parts) >= 4: + return parts[3] + raise ValueError(f"Invalid ARN format: {arn}") + + def _get_runtime_session_id(self, optional_params: dict) -> str: + """ + Get or generate runtime session ID (must be 33+ chars) + """ + session_id = optional_params.get("runtimeSessionId", None) + if session_id: + return session_id + + # Generate a session ID with 33+ characters + return f"litellm-session-{str(uuid.uuid4())}" + + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Transform the request to AgentCore format. + + Based on boto3's implementation: + - Session ID goes in header: X-Amzn-Bedrock-AgentCore-Runtime-Session-Id + - Qualifier goes as query parameter + - Only the payload goes in the request body + + Returns: + dict: Payload dict containing the prompt + """ + # Use the last message content as the prompt + prompt = convert_content_list_to_str(messages[-1]) + + # Create the payload - this is what goes in the body (raw JSON) + payload: dict = {"prompt": prompt} + + # Get or generate session ID - this goes in the header + runtime_session_id = self._get_runtime_session_id(optional_params) + headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] = runtime_session_id + + # The request data is the payload dict (will be JSON encoded by the HTTP handler) + # Qualifier will be handled as a query parameter in get_complete_url + + return payload + + def _extract_sse_json(self, line: str) -> Optional[Dict]: + """Extract and parse JSON from an SSE data line.""" + if not line.startswith('data:'): + return None + + json_str = line[5:].strip() + if not json_str: + return None + + try: + data = json.loads(json_str) + # Skip non-dict data (some lines contain JSON strings) + return data if isinstance(data, dict) else None + except json.JSONDecodeError: + verbose_logger.debug(f"Skipping non-JSON line: {line[:100]}") + return None + + def _extract_usage_from_event(self, event_data: Dict) -> Optional[AgentCoreUsage]: + """Extract usage information from event metadata.""" + event_payload = event_data.get("event") + if not event_payload: + return None + + metadata = event_payload.get("metadata") + if metadata and "usage" in metadata: + return metadata["usage"] # type: ignore + + return None + + def _extract_content_delta(self, event_data: Dict) -> Optional[str]: + """Extract text content from contentBlockDelta event.""" + event_payload = event_data.get("event") + if not event_payload: + return None + + content_block_delta = event_payload.get("contentBlockDelta") + if not content_block_delta: + return None + + delta = content_block_delta.get("delta", {}) + return delta.get("text") + + def _extract_content_from_message(self, message: AgentCoreMessage) -> str: + """ + Extract text content from message content blocks. + This works for both SSE messages and JSON responses. + """ + content_list = message.get("content", []) + if not isinstance(content_list, list): + return "" + + return "".join( + block["text"] + for block in content_list + if isinstance(block, dict) and "text" in block + ) + + def _calculate_usage( + self, model: str, messages: List[AllMessageValues], content: str + ) -> Optional[Usage]: + """ + Calculate token usage using LiteLLM's token counter. + + Args: + model: The model name + messages: Input messages + content: Response content + + Returns: + Usage object with calculated tokens, or None if calculation fails + """ + try: + from litellm.utils import token_counter + + prompt_tokens = token_counter(model=model, messages=messages) + completion_tokens = token_counter( + model=model, + text=content, + count_response_tokens=True + ) + total_tokens = prompt_tokens + completion_tokens + + verbose_logger.debug( + f"Calculated usage - prompt: {prompt_tokens}, " + f"completion: {completion_tokens}, total: {total_tokens}" + ) + + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + ) + except Exception as e: + verbose_logger.warning(f"Failed to calculate token usage: {str(e)}") + return None + + def _parse_json_response(self, response_json: dict) -> AgentCoreParsedResponse: + """ + Parse direct JSON response (non-streaming). + + JSON response structure: + { + "result": { + "role": "assistant", + "content": [{"text": "..."}] + } + } + """ + result = response_json.get("result", {}) + + # Extract content using the same helper as SSE parsing + content = self._extract_content_from_message(result) # type: ignore + + # JSON responses don't include usage data + return AgentCoreParsedResponse( + content=content, + usage=None, + final_message=result # type: ignore + ) + + def _get_parsed_response( + self, raw_response: httpx.Response + ) -> AgentCoreParsedResponse: + """ + Parse AgentCore response based on content type. + + Args: + raw_response: Raw HTTP response from AgentCore + + Returns: + AgentCoreParsedResponse: Parsed response data + """ + content_type = raw_response.headers.get("content-type", "").lower() + verbose_logger.debug(f"AgentCore response Content-Type: {content_type}") + + # Parse response based on content type + if "application/json" in content_type: + # Direct JSON response + verbose_logger.debug("Parsing JSON response") + response_json = raw_response.json() + verbose_logger.debug(f"Response JSON: {response_json}") + return self._parse_json_response(response_json) + else: + # SSE stream response (text/event-stream or default) + verbose_logger.debug("Parsing SSE stream response") + response_text = raw_response.text + verbose_logger.debug(f"AgentCore response (first 500 chars): {response_text[:500]}") + return self._parse_sse_stream(response_text) + + def _parse_sse_stream(self, response_text: str) -> AgentCoreParsedResponse: + """ + Parse Server-Sent Events (SSE) stream format. + Each line starts with 'data:' followed by JSON. + + Returns: + AgentCoreParsedResponse: Parsed response with content, usage, and message + """ + final_message: Optional[AgentCoreMessage] = None + usage_data: Optional[AgentCoreUsage] = None + content_blocks: List[str] = [] + + for line in response_text.strip().split('\n'): + line = line.strip() + if not line: + continue + + data = self._extract_sse_json(line) + if not data: + continue + + verbose_logger.debug(f"SSE event keys: {list(data.keys())}") + + # Check for final complete message + if "message" in data and isinstance(data["message"], dict): + final_message = data["message"] # type: ignore + verbose_logger.debug("Found final message") + + # Process event data + if "event" in data and isinstance(data["event"], dict): + event_payload = data["event"] + verbose_logger.debug(f"Event payload keys: {list(event_payload.keys())}") + + # Extract usage metadata + if usage := self._extract_usage_from_event(data): + usage_data = usage + verbose_logger.debug(f"Found usage data: {usage_data}") + + # Collect content deltas + if text := self._extract_content_delta(data): + content_blocks.append(text) + + # Build final content + content = ( + self._extract_content_from_message(final_message) + if final_message + else "".join(content_blocks) + ) + + verbose_logger.debug(f"Final usage_data: {usage_data}") + + return AgentCoreParsedResponse( + content=content, + usage=usage_data, + final_message=final_message + ) + + def get_streaming_response( + self, + model: str, + raw_response: httpx.Response, + ) -> AgentCoreSSEStreamIterator: + """ + Return a streaming iterator for SSE responses. + + Args: + model: The model name + raw_response: Raw HTTP response with streaming data + + Returns: + AgentCoreSSEStreamIterator: Iterator that yields ModelResponse chunks + """ + return AgentCoreSSEStreamIterator(response=raw_response, model=model) + + def get_sync_custom_stream_wrapper( + self, + model: str, + custom_llm_provider: str, + logging_obj: LiteLLMLoggingObj, + api_base: str, + headers: dict, + data: dict, + messages: list, + client: Optional[Union[HTTPHandler, "AsyncHTTPHandler"]] = None, + json_mode: Optional[bool] = None, + signed_json_body: Optional[bytes] = None, + ) -> CustomStreamWrapper: + """ + Get a CustomStreamWrapper for synchronous streaming. + + This is called when stream=True is passed to completion(). + """ + from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler, + _get_httpx_client, + ) + from litellm.utils import CustomStreamWrapper + + if client is None or not isinstance(client, HTTPHandler): + client = _get_httpx_client(params={}) + + # Make streaming request + response = client.post( + api_base, + headers=headers, + data=signed_json_body if signed_json_body else json.dumps(data), + stream=True, # THIS IS KEY - tells httpx to not buffer + logging_obj=logging_obj, + ) + + if response.status_code != 200: + raise BedrockError( + status_code=response.status_code, message=str(response.read()) + ) + + # Create iterator for SSE stream + completion_stream = self.get_streaming_response(model=model, raw_response=response) + + streaming_response = CustomStreamWrapper( + completion_stream=completion_stream, + model=model, + custom_llm_provider=custom_llm_provider, + logging_obj=logging_obj, + ) + + # LOGGING + logging_obj.post_call( + input=messages, + api_key="", + original_response="first stream response received", + additional_args={"complete_input_dict": data}, + ) + + return streaming_response + + @property + def has_custom_stream_wrapper(self) -> bool: + """Indicates that this config has custom streaming support.""" + return True + + @property + def supports_stream_param_in_request_body(self) -> bool: + """ + AgentCore does not allow passing `stream` in the request body. + Streaming is automatic based on the response format. + """ + return False + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: LiteLLMLoggingObj, + request_data: dict, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + encoding: Any, + api_key: Optional[str] = None, + json_mode: Optional[bool] = None, + ) -> ModelResponse: + """ + Transform the AgentCore response to LiteLLM ModelResponse format. + AgentCore can return either JSON or SSE (Server-Sent Events) stream responses. + + Note: For streaming responses, use get_streaming_response() instead. + """ + try: + # Parse the response based on content type (JSON or SSE) + parsed_data = self._get_parsed_response(raw_response) + + content = parsed_data["content"] + usage_data = parsed_data["usage"] + + verbose_logger.debug(f"Parsed content length: {len(content)}") + verbose_logger.debug(f"Usage data: {usage_data}") + + # Create the message + message = Message(content=content, role="assistant") + + # Create choices + choice = Choices(finish_reason="stop", index=0, message=message) + + # Update model response + model_response.choices = [choice] + model_response.model = model + + # Add usage information if available + # Note: AgentCore JSON responses don't include usage data + # SSE responses may include usage in metadata events + if usage_data: + usage = Usage( + prompt_tokens=usage_data.get("inputTokens", 0), + completion_tokens=usage_data.get("outputTokens", 0), + total_tokens=usage_data.get("totalTokens", 0), + ) + setattr(model_response, "usage", usage) + else: + # Calculate token usage using LiteLLM's token counter + verbose_logger.debug("No usage data from AgentCore - calculating tokens") + calculated_usage = self._calculate_usage(model, messages, content) + if calculated_usage: + setattr(model_response, "usage", calculated_usage) + + return model_response + + except Exception as e: + verbose_logger.error( + f"Error processing Bedrock AgentCore response: {str(e)}" + ) + raise BedrockError( + message=f"Error processing response: {str(e)}", + status_code=raw_response.status_code, + ) + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + return headers + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] + ) -> BaseLLMException: + return BedrockError(status_code=status_code, message=error_message) + + def should_fake_stream( + self, + model: Optional[str], + stream: Optional[bool], + custom_llm_provider: Optional[str] = None, + ) -> bool: + return True + diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 661817e4ce..d73860e1bc 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -445,17 +445,18 @@ class BedrockModelInfo(BaseLLMModelInfo): @staticmethod def get_bedrock_route( model: str, - ) -> Literal["converse", "invoke", "converse_like", "agent", "async_invoke"]: + ) -> Literal["converse", "invoke", "converse_like", "agent", "agentcore", "async_invoke"]: """ Get the bedrock route for the given model. """ route_mappings: Dict[ - str, Literal["invoke", "converse_like", "converse", "agent", "async_invoke"] + str, Literal["invoke", "converse_like", "converse", "agent", "agentcore", "async_invoke"] ] = { "invoke/": "invoke", "converse_like/": "converse_like", "converse/": "converse", "agent/": "agent", + "agentcore/": "agentcore", "async_invoke/": "async_invoke", } @@ -494,6 +495,13 @@ class BedrockModelInfo(BaseLLMModelInfo): """ return "agent/" in model + @staticmethod + def _explicit_agentcore_route(model: str) -> bool: + """ + Check if the model is an explicit agentcore route. + """ + return "agentcore/" in model + @staticmethod def _explicit_converse_like_route(model: str) -> bool: """ @@ -538,6 +546,65 @@ class BedrockModelInfo(BaseLLMModelInfo): return None +def get_bedrock_chat_config(model: str): + """ + Helper function to get the appropriate Bedrock chat config based on model and route. + + Args: + model: The model name/identifier + + Returns: + The appropriate Bedrock config class instance + """ + bedrock_route = BedrockModelInfo.get_bedrock_route(model) + bedrock_invoke_provider = litellm.BedrockLLM.get_bedrock_invoke_provider( + model=model + ) + base_model = BedrockModelInfo.get_base_model(model) + + # Handle explicit routes first + if bedrock_route == "converse" or bedrock_route == "converse_like": + return litellm.AmazonConverseConfig() + elif bedrock_route == "agent": + from litellm.llms.bedrock.chat.invoke_agent.transformation import ( + AmazonInvokeAgentConfig, + ) + return AmazonInvokeAgentConfig() + elif bedrock_route == "agentcore": + from litellm.llms.bedrock.chat.agentcore.transformation import ( + AmazonAgentCoreConfig, + ) + return AmazonAgentCoreConfig() + + # Handle provider-specific configs + if bedrock_invoke_provider == "amazon": + return litellm.AmazonTitanConfig() + elif bedrock_invoke_provider == "anthropic": + if ( + base_model + in litellm.AmazonAnthropicConfig.get_legacy_anthropic_model_names() + ): + return litellm.AmazonAnthropicConfig() + else: + return litellm.AmazonAnthropicClaudeConfig() + elif bedrock_invoke_provider == "meta" or bedrock_invoke_provider == "llama": + return litellm.AmazonLlamaConfig() + elif bedrock_invoke_provider == "ai21": + return litellm.AmazonAI21Config() + elif bedrock_invoke_provider == "cohere": + return litellm.AmazonCohereConfig() + elif bedrock_invoke_provider == "mistral": + return litellm.AmazonMistralConfig() + elif bedrock_invoke_provider == "deepseek_r1": + return litellm.AmazonDeepSeekR1Config() + elif bedrock_invoke_provider == "nova": + return litellm.AmazonInvokeNovaConfig() + elif bedrock_invoke_provider == "qwen3": + return litellm.AmazonQwen3Config() + else: + return litellm.AmazonInvokeConfig() + + class BedrockEventStreamDecoderBase: """ Base class for event stream decoding for Bedrock diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 1b029e824f..ce88f616e5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10069,6 +10069,96 @@ "supports_vision": true, "supports_web_search": true }, + "gemini-live-2.5-flash-preview-native-audio-09-2025": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "vertex_ai-language-models", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true + }, + "gemini/gemini-live-2.5-flash-preview-native-audio-09-2025": { + "cache_read_input_token_cost": 7.5e-08, + "input_cost_per_audio_token": 3e-06, + "input_cost_per_token": 3e-07, + "litellm_provider": "gemini", + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_images_per_prompt": 3000, + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_pdf_size_mb": 30, + "max_tokens": 65535, + "max_video_length": 1, + "max_videos_per_prompt": 10, + "mode": "chat", + "output_cost_per_audio_token": 1.2e-05, + "output_cost_per_token": 2e-06, + "source": "https://ai.google.dev/gemini-api/docs/pricing", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image", + "audio", + "video" + ], + "supported_output_modalities": [ + "text", + "audio" + ], + "supports_audio_input": true, + "supports_audio_output": true, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_url_context": true, + "supports_vision": true, + "supports_web_search": true + }, "gemini-2.5-flash-lite-preview-06-17": { "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, @@ -10203,96 +10293,6 @@ "supports_vision": true, "supports_web_search": true }, - "gemini-live-2.5-flash-preview-native-audio-09-2025": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "vertex_ai-language-models", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true - }, - "gemini/gemini-live-2.5-flash-preview-native-audio-09-2025": { - "cache_read_input_token_cost": 7.5e-08, - "input_cost_per_audio_token": 3e-06, - "input_cost_per_token": 3e-07, - "litellm_provider": "gemini", - "max_audio_length_hours": 8.4, - "max_audio_per_prompt": 1, - "max_images_per_prompt": 3000, - "max_input_tokens": 1048576, - "max_output_tokens": 65535, - "max_pdf_size_mb": 30, - "max_tokens": 65535, - "max_video_length": 1, - "max_videos_per_prompt": 10, - "mode": "chat", - "output_cost_per_audio_token": 1.2e-05, - "output_cost_per_token": 2.5e-06, - "source": "https://ai.google.dev/gemini-api/docs/pricing", - "supported_endpoints": [ - "/v1/chat/completions", - "/v1/completions" - ], - "supported_modalities": [ - "text", - "image", - "audio", - "video" - ], - "supported_output_modalities": [ - "text", - "audio" - ], - "supports_audio_input": true, - "supports_audio_output": true, - "supports_function_calling": true, - "supports_parallel_function_calling": true, - "supports_pdf_input": true, - "supports_prompt_caching": true, - "supports_response_schema": true, - "supports_system_messages": true, - "supports_tool_choice": true, - "supports_url_context": true, - "supports_vision": true, - "supports_web_search": true - }, "gemini-2.5-pro": { "cache_read_input_token_cost": 1.25e-07, "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 40229f4c5d..41df2e6c2a 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -9,6 +9,11 @@ model_list: model: bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0 aws_region_name: us-west-2 custom_llm_provider: bedrock + - model_name: bedrock/* + litellm_params: + model: bedrock/* + custom_llm_provider: bedrock + aws_region_name: us-west-2 # like MCPs/vector stores @@ -39,4 +44,5 @@ litellm_settings: s3_verify: False cache: True cache_params: - type: local \ No newline at end of file + type: local + drop_params: True \ No newline at end of file diff --git a/litellm/types/llms/bedrock_agentcore.py b/litellm/types/llms/bedrock_agentcore.py new file mode 100644 index 0000000000..49c3bfb2d5 --- /dev/null +++ b/litellm/types/llms/bedrock_agentcore.py @@ -0,0 +1,135 @@ +""" +Type definitions for AWS Bedrock AgentCore API. + +https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgentRuntime.html +""" + +from typing import Dict, List, Optional + +from typing_extensions import Literal, TypedDict + + +# Request Types +class AgentCoreRequestPayload(TypedDict): + """Payload for AgentCore request.""" + + prompt: str + + +class AgentCoreRequest(TypedDict, total=False): + """Complete request structure for AgentCore API (internal use).""" + + payload: str # JSON-encoded AgentCoreRequestPayload + + +# Response SSE Event Types +class AgentCoreMessageRole(TypedDict): + """Message role information.""" + + role: Literal["assistant"] + + +class AgentCoreMessageStart(TypedDict): + """Message start event.""" + + role: Literal["assistant"] + + +class AgentCoreContentBlockDelta(TypedDict): + """Content delta information.""" + + text: str + + +class AgentCoreContentBlockDeltaEvent(TypedDict): + """Content block delta event.""" + + delta: AgentCoreContentBlockDelta + contentBlockIndex: int + + +class AgentCoreContentBlockStop(TypedDict): + """Content block stop event.""" + + contentBlockIndex: int + + +class AgentCoreMessageStop(TypedDict): + """Message stop event.""" + + stopReason: Literal["end_turn", "max_tokens", "stop_sequence"] + + +class AgentCoreUsage(TypedDict): + """Token usage information.""" + + inputTokens: int + outputTokens: int + totalTokens: int + + +class AgentCoreMetrics(TypedDict): + """Response metrics.""" + + latencyMs: int + + +class AgentCoreMetadata(TypedDict): + """Metadata event payload.""" + + usage: AgentCoreUsage + metrics: AgentCoreMetrics + + +class AgentCoreEventPayload(TypedDict, total=False): + """Union payload for different event types.""" + + # messageStart event + messageStart: Optional[AgentCoreMessageStart] + + # contentBlockDelta event + contentBlockDelta: Optional[AgentCoreContentBlockDeltaEvent] + + # contentBlockStop event + contentBlockStop: Optional[AgentCoreContentBlockStop] + + # messageStop event + messageStop: Optional[AgentCoreMessageStop] + + # metadata event + metadata: Optional[AgentCoreMetadata] + + +class AgentCoreEvent(TypedDict, total=False): + """SSE event structure from AgentCore.""" + + event: Optional[AgentCoreEventPayload] + + +class AgentCoreContentBlock(TypedDict): + """Content block in final message.""" + + text: str + + +class AgentCoreMessage(TypedDict): + """Complete message structure.""" + + role: Literal["assistant"] + content: List[AgentCoreContentBlock] + + +class AgentCoreFinalMessage(TypedDict): + """Final message event containing complete response.""" + + message: AgentCoreMessage + + +# Response parsing result (internal use) +class AgentCoreParsedResponse(TypedDict): + """Parsed response from SSE stream.""" + + content: str + usage: Optional[AgentCoreUsage] + final_message: Optional[AgentCoreMessage] + diff --git a/litellm/utils.py b/litellm/utils.py index 8c03e9d505..6e8d6d1018 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7193,49 +7193,9 @@ class ProviderConfigManager: elif litellm.LlmProviders.MORPH == provider: return litellm.MorphChatConfig() elif litellm.LlmProviders.BEDROCK == provider: - bedrock_route = BedrockModelInfo.get_bedrock_route(model) - bedrock_invoke_provider = litellm.BedrockLLM.get_bedrock_invoke_provider( - model=model - ) + from litellm.llms.bedrock.common_utils import get_bedrock_chat_config - base_model = BedrockModelInfo.get_base_model(model) - - if bedrock_route == "converse" or bedrock_route == "converse_like": - return litellm.AmazonConverseConfig() - elif bedrock_route == "agent": - from litellm.llms.bedrock.chat.invoke_agent.transformation import ( - AmazonInvokeAgentConfig, - ) - - return AmazonInvokeAgentConfig() - elif bedrock_invoke_provider == "amazon": # amazon titan llms - return litellm.AmazonTitanConfig() - elif bedrock_invoke_provider == "anthropic": - if ( - base_model - in litellm.AmazonAnthropicConfig.get_legacy_anthropic_model_names() - ): - return litellm.AmazonAnthropicConfig() - else: - return litellm.AmazonAnthropicClaudeConfig() - elif ( - bedrock_invoke_provider == "meta" or bedrock_invoke_provider == "llama" - ): # amazon / meta llms - return litellm.AmazonLlamaConfig() - elif bedrock_invoke_provider == "ai21": # ai21 llms - return litellm.AmazonAI21Config() - elif bedrock_invoke_provider == "cohere": # cohere models on bedrock - return litellm.AmazonCohereConfig() - elif bedrock_invoke_provider == "mistral": # mistral models on bedrock - return litellm.AmazonMistralConfig() - elif bedrock_invoke_provider == "deepseek_r1": # deepseek models on bedrock - return litellm.AmazonDeepSeekR1Config() - elif bedrock_invoke_provider == "nova": - return litellm.AmazonInvokeNovaConfig() - elif bedrock_invoke_provider == "qwen3": - return litellm.AmazonQwen3Config() - else: - return litellm.AmazonInvokeConfig() + return get_bedrock_chat_config(model=model) elif litellm.LlmProviders.LITELLM_PROXY == provider: return litellm.LiteLLMProxyChatConfig() elif litellm.LlmProviders.OPENAI == provider: diff --git a/tests/llm_translation/test_bedrock_agentcore.py b/tests/llm_translation/test_bedrock_agentcore.py new file mode 100644 index 0000000000..6211fdbcdb --- /dev/null +++ b/tests/llm_translation/test_bedrock_agentcore.py @@ -0,0 +1,128 @@ +""" +Test Bedrock AgentCore integration +""" +import os +import sys +from dotenv import load_dotenv + +load_dotenv() + +sys.path.insert( + 0, os.path.abspath("../..") +) + +import litellm +from unittest.mock import MagicMock, patch +import pytest + +import pytest + +@pytest.mark.parametrize( + "model", [ + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_13sf6-cALnp38iZD", # non-streaming invocation + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", # streaming invocation + ] +) +def test_bedrock_agentcore_basic(model): + """ + Test AgentCore invocation parameterized by model + """ + litellm._turn_on_debug() + response = litellm.completion( + model=model, + messages=[{"role": "user", "content": "Explain machine learning in simple terms"}], + ) + print("response from agentcore=", response.model_dump_json(indent=4)) + # Assert that the message content has a response with some length + assert response.choices[0].message.content + assert len(response.choices[0].message.content) > 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "model", [ + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/non_stream_agent-mdfwS2DlAu", # non-streaming invocation + "bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", # streaming invocation + ] +) +async def test_bedrock_agentcore_with_streaming(model): + """ + Test AgentCore with streaming + """ + #litellm._turn_on_debug() + response = litellm.completion( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + messages=[ + { + "role": "user", + "content": "Explain machine learning in simple terms", + } + ], + stream=True, + ) + + for chunk in response: + print("chunk=", chunk) + + +def test_bedrock_agentcore_with_custom_params(): + """ + Test AgentCore request structure with custom parameters + """ + import json + + litellm._turn_on_debug() + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + + with patch.object(client, "post", return_value=MagicMock()) as mock_post: + try: + response = litellm.completion( + model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", + messages=[ + { + "role": "user", + "content": "Explain machine learning in simple terms", + } + ], + runtimeSessionId="litellm-test-session-id-12345678901234567890", + qualifier="DEFAULT", + client=client, + ) + except Exception as e: + print(f"Error: {e}") + + mock_post.assert_called_once() + call_kwargs = mock_post.call_args.kwargs + print(f"mock_post.call_args.kwargs: {call_kwargs}") + + # Verify URL structure - should include ARN and qualifier + assert "url" in call_kwargs + url = call_kwargs["url"] + print(f"URL: {url}") + assert "/runtimes/arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A888602223428%3Aruntime%2Fhosted_agent_r9jvp-3ySZuRHjLC/invocations" in url + assert "qualifier=DEFAULT" in url + + # Verify headers - session ID should be in header + assert "headers" in call_kwargs + headers = call_kwargs["headers"] + print(f"Headers: {headers}") + assert "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id" in headers + assert headers["X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"] == "litellm-test-session-id-12345678901234567890" + + # Verify the request body - should just be the payload + assert "data" in call_kwargs or "json" in call_kwargs + + # Parse the request data + if "data" in call_kwargs: + request_data = json.loads(call_kwargs["data"]) + else: + request_data = call_kwargs["json"] + + print(f"Request data: {json.dumps(request_data, indent=2)}") + + # Body should just contain the prompt + assert "prompt" in request_data + assert request_data["prompt"] == "Explain machine learning in simple terms" +