mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-13 03:59:12 +00:00
feat(strands_decider): call a Strands Decider on an AgentCore runtime ARN (#45678)
Backport of #45678 to rc/1.106.0.
Cherry-picked from ed9c31cfa8 (main).
This commit is contained in:
devin-ai-integration[bot]
mateo-berri
parent
3880fc285b
commit
a978acffb3
13 changed files with 1555 additions and 3 deletions
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -343,6 +343,7 @@ jobs:
|
|||
tests/unit/llms/snowflake
|
||||
tests/unit/llms/soniox
|
||||
tests/unit/llms/stability
|
||||
tests/unit/llms/strands_decider
|
||||
tests/unit/llms/tavily
|
||||
tests/unit/llms/tencent
|
||||
tests/unit/llms/tinyfish
|
||||
|
|
|
|||
|
|
@ -60,6 +60,7 @@ class _DecisionsCall:
|
|||
logging_obj: LiteLLMLoggingObj | None
|
||||
headers: Mapping[str, str]
|
||||
timeout: float | httpx.Timeout | None
|
||||
litellm_params: Mapping[str, object] = field(repr=False)
|
||||
|
||||
|
||||
def supported_decisions_providers() -> tuple[str, ...]:
|
||||
|
|
@ -256,6 +257,7 @@ def _prepare_call(
|
|||
logging_obj=logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
|
||||
headers=extra_headers or {},
|
||||
timeout=timeout,
|
||||
litellm_params=kwargs,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -339,6 +341,7 @@ async def adecisions(
|
|||
api_key=call.api_key,
|
||||
headers=call.headers,
|
||||
timeout=call.timeout,
|
||||
litellm_params=call.litellm_params,
|
||||
)
|
||||
except Exception as error:
|
||||
raise _map_upstream_exception(error, call) from error
|
||||
|
|
@ -384,6 +387,7 @@ def decisions(
|
|||
api_key=call.api_key,
|
||||
headers=call.headers,
|
||||
timeout=call.timeout,
|
||||
litellm_params=call.litellm_params,
|
||||
)
|
||||
except Exception as error:
|
||||
raise _map_upstream_exception(error, call) from error
|
||||
|
|
|
|||
|
|
@ -341,6 +341,20 @@ class BaseDecisionsConfig(ABC):
|
|||
def get_complete_url(self, api_base: str, model: str) -> str:
|
||||
return f"{api_base.rstrip('/').removesuffix('/v1')}{self.path}"
|
||||
|
||||
def signs_with_aws(self, api_base: str) -> bool:
|
||||
return False
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
url: str,
|
||||
api_base: str,
|
||||
body: Mapping[str, object],
|
||||
api_key: str | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[Mapping[str, str], bytes | None]:
|
||||
return headers, None
|
||||
|
||||
def transform_decisions_request(
|
||||
self,
|
||||
model: str,
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import json
|
|||
import ssl
|
||||
from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
from functools import lru_cache
|
||||
from functools import lru_cache, partial
|
||||
from types import MappingProxyType
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
|
|
@ -1517,6 +1517,7 @@ class BaseLLMHTTPHandler:
|
|||
api_key: str | None,
|
||||
headers: Mapping[str, str],
|
||||
timeout: float | httpx.Timeout | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
client: HTTPHandler | None = None,
|
||||
) -> DecisionsIRResponse:
|
||||
url, outbound_headers, data = self._prepare_decisions_request(
|
||||
|
|
@ -1528,10 +1529,23 @@ class BaseLLMHTTPHandler:
|
|||
api_key=api_key,
|
||||
headers=headers,
|
||||
)
|
||||
signed_headers, signed_body = provider_config.sign_request(
|
||||
headers=outbound_headers,
|
||||
url=url,
|
||||
api_base=api_base,
|
||||
body=data,
|
||||
api_key=api_key,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
sync_httpx_client: Final = client if client is not None else get_httpx_client()
|
||||
try:
|
||||
response: Final = sync_httpx_client.post(
|
||||
url, json=data, headers=outbound_headers, timeout=timeout, logging_obj=logging_obj
|
||||
url,
|
||||
content=signed_body,
|
||||
json=data if signed_body is None else None,
|
||||
headers=dict(signed_headers),
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except httpx.HTTPError as e:
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
|
|
@ -1551,6 +1565,7 @@ class BaseLLMHTTPHandler:
|
|||
api_key: str | None,
|
||||
headers: Mapping[str, str],
|
||||
timeout: float | httpx.Timeout | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
) -> DecisionsIRResponse:
|
||||
url, outbound_headers, data = self._prepare_decisions_request(
|
||||
|
|
@ -1562,6 +1577,18 @@ class BaseLLMHTTPHandler:
|
|||
api_key=api_key,
|
||||
headers=headers,
|
||||
)
|
||||
sign: Final = partial(
|
||||
provider_config.sign_request,
|
||||
headers=outbound_headers,
|
||||
url=url,
|
||||
api_base=api_base,
|
||||
body=data,
|
||||
api_key=api_key,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
signed_headers, signed_body = (
|
||||
await run_aws_signing(sign) if provider_config.signs_with_aws(api_base) else sign()
|
||||
)
|
||||
async_httpx_client: Final = (
|
||||
client
|
||||
if client is not None
|
||||
|
|
@ -1569,7 +1596,12 @@ class BaseLLMHTTPHandler:
|
|||
)
|
||||
try:
|
||||
response: Final = await async_httpx_client.post(
|
||||
url, json=data, headers=outbound_headers, timeout=timeout, logging_obj=logging_obj
|
||||
url,
|
||||
content=signed_body,
|
||||
json=data if signed_body is None else None,
|
||||
headers=dict(signed_headers),
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except httpx.HTTPError as e:
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,161 @@
|
|||
import hashlib
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import quote, unquote
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, sign_aws_json_post
|
||||
from litellm.types.decisions import DecisionsIRRequest, DecisionsIRResponse
|
||||
from litellm.types.llms.bedrock import AwsAuthParams
|
||||
|
||||
AGENTCORE_SESSION_HEADER: Final = "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id"
|
||||
_RESERVED_HEADER_PREFIXES: Final = ("x-amzn-bedrock-agentcore-runtime-", "x-amz-")
|
||||
_RUNTIME_ARN: Final = re.compile(
|
||||
r"\Aarn:aws(?:-[a-z]+)*:bedrock-agentcore:(?P<region>[a-z0-9-]+):[0-9]{12}:runtime/[A-Za-z0-9_-]+\Z"
|
||||
)
|
||||
_INVOCATIONS_PATH: Final = re.compile(r"\A/runtimes/(?P<arn>[^/]+)/invocations\Z")
|
||||
_RUNTIME_ERROR_STATUS: Final = MappingProxyType({"bad_request": 400, "invalid_request": 400, "loading": 503})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class AgentCoreRuntime:
|
||||
arn: str
|
||||
region: str
|
||||
|
||||
@property
|
||||
def invocations_url(self) -> str:
|
||||
host: Final = f"bedrock-agentcore.{self.region}.{get_aws_dns_suffix(self.region)}"
|
||||
return f"https://{host}/runtimes/{quote(self.arn, safe='')}/invocations"
|
||||
|
||||
@property
|
||||
def default_session_id(self) -> str:
|
||||
return f"litellm-decider-{hashlib.sha256(self.arn.encode()).hexdigest()}"
|
||||
|
||||
|
||||
def agentcore_runtime(api_base: str) -> AgentCoreRuntime | None:
|
||||
match: Final = _RUNTIME_ARN.match(api_base)
|
||||
return None if match is None else AgentCoreRuntime(arn=api_base, region=match["region"])
|
||||
|
||||
|
||||
def invoked_runtime(url: httpx.URL) -> AgentCoreRuntime | None:
|
||||
match: Final = _INVOCATIONS_PATH.match(url.raw_path.decode("ascii"))
|
||||
runtime: Final = None if match is None else agentcore_runtime(unquote(match["arn"]))
|
||||
return runtime if runtime is not None and runtime.invocations_url == str(url) else None
|
||||
|
||||
|
||||
class _RuntimeError(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
code: str
|
||||
message: str = ""
|
||||
|
||||
|
||||
class _RuntimeReply(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
error: _RuntimeError | None = None
|
||||
answers: object = None
|
||||
|
||||
|
||||
def _runtime_error(content: bytes) -> _RuntimeError | None:
|
||||
try:
|
||||
reply: Final = _RuntimeReply.model_validate_json(content)
|
||||
except ValidationError:
|
||||
return None
|
||||
return reply.error if reply.answers is None else None
|
||||
|
||||
|
||||
class _SessionParams(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
agentcore_runtime_session_id: str | None = None
|
||||
|
||||
|
||||
def _is_reserved_header(name: str) -> bool:
|
||||
lowered: Final = name.lower()
|
||||
return lowered == "host" or lowered.startswith(_RESERVED_HEADER_PREFIXES)
|
||||
|
||||
|
||||
def _runtime_headers(
|
||||
headers: Mapping[str, str], runtime: AgentCoreRuntime, litellm_params: Mapping[str, object]
|
||||
) -> Mapping[str, str]:
|
||||
dropped: Final = sorted(name for name in headers if _is_reserved_header(name))
|
||||
if dropped:
|
||||
verbose_logger.warning("Strands Decider: not forwarding reserved AgentCore header(s) %s", dropped)
|
||||
configured: Final = _SessionParams.model_validate(litellm_params).agentcore_runtime_session_id
|
||||
return {
|
||||
**{name: value for name, value in headers.items() if not _is_reserved_header(name)},
|
||||
AGENTCORE_SESSION_HEADER: runtime.default_session_id if configured is None else configured,
|
||||
}
|
||||
|
||||
|
||||
class StrandsDeciderDecisionsConfig(BaseDecisionsConfig):
|
||||
api_key_env = ("STRANDS_DECIDER_API_KEY",)
|
||||
api_base_env = ("STRANDS_DECIDER_API_BASE",)
|
||||
api_key_required = False
|
||||
|
||||
def __init__(self, aws: BaseAWSLLM | None = None) -> None:
|
||||
self._aws: Final = aws if aws is not None else BaseAWSLLM()
|
||||
|
||||
def get_complete_url(self, api_base: str, model: str) -> str:
|
||||
runtime: Final = agentcore_runtime(api_base)
|
||||
return super().get_complete_url(api_base, model) if runtime is None else runtime.invocations_url
|
||||
|
||||
def signs_with_aws(self, api_base: str) -> bool:
|
||||
return agentcore_runtime(api_base) is not None
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: Mapping[str, str],
|
||||
url: str,
|
||||
api_base: str,
|
||||
body: Mapping[str, object],
|
||||
api_key: str | None,
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[Mapping[str, str], bytes | None]:
|
||||
runtime: Final = agentcore_runtime(api_base)
|
||||
if runtime is None:
|
||||
return headers, None
|
||||
session_headers: Final = _runtime_headers(headers, runtime, litellm_params)
|
||||
if api_key is not None:
|
||||
return session_headers, None
|
||||
payload: Final = json.dumps(body)
|
||||
signed: Final = sign_aws_json_post(
|
||||
get_credentials=partial(
|
||||
self._aws.resolve_credentials, AwsAuthParams.model_validate(litellm_params), runtime.region
|
||||
),
|
||||
service_name="bedrock-agentcore",
|
||||
aws_region_name=runtime.region,
|
||||
url=url,
|
||||
body=payload,
|
||||
headers=session_headers,
|
||||
)
|
||||
return dict(signed.headers.items()), payload.encode()
|
||||
|
||||
def transform_decisions_response(
|
||||
self,
|
||||
model: str,
|
||||
custom_llm_provider: str,
|
||||
raw_response: httpx.Response,
|
||||
request: DecisionsIRRequest,
|
||||
) -> DecisionsIRResponse:
|
||||
runtime_error: Final = (
|
||||
None if invoked_runtime(raw_response.request.url) is None else _runtime_error(raw_response.content)
|
||||
)
|
||||
if runtime_error is not None:
|
||||
raise BaseLLMException(
|
||||
status_code=_RUNTIME_ERROR_STATUS.get(runtime_error.code, 500),
|
||||
message=f"Strands Decider runtime error '{runtime_error.code}': {runtime_error.message}",
|
||||
)
|
||||
return super().transform_decisions_response(model, custom_llm_provider, raw_response, request)
|
||||
|
|
|
|||
|
|
@ -411,6 +411,8 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
|
|||
"anthropic_workspace_id",
|
||||
"anthropic-workspace-id",
|
||||
"bedrock_tags",
|
||||
# Deployment-pinned AgentCore session: every new caller-chosen value provisions another microVM.
|
||||
"agentcore_runtime_session_id",
|
||||
# Provider-specific endpoint overrides that flow into the outbound
|
||||
# request via ``optional_params``. Same threat as ``api_base``:
|
||||
# ``s3_endpoint_url`` redirects Bedrock file uploads to attacker
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import socket
|
||||
import threading
|
||||
from collections.abc import Generator, Mapping
|
||||
from contextlib import contextmanager
|
||||
|
|
@ -69,3 +71,47 @@ def refusing_forward_proxy(port: int = 0) -> Generator[ForwardProxy, None, None]
|
|||
thread.join(timeout=6)
|
||||
assert not thread.is_alive(), "Owned forward proxy survived cleanup"
|
||||
server.server_close()
|
||||
|
||||
|
||||
def _pipe(source: socket.socket, sink: socket.socket) -> None:
|
||||
with contextlib.suppress(OSError):
|
||||
for chunk in iter(lambda: source.recv(65536), b""):
|
||||
sink.sendall(chunk)
|
||||
with contextlib.suppress(OSError):
|
||||
sink.shutdown(socket.SHUT_WR)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def tunnelling_forward_proxy(destination_port: int, port: int = 0) -> Generator[ForwardProxy, None, None]:
|
||||
received: Final[SimpleQueue[Tunnel]] = SimpleQueue()
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
protocol_version = "HTTP/1.1"
|
||||
rbufsize = 0
|
||||
timeout = 30
|
||||
request: socket.socket
|
||||
|
||||
def do_CONNECT(self) -> None:
|
||||
received.put(Tunnel(self.command, self.path, {name.lower(): value for name, value in self.headers.items()}))
|
||||
self.close_connection = True
|
||||
with socket.create_connection(("127.0.0.1", destination_port), timeout=30) as destination:
|
||||
self.send_response(200, "Connection established")
|
||||
self.end_headers()
|
||||
outbound: Final = threading.Thread(target=_pipe, args=(self.request, destination))
|
||||
outbound.start()
|
||||
_pipe(destination, self.request)
|
||||
outbound.join(timeout=35)
|
||||
|
||||
def log_message(self, format: str, *args: object) -> None:
|
||||
pass
|
||||
|
||||
with ThreadingHTTPServer(("127.0.0.1", port), Handler) as server:
|
||||
thread: Final = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.05})
|
||||
thread.start()
|
||||
try:
|
||||
yield ForwardProxy(f"http://127.0.0.1:{server.server_port}", server.server_port, received)
|
||||
finally:
|
||||
server.shutdown()
|
||||
thread.join(timeout=6)
|
||||
assert not thread.is_alive(), "Owned forward proxy survived cleanup"
|
||||
server.server_close()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,827 @@
|
|||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, replace
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
from urllib.parse import quote, unquote
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.forward_proxy import ForwardProxy, tunnelling_forward_proxy
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.sigv4 import encoded_path, signature
|
||||
from integration._support.tls import server_context, write_self_signed_cert
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.llms.strands_decider.decisions.transformation import StrandsDeciderDecisionsConfig
|
||||
|
||||
_ACCOUNT: Final = "123456789012"
|
||||
_ACCESS_KEY: Final = "AKIAAGENTCOREAUDIT01"
|
||||
_SECRET_KEY: Final = "synthetic-agentcore-audit-secret"
|
||||
_SIGNING: Final[Mapping[str, JsonValue]] = {"aws_access_key_id": _ACCESS_KEY, "aws_secret_access_key": _SECRET_KEY}
|
||||
_JWT: Final = "synthetic-agentcore-jwt"
|
||||
_ENV_JWT: Final = "synthetic-agentcore-env-jwt"
|
||||
_US_EAST: Final = "us-east-1"
|
||||
_EU_WEST: Final = "eu-west-1"
|
||||
_MODEL: Final = "strands_decider/systemone-decider"
|
||||
_BODY_MODEL: Final = "systemone-decider"
|
||||
_SESSION_HEADER: Final = "x-amzn-bedrock-agentcore-runtime-session-id"
|
||||
_CONFIGURED_SESSION: Final = "agentcore-audit-configured-session-0000000001"
|
||||
_STATE: Final[dict[str, JsonValue]] = {"ticket": "Refund requested twice for one order", "queue": "payments"}
|
||||
_QUESTIONS: Final[dict[str, JsonValue]] = {
|
||||
"duplicate": {"type": "noul", "instructions": "Is this a duplicate refund request?"},
|
||||
"urgency": {"type": "choice", "criteria": {"low": "can wait", "high": "money at risk"}},
|
||||
}
|
||||
_SDK_QUESTIONS: Final = TypeAdapter(dict[str, dict[str, object]]).validate_python(_QUESTIONS)
|
||||
_ANSWERS: Final[dict[str, JsonValue]] = {
|
||||
"duplicate": {"type": "noul", "noul": 0.88},
|
||||
"urgency": {"type": "choice", "choice": "high", "confidence": 0.75, "probabilities": {"low": 0.25, "high": 0.75}},
|
||||
}
|
||||
_USAGE: Final[dict[str, JsonValue]] = {"input_tokens": 211, "output_tokens": 2}
|
||||
_UPSTREAM_BODY: Final[dict[str, JsonValue]] = {"model": _BODY_MODEL, "state": _STATE, "questions": _QUESTIONS}
|
||||
_INPUT: Final = "Refund requested twice for one order"
|
||||
_OPENAI_QUESTIONS: Final[list[JsonValue]] = [
|
||||
{"type": "predicate", "name": "duplicate", "instructions": "Is this a duplicate refund request?"},
|
||||
{
|
||||
"type": "choice",
|
||||
"name": "urgency",
|
||||
"instructions": "How urgent is it?",
|
||||
"choices": [{"value": "low", "description": "can wait"}, {"value": "high", "description": "money at risk"}],
|
||||
},
|
||||
]
|
||||
_OPENAI_UPSTREAM_BODY: Final[dict[str, JsonValue]] = {
|
||||
"model": _BODY_MODEL,
|
||||
"state": _INPUT,
|
||||
"questions": {
|
||||
"duplicate": {"type": "noul", "instructions": "Is this a duplicate refund request?"},
|
||||
"urgency": {
|
||||
"type": "choice",
|
||||
"instructions": "How urgent is it?",
|
||||
"criteria": {"low": "can wait", "high": "money at risk"},
|
||||
},
|
||||
},
|
||||
}
|
||||
_OPENAI_ANSWER: Final[dict[str, JsonValue]] = {
|
||||
"model": _BODY_MODEL,
|
||||
"answers": [
|
||||
{"type": "predicate", "name": "duplicate", "probability": 0.88},
|
||||
{
|
||||
"type": "choice",
|
||||
"name": "urgency",
|
||||
"choice": "high",
|
||||
"probabilities": [{"value": "low", "probability": 0.25}, {"value": "high", "probability": 0.75}],
|
||||
"confidence": 0.75,
|
||||
},
|
||||
],
|
||||
"usage": {
|
||||
"input_tokens": 211,
|
||||
"input_tokens_details": {"cached_tokens": 0, "cache_write_tokens": 0},
|
||||
"output_tokens": 2,
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
"total_tokens": 213,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Route:
|
||||
path: str
|
||||
body: Mapping[str, JsonValue]
|
||||
upstream: Mapping[str, JsonValue]
|
||||
answer: Mapping[str, JsonValue]
|
||||
|
||||
|
||||
_SYSTEMONE: Final = _Route(
|
||||
"/v1/systemone",
|
||||
{"state": _STATE, "questions": _QUESTIONS},
|
||||
_UPSTREAM_BODY,
|
||||
{"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE},
|
||||
)
|
||||
_OPENAI: Final = _Route(
|
||||
"/v1/decisions", {"input": _INPUT, "questions": _OPENAI_QUESTIONS}, _OPENAI_UPSTREAM_BODY, _OPENAI_ANSWER
|
||||
)
|
||||
_ROUTES: Final = (_SYSTEMONE, _OPENAI, replace(_SYSTEMONE, path="/systemone"), replace(_OPENAI, path="/decisions"))
|
||||
_SPEND_QUERY: Final = (
|
||||
"SELECT status, call_type, custom_llm_provider, model_group, api_base, prompt_tokens, completion_tokens "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE request_id = %s'
|
||||
)
|
||||
_RUNTIME_TARGET: Final = re.compile(
|
||||
r"\A/runtimes/arn:aws:bedrock-agentcore:[a-z0-9-]+:[0-9]{12}:runtime/(?P<runtime>[A-Za-z0-9_-]+)/invocations\Z"
|
||||
)
|
||||
_ERROR_STATUS: Final = {"bad_request": 400, "invalid_request": 400, "loading": 503, "inference_error": 500}
|
||||
_HTTP_ERRORS: Final = {
|
||||
403: "AccessDeniedException",
|
||||
404: "ResourceNotFoundException",
|
||||
429: "ThrottlingException",
|
||||
}
|
||||
_MALFORMED_SESSIONS: Final[Mapping[str, JsonValue]] = {"int": 1234567, "list": ["a" * 40], "empty": ""}
|
||||
_LONG_SESSION: Final = "s" * 5120
|
||||
_NEAR_MISS_ARNS: Final = {
|
||||
"endpoint": f"arn:aws:bedrock-agentcore:us-east-1:{_ACCOUNT}:runtime/audit_qualified/runtime-endpoint/DEFAULT",
|
||||
"account": "arn:aws:bedrock-agentcore:us-east-1:12345:runtime/audit_short_account",
|
||||
"service": f"arn:aws:bedrock:us-east-1:{_ACCOUNT}:runtime/audit_other_service",
|
||||
"host": f"arn:aws:bedrock-agentcore:us-east-1.example.com:{_ACCOUNT}:runtime/audit_crafted_host",
|
||||
}
|
||||
_DEPLOYMENT_HEADERS: Final[dict[str, JsonValue]] = {
|
||||
"x-amz-target": "AmazonBedrockAgentCore.Forged",
|
||||
"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "deployment-forged-session-00000000000001",
|
||||
"Host": "bedrock-agentcore.us-west-2.amazonaws.com",
|
||||
"x-audit-trace": "deployment-trace",
|
||||
}
|
||||
_SIGNED: Final = "agentcore-audit-signed"
|
||||
_REGION: Final = "agentcore-audit-region"
|
||||
_SESSION: Final = "agentcore-audit-session"
|
||||
_LONG: Final = "agentcore-audit-session-long"
|
||||
_BEARER: Final = "agentcore-audit-bearer"
|
||||
_HEADERS: Final = "agentcore-audit-headers"
|
||||
_OPT_IN: Final = "agentcore-audit-opt-in"
|
||||
_HEALTH: Final = "agentcore-audit-health"
|
||||
_DROPPED: Final = "agentcore-audit-dropped"
|
||||
_OUTAGE: Final = "agentcore-audit-outage"
|
||||
_ENV: Final = "agentcore-audit-env"
|
||||
_UNSIGNED: Final = "agentcore-audit-unsigned"
|
||||
_SCRUBBED_PREFIXES: Final = ("AWS_", "STRANDS_DECIDER_")
|
||||
_SCRUBBED_NAMES: Final = frozenset(
|
||||
{"HTTPS_PROXY", "HTTP_PROXY", "ALL_PROXY", "NO_PROXY", "SSL_VERIFY", "SSL_CERT_FILE", "REQUESTS_CA_BUNDLE"}
|
||||
)
|
||||
_load_yaml: Final[Callable[[str], object]] = yaml.safe_load
|
||||
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_OWNED_PROXY_CELL_SECONDS: Final = 2 * max(30.0, float(os.environ.get("INTEGRATION_PROXY_READY_SECONDS") or 70)) + 120
|
||||
pytestmark: Final = pytest.mark.timeout(_OWNED_PROXY_CELL_SECONDS)
|
||||
|
||||
|
||||
def _host(region: str) -> str:
|
||||
return f"bedrock-agentcore.{region}.amazonaws.com"
|
||||
|
||||
|
||||
def _arn(runtime: str, region: str = _US_EAST) -> str:
|
||||
return f"arn:aws:bedrock-agentcore:{region}:{_ACCOUNT}:runtime/{runtime}"
|
||||
|
||||
|
||||
def _invocations_url(runtime: str, region: str = _US_EAST) -> str:
|
||||
return f"https://{_host(region)}/runtimes/{quote(_arn(runtime, region), safe='')}/invocations"
|
||||
|
||||
|
||||
def _default_session(runtime: str, region: str = _US_EAST) -> str:
|
||||
return f"litellm-decider-{hashlib.sha256(_arn(runtime, region).encode()).hexdigest()}"
|
||||
|
||||
|
||||
def _runtime_of(request: Request) -> str | None:
|
||||
match: Final = _RUNTIME_TARGET.match(unquote(request.target))
|
||||
return None if match is None else match["runtime"]
|
||||
|
||||
|
||||
def _json_reply(body: Mapping[str, JsonValue], status: int = 200, **headers: str) -> Reply:
|
||||
return Reply(status=status, body=json.dumps(body).encode(), headers=headers)
|
||||
|
||||
|
||||
def _respond(outage: threading.Event, request: Request) -> Reply:
|
||||
runtime: Final = _runtime_of(request) or ""
|
||||
if runtime == "audit_dropped" or (runtime == "audit_outage" and outage.is_set()):
|
||||
return Reply(drop_connection=True)
|
||||
if runtime.startswith("audit_error_"):
|
||||
code: Final = runtime.removeprefix("audit_error_")
|
||||
return _json_reply({"error": {"code": code, "message": f"runtime refused with {code}"}})
|
||||
if runtime.startswith("audit_status_"):
|
||||
status: Final = int(runtime.removeprefix("audit_status_"))
|
||||
return _json_reply(
|
||||
{"message": f"scripted {_HTTP_ERRORS[status]}"}, status, **{"x-amzn-errortype": _HTTP_ERRORS[status]}
|
||||
)
|
||||
if not runtime:
|
||||
return _json_reply({"message": f"unknown target {request.target}"}, 404)
|
||||
return _json_reply({"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Runtime:
|
||||
wire: Wire
|
||||
tunnel: ForwardProxy
|
||||
cert: Path
|
||||
outage: threading.Event
|
||||
|
||||
def reset(self) -> None:
|
||||
self.outage.clear()
|
||||
self.wire.drain()
|
||||
self.tunnel.drain()
|
||||
|
||||
def calls(self) -> tuple[Request, ...]:
|
||||
return self.wire.drain()
|
||||
|
||||
|
||||
def _calls_to(calls: Sequence[Request], runtime: str) -> tuple[Request, ...]:
|
||||
return tuple(call for call in calls if _runtime_of(call) == runtime)
|
||||
|
||||
|
||||
def _deployment(name: str, **litellm_params: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {"model": _MODEL, **litellm_params},
|
||||
"model_info": {"id": name, "mode": "evaluation"},
|
||||
}
|
||||
|
||||
|
||||
def _signed_deployment(name: str, runtime: str, region: str = _US_EAST, **extra: JsonValue) -> dict[str, JsonValue]:
|
||||
return _deployment(name, api_base=_arn(runtime, region), **_SIGNING, **extra)
|
||||
|
||||
|
||||
def _deployments() -> list[dict[str, JsonValue]]:
|
||||
return [
|
||||
_signed_deployment(_SIGNED, "audit_signed", aws_region_name=_US_EAST),
|
||||
_signed_deployment(_REGION, "audit_region", _EU_WEST, aws_region_name=_US_EAST),
|
||||
_signed_deployment(_SESSION, "audit_session", agentcore_runtime_session_id=_CONFIGURED_SESSION),
|
||||
_signed_deployment(_LONG, "audit_session_long", agentcore_runtime_session_id=_LONG_SESSION),
|
||||
_deployment(_BEARER, api_base=_arn("audit_bearer"), api_key=_JWT),
|
||||
_signed_deployment(_HEADERS, "audit_headers", extra_headers=_DEPLOYMENT_HEADERS),
|
||||
_signed_deployment(
|
||||
_OPT_IN, "audit_opt_in", configurable_clientside_auth_params=["agentcore_runtime_session_id"]
|
||||
),
|
||||
_signed_deployment(_HEALTH, "audit_health"),
|
||||
_signed_deployment(_DROPPED, "audit_dropped"),
|
||||
_signed_deployment(_OUTAGE, "audit_outage"),
|
||||
_deployment(_ENV, **_SIGNING),
|
||||
_deployment(_UNSIGNED, api_base=_arn("audit_unsigned")),
|
||||
*(_signed_deployment(f"agentcore-audit-error-{code}", f"audit_error_{code}") for code in _ERROR_STATUS),
|
||||
*(_signed_deployment(f"agentcore-audit-status-{status}", f"audit_status_{status}") for status in _HTTP_ERRORS),
|
||||
*(
|
||||
_signed_deployment(
|
||||
f"agentcore-audit-session-{label}", f"audit_session_{label}", agentcore_runtime_session_id=value
|
||||
)
|
||||
for label, value in _MALFORMED_SESSIONS.items()
|
||||
),
|
||||
*(
|
||||
_deployment(f"agentcore-audit-near-miss-{label}", api_base=value, **_SIGNING)
|
||||
for label, value in _NEAR_MISS_ARNS.items()
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _write_config(directory: Path) -> Path:
|
||||
base: Final = JSON_OBJECT.validate_python(_load_yaml(Path("tests/integration/proxy_config.yaml").read_text()))
|
||||
configuration: Final = {
|
||||
**base,
|
||||
"model_list": _deployments(),
|
||||
"router_settings": {"disable_cooldowns": True, "num_retries": 0},
|
||||
}
|
||||
path: Final = directory / "agentcore.yaml"
|
||||
path.write_text(yaml.safe_dump(configuration))
|
||||
return path
|
||||
|
||||
|
||||
def _scrubbed() -> tuple[str, ...]:
|
||||
return tuple(
|
||||
name for name in os.environ if name.upper().startswith(_SCRUBBED_PREFIXES) or name.upper() in _SCRUBBED_NAMES
|
||||
)
|
||||
|
||||
|
||||
def _isolated_aws(directory: Path) -> dict[str, str]:
|
||||
return {
|
||||
"AWS_EC2_METADATA_DISABLED": "true",
|
||||
"AWS_SHARED_CREDENTIALS_FILE": str(directory / "absent-credentials"),
|
||||
"AWS_CONFIG_FILE": str(directory / "absent-config"),
|
||||
}
|
||||
|
||||
|
||||
def _proxy_environment(runtime: _Runtime, directory: Path) -> dict[str, str]:
|
||||
return {
|
||||
"HTTPS_PROXY": runtime.tunnel.url,
|
||||
"NO_PROXY": "127.0.0.1,localhost",
|
||||
"SSL_CERT_FILE": str(runtime.cert),
|
||||
"STRANDS_DECIDER_API_BASE": _arn("audit_env"),
|
||||
**_isolated_aws(directory),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def module_gateway() -> Iterator[Gateway]:
|
||||
with gateway_from_environment() as value:
|
||||
yield value
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def runtime(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Runtime]:
|
||||
cert, key = write_self_signed_cert(tmp_path_factory.mktemp("agentcore-cert"), (_host(_US_EAST), _host(_EU_WEST)))
|
||||
outage: Final = threading.Event()
|
||||
with wire_server(partial(_respond, outage), tls=server_context(cert, key)) as wire:
|
||||
with tunnelling_forward_proxy(int(wire.url.rsplit(":", 1)[1])) as tunnel:
|
||||
yield _Runtime(wire, tunnel, cert, outage)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def agentcore(
|
||||
module_gateway: Gateway, runtime: _Runtime, tmp_path_factory: pytest.TempPathFactory
|
||||
) -> Iterator[Gateway]:
|
||||
directory: Final = tmp_path_factory.mktemp("agentcore-proxy")
|
||||
with owned_proxy(
|
||||
module_gateway,
|
||||
directory,
|
||||
_proxy_environment(runtime, directory),
|
||||
config=_write_config(directory),
|
||||
remove_environment=_scrubbed(),
|
||||
workers=2,
|
||||
) as candidate:
|
||||
yield candidate
|
||||
|
||||
|
||||
def _decide(gateway: Gateway, model: str, route: _Route = _SYSTEMONE, **extra: JsonValue) -> httpx.Response:
|
||||
return gateway.request("POST", route.path, {"model": model, **route.body, **extra})
|
||||
|
||||
|
||||
def _authorization_fields(request: Request) -> dict[str, str]:
|
||||
authorization: Final = request.headers.get("authorization", "")
|
||||
assert authorization.startswith("AWS4-HMAC-SHA256 "), dict(request.headers)
|
||||
return dict(part.split("=", 1) for part in authorization.removeprefix("AWS4-HMAC-SHA256 ").split(", "))
|
||||
|
||||
|
||||
def _assert_signed(request: Request, region: str = _US_EAST) -> None:
|
||||
fields: Final = _authorization_fields(request)
|
||||
access_key, scope = fields["Credential"].split("/", 1)
|
||||
assert access_key == _ACCESS_KEY, fields
|
||||
assert scope == f"{request.headers['x-amz-date'][:8]}/{region}/bedrock-agentcore/aws4_request", fields
|
||||
signed: Final = fields["SignedHeaders"]
|
||||
assert {"host", "x-amz-date", _SESSION_HEADER}.issubset(signed.split(";")), signed
|
||||
expected: Final = signature(
|
||||
"POST", encoded_path(request.target), request.headers, signed, request.body, _SECRET_KEY, scope
|
||||
)
|
||||
assert fields["Signature"] == expected[1], (request.target, fields)
|
||||
|
||||
|
||||
def _assert_answered(response: httpx.Response, model: str, route: _Route = _SYSTEMONE) -> None:
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == route.answer, response.text
|
||||
assert response.headers["x-litellm-model-group"] == model, dict(response.headers)
|
||||
|
||||
|
||||
def _assert_runtime_target(request: Request, runtime: str, region: str = _US_EAST) -> None:
|
||||
assert request.method == "POST", request.method
|
||||
assert request.target == f"/runtimes/{quote(_arn(runtime, region), safe='')}/invocations", request.target
|
||||
assert request.headers["host"] == _host(region), dict(request.headers)
|
||||
|
||||
|
||||
def _assert_runtime_call(request: Request, runtime: str, region: str = _US_EAST, route: _Route = _SYSTEMONE) -> None:
|
||||
_assert_runtime_target(request, runtime, region)
|
||||
assert json.loads(request.body) == route.upstream, request.body
|
||||
|
||||
|
||||
def _canonical(body: JsonValue | Mapping[str, JsonValue]) -> str:
|
||||
return json.dumps(body, sort_keys=True)
|
||||
|
||||
|
||||
def _spend_row(call_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = eventually(lambda: read_rows(_SPEND_QUERY, (call_id,)), lambda found: len(found) == 1, seconds=70)
|
||||
return rows[0]
|
||||
|
||||
|
||||
def _spend_rows(call_ids: Sequence[str]) -> list[dict[str, JsonValue]]:
|
||||
query: Final = (
|
||||
"SELECT request_id, status FROM \"LiteLLM_SpendLogs\" WHERE request_id = ANY(string_to_array(%s, ','))"
|
||||
)
|
||||
return eventually(
|
||||
lambda: read_rows(query, (",".join(call_ids),)), lambda found: len(found) >= len(call_ids), seconds=70
|
||||
)
|
||||
|
||||
|
||||
def _connect_targets(runtime: _Runtime) -> frozenset[str]:
|
||||
return frozenset(runtime.tunnel.targets())
|
||||
|
||||
|
||||
def test_an_arn_deployment_answers_on_every_decisions_route_signed_for_bedrock_agentcore(
|
||||
agentcore: Gateway, runtime: _Runtime
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
responses: Final = tuple(_decide(agentcore, _SIGNED, route) for route in _ROUTES)
|
||||
for response, route in zip(responses, _ROUTES, strict=True):
|
||||
_assert_answered(response, _SIGNED, route)
|
||||
assert _connect_targets(runtime) == {f"{_host(_US_EAST)}:443"}
|
||||
calls: Final = runtime.calls()
|
||||
assert len(calls) == len(_ROUTES), calls
|
||||
for call, route in zip(calls, _ROUTES, strict=True):
|
||||
_assert_runtime_call(call, "audit_signed", route=route)
|
||||
assert call.headers[_SESSION_HEADER] == _default_session("audit_signed"), dict(call.headers)
|
||||
_assert_signed(call)
|
||||
for response in responses:
|
||||
row = _spend_row(response.headers["x-litellm-call-id"])
|
||||
assert row == {
|
||||
"status": "success",
|
||||
"call_type": "adecisions",
|
||||
"custom_llm_provider": "strands_decider",
|
||||
"model_group": _SIGNED,
|
||||
"api_base": _invocations_url("audit_signed"),
|
||||
"prompt_tokens": 211,
|
||||
"completion_tokens": 2,
|
||||
}, row
|
||||
|
||||
|
||||
def test_the_arn_region_picks_the_host_and_signing_scope_over_aws_region_name(
|
||||
agentcore: Gateway, runtime: _Runtime
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
_assert_answered(_decide(agentcore, _REGION), _REGION)
|
||||
assert _connect_targets(runtime) == {f"{_host(_EU_WEST)}:443"}
|
||||
(call,) = runtime.calls()
|
||||
_assert_runtime_call(call, "audit_region", _EU_WEST)
|
||||
_assert_signed(call, _EU_WEST)
|
||||
|
||||
|
||||
def test_a_configured_session_id_is_sent_as_written_and_signed(agentcore: Gateway, runtime: _Runtime) -> None:
|
||||
runtime.reset()
|
||||
_assert_answered(_decide(agentcore, _SESSION), _SESSION)
|
||||
_assert_answered(_decide(agentcore, _LONG), _LONG)
|
||||
calls: Final = runtime.calls()
|
||||
(configured,) = _calls_to(calls, "audit_session")
|
||||
(long,) = _calls_to(calls, "audit_session_long")
|
||||
assert configured.headers[_SESSION_HEADER] == _CONFIGURED_SESSION, dict(configured.headers)
|
||||
assert long.headers[_SESSION_HEADER] == _LONG_SESSION
|
||||
_assert_signed(configured)
|
||||
_assert_signed(long)
|
||||
|
||||
|
||||
def test_a_deployment_api_key_sends_bearer_and_the_session_without_sigv4(agentcore: Gateway, runtime: _Runtime) -> None:
|
||||
runtime.reset()
|
||||
_assert_answered(_decide(agentcore, _BEARER), _BEARER)
|
||||
(call,) = runtime.calls()
|
||||
_assert_runtime_call(call, "audit_bearer")
|
||||
assert call.headers["authorization"] == f"Bearer {_JWT}", dict(call.headers)
|
||||
assert call.headers[_SESSION_HEADER] == _default_session("audit_bearer"), dict(call.headers)
|
||||
assert "x-amz-date" not in call.headers, dict(call.headers)
|
||||
|
||||
|
||||
def test_reserved_extra_headers_are_dropped_and_the_rest_are_kept(agentcore: Gateway, runtime: _Runtime) -> None:
|
||||
runtime.reset()
|
||||
_assert_answered(_decide(agentcore, _HEADERS), _HEADERS)
|
||||
caller_headers: Final[dict[str, JsonValue]] = {
|
||||
"x-amz-security-token": "caller-forged-token",
|
||||
"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "caller-forged-session-0000000000000000001",
|
||||
"x-audit-caller": "caller-trace",
|
||||
}
|
||||
_assert_answered(_decide(agentcore, _HEADERS, extra_headers=caller_headers), _HEADERS)
|
||||
deployment, caller = runtime.calls()
|
||||
for call in (deployment, caller):
|
||||
_assert_runtime_call(call, "audit_headers")
|
||||
assert call.headers[_SESSION_HEADER] == _default_session("audit_headers"), dict(call.headers)
|
||||
assert "x-amz-target" not in call.headers, dict(call.headers)
|
||||
assert "x-amz-security-token" not in call.headers, dict(call.headers)
|
||||
_assert_signed(call)
|
||||
assert deployment.headers["x-audit-trace"] == "deployment-trace", dict(deployment.headers)
|
||||
assert caller.headers["x-audit-caller"] == "caller-trace", dict(caller.headers)
|
||||
|
||||
|
||||
def test_an_opted_in_caller_session_id_becomes_the_header_and_stays_out_of_the_body(
|
||||
agentcore: Gateway, runtime: _Runtime
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
chosen: Final = f"agentcore-audit-caller-session-{uuid.uuid4().hex}"
|
||||
_assert_answered(_decide(agentcore, _OPT_IN, agentcore_runtime_session_id=chosen), _OPT_IN)
|
||||
_assert_answered(_decide(agentcore, _OPT_IN), _OPT_IN)
|
||||
first, second = runtime.calls()
|
||||
for call in (first, second):
|
||||
_assert_runtime_call(call, "audit_opt_in")
|
||||
_assert_signed(call)
|
||||
assert first.headers[_SESSION_HEADER] == chosen, dict(first.headers)
|
||||
assert second.headers[_SESSION_HEADER] == _default_session("audit_opt_in"), dict(second.headers)
|
||||
|
||||
|
||||
def test_an_arn_deployment_without_api_base_takes_the_runtime_from_the_environment(
|
||||
agentcore: Gateway, runtime: _Runtime
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
response: Final = _decide(agentcore, _ENV)
|
||||
_assert_answered(response, _ENV)
|
||||
(call,) = runtime.calls()
|
||||
_assert_runtime_call(call, "audit_env")
|
||||
_assert_signed(call)
|
||||
assert _spend_row(response.headers["x-litellm-call-id"])["api_base"] == _invocations_url("audit_env")
|
||||
|
||||
|
||||
def test_an_arn_deployment_added_through_the_management_api_answers_signed(
|
||||
agentcore: Gateway, runtime: _Runtime
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
with agentcore.scenario() as scenario:
|
||||
model: Final = scenario.model(
|
||||
model_info=None,
|
||||
model=_MODEL,
|
||||
api_base=_arn("audit_stored"),
|
||||
api_key=None,
|
||||
aws_access_key_id=_ACCESS_KEY,
|
||||
aws_secret_access_key=_SECRET_KEY,
|
||||
)
|
||||
_assert_answered(_decide(agentcore, model), model)
|
||||
(call,) = runtime.calls()
|
||||
_assert_runtime_call(call, "audit_stored")
|
||||
_assert_signed(call)
|
||||
|
||||
|
||||
def test_health_and_test_connection_probe_the_runtime_signed(agentcore: Gateway, runtime: _Runtime) -> None:
|
||||
runtime.reset()
|
||||
health: Final = agentcore.request("GET", "/health", params={"model": _HEALTH})
|
||||
assert health.status_code == 200, health.text
|
||||
assert (health.json()["healthy_count"], health.json()["unhealthy_count"]) == (1, 0), health.text
|
||||
(probe,) = runtime.calls()
|
||||
assert probe.target == f"/runtimes/{quote(_arn('audit_health'), safe='')}/invocations", probe.target
|
||||
_assert_signed(probe)
|
||||
connection: Final = agentcore.request(
|
||||
"POST",
|
||||
"/health/test_connection",
|
||||
{
|
||||
"litellm_params": {"model": _MODEL, "api_base": _arn("audit_test_connection"), **_SIGNING},
|
||||
"mode": "evaluation",
|
||||
},
|
||||
)
|
||||
assert connection.status_code == 200, connection.text
|
||||
assert connection.json()["status"] == "success", connection.text
|
||||
(tested,) = runtime.calls()
|
||||
assert _runtime_of(tested) == "audit_test_connection", tested.target
|
||||
_assert_signed(tested)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("code", "status"), tuple(_ERROR_STATUS.items()), ids=tuple(_ERROR_STATUS))
|
||||
def test_a_runtime_error_reply_maps_to_its_status_with_the_runtime_message(
|
||||
agentcore: Gateway, runtime: _Runtime, code: str, status: int
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
model: Final = f"agentcore-audit-error-{code}"
|
||||
response: Final = _decide(agentcore, model)
|
||||
assert response.status_code == status, response.text
|
||||
assert f"Strands Decider runtime error '{code}': runtime refused with {code}" in response.text, response.text
|
||||
assert len(_calls_to(runtime.calls(), f"audit_error_{code}")) == 1
|
||||
row: Final = _spend_row(response.headers["x-litellm-call-id"])
|
||||
assert (row["status"], row["model_group"]) == ("failure", model), row
|
||||
|
||||
|
||||
@pytest.mark.parametrize("status", tuple(_HTTP_ERRORS), ids=tuple(str(status) for status in _HTTP_ERRORS))
|
||||
def test_a_runtime_http_error_keeps_its_status_and_message(agentcore: Gateway, runtime: _Runtime, status: int) -> None:
|
||||
runtime.reset()
|
||||
model: Final = f"agentcore-audit-status-{status}"
|
||||
response: Final = _decide(agentcore, model)
|
||||
assert response.status_code == status, response.text
|
||||
assert f"scripted {_HTTP_ERRORS[status]}" in response.text, response.text
|
||||
assert len(_calls_to(runtime.calls(), f"audit_status_{status}")) == 1
|
||||
assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("label", tuple(_MALFORMED_SESSIONS))
|
||||
def test_a_malformed_configured_session_id_fails_only_its_own_deployment(
|
||||
agentcore: Gateway, runtime: _Runtime, label: str
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
model: Final = f"agentcore-audit-session-{label}"
|
||||
response: Final = _decide(agentcore, model)
|
||||
calls: Final = _calls_to(runtime.calls(), f"audit_session_{label}")
|
||||
if label == "empty":
|
||||
_assert_answered(response, model)
|
||||
(call,) = calls
|
||||
assert call.headers[_SESSION_HEADER] == "", dict(call.headers)
|
||||
else:
|
||||
assert response.status_code == 500, response.text
|
||||
assert "agentcore_runtime_session_id" in response.text, response.text
|
||||
assert calls == ()
|
||||
_assert_answered(_decide(agentcore, _SIGNED), _SIGNED)
|
||||
|
||||
|
||||
def test_a_deployment_without_aws_credentials_fails_before_dialing_the_runtime(
|
||||
agentcore: Gateway, runtime: _Runtime
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
response: Final = _decide(agentcore, _UNSIGNED)
|
||||
assert response.status_code == 500, response.text
|
||||
assert "credentials" in response.text.lower(), response.text
|
||||
assert runtime.calls() == ()
|
||||
assert _connect_targets(runtime) == frozenset()
|
||||
_assert_answered(_decide(agentcore, _SIGNED), _SIGNED)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("label", tuple(_NEAR_MISS_ARNS))
|
||||
def test_a_near_miss_arn_keeps_the_plain_api_base_path(agentcore: Gateway, runtime: _Runtime, label: str) -> None:
|
||||
runtime.reset()
|
||||
response: Final = _decide(agentcore, f"agentcore-audit-near-miss-{label}")
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert "litellm.APIConnectionError" in response.text, response.text
|
||||
assert _connect_targets(runtime) == frozenset()
|
||||
assert runtime.calls() == ()
|
||||
|
||||
|
||||
def test_a_dropped_connection_is_a_connection_error_after_one_attempt(agentcore: Gateway, runtime: _Runtime) -> None:
|
||||
runtime.reset()
|
||||
response: Final = _decide(agentcore, _DROPPED)
|
||||
assert response.status_code == 500, response.text
|
||||
assert "litellm.APIConnectionError" in response.text, response.text
|
||||
assert len(_calls_to(runtime.calls(), "audit_dropped")) == 1
|
||||
assert _spend_row(response.headers["x-litellm-call-id"])["status"] == "failure"
|
||||
|
||||
|
||||
def _burst_routes(size: int) -> tuple[_Route, ...]:
|
||||
return tuple(_ROUTES[index % len(_ROUTES)] for index in range(size))
|
||||
|
||||
|
||||
def _burst(gateway: Gateway, model: str, size: int) -> tuple[httpx.Response, ...]:
|
||||
with ThreadPoolExecutor(max_workers=size) as pool:
|
||||
return tuple(pool.map(partial(_decide, gateway, model), _burst_routes(size)))
|
||||
|
||||
|
||||
def _health_counts(gateway: Gateway, model: str) -> tuple[int, int, int]:
|
||||
response: Final = gateway.request("GET", "/health", params={"model": model})
|
||||
body: Final = JSON_OBJECT.validate_json(response.content)
|
||||
healthy: Final = body["healthy_count"]
|
||||
unhealthy: Final = body["unhealthy_count"]
|
||||
assert isinstance(healthy, int) and isinstance(unhealthy, int), body
|
||||
return response.status_code, healthy, unhealthy
|
||||
|
||||
|
||||
def test_a_runtime_outage_mid_burst_fails_each_call_once_and_recovers_signed(
|
||||
agentcore: Gateway, runtime: _Runtime
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
runtime.outage.set()
|
||||
failed: Final = _burst(agentcore, _OUTAGE, 10)
|
||||
assert [response.status_code for response in failed] == [500] * 10, [response.text for response in failed]
|
||||
assert _health_counts(agentcore, _OUTAGE) == (503, 0, 1)
|
||||
runtime.outage.clear()
|
||||
runtime.calls()
|
||||
served: Final = _burst(agentcore, _OUTAGE, 20)
|
||||
for response, route in zip(served, _burst_routes(20), strict=True):
|
||||
_assert_answered(response, _OUTAGE, route)
|
||||
calls: Final = runtime.calls()
|
||||
assert len(calls) == 20, calls
|
||||
for call in calls:
|
||||
_assert_runtime_target(call, "audit_outage")
|
||||
_assert_signed(call)
|
||||
assert sorted(_canonical(_JSON.validate_json(call.body)) for call in calls) == sorted(
|
||||
_canonical(route.upstream) for route in _burst_routes(20)
|
||||
)
|
||||
assert _health_counts(agentcore, _OUTAGE) == (200, 1, 0)
|
||||
expected: Final = {
|
||||
**{response.headers["x-litellm-call-id"]: "failure" for response in failed},
|
||||
**{response.headers["x-litellm-call-id"]: "success" for response in served},
|
||||
}
|
||||
assert len(expected) == 30, expected
|
||||
rows: Final = _spend_rows(tuple(expected))
|
||||
assert sorted((str(row["request_id"]), row["status"]) for row in rows) == sorted(expected.items()), rows
|
||||
|
||||
|
||||
async def test_sdk_sync_and_async_calls_reach_the_runtime_signed(
|
||||
runtime: _Runtime, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
|
||||
) -> None:
|
||||
runtime.reset()
|
||||
for name in _scrubbed():
|
||||
monkeypatch.delenv(name)
|
||||
for name, value in {
|
||||
"HTTPS_PROXY": runtime.tunnel.url,
|
||||
"NO_PROXY": "127.0.0.1,localhost",
|
||||
"SSL_CERT_FILE": str(runtime.cert),
|
||||
**_isolated_aws(tmp_path),
|
||||
}.items():
|
||||
monkeypatch.setenv(name, value)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
try:
|
||||
synchronous: Final = litellm.decisions(
|
||||
model=_MODEL,
|
||||
state=_STATE,
|
||||
questions=_SDK_QUESTIONS,
|
||||
api_base=_arn("audit_sdk"),
|
||||
aws_access_key_id=_ACCESS_KEY,
|
||||
aws_secret_access_key=_SECRET_KEY,
|
||||
)
|
||||
asynchronous: Final = await litellm.adecisions(
|
||||
model=_MODEL,
|
||||
state=_STATE,
|
||||
questions=_SDK_QUESTIONS,
|
||||
api_base=_arn("audit_sdk"),
|
||||
aws_access_key_id=_ACCESS_KEY,
|
||||
aws_secret_access_key=_SECRET_KEY,
|
||||
)
|
||||
monkeypatch.setenv("STRANDS_DECIDER_API_KEY", _ENV_JWT)
|
||||
bearer: Final = await litellm.adecisions(
|
||||
model=_MODEL, state=_STATE, questions=_SDK_QUESTIONS, api_base=_arn("audit_sdk_env_key")
|
||||
)
|
||||
finally:
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
for response in (synchronous, asynchronous, bearer):
|
||||
assert response.model_dump(mode="json") == {"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE}
|
||||
calls: Final = runtime.calls()
|
||||
signed: Final = _calls_to(calls, "audit_sdk")
|
||||
assert len(signed) == 2, calls
|
||||
for call in signed:
|
||||
_assert_runtime_call(call, "audit_sdk")
|
||||
assert call.headers[_SESSION_HEADER] == _default_session("audit_sdk"), dict(call.headers)
|
||||
_assert_signed(call)
|
||||
(env_key,) = _calls_to(calls, "audit_sdk_env_key")
|
||||
assert env_key.headers["authorization"] == f"Bearer {_ENV_JWT}", dict(env_key.headers)
|
||||
assert "x-amz-date" not in env_key.headers, dict(env_key.headers)
|
||||
|
||||
|
||||
def _answering(request: Request) -> Reply:
|
||||
return _json_reply({"model": _BODY_MODEL, "answers": _ANSWERS, "usage": _USAGE})
|
||||
|
||||
|
||||
def _envelope(request: Request) -> Reply:
|
||||
return _json_reply({"error": {"code": "bad_request", "message": "plain server refused"}})
|
||||
|
||||
|
||||
def _plain_model(scenario: Scenario, wire: Wire, **extra: JsonValue) -> str:
|
||||
return scenario.model(model_info=None, model=_MODEL, api_base=wire.url, api_key=None, **extra)
|
||||
|
||||
|
||||
def test_a_request_body_session_id_is_refused_on_every_route_without_an_upstream_call(gateway: Gateway) -> None:
|
||||
with wire_server(_answering) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _plain_model(scenario, wire)
|
||||
responses: Final = tuple(
|
||||
_decide(gateway, model, route, agentcore_runtime_session_id=f"caller-{uuid.uuid4().hex}")
|
||||
for route in _ROUTES
|
||||
)
|
||||
for response in responses:
|
||||
assert response.status_code == 401, response.text
|
||||
assert "agentcore_runtime_session_id is not allowed in request body" in response.text, response.text
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def test_a_plain_api_base_opted_into_the_session_param_sends_no_session_header_or_body_field(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
with wire_server(_answering) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _plain_model(
|
||||
scenario, wire, configurable_clientside_auth_params=["agentcore_runtime_session_id"]
|
||||
)
|
||||
_assert_answered(_decide(gateway, model, agentcore_runtime_session_id=f"caller-{uuid.uuid4().hex}"), model)
|
||||
(call,) = wire.drain()
|
||||
assert call.target == "/v1/systemone", call.target
|
||||
assert _SESSION_HEADER not in call.headers, dict(call.headers)
|
||||
assert json.loads(call.body) == _UPSTREAM_BODY, call.body
|
||||
|
||||
|
||||
def test_a_plain_api_base_error_envelope_keeps_the_unexpected_response_server_error(gateway: Gateway) -> None:
|
||||
with wire_server(_envelope) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _plain_model(scenario, wire)
|
||||
response: Final = _decide(gateway, model)
|
||||
assert response.status_code == 500, response.text
|
||||
assert "Strands Decider runtime error" not in response.text, response.text
|
||||
assert len(wire.drain()) == 1
|
||||
|
||||
|
||||
def test_a_plain_api_base_forwards_amazon_shaped_extra_headers_unchanged(gateway: Gateway) -> None:
|
||||
headers: Final[dict[str, JsonValue]] = {
|
||||
"x-amz-target": "plain-target",
|
||||
"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": "plain-session-kept",
|
||||
"x-audit-trace": "plain-trace",
|
||||
}
|
||||
with wire_server(_answering) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _plain_model(scenario, wire, extra_headers=headers)
|
||||
_assert_answered(_decide(gateway, model), model)
|
||||
(call,) = wire.drain()
|
||||
assert (call.headers["x-amz-target"], call.headers[_SESSION_HEADER], call.headers["x-audit-trace"]) == (
|
||||
"plain-target",
|
||||
"plain-session-kept",
|
||||
"plain-trace",
|
||||
), dict(call.headers)
|
||||
assert "authorization" not in call.headers, dict(call.headers)
|
||||
assert json.loads(call.body) == _UPSTREAM_BODY, call.body
|
||||
|
||||
|
||||
def _worst_case_api_base() -> str:
|
||||
return "arn:aws" + "-a" * 131072 + "!"
|
||||
|
||||
|
||||
def test_a_worst_case_runtime_arn_pattern_input_is_classified_in_linear_time(
|
||||
gateway: Gateway, record_property: Callable[[str, object], None]
|
||||
) -> None:
|
||||
worst: Final = _worst_case_api_base()
|
||||
started: Final = time.perf_counter()
|
||||
url: Final = StrandsDeciderDecisionsConfig().get_complete_url(worst, _BODY_MODEL)
|
||||
classified: Final = time.perf_counter() - started
|
||||
record_property("classifier_seconds", f"{classified:.6f}")
|
||||
assert url == f"{worst}/v1/systemone"
|
||||
assert classified < 0.5, classified
|
||||
with wire_server(_answering) as wire, gateway.scenario() as scenario:
|
||||
model: Final = _plain_model(scenario, wire, configurable_clientside_auth_params=["api_base"])
|
||||
with ThreadPoolExecutor(max_workers=2) as pool:
|
||||
timed: Final = pool.submit(_timed_decide, gateway, model, worst)
|
||||
liveliness_started: Final = time.perf_counter()
|
||||
liveliness: Final = gateway.request("GET", "/health/liveliness")
|
||||
liveliness_seconds: Final = time.perf_counter() - liveliness_started
|
||||
response, decide_seconds = timed.result()
|
||||
record_property("decide_seconds", f"{decide_seconds:.6f}")
|
||||
record_property("liveliness_seconds", f"{liveliness_seconds:.6f}")
|
||||
assert liveliness.status_code == 200, liveliness.text
|
||||
assert 500 <= response.status_code < 600, response.text
|
||||
assert "litellm.APIConnectionError" in response.text, response.text[:2000]
|
||||
assert decide_seconds < 5, decide_seconds
|
||||
assert liveliness_seconds < 2, liveliness_seconds
|
||||
assert wire.drain() == ()
|
||||
|
||||
|
||||
def _timed_decide(gateway: Gateway, model: str, api_base: str) -> tuple[httpx.Response, float]:
|
||||
started: Final = time.perf_counter()
|
||||
response: Final = _decide(gateway, model, api_base=api_base)
|
||||
return response, time.perf_counter() - started
|
||||
|
|
@ -2,6 +2,8 @@ from __future__ import annotations
|
|||
|
||||
import asyncio
|
||||
import json
|
||||
import threading
|
||||
from concurrent.futures import wait
|
||||
from collections.abc import Mapping
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
|
@ -9,11 +11,16 @@ from typing import Final
|
|||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
import litellm
|
||||
from litellm.cost_calculator import get_response_cost_from_hidden_params
|
||||
from litellm.constants import AWS_SIGNING_MAX_THREADS
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.bedrock.base_aws_llm import AWS_SIGNING_EXECUTOR
|
||||
from litellm.types.decisions import (
|
||||
ChoiceAnswer,
|
||||
DecisionsResponse,
|
||||
|
|
@ -698,6 +705,33 @@ async def test_strands_decider_without_key_preserves_response_extras(
|
|||
assert severity.legend == {"0": "none", "1": "low", "2": "high"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strands_decider_plain_api_base_answers_while_every_aws_signing_thread_is_busy(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
|
||||
respx_mock.post("https://strands.example/v1/systemone").respond(json=_STRANDS_RESPONSE)
|
||||
release: Final = threading.Event()
|
||||
busy: Final = tuple(AWS_SIGNING_EXECUTOR.submit(release.wait) for _ in range(AWS_SIGNING_MAX_THREADS))
|
||||
|
||||
try:
|
||||
response: Final = await asyncio.wait_for(
|
||||
litellm.adecisions(
|
||||
model="strands_decider/strands-decider-2B-hobson-v19",
|
||||
state="review",
|
||||
questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
|
||||
api_base="https://strands.example",
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
finally:
|
||||
release.set()
|
||||
wait(busy)
|
||||
|
||||
assert isinstance(response.answers["severity"], ScoreAnswer)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strands_decider_uses_key_from_matching_environment_base(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
|
|
@ -750,6 +784,172 @@ async def test_strands_decider_provider_resolution_and_router_dispatch(
|
|||
assert response.model == _STRANDS_RESPONSE["model"]
|
||||
|
||||
|
||||
_RUNTIME_ARN: Final = "arn:aws:bedrock-agentcore:us-east-1:123456789012:runtime/qa_decider-AbC123xyz0"
|
||||
_RUNTIME_URL: Final = (
|
||||
"https://bedrock-agentcore.us-east-1.amazonaws.com/runtimes/"
|
||||
"arn%3Aaws%3Abedrock-agentcore%3Aus-east-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0/invocations"
|
||||
)
|
||||
_RUNTIME_CREDENTIALS: Final = Credentials(access_key="AKIDEXAMPLE", secret_key="example-secret")
|
||||
_SESSION_HEADER: Final = "x-amzn-bedrock-agentcore-runtime-session-id"
|
||||
|
||||
|
||||
def _sigv4_verifies(request: httpx.Request) -> bool:
|
||||
authorization: Final = request.headers["authorization"]
|
||||
signed_names: Final = authorization.split("SignedHeaders=")[1].split(",")[0].split(";")
|
||||
replayed: Final = AWSRequest(
|
||||
method=request.method,
|
||||
url=str(request.url),
|
||||
data=request.content,
|
||||
headers={name: request.headers[name] for name in signed_names},
|
||||
)
|
||||
replayed.context["timestamp"] = request.headers["x-amz-date"]
|
||||
auth: Final = SigV4Auth(_RUNTIME_CREDENTIALS, "bedrock-agentcore", "us-east-1")
|
||||
string_to_sign: Final = auth.string_to_sign(replayed, auth.canonical_request(replayed))
|
||||
return authorization.endswith(f"Signature={auth.signature(string_to_sign, replayed)}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_async", [True, False])
|
||||
@pytest.mark.asyncio
|
||||
async def test_strands_decider_runtime_arn_sends_a_sigv4_signed_invoke_agent_runtime_call(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
use_async: bool,
|
||||
) -> None:
|
||||
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE)
|
||||
call_args: Final[Mapping[str, object]] = {
|
||||
"model": "strands_decider/strands-decider-2B-hobson-v19",
|
||||
"state": "review",
|
||||
"questions": {"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
|
||||
"api_base": _RUNTIME_ARN,
|
||||
"aws_access_key_id": _RUNTIME_CREDENTIALS.access_key,
|
||||
"aws_secret_access_key": _RUNTIME_CREDENTIALS.secret_key,
|
||||
}
|
||||
|
||||
response: Final = await litellm.adecisions(**call_args) if use_async else litellm.decisions(**call_args)
|
||||
|
||||
assert route.call_count == 1
|
||||
sent: Final = route.calls[0].request
|
||||
assert _sigv4_verifies(sent)
|
||||
assert json.loads(sent.content)["model"] == "strands-decider-2B-hobson-v19"
|
||||
assert sent.headers[_SESSION_HEADER].startswith("litellm-decider-")
|
||||
assert response.model == _STRANDS_RESPONSE["model"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strands_decider_runtime_arn_router_deployment_answers_openai_format_calls(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "strands",
|
||||
"litellm_params": {
|
||||
"model": "strands_decider/strands-decider-2B-hobson-v19",
|
||||
"api_base": _RUNTIME_ARN,
|
||||
"aws_access_key_id": _RUNTIME_CREDENTIALS.access_key,
|
||||
"aws_secret_access_key": _RUNTIME_CREDENTIALS.secret_key,
|
||||
"aws_region_name": "eu-west-1",
|
||||
"agentcore_runtime_session_id": "pooled-session-0123456789abcdef0123456",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE)
|
||||
|
||||
response: Final = await router.adecisions(
|
||||
model="strands",
|
||||
input="review",
|
||||
extra_headers={_SESSION_HEADER: "caller-session-0123456789abcdef01234567"},
|
||||
questions=[
|
||||
{
|
||||
"type": "score",
|
||||
"name": "severity",
|
||||
"instructions": "How bad?",
|
||||
"levels": [{"label": "none"}, {"label": "low"}, {"label": "high"}],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
sent: Final = route.calls[0].request
|
||||
assert _sigv4_verifies(sent)
|
||||
assert sent.headers[_SESSION_HEADER] == "pooled-session-0123456789abcdef0123456"
|
||||
assert isinstance(response, OpenAIDecisionResponse)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strands_decider_runtime_arn_with_api_key_sends_the_bearer_token_unsigned(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
|
||||
route: Final = respx_mock.post(_RUNTIME_URL).respond(json=_STRANDS_RESPONSE)
|
||||
|
||||
await litellm.adecisions(
|
||||
model="strands_decider/strands-decider-2B-hobson-v19",
|
||||
state="review",
|
||||
questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
|
||||
api_base=_RUNTIME_ARN,
|
||||
api_key="runtime-jwt",
|
||||
)
|
||||
|
||||
sent: Final = route.calls[0].request
|
||||
assert sent.headers["authorization"] == "Bearer runtime-jwt"
|
||||
assert "x-amz-date" not in sent.headers
|
||||
assert _SESSION_HEADER in sent.headers
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("code", "error_class"),
|
||||
[
|
||||
("bad_request", litellm.BadRequestError),
|
||||
("loading", litellm.ServiceUnavailableError),
|
||||
("inference_error", litellm.InternalServerError),
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_strands_decider_runtime_error_envelope_raises_instead_of_answering(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
code: str,
|
||||
error_class: type[Exception],
|
||||
) -> None:
|
||||
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
|
||||
respx_mock.post(_RUNTIME_URL).respond(json={"error": {"code": code, "message": "at most 16 questions"}})
|
||||
|
||||
with pytest.raises(error_class, match="at most 16 questions"):
|
||||
await litellm.adecisions(
|
||||
model="strands_decider/strands-decider-2B-hobson-v19",
|
||||
state="review",
|
||||
questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
|
||||
api_base=_RUNTIME_ARN,
|
||||
aws_access_key_id=_RUNTIME_CREDENTIALS.access_key,
|
||||
aws_secret_access_key=_RUNTIME_CREDENTIALS.secret_key,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_strands_decider_plain_api_base_error_envelope_stays_an_unexpected_response(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
respx_mock: respx.MockRouter,
|
||||
) -> None:
|
||||
monkeypatch.delenv("STRANDS_DECIDER_API_KEY", raising=False)
|
||||
respx_mock.post("https://strands.example/v1/systemone").respond(
|
||||
json={"error": {"code": "bad_request", "message": "at most 16 questions"}}
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.InternalServerError, match="returned an unexpected response"):
|
||||
await litellm.adecisions(
|
||||
model="strands_decider/strands-decider-2B-hobson-v19",
|
||||
state="review",
|
||||
questions={"severity": {"type": "score", "criteria": ["none", "low", "high"]}},
|
||||
api_base="https://strands.example",
|
||||
)
|
||||
|
||||
|
||||
_VLLM_RESPONSE: Final[Mapping[str, object]] = {
|
||||
"id": "systemone-4af3d2c1",
|
||||
"object": "structured_decision",
|
||||
|
|
|
|||
0
tests/unit/llms/strands_decider/__init__.py
Normal file
0
tests/unit/llms/strands_decider/__init__.py
Normal file
0
tests/unit/llms/strands_decider/decisions/__init__.py
Normal file
0
tests/unit/llms/strands_decider/decisions/__init__.py
Normal file
|
|
@ -0,0 +1,256 @@
|
|||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from botocore.credentials import Credentials
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.decisions.transformation import systemone_request_to_ir
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.strands_decider.decisions.transformation import (
|
||||
AGENTCORE_SESSION_HEADER,
|
||||
StrandsDeciderDecisionsConfig,
|
||||
)
|
||||
from litellm.types.decisions import DecisionsIRResponse, DecisionsRequestBody
|
||||
from litellm.types.llms.bedrock import AwsAuthParams
|
||||
|
||||
_ARN: Final = "arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0"
|
||||
_BODY: Final[Mapping[str, object]] = {
|
||||
"model": "strands-decider-2B-hobson-v21",
|
||||
"state": "My order arrived broken.",
|
||||
"questions": {"is_complaint": {"type": "noul", "instructions": "Is this a complaint?"}},
|
||||
}
|
||||
_IR_REQUEST: Final = systemone_request_to_ir(
|
||||
TypeAdapter(DecisionsRequestBody).validate_python({"state": _BODY["state"], "questions": _BODY["questions"]})
|
||||
)
|
||||
|
||||
|
||||
class _RecordingAWS(BaseAWSLLM):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.resolved: tuple[tuple[AwsAuthParams, str | None], ...] = ()
|
||||
|
||||
def resolve_credentials(self, auth_params: AwsAuthParams, aws_region_name: str | None) -> Credentials:
|
||||
self.resolved = (*self.resolved, (auth_params, aws_region_name))
|
||||
return Credentials(access_key="AKIDEXAMPLE", secret_key="example-secret")
|
||||
|
||||
|
||||
def _sign(
|
||||
config: StrandsDeciderDecisionsConfig,
|
||||
api_base: str,
|
||||
headers: Mapping[str, str],
|
||||
api_key: str | None = None,
|
||||
litellm_params: Mapping[str, object] | None = None,
|
||||
) -> tuple[Mapping[str, str], bytes | None]:
|
||||
return config.sign_request(
|
||||
headers=headers,
|
||||
url=config.get_complete_url(api_base=api_base, model="strands-decider-2B-hobson-v21"),
|
||||
api_base=api_base,
|
||||
body=_BODY,
|
||||
api_key=api_key,
|
||||
litellm_params=litellm_params or {},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("arn", "expected_url"),
|
||||
[
|
||||
(
|
||||
_ARN,
|
||||
"https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/"
|
||||
"arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0/invocations",
|
||||
),
|
||||
(
|
||||
"arn:aws-cn:bedrock-agentcore:cn-north-1:123456789012:runtime/qa_decider-AbC123xyz0",
|
||||
"https://bedrock-agentcore.cn-north-1.amazonaws.com.cn/runtimes/"
|
||||
"arn%3Aaws-cn%3Abedrock-agentcore%3Acn-north-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0"
|
||||
"/invocations",
|
||||
),
|
||||
(
|
||||
"arn:aws-us-gov:bedrock-agentcore:us-gov-west-1:123456789012:runtime/qa_decider-AbC123xyz0",
|
||||
"https://bedrock-agentcore.us-gov-west-1.amazonaws.com/runtimes/"
|
||||
"arn%3Aaws-us-gov%3Abedrock-agentcore%3Aus-gov-west-1%3A123456789012%3Aruntime%2Fqa_decider-AbC123xyz0"
|
||||
"/invocations",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_runtime_arn_targets_the_regional_invoke_agent_runtime_url(arn: str, expected_url: str) -> None:
|
||||
assert StrandsDeciderDecisionsConfig().get_complete_url(api_base=arn, model="any") == expected_url
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"api_base",
|
||||
[
|
||||
"https://strands.example",
|
||||
"arn:aws:bedrock-agentcore:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0/runtime-endpoint/DEFAULT",
|
||||
"arn:aws:bedrock-agentcore:evil.example/x:123456789012:runtime/qa_decider-AbC123xyz0",
|
||||
"arn:aws:bedrock:us-west-2:123456789012:runtime/qa_decider-AbC123xyz0",
|
||||
],
|
||||
)
|
||||
def test_other_api_bases_keep_the_plain_systemone_request(api_base: str) -> None:
|
||||
aws: Final = _RecordingAWS()
|
||||
config: Final = StrandsDeciderDecisionsConfig(aws=aws)
|
||||
headers: Final = {"Content-Type": "application/json"}
|
||||
|
||||
assert config.get_complete_url(api_base=api_base, model="any") == f"{api_base}/v1/systemone"
|
||||
assert _sign(config, api_base, headers) == (headers, None)
|
||||
assert aws.resolved == ()
|
||||
|
||||
|
||||
def test_sign_request_signs_the_body_for_the_arn_region_with_deployment_credentials() -> None:
|
||||
aws: Final = _RecordingAWS()
|
||||
config: Final = StrandsDeciderDecisionsConfig(aws=aws)
|
||||
|
||||
signed_headers, signed_body = _sign(
|
||||
config,
|
||||
_ARN,
|
||||
{"Content-Type": "application/json"},
|
||||
litellm_params={"aws_profile_name": "qa", "aws_role_name": "decider-role", "aws_region_name": "us-east-1"},
|
||||
)
|
||||
|
||||
assert signed_body == json.dumps(_BODY).encode()
|
||||
assert len(aws.resolved) == 1
|
||||
auth_params, region = aws.resolved[0]
|
||||
assert (auth_params.aws_profile_name, auth_params.aws_role_name, region) == ("qa", "decider-role", "us-west-2")
|
||||
authorization: Final = signed_headers["Authorization"]
|
||||
assert authorization.startswith("AWS4-HMAC-SHA256 Credential=AKIDEXAMPLE/")
|
||||
assert "/us-west-2/bedrock-agentcore/aws4_request" in authorization
|
||||
assert AGENTCORE_SESSION_HEADER.lower() in authorization
|
||||
|
||||
|
||||
def test_default_session_id_is_stable_per_runtime_and_within_the_agentcore_length_bounds() -> None:
|
||||
other_arn: Final = _ARN.replace("AbC123xyz0", "ZyX987cba1")
|
||||
|
||||
first, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {})
|
||||
again, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {})
|
||||
other, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), other_arn, {})
|
||||
|
||||
session_id: Final = first[AGENTCORE_SESSION_HEADER]
|
||||
assert session_id == again[AGENTCORE_SESSION_HEADER]
|
||||
assert session_id != other[AGENTCORE_SESSION_HEADER]
|
||||
assert 33 <= len(session_id) <= 256
|
||||
|
||||
|
||||
_CALLER_RESERVED_HEADERS: Final = {
|
||||
AGENTCORE_SESSION_HEADER.lower(): "caller-session-0123456789abcdef0123456789",
|
||||
"X-Amzn-Bedrock-AgentCore-Runtime-User-Id": "someone-else",
|
||||
"X-Amz-Security-Token": "caller-token",
|
||||
"Host": "evil.example",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("api_key", [None, "runtime-jwt"])
|
||||
def test_caller_reserved_agentcore_headers_never_reach_the_runtime(api_key: str | None) -> None:
|
||||
signed_headers, _ = _sign(
|
||||
StrandsDeciderDecisionsConfig(aws=_RecordingAWS()),
|
||||
_ARN,
|
||||
{**_CALLER_RESERVED_HEADERS, "X-Request-Source": "batch"},
|
||||
api_key=api_key,
|
||||
)
|
||||
default_headers, _ = _sign(StrandsDeciderDecisionsConfig(aws=_RecordingAWS()), _ARN, {}, api_key=api_key)
|
||||
|
||||
forwarded: Final = {name.lower(): value for name, value in signed_headers.items()}
|
||||
assert forwarded[AGENTCORE_SESSION_HEADER.lower()] == default_headers[AGENTCORE_SESSION_HEADER]
|
||||
assert not any(value in forwarded.values() for value in _CALLER_RESERVED_HEADERS.values())
|
||||
assert "x-amzn-bedrock-agentcore-runtime-user-id" not in forwarded
|
||||
assert "host" not in forwarded
|
||||
assert forwarded["x-request-source"] == "batch"
|
||||
|
||||
|
||||
def test_deployment_session_param_picks_the_session_over_the_default_and_any_caller_header() -> None:
|
||||
replica_session: Final = "strands-decider-replica-a-0000000000"
|
||||
|
||||
signed_headers, _ = _sign(
|
||||
StrandsDeciderDecisionsConfig(aws=_RecordingAWS()),
|
||||
_ARN,
|
||||
_CALLER_RESERVED_HEADERS,
|
||||
litellm_params={"agentcore_runtime_session_id": replica_session},
|
||||
)
|
||||
|
||||
session_headers: Final = [
|
||||
value for name, value in signed_headers.items() if name.lower() == "x-amzn-bedrock-agentcore-runtime-session-id"
|
||||
]
|
||||
assert session_headers == [replica_session]
|
||||
assert AGENTCORE_SESSION_HEADER.lower() in signed_headers["Authorization"]
|
||||
|
||||
|
||||
def test_api_key_on_a_runtime_arn_sends_the_bearer_token_without_sigv4() -> None:
|
||||
aws: Final = _RecordingAWS()
|
||||
|
||||
signed_headers, signed_body = _sign(
|
||||
StrandsDeciderDecisionsConfig(aws=aws), _ARN, {"Authorization": "Bearer runtime-jwt"}, api_key="runtime-jwt"
|
||||
)
|
||||
|
||||
assert signed_body is None
|
||||
assert signed_headers["Authorization"] == "Bearer runtime-jwt"
|
||||
assert AGENTCORE_SESSION_HEADER in signed_headers
|
||||
assert aws.resolved == ()
|
||||
|
||||
|
||||
def _response(url: str, payload: object) -> httpx.Response:
|
||||
return httpx.Response(200, json=payload, request=httpx.Request("POST", url))
|
||||
|
||||
|
||||
def _transform(api_base: str, payload: object) -> DecisionsIRResponse:
|
||||
config: Final = StrandsDeciderDecisionsConfig()
|
||||
url: Final = config.get_complete_url(api_base=api_base, model="strands-decider-2B-hobson-v21")
|
||||
return config.transform_decisions_response(
|
||||
model="strands-decider-2B-hobson-v21",
|
||||
custom_llm_provider="strands_decider",
|
||||
raw_response=_response(url, payload),
|
||||
request=_IR_REQUEST,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("code", "status_code"),
|
||||
[("bad_request", 400), ("invalid_request", 400), ("loading", 503), ("inference_error", 500), ("new_code", 500)],
|
||||
)
|
||||
def test_runtime_error_envelope_raises_with_the_mapped_status(code: str, status_code: int) -> None:
|
||||
with pytest.raises(BaseLLMException) as raised:
|
||||
_transform(_ARN, {"error": {"code": code, "message": "at most 16 questions"}})
|
||||
|
||||
assert raised.value.status_code == status_code
|
||||
assert code in raised.value.message
|
||||
assert "at most 16 questions" in raised.value.message
|
||||
|
||||
|
||||
@pytest.mark.parametrize("code", ["bad_request", "loading", "inference_error"])
|
||||
def test_plain_api_base_error_envelope_keeps_the_unexpected_response_error(code: str) -> None:
|
||||
with pytest.raises(BaseLLMException) as raised:
|
||||
_transform("https://strands.example", {"error": {"code": code, "message": "at most 16 questions"}})
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
assert raised.value.message.startswith("Decisions provider 'strands_decider' returned an unexpected response")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url",
|
||||
[
|
||||
"https://strands.example/runtimes/arn%3Aaws%3Abedrock-agentcore%3Aus-west-2%3A123456789012%3Aruntime"
|
||||
"%2Fqa_decider-AbC123xyz0/invocations",
|
||||
"https://bedrock-agentcore.us-west-2.amazonaws.com/runtimes/not-an-arn/invocations",
|
||||
],
|
||||
)
|
||||
def test_error_envelope_from_a_url_that_is_not_the_runtime_keeps_the_unexpected_response_error(url: str) -> None:
|
||||
with pytest.raises(BaseLLMException) as raised:
|
||||
StrandsDeciderDecisionsConfig().transform_decisions_response(
|
||||
model="strands-decider-2B-hobson-v21",
|
||||
custom_llm_provider="strands_decider",
|
||||
raw_response=_response(url, {"error": {"code": "bad_request", "message": "at most 16 questions"}}),
|
||||
request=_IR_REQUEST,
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
|
||||
|
||||
@pytest.mark.parametrize("payload", [["not", "an", "object"], {"error": "unstructured"}])
|
||||
def test_runtime_reply_that_is_not_an_error_envelope_keeps_the_unexpected_response_error(payload: object) -> None:
|
||||
with pytest.raises(BaseLLMException) as raised:
|
||||
_transform(_ARN, payload)
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
assert raised.value.message.startswith("Decisions provider 'strands_decider' returned an unexpected response")
|
||||
|
|
@ -2854,6 +2854,22 @@ class TestIsRequestBodySafeBlocksBedrockProjectOverride:
|
|||
)
|
||||
|
||||
|
||||
def test_agentcore_runtime_session_id_in_request_body_is_rejected_without_admin_opt_in():
|
||||
body: Final = {"model": "strands-decider", "agentcore_runtime_session_id": "caller-session-0123456789abcdef0123"}
|
||||
|
||||
with pytest.raises(ValueError, match="agentcore_runtime_session_id"):
|
||||
is_request_body_safe(request_body=body, general_settings={}, llm_router=None, model="strands-decider")
|
||||
assert (
|
||||
is_request_body_safe(
|
||||
request_body=body,
|
||||
general_settings={"allow_client_side_credentials": True},
|
||||
llm_router=None,
|
||||
model="strands-decider",
|
||||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
class TestIsRequestBodySafeBlocksClaudePlatformWorkspaceOverride:
|
||||
@pytest.mark.parametrize(
|
||||
"alias", ["workspace_id", "aws_workspace_id", "anthropic_workspace_id", "anthropic-workspace-id"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue