diff --git a/.github/workflows/auto_update_price_and_context_window.yml b/.github/workflows/auto_update_price_and_context_window.yml index 60e89936219..8265e0d09c5 100644 --- a/.github/workflows/auto_update_price_and_context_window.yml +++ b/.github/workflows/auto_update_price_and_context_window.yml @@ -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' diff --git a/.github/workflows/check_duplicate_issues.yml b/.github/workflows/check_duplicate_issues.yml index 289d78880ad..539290bfeaf 100644 --- a/.github/workflows/check_duplicate_issues.yml +++ b/.github/workflows/check_duplicate_issues.yml @@ -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' diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index e86fca17c7a..0e49b3af138 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -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 diff --git a/.github/workflows/codspeed.yml b/.github/workflows/codspeed.yml index 52d64addea9..749242b1cd6 100644 --- a/.github/workflows/codspeed.yml +++ b/.github/workflows/codspeed.yml @@ -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 diff --git a/.github/workflows/create_daily_staging_branch.yml b/.github/workflows/create_daily_staging_branch.yml index 424d8de0a41..0df5f4f92ea 100644 --- a/.github/workflows/create_daily_staging_branch.yml +++ b/.github/workflows/create_daily_staging_branch.yml @@ -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: diff --git a/.github/workflows/helm_unit_test.yml b/.github/workflows/helm_unit_test.yml index 06836b1d1cd..416523f241b 100644 --- a/.github/workflows/helm_unit_test.yml +++ b/.github/workflows/helm_unit_test.yml @@ -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 diff --git a/.github/workflows/issue-keyword-labeler.yml b/.github/workflows/issue-keyword-labeler.yml index 7e2693209b6..59b8fd9cf9f 100644 --- a/.github/workflows/issue-keyword-labeler.yml +++ b/.github/workflows/issue-keyword-labeler.yml @@ -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 diff --git a/.github/workflows/llm-translation-testing.yml b/.github/workflows/llm-translation-testing.yml index 922013c4b54..a6b643dd92e 100644 --- a/.github/workflows/llm-translation-testing.yml +++ b/.github/workflows/llm-translation-testing.yml @@ -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 diff --git a/.github/workflows/read_pyproject_version.yml b/.github/workflows/read_pyproject_version.yml index 04b4a38ce19..a9d16ca4413 100644 --- a/.github/workflows/read_pyproject_version.yml +++ b/.github/workflows/read_pyproject_version.yml @@ -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 diff --git a/.github/workflows/run_observatory_tests.yml b/.github/workflows/run_observatory_tests.yml index a25b96766d7..b0706b7b716 100644 --- a/.github/workflows/run_observatory_tests.yml +++ b/.github/workflows/run_observatory_tests.yml @@ -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 diff --git a/.github/workflows/scan_duplicate_issues.yml b/.github/workflows/scan_duplicate_issues.yml index 222ff11f304..6c88a54554e 100644 --- a/.github/workflows/scan_duplicate_issues.yml +++ b/.github/workflows/scan_duplicate_issues.yml @@ -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 diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index c905bb12312..e7ce48b9fbd 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -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' diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 5bb85716a17..1424088eaa2 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -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 diff --git a/.github/workflows/test-litellm-matrix.yml b/.github/workflows/test-litellm-matrix.yml index dafabe4d83e..37c55fc4bf0 100644 --- a/.github/workflows/test-litellm-matrix.yml +++ b/.github/workflows/test-litellm-matrix.yml @@ -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 diff --git a/.github/workflows/test-litellm-ui-build.yml b/.github/workflows/test-litellm-ui-build.yml index bef568298e0..4ac268e4e7a 100644 --- a/.github/workflows/test-litellm-ui-build.yml +++ b/.github/workflows/test-litellm-ui-build.yml @@ -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 diff --git a/.github/workflows/test-litellm.yml b/.github/workflows/test-litellm.yml index 0c040b3ebe7..57051e9fd6c 100644 --- a/.github/workflows/test-litellm.yml +++ b/.github/workflows/test-litellm.yml @@ -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: | diff --git a/.github/workflows/test-mcp.yml b/.github/workflows/test-mcp.yml index 1b228ab76bb..3b8f2920aba 100644 --- a/.github/workflows/test-mcp.yml +++ b/.github/workflows/test-mcp.yml @@ -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: | diff --git a/.github/workflows/test-model-map.yaml b/.github/workflows/test-model-map.yaml index 429f9e1ce0a..2874a001e5f 100644 --- a/.github/workflows/test-model-map.yaml +++ b/.github/workflows/test-model-map.yaml @@ -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: | diff --git a/.github/workflows/test-proxy-e2e-azure-batches.yml b/.github/workflows/test-proxy-e2e-azure-batches.yml index 7cbbe0b338f..b579b37de9c 100644 --- a/.github/workflows/test-proxy-e2e-azure-batches.yml +++ b/.github/workflows/test-proxy-e2e-azure-batches.yml @@ -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 diff --git a/.github/workflows/test_server_root_path.yml b/.github/workflows/test_server_root_path.yml index 47636ce8e92..647d81849c6 100644 --- a/.github/workflows/test_server_root_path.yml +++ b/.github/workflows/test_server_root_path.yml @@ -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 diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py b/enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py index 8824f4c02de..7353b995d2a 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py @@ -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 \ No newline at end of file + verbose_logger.warning( + f"Disabling callbacks using request headers is an enterprise feature. {CommonProxyErrors.not_premium_user.value}" + ) + return False diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py index 7a77898b160..b6a5918f8db 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py @@ -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." diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py index 8fc2d66d531..2dc158a3cfb 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py @@ -79,4 +79,4 @@ class SendGridEmailLogger(BaseEmailLogger): verbose_logger.debug( f"SendGrid response status={response.status_code}, body={response.text}" ) - return \ No newline at end of file + return diff --git a/enterprise/litellm_enterprise/proxy/auth/__init__.py b/enterprise/litellm_enterprise/proxy/auth/__init__.py index f67826ca7fa..dc70b57ab55 100644 --- a/enterprise/litellm_enterprise/proxy/auth/__init__.py +++ b/enterprise/litellm_enterprise/proxy/auth/__init__.py @@ -7,4 +7,4 @@ including custom SSO handlers and advanced authentication features. from .custom_sso_handler import EnterpriseCustomSSOHandler -__all__ = ["EnterpriseCustomSSOHandler"] \ No newline at end of file +__all__ = ["EnterpriseCustomSSOHandler"] diff --git a/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py b/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py index a3682320387..1c74ca3c49c 100644 --- a/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py +++ b/enterprise/litellm_enterprise/proxy/auth/custom_sso_handler.py @@ -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, - ) \ No newline at end of file + ) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 356f6ecd4b5..e4ae1536865 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -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], diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 54fbc7abcc5..997539c7325 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -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" ) - diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 60c564072a0..7dd982e4e05 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -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}" ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py index 254d816039c..70634537c55 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_vector_stores.py @@ -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,;target_model_names,;resource_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 """ diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/key_management_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/key_management_endpoints.py index 794568b210b..9bb0a4200bf 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/key_management_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/key_management_endpoints.py @@ -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 diff --git a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py index 5e799599862..cf8c38719d7 100644 --- a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py +++ b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py @@ -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 diff --git a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py index 380b0a6facb..d9d5a989abb 100644 --- a/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py +++ b/enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py @@ -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() \ No newline at end of file + return cls().to_dict() diff --git a/litellm/integrations/langfuse/__init__.py b/litellm/integrations/langfuse/__init__.py new file mode 100644 index 00000000000..d48ae2675ab --- /dev/null +++ b/litellm/integrations/langfuse/__init__.py @@ -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"] diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index bea027aa63d..701ff0244b9 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -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, diff --git a/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_init.py b/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_init.py new file mode 100644 index 00000000000..fd43cdf3768 --- /dev/null +++ b/tests/test_litellm/integrations/langfuse/test_langfuse_prompt_init.py @@ -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"