scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

Merge ade7a64ce4 into 3e90565ca3

Browse source
This commit is contained in:
Thiago Riemma Carbonera 2026-04-04 17:08:14 -07:00 • committed by GitHub
commit 8bd84bdca3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
35 changed files with 512 additions and 412 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
.github/workflows/auto_update_price_and_context_window.yml vendored
View file

@ -5,18 +5,12 @@ on:
- cron: "0 0 * * 0" # Run every Sundays at midnight
#- cron: "0 0 * * *" # Run daily at midnight
permissions:
contents: write
pull-requests: write
jobs:
auto_update_price_and_context_window:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Install Dependencies
run: |
pip install 'aiohttp==3.13.3'

1
.github/workflows/check_duplicate_issues.yml vendored
View file

@ -33,7 +33,6 @@ jobs:
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: .github/scripts
persist-credentials: false
- name: Set up Python
if: github.event.action == 'opened'

6
.github/workflows/codeql.yml vendored
View file

@ -6,8 +6,8 @@ on:
pull_request:
branches: [main]
schedule:
# Run daily at 04:00 UTC
- cron: "0 4 * * *"
# Run weekly on Sundays at 04:00 UTC
- cron: "0 4 * * 0"
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
@ -39,8 +39,6 @@ jobs:
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Initialize CodeQL
uses: github/codeql-action/init@ebcb5b36ded6beda4ceefea6a8bc4cc885255bb3 # v3

2
.github/workflows/codspeed.yml vendored
View file

@ -26,8 +26,6 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0

6
.github/workflows/create_daily_staging_branch.yml vendored
View file

@ -9,15 +9,12 @@ jobs:
create-staging-branch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Create daily staging branch
env:
@ -49,15 +46,12 @@ jobs:
create-internal-dev-branch:
if: github.repository == 'BerriAI/litellm'
runs-on: ubuntu-latest
permissions:
contents: write
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Create internal dev branch
env:

5
.github/workflows/helm_unit_test.yml vendored
View file

@ -6,17 +6,12 @@ on:
branches:
- main
permissions:
contents: read
jobs:
unit-test:
runs-on: ubuntu-latest
steps:
- name: Checkout
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Helm 3.11.1
uses: azure/setup-helm@1a275c3b69536ee54be43f2070a358922e12c8d4 # v4.3.1

2
.github/workflows/issue-keyword-labeler.yml vendored
View file

@ -14,8 +14,6 @@ jobs:
steps:
- name: Checkout code
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Scan for provider keywords
id: scan

15
.github/workflows/llm-translation-testing.yml vendored
View file

@ -11,9 +11,6 @@ on:
tags:
- "v*-rc*" # Triggers on release candidate tags like v1.0.0-rc1
permissions:
contents: read
jobs:
run-llm-translation-tests:
runs-on: ubuntu-latest
@ -23,7 +20,6 @@ jobs:
- name: Checkout code
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
ref: ${{ github.event.inputs.release_candidate_tag || github.ref }}
- name: Set up Python
@ -37,8 +33,8 @@ jobs:
poetry config virtualenvs.create true
poetry config virtualenvs.in-project true
- name: Restore Poetry dependencies cache
uses: actions/cache/restore@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.0.0
- name: Cache Poetry dependencies
uses: actions/cache@0057852bfaa89a56745cba8c7296529d2fc39830 # v4.0.0
with:
path: |
~/.cache/pypoetry
@ -64,12 +60,11 @@ jobs:
AZURE_API_KEY: ${{ secrets.AZURE_API_KEY }}
AZURE_API_BASE: ${{ secrets.AZURE_API_BASE }}
AZURE_API_VERSION: ${{ secrets.AZURE_API_VERSION }}
RC_TAG: ${{ github.event.inputs.release_candidate_tag || github.ref_name }}
COMMIT_SHA: ${{ github.sha }}
# Add other API keys as needed
run: |
python .github/workflows/run_llm_translation_tests.py \
--tag "$RC_TAG" \
--commit "$COMMIT_SHA" \
--tag "${{ github.event.inputs.release_candidate_tag || github.ref_name }}" \
--commit "${{ github.sha }}" \
|| true # Continue even if tests fail
- name: Display test summary

5
.github/workflows/read_pyproject_version.yml vendored
View file

@ -5,9 +5,6 @@ on:
branches:
- main # Change this to the default branch of your repository
permissions:
contents: read
jobs:
read-version:
runs-on: ubuntu-latest
@ -15,8 +12,6 @@ jobs:
steps:
- name: Checkout code
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Read version from pyproject.toml
id: read-version

11
.github/workflows/run_observatory_tests.yml vendored
View file

@ -34,8 +34,6 @@ jobs:
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Validate tag input
env:
@ -51,12 +49,11 @@ jobs:
TAG: ${{ inputs.tag }}
AZURE_API_KEY: ${{ secrets.AZURE_API_KEY }}
AZURE_API_BASE: ${{ secrets.AZURE_API_BASE }}
WORKSPACE: ${{ github.workspace }}
run: |
docker run -d \
--name litellm-rc \
-p 4000:4000 \
-v "${WORKSPACE}/.github/observatory/litellm_config.yaml:/app/config.yaml" \
-v "${{ github.workspace }}/.github/observatory/litellm_config.yaml:/app/config.yaml" \
-e LITELLM_MASTER_KEY="${LITELLM_MASTER_KEY}" \
-e AZURE_API_KEY="${AZURE_API_KEY}" \
-e AZURE_API_BASE="${AZURE_API_BASE}" \
@ -107,11 +104,11 @@ jobs:
- name: Verify tunnel connectivity
run: |
echo "Testing tunnel at ${TUNNEL_URL}..."
echo "Testing tunnel at ${{ env.TUNNEL_URL }}..."
# Quick tunnels need time for DNS propagation; retry to avoid
# transient NXDOMAIN (curl exit code 6) on first attempt.
for i in $(seq 1 10); do
if curl -sf "${TUNNEL_URL}/health/liveliness" > /dev/null 2>&1; then
if curl -sf "${{ env.TUNNEL_URL }}/health/liveliness" > /dev/null 2>&1; then
echo "Tunnel is working (attempt $i)"
exit 0
fi
@ -225,5 +222,5 @@ jobs:
- name: Cleanup
if: always()
run: |
kill "$CLOUDFLARED_PID" 2>/dev/null || true
kill "${{ env.CLOUDFLARED_PID }}" 2>/dev/null || true
docker rm -f litellm-rc 2>/dev/null || true

1
.github/workflows/scan_duplicate_issues.yml vendored
View file

@ -24,7 +24,6 @@ jobs:
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
sparse-checkout: .github/scripts
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0

4
.github/workflows/stale.yml vendored
View file

@ -5,10 +5,6 @@ on:
- cron: "0 0 * * *" # Runs daily at midnight UTC
workflow_dispatch:
permissions:
issues: write
pull-requests: write
jobs:
stale:
if: github.repository == 'BerriAI/litellm'

5
.github/workflows/test-linting.yml vendored
View file

@ -4,9 +4,6 @@ on:
pull_request:
branches: [main]
permissions:
contents: read
jobs:
lint:
runs-on: ubuntu-latest
@ -17,7 +14,6 @@ jobs:
with:
fetch-depth: 0
clean: true
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
@ -91,7 +87,6 @@ jobs:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
fetch-depth: 0
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0

50
.github/workflows/test-litellm-matrix.yml vendored
View file

@ -4,9 +4,6 @@ on:
pull_request:
branches: [main]
permissions:
contents: read
# Cancel in-progress runs for the same PR
concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
@ -121,8 +118,6 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
@ -156,10 +151,7 @@ jobs:
poetry run pip install --force-reinstall --no-deps -e enterprise/
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
poetry run pip install nodejs-wheel-binaries==24.13.1
poetry run prisma generate --schema litellm/proxy/schema.prisma
- name: Run tests - ${{ matrix.test-group.name }}
@ -171,44 +163,4 @@ jobs:
--reruns ${{ matrix.test-group.reruns }} \
--reruns-delay 1 \
--dist=loadscope \
--durations=20 \
--cov=litellm \
--cov-report=xml:coverage-${{ matrix.test-group.name }}.xml \
--cov-config=pyproject.toml
- name: Save coverage report
if: always()
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: coverage-${{ matrix.test-group.name }}
path: coverage-${{ matrix.test-group.name }}.xml
retention-days: 1
upload-coverage:
name: Upload coverage to Codecov
needs: test
if: always()
runs-on: ubuntu-latest
permissions:
contents: read
id-token: write # Required for OIDC tokenless upload
pull-requests: write # Required for Codecov PR comments
steps:
- name: Checkout code
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
- name: Download all coverage reports
uses: actions/download-artifact@95815c38cf2ff2164869cbab79da8d1f422bc89e # v4.2.1
with:
pattern: coverage-*
path: coverage-reports
merge-multiple: true
- name: Upload to Codecov
uses: codecov/codecov-action@aa56896cf108bd10b5eb883cd1d24196da57f695 # v5.5.4
with:
use_oidc: true
directory: coverage-reports
root_dir: ${{ github.workspace }}
fail_ci_if_error: false
--durations=20

4
.github/workflows/test-litellm-ui-build.yml vendored
View file

@ -17,8 +17,6 @@ jobs:
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Setup Node.js
uses: actions/setup-node@a0853c24544627f65ddf259abe73b1d18a591444 # v5.0
@ -28,7 +26,7 @@ jobs:
cache-dependency-path: ui/litellm-dashboard/package-lock.json
- name: Install dependencies
run: npm ci
run: npm install
- name: Build
run: npm run build

5
.github/workflows/test-litellm.yml vendored
View file

@ -8,9 +8,6 @@ on:
# pull_request:
# branches: [ main ]
permissions:
contents: read
jobs:
test:
runs-on: ubuntu-latest
@ -18,8 +15,6 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Thank You Message
run: |

5
.github/workflows/test-mcp.yml vendored
View file

@ -4,9 +4,6 @@ on:
pull_request:
branches: [main]
permissions:
contents: read
jobs:
test:
runs-on: ubuntu-latest
@ -14,8 +11,6 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Thank You Message
run: |

5
.github/workflows/test-model-map.yaml vendored
View file

@ -4,16 +4,11 @@ on:
pull_request:
branches: [main]
permissions:
contents: read
jobs:
validate-model-prices-json:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Validate model_prices_and_context_window.json
run: |

8
.github/workflows/test-proxy-e2e-azure-batches.yml vendored
View file

@ -9,9 +9,6 @@ concurrency:
group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
proxy_e2e_azure_batches_tests:
runs-on: ubuntu-latest
@ -34,8 +31,6 @@ jobs:
steps:
- uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Python
uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
@ -68,10 +63,7 @@ jobs:
poetry run pip install --force-reinstall --no-deps -e enterprise/
- name: Generate Prisma client
env:
PRISMA_BINARY_CACHE_DIR: ${{ runner.temp }}/prisma-cache
run: |
poetry run pip install nodejs-wheel-binaries==24.13.1
poetry run prisma generate --schema litellm/proxy/schema.prisma
- name: Run Prisma migrations

2
.github/workflows/test_server_root_path.yml vendored
View file

@ -18,8 +18,6 @@ jobs:
steps:
- name: Checkout repository
uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0
with:
persist-credentials: false
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3.12

141
enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py
View file

@ -14,53 +14,74 @@ from litellm.types.utils import StandardCallbackDynamicParams
class EnterpriseCallbackControls:
@staticmethod
def is_callback_disabled_dynamically(
callback: litellm.CALLBACK_TYPES,
litellm_params: dict,
standard_callback_dynamic_params: StandardCallbackDynamicParams
) -> bool:
"""
Check if a callback is disabled via the x-litellm-disable-callbacks header or via `litellm_disabled_callbacks` in standard_callback_dynamic_params.
Args:
callback: The callback to check (can be string, CustomLogger instance, or callable)
litellm_params: Parameters containing proxy server request info
Returns:
bool: True if the callback should be disabled, False otherwise
"""
from litellm.litellm_core_utils.custom_logger_registry import (
CustomLoggerRegistry,
)
callback: litellm.CALLBACK_TYPES,
litellm_params: dict,
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> bool:
"""
Check if a callback is disabled via the x-litellm-disable-callbacks header or via `litellm_disabled_callbacks` in standard_callback_dynamic_params.
Args:
callback: The callback to check (can be string, CustomLogger instance, or callable)
litellm_params: Parameters containing proxy server request info
Returns:
bool: True if the callback should be disabled, False otherwise
"""
from litellm.litellm_core_utils.custom_logger_registry import (
CustomLoggerRegistry,
)
try:
disabled_callbacks = EnterpriseCallbackControls.get_disabled_callbacks(
litellm_params, standard_callback_dynamic_params
)
verbose_logger.debug(
f"Dynamically disabled callbacks from {X_LITELLM_DISABLE_CALLBACKS}: {disabled_callbacks}"
)
verbose_logger.debug(
f"Checking if {callback} is disabled via headers. Disable callbacks from headers: {disabled_callbacks}"
)
if disabled_callbacks is not None:
#########################################################
# premium user check
#########################################################
if (
not EnterpriseCallbackControls._should_allow_dynamic_callback_disabling()
):
return False
#########################################################
if isinstance(callback, str):
if callback.lower() in disabled_callbacks:
verbose_logger.debug(
f"Not logging to {callback} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}"
)
return True
elif isinstance(callback, CustomLogger):
# get the string name of the callback
callback_str = (
CustomLoggerRegistry.get_callback_str_from_class_type(
callback.__class__
)
)
if (
callback_str is not None
and callback_str.lower() in disabled_callbacks
):
verbose_logger.debug(
f"Not logging to {callback_str} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}"
)
return True
return False
except Exception as e:
verbose_logger.debug(f"Error checking disabled callbacks header: {str(e)}")
return False
try:
disabled_callbacks = EnterpriseCallbackControls.get_disabled_callbacks(litellm_params, standard_callback_dynamic_params)
verbose_logger.debug(f"Dynamically disabled callbacks from {X_LITELLM_DISABLE_CALLBACKS}: {disabled_callbacks}")
verbose_logger.debug(f"Checking if {callback} is disabled via headers. Disable callbacks from headers: {disabled_callbacks}")
if disabled_callbacks is not None:
#########################################################
# premium user check
#########################################################
if not EnterpriseCallbackControls._should_allow_dynamic_callback_disabling():
return False
#########################################################
if isinstance(callback, str):
if callback.lower() in disabled_callbacks:
verbose_logger.debug(f"Not logging to {callback} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}")
return True
elif isinstance(callback, CustomLogger):
# get the string name of the callback
callback_str = CustomLoggerRegistry.get_callback_str_from_class_type(callback.__class__)
if callback_str is not None and callback_str.lower() in disabled_callbacks:
verbose_logger.debug(f"Not logging to {callback_str} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}")
return True
return False
except Exception as e:
verbose_logger.debug(
f"Error checking disabled callbacks header: {str(e)}"
)
return False
@staticmethod
def get_disabled_callbacks(litellm_params: dict, standard_callback_dynamic_params: StandardCallbackDynamicParams) -> Optional[List[str]]:
def get_disabled_callbacks(
litellm_params: dict,
standard_callback_dynamic_params: StandardCallbackDynamicParams,
) -> Optional[List[str]]:
"""
Get the disabled callbacks from the standard callback dynamic params.
"""
@ -71,18 +92,24 @@ class EnterpriseCallbackControls:
request_headers = get_proxy_server_request_headers(litellm_params)
disabled_callbacks = request_headers.get(X_LITELLM_DISABLE_CALLBACKS, None)
if disabled_callbacks is not None:
disabled_callbacks = set([cb.strip().lower() for cb in disabled_callbacks.split(",")])
disabled_callbacks = set(
[cb.strip().lower() for cb in disabled_callbacks.split(",")]
)
return list(disabled_callbacks)
#########################################################
# check if disabled via request body
#########################################################
if standard_callback_dynamic_params.get("litellm_disabled_callbacks", None) is not None:
return standard_callback_dynamic_params.get("litellm_disabled_callbacks", None)
if (
standard_callback_dynamic_params.get("litellm_disabled_callbacks", None)
is not None
):
return standard_callback_dynamic_params.get(
"litellm_disabled_callbacks", None
)
return None
@staticmethod
def _should_allow_dynamic_callback_disabling():
import litellm
@ -90,10 +117,14 @@ class EnterpriseCallbackControls:
# Check if admin has disabled this feature
if litellm.allow_dynamic_callback_disabling is not True:
verbose_logger.debug("Dynamic callback disabling is disabled by admin via litellm.allow_dynamic_callback_disabling")
verbose_logger.debug(
"Dynamic callback disabling is disabled by admin via litellm.allow_dynamic_callback_disabling"
)
return False
if premium_user:
return True
verbose_logger.warning(f"Disabling callbacks using request headers is an enterprise feature. {CommonProxyErrors.not_premium_user.value}")
return False
verbose_logger.warning(
f"Disabling callbacks using request headers is an enterprise feature. {CommonProxyErrors.not_premium_user.value}"
)
return False

6
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
View file

@ -612,9 +612,9 @@ class BaseEmailLogger(CustomLogger):
else "LiteLLM Notification"
)
recipient_email: Optional[str] = (
user_email or await self._lookup_user_email_from_db(user_id=user_id)
)
recipient_email: Optional[
str
] = user_email or await self._lookup_user_email_from_db(user_id=user_id)
if recipient_email is None:
raise ValueError(
f"User email not found for user_id: {user_id}. User email is required to send email."

2
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py
View file

@ -79,4 +79,4 @@ class SendGridEmailLogger(BaseEmailLogger):
verbose_logger.debug(
f"SendGrid response status={response.status_code}, body={response.text}"
)
return
return

2
enterprise/litellm_enterprise/proxy/auth/__init__.py
View file

@ -7,4 +7,4 @@ including custom SSO handlers and advanced authentication features.
from .custom_sso_handler import EnterpriseCustomSSOHandler
__all__ = ["EnterpriseCustomSSOHandler"]
__all__ = ["EnterpriseCustomSSOHandler"]

37
enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py
View file

@ -26,12 +26,12 @@ from litellm.proxy.management_endpoints.types import CustomOpenID
class EnterpriseCustomSSOHandler:
"""
Enterprise Custom SSO Handler for LiteLLM Proxy
This class provides methods for handling custom SSO authentication flows
where users can implement their own authentication logic by processing
request headers and returning user information in OpenID format.
"""
@staticmethod
async def handle_custom_ui_sso_sign_in(
request: Request,
@ -40,16 +40,16 @@ class EnterpriseCustomSSOHandler:
Allow a user to execute their custom code to parse incoming request headers and return a OpenID object
Use this when you have an OAuth proxy in front of LiteLLM (where the OAuth proxy has already authenticated the user)
Args:
request: The FastAPI request object containing headers and other request data
Returns:
RedirectResponse: Redirect response that sends the user to the LiteLLM UI with authentication token
Raises:
ValueError: If custom_ui_sso_sign_in_handler is not configured
Example:
This method is typically called when a user has already been authenticated by an
external OAuth proxy and the proxy has added custom headers containing user information.
@ -63,24 +63,31 @@ class EnterpriseCustomSSOHandler:
premium_user,
user_custom_ui_sso_sign_in_handler,
)
if premium_user is not True:
raise ValueError(CommonProxyErrors.not_premium_user.value)
if user_custom_ui_sso_sign_in_handler is None:
raise ValueError("custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings.")
custom_sso_login_handler = cast(CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler)
openid_response: OpenID = await custom_sso_login_handler.handle_custom_ui_sso_sign_in(
request=request,
raise ValueError(
"custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings."
)
custom_sso_login_handler = cast(
CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler
)
openid_response: OpenID = (
await custom_sso_login_handler.handle_custom_ui_sso_sign_in(
request=request,
)
)
# Import here to avoid circular imports
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
result=openid_response,
request=request,
received_response=None,
generic_client_id=None,
ui_access_mode=None,
)
)

65
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View file

@ -53,7 +53,9 @@ class CheckBatchCost:
"user_api_key_alias": getattr(user_row, "user_alias", None),
}
except Exception as e:
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
verbose_proxy_logger.error(
f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}"
)
return {}
async def _cleanup_stale_managed_objects(self) -> None:
@ -62,11 +64,22 @@ class CheckBatchCost:
in non-terminal states as 'stale_expired'. These will never complete and
should not be polled.
"""
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
cutoff = datetime.now(timezone.utc) - timedelta(
days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS
)
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={
"file_purpose": "batch",
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
"status": {
"not_in": [
"completed",
"complete",
"failed",
"expired",
"cancelled",
"stale_expired",
]
},
"created_at": {"lt": cutoff},
},
data={"status": "stale_expired"},
@ -161,7 +174,11 @@ class CheckBatchCost:
order={"created_at": "asc"},
)
except Exception as query_err:
if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower():
if (
"batch_processed" not in str(query_err).lower()
and "unknown column" not in str(query_err).lower()
and "does not exist" not in str(query_err).lower()
):
raise
# Permanent schema gap — cache the result so future cycles skip straight to fallback
self._has_batch_processed_column = False
@ -220,10 +237,7 @@ class CheckBatchCost:
continue
## RETRIEVE THE BATCH JOB OUTPUT FILE
if (
response.status == "completed"
and response.output_file_id is not None
):
if response.status == "completed" and response.output_file_id is not None:
verbose_proxy_logger.info(
f"Batch ID: {batch_id} is complete, tracking cost and usage"
)
@ -250,20 +264,25 @@ class CheckBatchCost:
decoded = _is_base64_encoded_unified_file_id(raw_output_file_id)
if decoded:
try:
raw_output_file_id = decoded.split("llm_output_file_id,")[1].split(";")[0]
raw_output_file_id = decoded.split("llm_output_file_id,")[
1
].split(";")[0]
except (IndexError, AttributeError):
pass
credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
credentials = (
self.llm_router.get_deployment_credentials_with_provider(model_id)
or {}
)
_file_content = await afile_content(
file_id=raw_output_file_id,
**credentials,
)
# Access content - handle both direct attribute and method call
if hasattr(_file_content, 'content'):
if hasattr(_file_content, "content"):
content_bytes = _file_content.content # type: ignore[union-attr]
elif hasattr(_file_content, 'read'):
elif hasattr(_file_content, "read"):
content_bytes = await _file_content.read() # type: ignore[misc]
else:
content_bytes = _file_content # type: ignore[assignment]
@ -302,14 +321,20 @@ class CheckBatchCost:
# Pass deployment model_info so custom batch pricing
# (input_cost_per_token_batches etc.) is used for cost calc
deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {}
batch_cost, batch_usage, batch_models = (
await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=llm_provider, # type: ignore
model_name=model_name,
model_info=deployment_model_info, # type: ignore[arg-type]
)
deployment_model_info = (
deployment_info.model_info.model_dump()
if deployment_info.model_info
else {}
)
(
batch_cost,
batch_usage,
batch_models,
) = await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=llm_provider, # type: ignore
model_name=model_name,
model_info=deployment_model_info, # type: ignore[arg-type]
)
logging_obj = LiteLLMLogging(
model=batch_models[0],

42
enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py
View file

@ -38,11 +38,22 @@ class CheckResponsesCost:
in non-terminal states as 'stale_expired'. These will never complete and
should not be polled.
"""
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
cutoff = datetime.now(timezone.utc) - timedelta(
days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS
)
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
where={
"file_purpose": "response",
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
"status": {
"not_in": [
"completed",
"complete",
"failed",
"expired",
"cancelled",
"stale_expired",
]
},
"created_at": {"lt": cutoff},
},
data={"status": "stale_expired"},
@ -76,7 +87,7 @@ class CheckResponsesCost:
take=MAX_OBJECTS_PER_POLL_CYCLE,
order={"created_at": "asc"},
)
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
completed_jobs = []
@ -91,29 +102,35 @@ class CheckResponsesCost:
# Get the stored response object to extract model information
stored_response = job.file_object
model_name = stored_response.get("model", None)
# Decrypt the response ID
responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
(
responses_id_security,
_,
_,
) = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
# Prepare metadata with model information for cost tracking
litellm_metadata = {
"user_api_key_user_id": job.created_by or "default-user-id",
}
# Add model information if available
if model_name:
litellm_metadata["model"] = model_name
litellm_metadata["model_group"] = model_name # Use same value for model_group
litellm_metadata[
"model_group"
] = model_name # Use same value for model_group
response = await litellm.aget_responses(
response_id=responses_id_security,
litellm_metadata=litellm_metadata,
)
verbose_proxy_logger.debug(
f"Response {unified_object_id} status: {response.status}, model: {model_name}"
)
except Exception as e:
verbose_proxy_logger.info(
f"Skipping job {unified_object_id} due to error: {e}"
@ -126,7 +143,7 @@ class CheckResponsesCost:
f"Response {unified_object_id} is complete. Cost automatically tracked by aget_responses."
)
completed_jobs.append(job)
elif response.status in ["failed", "cancelled"]:
verbose_proxy_logger.info(
f"Response {unified_object_id} has status {response.status}, marking as complete"
@ -142,4 +159,3 @@ class CheckResponsesCost:
verbose_proxy_logger.info(
f"Marked {len(completed_jobs)} response jobs as completed"
)

249
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -125,7 +125,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
db_data["storage_backend"] = hidden_params["storage_backend"]
if "storage_url" in hidden_params:
db_data["storage_url"] = hidden_params["storage_url"]
verbose_logger.debug(
f"Storage metadata: storage_backend={db_data.get('storage_backend')}, "
f"storage_url={db_data.get('storage_url')}"
@ -285,28 +285,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
raise Exception(
"Filtering by 'target_model_names' is not supported when using managed batches."
)
where_clause: Dict[str, Any] = {"file_purpose": "batch"}
# Filter by user who created the batch
if user_api_key_dict.user_id:
where_clause["created_by"] = user_api_key_dict.user_id
if after:
where_clause["id"] = {"gt": after}
# Fetch more than needed to allow for post-fetch filtering
fetch_limit = limit or 20
if target_model_names:
# Fetch extra to account for filtering
fetch_limit = max(fetch_limit * 3, 100)
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
where=where_clause,
take=fetch_limit,
order={"created_at": "desc"},
)
batch_objects: List[LiteLLMBatch] = []
for batch in batches:
try:
@ -314,7 +314,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if len(batch_objects) >= (limit or 20):
break
batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object
batch_data = (
json.loads(batch.file_object)
if isinstance(batch.file_object, str)
else batch.file_object
)
batch_obj = LiteLLMBatch(**batch_data)
batch_obj.id = batch.unified_object_id
batch_objects.append(batch_obj)
@ -324,7 +328,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
f"Failed to parse batch object {batch.unified_object_id}: {e}"
)
continue
return {
"object": "list",
"data": batch_objects,
@ -377,11 +381,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"""
Check if the user has access to a list of file IDs.
Only checks managed (unified) file IDs.
Args:
file_ids: List of file IDs to check access for
user_api_key_dict: User API key authentication details
Raises:
HTTPException: If user doesn't have access to any of the files
"""
@ -419,10 +423,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
### HANDLE TRANSFORMATIONS ###
# Check both completion and acompletion call types
is_completion_call = (
call_type == CallTypes.completion.value
call_type == CallTypes.completion.value
or call_type == CallTypes.acompletion.value
)
if is_completion_call:
messages = data.get("messages")
model = data.get("model", "")
@ -431,22 +435,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if file_ids:
# Check user has access to all managed files
await self.check_file_ids_access(file_ids, user_api_key_dict)
# Check if any files are stored in storage backends and need base64 conversion
# This is needed for Vertex AI/Gemini which requires base64 content
is_vertex_ai = model and ("vertex_ai" in model or "gemini" in model.lower())
is_vertex_ai = model and (
"vertex_ai" in model or "gemini" in model.lower()
)
if is_vertex_ai:
await self._convert_storage_files_to_base64(
messages=messages,
file_ids=file_ids,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
model_file_id_mapping = await self.get_model_file_id_mapping(
file_ids, user_api_key_dict.parent_otel_span
)
data["model_file_id_mapping"] = model_file_id_mapping
elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value:
elif (
call_type == CallTypes.aresponses.value
or call_type == CallTypes.responses.value
):
# Handle managed files in responses API input and tools
file_ids = []
@ -611,7 +620,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if model_id is None:
model_id = cast(
Optional[str],
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
kwargs.get("litellm_metadata", {})
.get("model_info", {})
.get("id", None),
)
mapped_file_id: Optional[str] = None
if input_file_id and model_file_id_mapping and model_id:
@ -648,7 +659,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
) -> List[str]:
"""
Gets file ids from responses API input.
The input can be:
- A string (no files)
- A list of input items, where each item can have:
@ -656,32 +667,35 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
- content: a list that can contain items with type: "input_file" and file_id
"""
file_ids: List[str] = []
if isinstance(input, str):
return file_ids
if not isinstance(input, list):
return file_ids
for item in input:
if not isinstance(item, dict):
continue
# Check for direct input_file type
if item.get("type") == "input_file":
file_id = item.get("file_id")
if file_id:
file_ids.append(file_id)
# Check for input_file in content array
content = item.get("content")
if isinstance(content, list):
for content_item in content:
if isinstance(content_item, dict) and content_item.get("type") == "input_file":
if (
isinstance(content_item, dict)
and content_item.get("type") == "input_file"
):
file_id = content_item.get("file_id")
if file_id:
file_ids.append(file_id)
return file_ids
def get_file_ids_from_responses_tools(
@ -689,7 +703,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
) -> List[str]:
"""
Gets file ids from responses API tools parameter.
The tools can contain code_interpreter with container.file_ids:
[
{
@ -699,14 +713,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
]
"""
file_ids: List[str] = []
if not isinstance(tools, list):
return file_ids
for tool in tools:
if not isinstance(tool, dict):
continue
# Check for code_interpreter with container file_ids
if tool.get("type") == "code_interpreter":
container = tool.get("container")
@ -716,7 +730,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
for file_id in container_file_ids:
if isinstance(file_id, str):
file_ids.append(file_id)
return file_ids
def get_vector_store_ids_from_file_search_tools(
@ -1073,16 +1087,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
model_name=resolved_model_name,
)
setattr(response, file_attr, unified_file_id)
# Use llm_router credentials when available. Without credentials,
# Azure and other auth-required providers return 500/401.
file_object = None
try:
# Import module and use getattr for better testability with mocks
import litellm.proxy.proxy_server as proxy_server_module
_llm_router = getattr(proxy_server_module, 'llm_router', None)
_llm_router = getattr(
proxy_server_module, "llm_router", None
)
if _llm_router is not None and model_id:
_creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {}
_creds = (
_llm_router.get_deployment_credentials_with_provider(
model_id
)
or {}
)
file_object = await litellm.afile_retrieve(
file_id=original_file_id,
**_creds,
@ -1099,7 +1121,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
verbose_logger.warning(
f"Failed to retrieve file object for {file_attr}={original_file_id}: {str(e)}. Storing with None and will fetch on-demand."
)
await self.store_unified_file_id(
file_id=unified_file_id,
file_object=file_object,
@ -1199,7 +1221,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
# Case 1 : This is not a managed file
if not stored_file_object:
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
# Case 2: Managed file and the file object exists in the database
# The stored file_object has the raw provider ID. Replace with the unified ID
# so callers see a consistent ID (matching Case 3 which does response.id = file_id).
@ -1217,13 +1239,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
try:
model_id, model_file_id = next(iter(stored_file_object.model_mappings.items()))
credentials = llm_router.get_deployment_credentials_with_provider(model_id) or {}
response = await litellm.afile_retrieve(file_id=model_file_id, **credentials)
model_id, model_file_id = next(
iter(stored_file_object.model_mappings.items())
)
credentials = (
llm_router.get_deployment_credentials_with_provider(model_id) or {}
)
response = await litellm.afile_retrieve(
file_id=model_file_id, **credentials
)
response.id = file_id # Replace with unified ID
return response
except Exception as e:
raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e
raise Exception(
f"Failed to retrieve file {file_id} from provider: {str(e)}"
) from e
async def afile_list(
self,
@ -1245,19 +1275,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
import litellm.proxy.proxy_server as proxy_server_module
# Check if the scheduler has the batch cost checking job registered
scheduler = getattr(proxy_server_module, 'scheduler', None)
scheduler = getattr(proxy_server_module, "scheduler", None)
if scheduler is None:
return False
# Check if the check_batch_cost_job exists in the scheduler
try:
job = scheduler.get_job('check_batch_cost_job')
job = scheduler.get_job("check_batch_cost_job")
if job is not None:
return True
except Exception:
# Job not found or scheduler doesn't support get_job
pass
return False
except Exception as e:
verbose_logger.warning(
@ -1265,28 +1295,26 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
)
return False
async def _get_batches_referencing_file(
self, file_id: str
) -> List[Dict[str, Any]]:
async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, Any]]:
"""
Find batches that reference this file and still need cost tracking.
Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost.
Args:
file_id: The unified file ID to check
Returns:
List of batch objects referencing this file in non-terminal state
(max 10 for error message display)
"""
# Prepare list of file IDs to check (both unified and provider IDs)
file_ids_to_check = [file_id]
# Get model-specific file IDs for this unified file ID if it's a managed file
try:
model_file_id_mapping = await self.get_model_file_id_mapping(
[file_id], litellm_parent_otel_span=None
)
if model_file_id_mapping and file_id in model_file_id_mapping:
# Add all provider file IDs for this unified file
provider_file_ids = list(model_file_id_mapping[file_id].values())
@ -1296,59 +1324,67 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
f"Could not get model file ID mapping for {file_id}: {e}. "
f"Will only check unified file ID."
)
MAX_MATCHES_TO_RETURN = 10
MAX_MATCHES_TO_RETURN = 10
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
where={
"file_purpose": "batch",
"batch_processed": False,
"status": {"not_in": ["failed", "expired", "cancelled"]}
"status": {"not_in": ["failed", "expired", "cancelled"]},
},
take=MAX_MATCHES_TO_RETURN,
order={"created_at": "desc"},
)
referencing_batches = []
for batch in batches:
try:
# Parse the batch file_object to check for file references
batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object
batch_data = (
json.loads(batch.file_object)
if isinstance(batch.file_object, str)
else batch.file_object
)
# Extract file IDs from batch
# Batches typically reference the unified file ID in input_file_id
# Output and error files are generated by the provider
input_file_id = batch_data.get("input_file_id")
output_file_id = batch_data.get("output_file_id")
error_file_id = batch_data.get("error_file_id")
referenced_file_ids = [fid for fid in [input_file_id, output_file_id, error_file_id] if fid]
referenced_file_ids = [
fid for fid in [input_file_id, output_file_id, error_file_id] if fid
]
# Check if any referenced file ID matches the file we're trying to delete
if any(ref_id in file_ids_to_check for ref_id in referenced_file_ids):
referencing_batches.append({
"batch_id": batch.unified_object_id,
"status": batch.status,
"created_at": batch.created_at,
})
referencing_batches.append(
{
"batch_id": batch.unified_object_id,
"status": batch.status,
"created_at": batch.created_at,
}
)
except Exception as e:
verbose_logger.warning(
f"Error parsing batch object {batch.unified_object_id}: {e}"
)
continue
return referencing_batches
async def _check_file_deletion_allowed(self, file_id: str) -> None:
"""
Check if file deletion should be blocked due to batch references.
Blocks deletion if:
1. File is referenced by any batch in non-terminal state, AND
2. Batch polling is configured (user wants cost tracking)
Args:
file_id: The unified file ID to check
Raises:
HTTPException: If file deletion should be blocked
"""
@ -1356,34 +1392,40 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if not self._is_batch_polling_enabled():
# Batch polling not configured, allow deletion
return
# Check if file is referenced by any non-terminal batches
referencing_batches = await self._get_batches_referencing_file(file_id)
if referencing_batches:
# File is referenced by non-terminal batches and polling is enabled
MAX_BATCHES_IN_ERROR = 5 # Limit batches shown in error message for readability
MAX_BATCHES_IN_ERROR = (
5 # Limit batches shown in error message for readability
)
# Show up to MAX_BATCHES_IN_ERROR in the error message
batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR]
batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show]
batch_statuses = [
f"{b['batch_id']}: {b['status']}" for b in batches_to_show
]
# Determine the count message
count_message = f"{len(referencing_batches)}"
if len(referencing_batches) >= 10: # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file
if (
len(referencing_batches) >= 10
): # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file
count_message = "10+"
error_message = (
f"Cannot delete file {file_id}. "
f"The file is referenced by {count_message} batch(es) in non-terminal state"
)
# Add specific batch details if not too many
if len(referencing_batches) <= MAX_BATCHES_IN_ERROR:
error_message += f": {', '.join(batch_statuses)}. "
else:
error_message += f" (showing {MAX_BATCHES_IN_ERROR} most recent): {', '.join(batch_statuses)}. "
error_message += (
f"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. "
f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)."
@ -1406,7 +1448,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
llm_router: Router,
**data: Dict,
) -> OpenAIFileObject:
# Check if file deletion should be blocked due to batch references
await self._check_file_deletion_allowed(file_id)
@ -1419,7 +1460,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
if specific_model_file_id_mapping:
# Remove conflicting keys from data to avoid duplicate keyword arguments
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
filtered_data = {
k: v for k, v in data.items() if k not in ("model", "file_id")
}
for model_id, model_file_id in specific_model_file_id_mapping.items():
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore
@ -1480,7 +1523,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
) -> None:
"""
Convert files stored in storage backends to base64 format for Vertex AI/Gemini.
This method checks if any managed files are stored in storage backends,
downloads them, and converts them to base64 format in the messages.
"""
@ -1488,29 +1531,29 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
for file_id in file_ids:
# Check if this is a base64 encoded unified file ID
decoded_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
if not decoded_unified_file_id:
continue
# Check database for storage backend info
# IMPORTANT: The database stores the base64 encoded unified_file_id (not the decoded version)
# So we query with the original file_id (which is base64 encoded)
db_file = await self.prisma_client.db.litellm_managedfiletable.find_first(
where={"unified_file_id": file_id}
)
if not db_file or not db_file.storage_backend or not db_file.storage_url:
continue
# File is stored in a storage backend, download and convert to base64
try:
from litellm.llms.base_llm.files.storage_backend_factory import (
get_storage_backend,
)
storage_backend_name = db_file.storage_backend
storage_url = db_file.storage_url
# Get storage backend (uses same env vars as callback)
try:
storage_backend = get_storage_backend(storage_backend_name)
@ -1519,18 +1562,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
f"Storage backend '{storage_backend_name}' error for file {file_id}: {str(e)}"
)
continue
file_content = await storage_backend.download_file(storage_url)
# Determine content type from file object
content_type = self._get_content_type_from_file_object(db_file.file_object)
content_type = self._get_content_type_from_file_object(
db_file.file_object
)
# Convert to base64
base64_data = base64.b64encode(file_content).decode("utf-8")
base64_data_uri = f"data:{content_type};base64,{base64_data}"
# Update messages to use base64 instead of file_id
self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type)
self._update_messages_with_base64_data(
messages, file_id, base64_data_uri, content_type
)
except Exception as e:
verbose_logger.exception(
f"Error converting file {file_id} from storage backend to base64: {str(e)}"
@ -1541,21 +1588,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
def _get_content_type_from_file_object(self, file_object: Optional[Any]) -> str:
"""
Determine content type from file object.
Uses the MIME type utility for consistent detection and normalization.
Args:
file_object: The file object from the database (can be dict, JSON string, or None)
Returns:
str: MIME type (defaults to "application/octet-stream" if cannot be determined)
"""
# Use utility function for detection
content_type = get_content_type_from_file_object(file_object)
# Normalize for Gemini/Vertex AI (requires image/jpeg, not image/jpg)
content_type = normalize_mime_type_for_provider(content_type, provider="gemini")
return content_type
def _update_messages_with_base64_data(
@ -1567,7 +1614,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
) -> None:
"""
Update messages to replace file_id with base64 data URI.
Args:
messages: List of messages to update
file_id: The file ID to replace
@ -1582,7 +1629,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
if element.get("type") == "file":
file_element = cast(ChatCompletionFileObject, element)
file_element_file = file_element.get("file", {})
if file_element_file.get("file_id") == file_id:
# Replace file_id with base64 data
file_element_file["file_data"] = base64_data_uri
@ -1590,7 +1637,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
file_element_file["format"] = content_type
# Remove file_id to ensure only file_data is used
file_element_file.pop("file_id", None)
verbose_logger.debug(
f"Converted file {file_id} from storage backend to base64 with format {content_type}"
)

92
enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py
View file

@ -41,7 +41,7 @@ class _PROXY_LiteLLMManagedVectorStores(
):
"""
Managed vector stores with target_model_names support.
This class provides functionality to:
- Create vector stores across multiple models
- Retrieve vector stores by unified ID
@ -77,14 +77,14 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> str:
"""
Generate the format string for the unified vector store ID.
Format:
litellm_proxy:vector_store;unified_id,<uuid>;target_model_names,<models>;resource_id,<vs_id>;model_id,<model_id>
"""
# VectorStoreCreateResponse is a TypedDict, so resource_object is a dictionary
# Extract provider resource ID from the response
provider_resource_id = resource_object.get("id", "")
# Model ID is stored in hidden params if the response object supports it
# For TypedDict responses, we need to check if _hidden_params was added
hidden_params: Dict[str, Any] = {}
@ -109,20 +109,18 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> VectorStoreCreateResponse:
"""
Create a vector store for a specific model.
Args:
llm_router: LiteLLM router instance
model: Model name to create vector store for
request_data: Request data for vector store creation
litellm_parent_otel_span: OpenTelemetry span for tracing
Returns:
VectorStoreCreateResponse from the provider
"""
# Use the router to create the vector store
response = await llm_router.avector_store_create(
model=model, **request_data
)
response = await llm_router.avector_store_create(model=model, **request_data)
return response
# ============================================================================
@ -139,14 +137,14 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> VectorStoreCreateResponse:
"""
Create a vector store across multiple models.
Args:
create_request: Vector store creation request parameters
llm_router: LiteLLM router instance
target_model_names_list: List of target model names
litellm_parent_otel_span: OpenTelemetry span for tracing
user_api_key_dict: User API key authentication details
Returns:
VectorStoreCreateResponse with unified ID
"""
@ -196,7 +194,7 @@ class _PROXY_LiteLLMManagedVectorStores(
# VectorStoreCreateResponse is a TypedDict, so we need to create a new dict with the unified ID
response = responses[0].copy()
response["id"] = unified_id
verbose_logger.info(
f"Successfully created managed vector store with unified ID: {unified_id}"
)
@ -212,13 +210,13 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> Dict[str, Any]:
"""
List vector stores created by a user.
Args:
user_api_key_dict: User API key authentication details
limit: Maximum number of vector stores to return
after: Cursor for pagination
order: Sort order ('asc' or 'desc')
Returns:
Dictionary with list of vector stores and pagination info
"""
@ -238,23 +236,23 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> bool:
"""
Check if user has access to a vector store.
Args:
vector_store_id: The unified vector store ID
user_api_key_dict: User API key authentication details
Returns:
True if user has access, False otherwise
"""
is_unified_id = is_base64_encoded_unified_id(vector_store_id)
if is_unified_id:
# Check access for managed vector store
return await self.can_user_access_unified_resource_id(
vector_store_id,
user_api_key_dict,
)
# Not a managed vector store, allow access
return True
@ -263,24 +261,22 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> bool:
"""
Check if user has access to a managed vector store in request data.
Args:
data: Request data containing vector_store_id
user_api_key_dict: User API key authentication details
Returns:
True if this is a managed vector store and user has access
Raises:
HTTPException: If user doesn't have access
"""
vector_store_id = cast(Optional[str], data.get("vector_store_id"))
is_unified_id = (
is_base64_encoded_unified_id(vector_store_id)
if vector_store_id
else False
is_base64_encoded_unified_id(vector_store_id) if vector_store_id else False
)
if is_unified_id and vector_store_id:
if await self.can_user_access_unified_resource_id(
vector_store_id, user_api_key_dict
@ -291,7 +287,7 @@ class _PROXY_LiteLLMManagedVectorStores(
status_code=403,
detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}",
)
return False
# ============================================================================
@ -307,18 +303,18 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> Union[Exception, str, Dict, None]:
"""
Pre-call hook to handle vector store operations.
This hook intercepts vector store requests and:
- Validates access for managed vector stores
- Transforms unified IDs to provider-specific IDs
- Adds model routing information
Args:
user_api_key_dict: User API key authentication details
cache: Cache instance
data: Request data
call_type: Type of call being made
Returns:
Modified request data or None
"""
@ -330,40 +326,40 @@ class _PROXY_LiteLLMManagedVectorStores(
# Handle vector store search operations
if call_type == "avector_store_search":
vector_store_id = data.get("vector_store_id")
if vector_store_id:
# Check if it's a managed vector store ID
decoded_id = is_base64_encoded_unified_id(vector_store_id)
if decoded_id:
verbose_logger.debug(
f"Processing managed vector store search: {vector_store_id}"
)
# Check access
has_access = await self.can_user_access_unified_resource_id(
vector_store_id, user_api_key_dict
)
if not has_access:
raise HTTPException(
status_code=403,
detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}",
)
# Parse the unified ID to extract components
parsed_id = parse_unified_id(vector_store_id)
if parsed_id:
# Extract the model ID and provider resource ID
model_id = parsed_id.get("model_id")
provider_resource_id = parsed_id.get("provider_resource_id")
target_model_names = parsed_id.get("target_model_names", [])
verbose_logger.debug(
f"Decoded vector store - model_id: {model_id}, provider_resource_id: {provider_resource_id}, target_model_names: {target_model_names}"
)
# Determine which model to use for routing
# Priority: model_id (deployment ID) > first target_model_name
routing_model = None
@ -371,28 +367,28 @@ class _PROXY_LiteLLMManagedVectorStores(
routing_model = model_id
elif target_model_names and len(target_model_names) > 0:
routing_model = target_model_names[0]
# Set the model for routing
if routing_model:
data["model"] = routing_model
verbose_logger.info(
f"Routing vector store search to model: {routing_model}"
)
# Replace the unified ID with the provider-specific ID
if provider_resource_id:
data["vector_store_id"] = provider_resource_id
verbose_logger.debug(
f"Replaced unified ID with provider resource ID: {provider_resource_id}"
)
# Handle vector store retrieve/delete operations
elif call_type in ("avector_store_retrieve", "avector_store_delete"):
await self.check_managed_vector_store_access(data, user_api_key_dict)
# If it's a managed vector store, we'll handle it in the endpoint
# No need to transform here as the endpoint will route to the hook
return data
# ============================================================================
@ -407,15 +403,15 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> Any:
"""
Post-call hook to transform responses.
This hook can be used to transform responses if needed.
For now, it just passes through the response.
Args:
data: Request data
user_api_key_dict: User API key authentication details
response: Response from the provider
Returns:
Potentially modified response
"""
@ -436,21 +432,21 @@ class _PROXY_LiteLLMManagedVectorStores(
) -> List[Dict]:
"""
Filter deployments based on vector store availability.
This is used by the router to select only deployments that have
the vector store available.
Note: This method signature is a compromise between CustomLogger and BaseManagedResource
parent classes which have incompatible signatures. The type: ignore[override] is necessary
due to this multiple inheritance conflict.
Args:
model: Model name
healthy_deployments: List of healthy deployments
messages: Messages (unused for vector stores, required by CustomLogger interface)
request_kwargs: Request kwargs containing vector_store_id and mappings
parent_otel_span: OpenTelemetry span for tracing
Returns:
Filtered list of deployments
"""

1
enterprise/litellm_enterprise/proxy/management_endpoints/key_management_endpoints.py
View file

@ -36,7 +36,6 @@ def apply_enterprise_key_management_params(
data: GenerateKeyRequest,
team_table: Optional[LiteLLM_TeamTable],
) -> GenerateKeyRequest:
data = add_team_member_key_duration(team_table, data)
data = add_team_organization_id(team_table, data)
return data

6
enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py
View file

@ -147,12 +147,12 @@ async def list_vector_stores(
vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db(
prisma_client=prisma_client
)
# Also clean up in-memory registry to remove any deleted vector stores
if litellm.vector_store_registry is not None:
db_vector_store_ids = {
vs.get("vector_store_id")
for vs in vector_stores_from_db
vs.get("vector_store_id")
for vs in vector_stores_from_db
if vs.get("vector_store_id")
}
# Remove any in-memory vector stores that no longer exist in database

12
enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py
View file

@ -39,15 +39,23 @@ class EmailEvent(str, enum.Enum):
soft_budget_crossed = "Soft Budget Crossed"
max_budget_alert = "Max Budget Alert"
class EmailEventSettings(BaseModel):
event: EmailEvent
enabled: bool
class EmailEventSettingsUpdateRequest(BaseModel):
settings: List[EmailEventSettings]
class EmailEventSettingsResponse(BaseModel):
settings: List[EmailEventSettings]
class DefaultEmailSettings(BaseModel):
"""Default settings for email events"""
settings: Dict[EmailEvent, bool] = Field(
default_factory=lambda: {
EmailEvent.virtual_key_created: True, # On by default
@ -57,10 +65,12 @@ class DefaultEmailSettings(BaseModel):
EmailEvent.max_budget_alert: True, # On by default
}
)
def to_dict(self) -> Dict[str, bool]:
"""Convert to dictionary with string keys for storage"""
return {event.value: enabled for event, enabled in self.settings.items()}
@classmethod
def get_defaults(cls) -> Dict[str, bool]:
"""Get the default settings as a dictionary with string keys"""
return cls().to_dict()
return cls().to_dict()

21
litellm/integrations/langfuse/__init__.py Normal file
View file

@ -0,0 +1,21 @@
"""Langfuse integration for LiteLLM Prompt Management."""
from .langfuse_prompt_management import LangfusePromptManagement
def initialize_prompt(litellm_params, prompt_spec):
"""
Initialization function that prompt_registry.py will call.
"""
return LangfusePromptManagement(
langfuse_public_key=getattr(litellm_params, "langfuse_public_key", None),
langfuse_secret=getattr(litellm_params, "langfuse_secret", None),
langfuse_host=getattr(litellm_params, "langfuse_host", None),
)
prompt_initializer_registry = {
"langfuse": initialize_prompt,
}
__all__ = ["initialize_prompt", "prompt_initializer_registry"]

67
litellm/integrations/langfuse/langfuse_prompt_management.py
View file

@ -9,8 +9,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Uni
from packaging.version import Version
from typing_extensions import TypeAlias
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.prompt_management_base import PromptManagementClient
from ..custom_prompt_management import CustomPromptManagement
from litellm.litellm_core_utils.asyncify import run_async_function
from litellm.types.llms.openai import AllMessageValues, ChatCompletionSystemMessage
from litellm.types.prompts.init_prompts import PromptSpec
@ -19,7 +19,6 @@ from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPa
from ...litellm_core_utils.specialty_caches.dynamic_logging_cache import (
DynamicLoggingCache,
)
from ..prompt_management_base import PromptManagementBase
from .langfuse import LangFuseLogger
from .langfuse_handler import LangFuseHandler
@ -108,7 +107,7 @@ def langfuse_client_init(
return client
class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogger):
class LangfusePromptManagement(CustomPromptManagement, LangFuseLogger):
def __init__(
self,
langfuse_public_key=None,
@ -175,6 +174,47 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
optional_params[k] = v
return optional_params
def get_chat_completion_prompt(
self,
model: str,
messages: List[AllMessageValues],
non_default_params: dict,
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_spec: Optional[PromptSpec] = None,
prompt_label: Optional[str] = None,
prompt_version: Optional[int] = None,
ignore_prompt_manager_model: Optional[bool] = False,
ignore_prompt_manager_optional_params: Optional[bool] = False,
tools: Optional[List[Dict]] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Overrides the base method to ensure Langfuse-specific prompt management is triggered.
"""
if prompt_id is None and not model.startswith("langfuse/"):
return model, messages, non_default_params
prompt_template = self.compile_prompt(
prompt_id=prompt_id or model.replace("langfuse/", ""),
prompt_variables=prompt_variables,
client_messages=messages,
dynamic_callback_params=dynamic_callback_params,
prompt_label=prompt_label,
prompt_version=prompt_version,
prompt_spec=prompt_spec,
)
return self.post_compile_prompt_processing(
prompt_template=prompt_template,
messages=messages,
non_default_params=non_default_params,
model=model,
ignore_prompt_manager_model=ignore_prompt_manager_model,
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
)
async def async_get_chat_completion_prompt(
self,
model: str,
@ -212,18 +252,15 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
dynamic_callback_params: StandardCallbackDynamicParams,
) -> bool:
if prompt_id is None:
return False
langfuse_client = langfuse_client_init(
langfuse_public_key=dynamic_callback_params.get("langfuse_public_key"),
langfuse_secret=dynamic_callback_params.get("langfuse_secret"),
langfuse_secret_key=dynamic_callback_params.get("langfuse_secret_key"),
langfuse_host=dynamic_callback_params.get("langfuse_host"),
)
langfuse_prompt_client = self._get_prompt_from_id(
langfuse_prompt_id=prompt_id,
langfuse_client=langfuse_client,
)
return langfuse_prompt_client is not None
return True
if (
prompt_spec is not None
and getattr(prompt_spec, "prompt_id", None) is not None
):
return True
return False
def _compile_prompt_helper(
self,

33
tests/test_litellm/integrations/langfuse/test_langfuse_prompt_init.py Normal file
View file

@ -0,0 +1,33 @@
"""
Tests whether Langfuse is dynamically discovered and can be initialized without real network calls.
"""
from unittest.mock import patch
from litellm.proxy.prompts.prompt_registry import get_prompt_initializer_from_integrations
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
def test_langfuse_discovery_and_init():
"""
Validates the fix in __init__.py by ensuring Langfuse is registered and initialized (mocked).
"""
# 1. Verify if the integration was registered in the global dictionary
registry = get_prompt_initializer_from_integrations()
assert "langfuse" in registry, "Error: Langfuse wasn't discovered by the prompt_registry!"
# 2. Test initialization using MOCK to avoid real network calls
# Patch the function that creates the real Langfuse client
with patch("litellm.integrations.langfuse.langfuse_prompt_management.langfuse_client_init") as mocked_client:
mocked_client.return_value = None # No real client needed for this test
init_func = registry["langfuse"]
params = PromptLiteLLMParams(
prompt_integration="langfuse",
langfuse_public_key="test-key",
langfuse_secret="test-secret"
)
spec = PromptSpec(prompt_id="test-id", litellm_params=params)
# Initialize the class (this would call the network, but is now mocked)
obj = init_func(params, spec)
assert obj is not None, "Initiation function returned None!"
assert obj.integration_name == "langfuse"