scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

ci: drop aws partition hardcode gate

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-17 20:16:07 +00:00
parent 6a9ae2bba2
commit fc35b78eb4
3 changed files with 1 additions and 258 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
.github/workflows/test-code-quality.yml vendored
View file

@ -146,9 +146,6 @@ jobs:
- name: check_migrations_no_data_rewrites
run: uv run --no-sync python ./tests/code_coverage_tests/check_migrations_no_data_rewrites.py
- name: check_aws_partition_hardcodes
run: uv run --no-sync python ./tests/code_coverage_tests/check_aws_partition_hardcodes.py
- name: memory_test
run: uv run --no-sync python ./tests/code_coverage_tests/memory_test.py

140
tests/code_coverage_tests/check_aws_partition_hardcodes.py
View file

@ -1,140 +0,0 @@
#!/usr/bin/env python3
"""Ban hardcoded commercial-partition AWS hosts and ARN prefixes under `litellm/`.
An endpoint or ARN built with a literal `amazonaws.com` or `arn:aws:` works in every
commercial region and breaks only for GovCloud (`us-gov-*`, `arn:aws-us-gov:`) and
China (`amazonaws.com.cn`, `arn:aws-cn:`) deployments, so the failure never shows up
in CI or on a developer laptop. `litellm/litellm_core_utils/aws_partition.py` derives
both from the region and is the only place those literals belong. Build hosts with
`get_aws_dns_suffix(region)` and ARNs with `get_aws_arn_prefix(region)`.
Every string constant in every `litellm/**/*.py` file is scanned, including the
literal parts of f-strings and the strings inside `.format()` calls and
concatenations. Docstrings and comments are not, since they never reach a request.
`amazonaws.com.cn` passes because it is already the China partition.
`ALLOWED` holds the (file, token, count) triples that are text rather than a request
target: a hosted logo, an IAM service principal, and hostnames quoted as examples
inside error messages and field descriptions. An entry only covers that many
occurrences of that exact token in that exact file, so a second copy of an allowed
literal is still caught, and an entry whose token is gone or whose count has changed
fails the check so the set only shrinks.
"""
from __future__ import annotations
import ast
import re
import sys
from collections import Counter
from pathlib import Path
from types import MappingProxyType
from typing import Final, NamedTuple
REPO_ROOT: Final = Path(__file__).resolve().parents[2]
SCAN_ROOT: Final = REPO_ROOT / "litellm"
PARTITION_HELPER: Final = "litellm/litellm_core_utils/aws_partition.py"
COMMERCIAL_TOKEN: Final = re.compile(r"[A-Za-z0-9.-]*amazonaws\.com(?!\.cn)|arn:aws:[A-Za-z0-9:/_.*-]*")
class Allowance(NamedTuple):
file: str
token: str
occurrences: int
ALLOWED: Final = frozenset(
{
Allowance("litellm/integrations/email_alerting.py", "litellm-listing.s3.amazonaws.com", 1),
Allowance("litellm/types/integrations/slack_alerting.py", "litellm-listing.s3.amazonaws.com", 1),
Allowance("litellm/rag/ingestion/bedrock_ingestion.py", "bedrock.amazonaws.com", 1),
Allowance(
"litellm/llms/bedrock/chat/agentcore/transformation.py",
"arn:aws:bedrock-agentcore:region:account:runtime/runtime_id",
1,
),
Allowance("litellm/llms/bedrock/search/transformation.py", ".amazonaws.com", 1),
Allowance(
"litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py",
"bucket.s3.amazonaws.com",
1,
),
Allowance("litellm/types/proxy/claude_code_endpoints.py", "bucket.s3.amazonaws.com", 1),
}
)
ALLOWED_COUNTS: Final = MappingProxyType({(entry.file, entry.token): entry.occurrences for entry in ALLOWED})
class Hit(NamedTuple):
file: str
line: int
token: str
def _docstring_ids(tree: ast.Module) -> frozenset[int]:
return frozenset(
id(statement.value)
for node in ast.walk(tree)
if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef))
for statement in node.body
if isinstance(statement, ast.Expr)
and isinstance(statement.value, ast.Constant)
and isinstance(statement.value.value, str)
)
def _hits_in_file(path: Path) -> tuple[Hit, ...]:
tree: Final = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
docstrings: Final = _docstring_ids(tree)
relative: Final = path.relative_to(REPO_ROOT).as_posix()
return tuple(
Hit(relative, node.lineno, match.group(0))
for node in ast.walk(tree)
if isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstrings
for match in COMMERCIAL_TOKEN.finditer(node.value)
)
def find_hits(scan_root: Path) -> tuple[Hit, ...]:
return tuple(
hit
for path in sorted(scan_root.rglob("*.py"))
if path.relative_to(REPO_ROOT).as_posix() != PARTITION_HELPER
for hit in _hits_in_file(path)
)
def _violation_message(hit: Hit, found: int) -> str:
allowed: Final = ALLOWED_COUNTS.get((hit.file, hit.token))
if allowed is None:
return f"{hit.file}:{hit.line}: hardcoded commercial AWS partition literal {hit.token!r}"
return (
f"{hit.file}:{hit.line}: {hit.token!r} appears {found} times but ALLOWED covers {allowed}; "
"build it from the region helper or update the count"
)
def main() -> int:
hits: Final = find_hits(SCAN_ROOT)
counts: Final = MappingProxyType(Counter((hit.file, hit.token) for hit in hits))
violations: Final = tuple(
sorted(hit for hit in hits if Allowance(hit.file, hit.token, counts[hit.file, hit.token]) not in ALLOWED)
)
stale: Final = tuple(entry for entry in sorted(ALLOWED) if (entry.file, entry.token) not in counts)
for hit in violations:
print(_violation_message(hit, counts[hit.file, hit.token]))
for allowance in stale:
print(f"{allowance.file}: ALLOWED entry {allowance.token!r} no longer matches anything, remove it")
if violations or stale:
print(
"\nBuild AWS hosts with get_aws_dns_suffix(region) and ARNs with get_aws_arn_prefix(region) "
"from litellm/litellm_core_utils/aws_partition.py so GovCloud and China regions resolve."
)
return 1
print(f"No hardcoded commercial AWS partition literals outside {PARTITION_HELPER}.")
return 0
if __name__ == "__main__":
sys.exit(main())

116
tests/test_litellm/litellm_core_utils/test_aws_partition.py
View file

@ -1,11 +1,9 @@
import ast
from pathlib import Path
from types import MappingProxyType
from typing import Final
from urllib.parse import unquote, urlparse
from urllib.parse import urlparse
import pytest
from botocore.credentials import Credentials
import litellm
from litellm.integrations.s3_v2 import S3Logger
@ -22,20 +20,8 @@ from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToS
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
from litellm.llms.bedrock.chat.agentcore.transformation import AmazonAgentCoreConfig
from litellm.llms.bedrock.chat.invoke_agent.transformation import AmazonInvokeAgentConfig
from litellm.llms.bedrock.common_utils import init_bedrock_client
from litellm.llms.bedrock.files.transformation import BedrockFilesConfig
from litellm.llms.bedrock.rerank.handler import BedrockRerankHandler
from litellm.llms.bedrock.vector_stores.transformation import BedrockVectorStoreConfig
from litellm.llms.sagemaker.chat.transformation import SagemakerChatConfig
from litellm.llms.sagemaker.completion.handler import SagemakerLLM
from litellm.proxy.auth.rds_iam_token import init_rds_client
from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2
STATIC_AWS_CREDENTIALS: Final = MappingProxyType(
{"aws_access_key_id": "test-key", "aws_secret_access_key": "test-secret"}
)
@pytest.mark.parametrize(
@ -120,48 +106,6 @@ def _s3_object_url(region: str) -> str:
return logger._build_object_url("2025-01-01/key.json")
def _bedrock_job_arn(region: str) -> str:
return f"{get_aws_arn_prefix(region)}bedrock:{region}:111122223333:model-invocation-job/abc1234567"
def _bedrock_files_upload_url(region: str) -> str:
return BedrockFilesConfig().get_complete_file_url(
api_base=None,
api_key=None,
model="amazon.nova-pro-v1:0",
optional_params={},
litellm_params={"s3_bucket_name": "batch-bucket", "s3_region_name": region},
data={"file": ("batch.jsonl", b"{}", "application/jsonl"), "purpose": "batch"},
)
def _bedrock_files_download_url(region: str) -> str:
return (
BedrockFilesConfig()
._s3_request_target(optional_params={}, litellm_params={"s3_region_name": region})
.endpoint_url
)
def _bedrock_guardrail_url(region: str) -> str:
guardrail = BedrockGuardrail(guardrailIdentifier="guardrail-id", guardrailVersion="1")
return guardrail._prepare_request(
credentials=Credentials("test-key", "test-secret"),
data={"source": "INPUT", "content": []},
optional_params={},
aws_region_name=region,
).url
def _secrets_manager_url(region: str) -> str:
endpoint_url, _headers, _body = AWSSecretsManagerV2(aws_region_name=region)._prepare_request(
action="GetSecretValue",
secret_name="my-secret",
optional_params=dict(STATIC_AWS_CREDENTIALS),
)
return endpoint_url
ENDPOINT_BUILDERS: Final = {
"bedrock_runtime_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("runtime", region),
"bedrock_agent_default": lambda region: BaseAWSLLM()._select_default_endpoint_url("agent", region),
@ -180,13 +124,6 @@ ENDPOINT_BUILDERS: Final = {
litellm_params={},
data={"input_file_id": "s3://bucket/key.jsonl"},
),
"bedrock_batches_retrieve": lambda region: BedrockBatchesConfig().transform_retrieve_batch_request(
batch_id=_bedrock_job_arn(region),
optional_params=dict(STATIC_AWS_CREDENTIALS),
litellm_params={},
)["url"],
"bedrock_files_upload": _bedrock_files_upload_url,
"bedrock_files_download": _bedrock_files_download_url,
"bedrock_agentcore_invoke": lambda region: AmazonAgentCoreConfig().get_complete_url(
api_base=None,
api_key=None,
@ -194,32 +131,6 @@ ENDPOINT_BUILDERS: Final = {
optional_params={},
litellm_params={},
),
"bedrock_invoke_agent": lambda region: AmazonInvokeAgentConfig().get_complete_url(
api_base=None,
api_key=None,
model="agent/AGENT123/ALIAS456",
optional_params={"aws_region_name": region},
litellm_params={},
),
"bedrock_guardrail_apply": _bedrock_guardrail_url,
"bedrock_rerank": lambda region: BedrockRerankHandler()._prepare_request(
model="amazon.rerank-v1:0",
api_base=None,
extra_headers=None,
data={"queries": [], "sources": []},
optional_params={"aws_region_name": region, **STATIC_AWS_CREDENTIALS},
)["endpoint_url"],
"bedrock_knowledgebase_search": lambda region: BedrockVectorStoreConfig().get_complete_url(
api_base=None, litellm_params={"aws_region_name": region}
),
"secrets_manager": _secrets_manager_url,
"rds_iam_client": lambda region: (
init_rds_client(
aws_region_name=region,
aws_access_key_id="test-key",
aws_secret_access_key="test-secret",
).meta.endpoint_url
),
"polly": lambda region: AWSPollyTextToSpeechConfig().get_complete_url(
model="polly/neural",
api_base=None,
@ -241,19 +152,6 @@ ENDPOINT_BUILDERS: Final = {
litellm_params={},
stream=True,
),
"sagemaker_completion": lambda region: (
SagemakerLLM()
._prepare_request(
credentials=Credentials("test-key", "test-secret"),
model="my-endpoint",
data={},
messages=[],
litellm_params={},
optional_params={},
aws_region_name=region,
)
.url
),
"s3_object_url": _s3_object_url,
}
@ -284,18 +182,6 @@ def test_every_endpoint_builder_keeps_amazonaws_com_outside_cn(builder_name: str
assert hostname.endswith(".amazonaws.com"), url
@pytest.mark.parametrize("region", ["us-gov-west-1", "us-gov-east-1"])
@pytest.mark.parametrize("builder_name", sorted(ENDPOINT_BUILDERS))
def test_every_endpoint_builder_respects_us_gov_partition(builder_name: str, region: str) -> None:
url = unquote(ENDPOINT_BUILDERS[builder_name](region))
hostname = urlparse(url).hostname
assert hostname is not None
assert hostname.endswith(f".{region}.amazonaws.com"), url
assert "arn:aws:" not in url, url
if "arn:" in url:
assert "arn:aws-us-gov:" in url, url
def _fstring_literal_offenders(needle: str) -> list[str]:
litellm_root = Path(litellm.__file__).parent
return [