mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-17 00:26:01 +00:00
[Feat] Add Bedrock Agentcore as a provider on LiteLLM Python SDK and LiteLLM AI Gateway (#16252)
* add agentcore in get_bedrock_route * add AmazonAgentCoreConfig * fix get_runtime_endpoint * init AmazonAgentCoreConfig * add get_bedrock_chat_config * get_bedrock_chat_config * add AmazonAgentCoreConfig * fix get_complete_url * refactor transform response * test agentcore * test_bedrock_agentcore_with_streaming * fix _parse_json_response * fix _calculate_usage * test_bedrock_agentcore_basic * add AgentCoreSSEStreamIterator * add native streaming for agentcore * test_bedrock_agentcore_with_streaming * test_bedrock_agentcore_basic * add agentcore * _calculate_usage * fix linting
This commit is contained in:
Ishaan Jaff
@@ -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
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="config-yaml" label="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
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
#### 2. Start the LiteLLM Proxy
|
||||
|
||||
```bash showLineNumbers title="Start LiteLLM Proxy"
|
||||
litellm --config config.yaml
|
||||
```
|
||||
|
||||
#### 3. Make requests to your AgentCore runtimes
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="curl" label="Curl">
|
||||
|
||||
```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
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
|
||||
<TabItem value="openai-sdk" label="OpenAI Python SDK">
|
||||
|
||||
```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="")
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## Provider-specific Parameters
|
||||
|
||||
AgentCore supports additional parameters that can be passed to customize the runtime invocation.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="sdk" label="SDK">
|
||||
|
||||
```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
|
||||
)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="proxy" label="Proxy">
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
### 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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -0,0 +1,4 @@
|
||||
from .transformation import AmazonAgentCoreConfig
|
||||
|
||||
__all__ = ["AmazonAgentCoreConfig"]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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=<value>
|
||||
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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
type: local
|
||||
drop_params: True
|
||||
@@ -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]
|
||||
|
||||
+2
-42
@@ -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:
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user