scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

adding oidc to apply gaurdrail

Browse source
This commit is contained in:
shivam 2026-03-20 15:47:27 -07:00
parent 50f88c8642
commit 7acbc1729d
2 changed files with 24 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/constants.py
View file

@ -359,6 +359,13 @@ NON_LLM_CONNECTION_TIMEOUT = int(
MAX_EXCEPTION_MESSAGE_LENGTH = int(os.getenv("MAX_EXCEPTION_MESSAGE_LENGTH", 2000))
MAX_STRING_LENGTH_PROMPT_IN_DB = int(os.getenv("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048))
BEDROCK_MAX_POLICY_SIZE = int(os.getenv("BEDROCK_MAX_POLICY_SIZE", 75))
# Actions allowed in the OIDC session policy when assuming role via AssumeRoleWithWebIdentity.
# Includes model invocation and guardrail apply for Bedrock.
BEDROCK_OIDC_SESSION_POLICY_ACTIONS = (
"bedrock:InvokeModel",
"bedrock:InvokeModelWithResponseStream",
"bedrock:ApplyGuardrail",
)
BEDROCK_MIN_THINKING_BUDGET_TOKENS = int(
os.getenv("BEDROCK_MIN_THINKING_BUDGET_TOKENS", 1024)
)

18
litellm/llms/bedrock/base_aws_llm.py
View file

@ -25,6 +25,7 @@ from litellm.constants import (
BEDROCK_EMBEDDING_PROVIDERS_LITERAL,
BEDROCK_INVOKE_PROVIDERS_LITERAL,
BEDROCK_MAX_POLICY_SIZE,
BEDROCK_OIDC_SESSION_POLICY_ACTIONS,
)
from litellm.litellm_core_utils.dd_tracing import tracer
from litellm.secret_managers.main import get_secret, get_secret_str
@ -695,12 +696,27 @@ class BaseAWSLLM:
# https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html
# https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html
session_policy = {
"Version": "2012-10-17",
"Statement": [
{
"Sid": "BedrockLiteLLM",
"Effect": "Allow",
"Action": list(BEDROCK_OIDC_SESSION_POLICY_ACTIONS),
"Resource": "*",
"Condition": {
"Bool": {"aws:SecureTransport": "true"},
"StringLike": {"aws:UserAgent": "litellm/*"},
},
}
],
}
assume_role_params = {
"RoleArn": aws_role_name,
"RoleSessionName": aws_session_name,
"WebIdentityToken": oidc_token,
"DurationSeconds": 3600,
"Policy": '{"Version":"2012-10-17","Statement":[{"Sid":"BedrockLiteLLM","Effect":"Allow","Action":["bedrock:InvokeModel","bedrock:InvokeModelWithResponseStream"],"Resource":"*","Condition":{"Bool":{"aws:SecureTransport":"true"},"StringLike":{"aws:UserAgent":"litellm/*"}}}]}',
"Policy": json.dumps(session_policy),
}
# Add ExternalId parameter if provided