diff --git a/.circleci/config.yml b/.circleci/config.yml index e30dc02b2ab..867accaf05a 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1980,6 +1980,7 @@ jobs: - run: python ./tests/code_coverage_tests/check_unsafe_enterprise_import.py - run: python ./tests/code_coverage_tests/ban_copy_deepcopy_kwargs.py - run: python ./tests/code_coverage_tests/check_fastuuid_usage.py + - run: python ./tests/code_coverage_tests/memory_test.py - run: helm lint ./deploy/charts/litellm-helm db_migration_disable_update_check: @@ -2008,10 +2009,13 @@ jobs: pip install "pytest-asyncio==0.21.1" pip install aiohttp pip install apscheduler + - attach_workspace: + at: ~/project - run: - name: Build Docker image + name: Load Docker Database Image command: | - docker build -t myapp . -f ./docker/Dockerfile.database + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database - run: name: Run Docker container command: | @@ -2024,7 +2028,7 @@ jobs: -v $(pwd)/litellm/proxy/example_config_yaml/bad_schema.prisma:/app/litellm/proxy/schema.prisma \ -v $(pwd)/litellm/proxy/example_config_yaml/disable_schema_update.yaml:/app/config.yaml \ --name my-app \ - myapp:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 - run: @@ -2276,9 +2280,13 @@ jobs: - run: name: Wait for PostgreSQL to be ready command: dockerize -wait tcp://localhost:5432 -timeout 1m + - attach_workspace: + at: ~/project - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + name: Load Docker Database Image + command: | + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database - run: name: Run Docker container command: | @@ -2313,7 +2321,7 @@ jobs: --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/oai_misc_config.yaml:/app/config.yaml \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 \ --detailed_debug \ @@ -2416,9 +2424,13 @@ jobs: - run: name: Wait for PostgreSQL to be ready command: dockerize -wait tcp://localhost:5432 -timeout 1m + - attach_workspace: + at: ~/project - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + name: Load Docker Database Image + command: | + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database - run: name: Run Docker container # intentionally give bad redis credentials here @@ -2451,7 +2463,7 @@ jobs: --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/otel_test_config.yaml:/app/config.yaml \ -v $(pwd)/litellm/proxy/example_config_yaml/custom_guardrail.py:/app/custom_guardrail.py \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 \ --detailed_debug \ @@ -2502,7 +2514,7 @@ jobs: --add-host host.docker.internal:host-gateway \ --name my-app-3 \ -v $(pwd)/litellm/proxy/example_config_yaml/enterprise_config.yaml:/app/config.yaml \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 \ --detailed_debug @@ -2577,9 +2589,13 @@ jobs: - run: name: Wait for PostgreSQL to be ready command: dockerize -wait tcp://localhost:5432 -timeout 1m + - attach_workspace: + at: ~/project - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + name: Load Docker Database Image + command: | + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database - run: name: Run Docker container # intentionally give bad redis credentials here @@ -2603,7 +2619,7 @@ jobs: --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/spend_tracking_config.yaml:/app/config.yaml \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 \ --detailed_debug \ @@ -2690,9 +2706,13 @@ jobs: - run: name: Wait for PostgreSQL to be ready command: dockerize -wait tcp://localhost:5432 -timeout 1m + - attach_workspace: + at: ~/project - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + name: Load Docker Database Image + command: | + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database - run: name: Run Docker container 1 # intentionally give bad redis credentials here @@ -2712,7 +2732,7 @@ jobs: --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml:/app/config.yaml \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 \ --detailed_debug \ @@ -2733,7 +2753,7 @@ jobs: --add-host host.docker.internal:host-gateway \ --name my-app-2 \ -v $(pwd)/litellm/proxy/example_config_yaml/multi_instance_simple_config.yaml:/app/config.yaml \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4001 \ --detailed_debug @@ -2826,9 +2846,13 @@ jobs: - run: name: Wait for PostgreSQL to be ready command: dockerize -wait tcp://localhost:5432 -timeout 1m + - attach_workspace: + at: ~/project - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + name: Load Docker Database Image + command: | + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database - run: name: Run Docker container # intentionally give bad redis credentials here @@ -2843,7 +2867,7 @@ jobs: --add-host host.docker.internal:host-gateway \ --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 \ --detailed_debug \ @@ -3058,10 +3082,13 @@ jobs: - run: name: Wait for PostgreSQL to be ready command: dockerize -wait tcp://localhost:5432 -timeout 1m - # Run pytest and generate JUnit XML report + - attach_workspace: + at: ~/project - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . + name: Load Docker Database Image + command: | + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database - run: name: Run Docker container command: | @@ -3083,7 +3110,7 @@ jobs: --name my-app \ -v $(pwd)/litellm/proxy/example_config_yaml/pass_through_config.yaml:/app/config.yaml \ -v $(pwd)/litellm/proxy/example_config_yaml/custom_auth_basic.py:/app/custom_auth_basic.py \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 \ --detailed_debug \ @@ -3421,6 +3448,37 @@ jobs: --coverage.reporter=html \ --coverage.reportsDirectory=coverage/html + build_docker_database_image: + machine: + image: ubuntu-2204:2023.10.1 + resource_class: xlarge + working_directory: ~/project + steps: + - checkout + + - run: + name: Upgrade Docker + command: | + curl -fsSL https://get.docker.com | sh + docker version + + - run: + name: Build Docker image + command: | + docker build \ + -t litellm-docker-database:ci \ + -f docker/Dockerfile.database . + + - run: + name: Save Docker image to workspace root + command: | + docker save litellm-docker-database:ci | gzip > litellm-docker-database.tar.gz + + - persist_to_workspace: + root: . + paths: + - litellm-docker-database.tar.gz + e2e_ui_testing: machine: image: ubuntu-2204:2023.10.1 @@ -3432,54 +3490,18 @@ jobs: - attach_workspace: at: ~/project - run: - name: Upgrade Docker to v24.x (API 1.44+) + name: Load Docker Database Image command: | - curl -fsSL https://get.docker.com | sh - sudo usermod -aG docker $USER - docker version - - run: - name: Install Python 3.9 - command: | - curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh - bash miniconda.sh -b -p $HOME/miniconda - export PATH="$HOME/miniconda/bin:$PATH" - conda init bash - source ~/.bashrc - conda create -n myenv python=3.9 -y - conda activate myenv - python --version + gunzip -c litellm-docker-database.tar.gz | docker load + docker images | grep litellm-docker-database - run: name: Install Dependencies command: | npm install -D @playwright/test - npm install @google-cloud/vertexai - pip install "pytest==7.3.1" - pip install "pytest-retry==1.6.3" - pip install "pytest-asyncio==0.21.1" - pip install aiohttp - pip install "openai==1.100.1" - python -m pip install --upgrade pip - pip install "pydantic==2.10.2" - pip install "pytest==7.3.1" - pip install "pytest-mock==3.12.0" - pip install "pytest-asyncio==0.21.1" - pip install "mypy==1.18.2" - pip install pyarrow - pip install numpydoc - pip install prisma - pip install fastapi - pip install jsonschema - pip install "httpx==0.24.1" - pip install "anyio==3.7.1" - pip install "asyncio==3.4.3" - run: name: Install Playwright Browsers command: | npx playwright install - - - run: - name: Build Docker image - command: docker build -t my-app:latest -f ./docker/Dockerfile.database . - run: name: Run Docker container command: | @@ -3491,9 +3513,9 @@ jobs: -e UI_USERNAME="admin" \ -e UI_PASSWORD="gm" \ -e LITELLM_LICENSE=$LITELLM_LICENSE \ - --name my-app \ + --name litellm-docker-database \ -v $(pwd)/litellm/proxy/example_config_yaml/simple_config.yaml:/app/config.yaml \ - my-app:latest \ + litellm-docker-database:ci \ --config /app/config.yaml \ --port 4000 \ --detailed_debug @@ -3507,7 +3529,7 @@ jobs: sudo rm dockerize-linux-amd64-v0.6.1.tar.gz - run: name: Start outputting logs - command: docker logs -f my-app + command: docker logs -f litellm-docker-database background: true - run: name: Wait for app to be ready @@ -3515,7 +3537,10 @@ jobs: - run: name: Run Playwright Tests command: | - npx playwright test e2e_ui_tests/ --reporter=html --output=test-results + npx playwright test \ + --config ui/litellm-dashboard/e2e_tests/playwright.config.ts \ + --reporter=html \ + --output=test-results no_output_timeout: 120m - store_artifacts: path: test-results @@ -3705,9 +3730,16 @@ workflows: only: - main - /litellm_.*/ + - build_docker_database_image: + filters: + branches: + only: + - main + - /litellm_.*/ - e2e_ui_testing: requires: - ui_build + - build_docker_database_image filters: branches: only: @@ -3720,30 +3752,40 @@ workflows: - main - /litellm_.*/ - e2e_openai_endpoints: + requires: + - build_docker_database_image filters: branches: only: - main - /litellm_.*/ - proxy_logging_guardrails_model_info_tests: + requires: + - build_docker_database_image filters: branches: only: - main - /litellm_.*/ - proxy_spend_accuracy_tests: + requires: + - build_docker_database_image filters: branches: only: - main - /litellm_.*/ - proxy_multi_instance_tests: + requires: + - build_docker_database_image filters: branches: only: - main - /litellm_.*/ - proxy_store_model_in_db_tests: + requires: + - build_docker_database_image filters: branches: only: @@ -3756,6 +3798,8 @@ workflows: - main - /litellm_.*/ - proxy_pass_through_endpoint_tests: + requires: + - build_docker_database_image filters: branches: only: @@ -3894,6 +3938,8 @@ workflows: - litellm_assistants_api_testing - auth_ui_unit_tests - db_migration_disable_update_check: + requires: + - build_docker_database_image filters: branches: only: @@ -3973,4 +4019,4 @@ workflows: - proxy_pass_through_endpoint_tests - check_code_and_doc_quality - publish_proxy_extras - - guardrails_testing \ No newline at end of file + - guardrails_testing diff --git a/.gitguardian.yaml b/.gitguardian.yaml index af8f2489eec..1eeec0677af 100644 --- a/.gitguardian.yaml +++ b/.gitguardian.yaml @@ -84,6 +84,10 @@ secret: - name: Langfuse test credentials in test_completion match: c39310f68cc3d3e22f7b298bb6353c4f45759adcc37080d8b7f4e535d3cfd7f4 + # Test password "sk-1234" in e2e test fixtures - test fixture, not a real secret + - name: Test password in e2e test fixtures + match: ce32b547202e209ec1dd50107b64be4cfcf2eb15c3b4f8e9dc611ef747af634f + # === Preventive patterns for test keys (pattern-based) === # Test API keys (124 instances across 45 files) @@ -102,3 +106,6 @@ secret: - name: Test API key patterns match: test-api-key + - name: Short fake sk keys (1–9 digits only) + match: \bsk-\d{1,9}\b + diff --git a/.gitignore b/.gitignore index 8196d1d9f24..fafacd874a0 100644 --- a/.gitignore +++ b/.gitignore @@ -101,3 +101,7 @@ tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py litellm/proxy/_experimental/out/guardrails/index.html scripts/test_vertex_ai_search.py LAZY_LOADING_IMPROVEMENTS.md +**/test-results +**/playwright-report +**/*.storageState.json +**/coverage \ No newline at end of file diff --git a/ci_cd/security_scans.sh b/ci_cd/security_scans.sh index 0036a304417..276f5abe330 100755 --- a/ci_cd/security_scans.sh +++ b/ci_cd/security_scans.sh @@ -128,6 +128,7 @@ run_grype_scans() { "GHSA-5j98-mcp5-4vw2" "CVE-2025-13836" # Python 3.13 HTTP response reading OOM/DoS - no fix available in base image "CVE-2025-12084" # Python 3.13 xml.dom.minidom quadratic algorithm - no fix available in base image + "CVE-2025-60876" # BusyBox wget HTTP request splitting - no fix available in Chainguard Wolfi base image ) # Build JSON array of allowlisted CVE IDs for jq diff --git a/deploy/charts/litellm-helm/templates/deployment.yaml b/deploy/charts/litellm-helm/templates/deployment.yaml index 0dab2ec40e0..19fa0479091 100644 --- a/deploy/charts/litellm-helm/templates/deployment.yaml +++ b/deploy/charts/litellm-helm/templates/deployment.yaml @@ -182,6 +182,10 @@ spec: {{- with .Values.volumeMounts }} {{- toYaml . | nindent 12 }} {{- end }} + {{- with .Values.lifecycle }} + lifecycle: + {{- toYaml . | nindent 12 }} + {{- end }} {{- with .Values.extraContainers }} {{- toYaml . | nindent 8 }} {{- end }} diff --git a/deploy/charts/litellm-helm/tests/deployment_tests.yaml b/deploy/charts/litellm-helm/tests/deployment_tests.yaml index f9c83966696..182a2362392 100644 --- a/deploy/charts/litellm-helm/tests/deployment_tests.yaml +++ b/deploy/charts/litellm-helm/tests/deployment_tests.yaml @@ -136,4 +136,26 @@ tests: path: spec.template.spec.containers[0].volumeMounts content: name: litellm-config - mountPath: /etc/litellm/ \ No newline at end of file + mountPath: /etc/litellm/ + - it: should work with lifecycle hooks + template: deployment.yaml + set: + lifecycle: + preStop: + exec: + command: + - /bin/sh + - -c + - echo "Container stopping" + asserts: + - exists: + path: spec.template.spec.containers[0].lifecycle + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[0] + value: /bin/sh + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[1] + value: -c + - equal: + path: spec.template.spec.containers[0].lifecycle.preStop.exec.command[2] + value: echo "Container stopping" \ No newline at end of file diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index a70e3d24188..8c11b8fd654 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -110,6 +110,22 @@ For stdio MCP servers, select "Standard Input/Output (stdio)" as the transport t

+### OAuth Configuration & Overrides + +LiteLLM attempts [OAuth 2.0 Authorization Server Discovery](https://datatracker.ietf.org/doc/html/rfc8414) by default. When you create an MCP server in the UI and set `Authentication: OAuth`, LiteLLM will locate the provider metadata, dynamically register a client, and perform PKCE-based authorization without you providing any additional details. + +**Customize the OAuth flow when needed:** + + + +- **Provide explicit client credentials** – If the MCP provider does not offer dynamic client registration or you prefer to manage the client yourself, fill in `client_id`, `client_secret`, and the desired `scopes`. +- **Override discovery URLs** – In some environments, LiteLLM might not be able to reach the provider's metadata endpoints. Use the optional `authorization_url`, `token_url`, and `registration_url` fields to point LiteLLM directly to the correct endpoints. + +
+ ### Static Headers Sometimes your MCP server needs specific headers on every request. Maybe it's an API key, maybe it's a custom header the server expects. Instead of configuring auth, you can just set them directly. diff --git a/docs/my-website/docs/observability/generic_api.md b/docs/my-website/docs/observability/generic_api.md index 2d1a24c317b..93a0762591a 100644 --- a/docs/my-website/docs/observability/generic_api.md +++ b/docs/my-website/docs/observability/generic_api.md @@ -47,6 +47,7 @@ callback_settings: | `endpoint` | string | Yes | HTTP endpoint to send logs to | | `headers` | dict | No | Custom headers for the request | | `event_types` | list | No | Filter events: `llm_api_success`, `llm_api_failure`. Defaults to all events. | +| `log_format` | string | No | Output format: `json_array` (default), `ndjson`, or `single`. Controls how logs are batched and sent. | ## Pre-configured Callbacks @@ -107,4 +108,62 @@ callback_settings: flush_interval: 60 # seconds, default: 60 ``` +## Log Format Options + +Control how logs are formatted and sent to your endpoint. + +### JSON Array (Default) + +```yaml +callback_settings: + my_api: + callback_type: generic_api + endpoint: https://your-endpoint.com + log_format: json_array # default if not specified +``` + +Sends all logs in a batch as a single JSON array `[{log1}, {log2}, ...]`. This is the default behavior and maintains backward compatibility. + +**When to use**: Most HTTP endpoints expecting batched JSON data. + +### NDJSON (Newline-Delimited JSON) + +```yaml +callback_settings: + my_api: + callback_type: generic_api + endpoint: https://your-endpoint.com + log_format: ndjson +``` + +Sends logs as newline-delimited JSON (one record per line): +``` +{log1} +{log2} +{log3} +``` + +**When to use**: Log aggregation services like Sumo Logic, Splunk, or Datadog that support field extraction on individual records. + +**Benefits**: +- Each log is ingested as a separate message +- Field Extraction Rules work at ingest time +- Better parsing and querying performance + +### Single + +```yaml +callback_settings: + my_api: + callback_type: generic_api + endpoint: https://your-endpoint.com + log_format: single +``` + +Sends each log as an individual HTTP request in parallel when the batch is flushed. + +**When to use**: Endpoints that expect individual records, or when you need maximum compatibility. + +**Note**: This mode sends N HTTP requests per batch (more overhead). Consider using `ndjson` instead if your endpoint supports it. + diff --git a/docs/my-website/docs/observability/levo_integration.md b/docs/my-website/docs/observability/levo_integration.md new file mode 100644 index 00000000000..3e46cf6b921 --- /dev/null +++ b/docs/my-website/docs/observability/levo_integration.md @@ -0,0 +1,162 @@ +--- +sidebar_label: Levo AI +--- + +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Levo AI + +
+
+ +
+
+ +
+
+ +[Levo](https://levo.ai/) is an AI observability and compliance platform that provides comprehensive monitoring, analysis, and compliance tracking for LLM applications. + +## Quick Start + +Send all your LLM requests and responses to Levo for monitoring and analysis using LiteLLM's built-in Levo integration. + +### What You'll Get + +- **Complete visibility** into all LLM API calls across all providers +- **Request and response data** including prompts, completions, and metadata +- **Usage and cost tracking** with token counts and cost breakdowns +- **Error monitoring** and performance metrics +- **Compliance tracking** for audit and governance + +### Setup Steps + +**1. Install OpenTelemetry dependencies:** + +```bash +pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-exporter-otlp-proto-grpc +``` + +**2. Enable Levo callback in your LiteLLM config:** + +Add to your `litellm_config.yaml`: + +```yaml +litellm_settings: + callbacks: ["levo"] +``` + +**3. Configure environment variables:** + +[Contact Levo support](mailto:support@levo.ai) to get your collector endpoint URL, API key, organization ID, and workspace ID. + +Set these required environment variables: + +```bash +export LEVOAI_API_KEY="" +export LEVOAI_ORG_ID="" +export LEVOAI_WORKSPACE_ID="" +export LEVOAI_COLLECTOR_URL="" +``` + +**Note:** The collector URL should be the full endpoint URL provided by Levo support. It will be used exactly as provided. + +**4. Start LiteLLM:** + +```bash +litellm --config config.yaml +``` + +**5. Make requests - they'll automatically be sent to Levo!** + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --data '{ + "model": "gpt-3.5-turbo", + "messages": [ + { + "role": "user", + "content": "Hello, this is a test message" + } + ] + }' +``` + +## What Data is Captured + +| Feature | Details | +|---------|---------| +| **What is logged** | OpenTelemetry Trace Data (OTLP format) | +| **Events** | Success + Failure | +| **Format** | OTLP (OpenTelemetry Protocol) | +| **Headers** | Automatically includes `Authorization: Bearer {LEVOAI_API_KEY}`, `x-levo-organization-id`, and `x-levo-workspace-id` | + +## Configuration Reference + +### Required Environment Variables + +| Variable | Description | Example | +|----------|-------------|---------| +| `LEVOAI_API_KEY` | Your Levo API key | `levo_abc123...` | +| `LEVOAI_ORG_ID` | Your Levo organization ID | `org-123456` | +| `LEVOAI_WORKSPACE_ID` | Your Levo workspace ID | `workspace-789` | +| `LEVOAI_COLLECTOR_URL` | Full collector endpoint URL from Levo support | `https://collector.levo.ai/v1/traces` | + +### Optional Environment Variables + +| Variable | Description | Default | +|----------|-------------|---------| +| `LEVOAI_ENV_NAME` | Environment name for tagging traces | `None` | + +**Note:** The collector URL is used exactly as provided by Levo support. No path manipulation is performed. + +## Troubleshooting + +### Not seeing traces in Levo? + +1. **Verify Levo callback is enabled**: Check LiteLLM startup logs for `initializing callbacks=['levo']` + +2. **Check required environment variables**: Ensure all required variables are set: + ```bash + echo $LEVOAI_API_KEY + echo $LEVOAI_ORG_ID + echo $LEVOAI_WORKSPACE_ID + echo $LEVOAI_COLLECTOR_URL + ``` + +3. **Verify collector connectivity**: Test if your collector is reachable: + ```bash + curl /health + ``` + +4. **Check for initialization errors**: Look for errors in LiteLLM startup logs. Common issues: + - Missing OpenTelemetry packages: Install with `pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-exporter-otlp-proto-grpc` + - Missing required environment variables: All four required variables must be set + - Invalid collector URL: Ensure the URL is correct and reachable + +5. **Enable debug logging**: + ```bash + export LITELLM_LOG="DEBUG" + ``` + +6. **Wait for async export**: OTLP sends traces asynchronously. Wait 10-15 seconds after making requests before checking Levo. + +### Common Errors + +**Error: "LEVOAI_COLLECTOR_URL environment variable is required"** +- Solution: Set the `LEVOAI_COLLECTOR_URL` environment variable with your collector endpoint URL from Levo support. + +**Error: "No module named 'opentelemetry'"** +- Solution: Install OpenTelemetry packages: `pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-exporter-otlp-proto-grpc` + +## Additional Resources + +- [Levo Documentation](https://docs.levo.ai) +- [OpenTelemetry Specification](https://opentelemetry.io/docs/specs/otel/) + +## Need Help? + +For issues or questions about the Levo integration with LiteLLM, please [contact Levo support](mailto:support@levo.ai) or open an issue on the [LiteLLM GitHub repository](https://github.com/BerriAI/litellm/issues). diff --git a/docs/my-website/docs/observability/opentelemetry_integration.md b/docs/my-website/docs/observability/opentelemetry_integration.md index 2b3cf1313ba..b6eff231620 100644 --- a/docs/my-website/docs/observability/opentelemetry_integration.md +++ b/docs/my-website/docs/observability/opentelemetry_integration.md @@ -4,7 +4,7 @@ import TabItem from '@theme/TabItem'; # OpenTelemetry - Tracing LLMs with any observability tool -OpenTelemetry is a CNCF standard for observability. It connects to any observability tool, such as Jaeger, Zipkin, Datadog, New Relic, Traceloop and others. +OpenTelemetry is a CNCF standard for observability. It connects to any observability tool, such as Jaeger, Zipkin, Datadog, New Relic, Traceloop, Levo AI and others. @@ -12,7 +12,9 @@ OpenTelemetry is a CNCF standard for observability. It connects to any observabi From v1.81.0, the request/response will be set as attributes on the parent "Received Proxy Server Request" span by default. This allows you to see the request/response in the parent span in your observability tool. -To use the older behavior with nested "litellm_request" spans, set the following environment variable: +**Note:** When making multiple LLM calls within an external OTEL span context, the last call's attributes will overwrite previous calls' attributes on the parent span. + +To use the older behavior with nested "litellm_request" spans (which creates separate spans for each call), set the following environment variable: ```shell USE_OTEL_LITELLM_REQUEST_SPAN=true diff --git a/docs/my-website/docs/observability/sumologic_integration.md b/docs/my-website/docs/observability/sumologic_integration.md index d0894146e4c..c30ee94dad4 100644 --- a/docs/my-website/docs/observability/sumologic_integration.md +++ b/docs/my-website/docs/observability/sumologic_integration.md @@ -148,6 +148,51 @@ Example payload: ## Advanced Configuration +### Log Format + +The Sumo Logic integration uses **NDJSON (newline-delimited JSON)** format by default. This format is optimal for Sumo Logic's parsing capabilities and allows Field Extraction Rules to work at ingest time. + +#### NDJSON Format + +Each log entry is sent as a separate line in the HTTP request: +``` +{"id":"chatcmpl-1","model":"gpt-3.5-turbo","response_cost":0.0001,...} +{"id":"chatcmpl-2","model":"gpt-4","response_cost":0.0003,...} +{"id":"chatcmpl-3","model":"gpt-3.5-turbo","response_cost":0.0001,...} +``` + +#### Benefits for Field Extraction Rules (FERs) + +With NDJSON format, you can create Field Extraction Rules directly: + +``` +_sourceCategory=litellm/logs +| json field=_raw "model", "response_cost", "user" as model, cost, user +``` + +**Before NDJSON** (with JSON array format): +- Required `parse regex ... multi` workaround +- FERs couldn't parse at ingest time +- Query-time parsing impacted dashboard performance + +**After NDJSON**: +- ✅ FERs parse fields at ingest time +- ✅ No query-time workarounds needed +- ✅ Better dashboard performance +- ✅ Simpler query syntax + +#### Changing the Log Format (Advanced) + +If you need to change the log format (not recommended for Sumo Logic): + +```yaml +callback_settings: + sumologic: + callback_type: generic_api + callback_name: sumologic + log_format: json_array # Override to use JSON array instead +``` + ### Batching Settings Control how LiteLLM batches logs before sending to Sumo Logic: diff --git a/docs/my-website/docs/providers/anthropic.md b/docs/my-website/docs/providers/anthropic.md index bcfb698a0f8..cae8657f1a0 100644 --- a/docs/my-website/docs/providers/anthropic.md +++ b/docs/my-website/docs/providers/anthropic.md @@ -444,7 +444,7 @@ Here's what a sample Raw Request from LiteLLM for Anthropic Context Caching look POST Request Sent from LiteLLM: curl -X POST \ https://api.anthropic.com/v1/messages \ --H 'accept: application/json' -H 'anthropic-version: 2023-06-01' -H 'content-type: application/json' -H 'x-api-key: sk-...' -H 'anthropic-beta: prompt-caching-2024-07-31' \ +-H 'accept: application/json' -H 'anthropic-version: 2023-06-01' -H 'content-type: application/json' -H 'x-api-key: sk-...' \ -d '{'model': 'claude-3-5-sonnet-20240620', [ { "role": "user", @@ -472,6 +472,8 @@ https://api.anthropic.com/v1/messages \ "max_tokens": 10 }' ``` + +**Note:** Anthropic no longer requires the `anthropic-beta: prompt-caching-2024-07-31` header. Prompt caching now works automatically when you use `cache_control` in your messages. ::: ### Caching - Large Context Caching diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index 122554fe8a4..f1eed4b4d52 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -2208,6 +2208,53 @@ response = completion( | `aws_role_name` | `RoleArn` | The Amazon Resource Name (ARN) of the role to assume | [AssumeRole API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role) | | `aws_session_name` | `RoleSessionName` | An identifier for the assumed role session | [AssumeRole API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts.html#STS.Client.assume_role) | +### IAM Roles Anywhere (On-Premise / External Workloads) + +[IAM Roles Anywhere](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/introduction.html) extends IAM roles to workloads **outside of AWS** (on-premise servers, edge devices, other clouds). It uses the same STS mechanism as regular IAM roles but authenticates via X.509 certificates instead of AWS credentials. + +**Setup**: Configure the [AWS Signing Helper](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/credential-helper.html) as a credential process in `~/.aws/config`: + +```ini +[profile litellm-roles-anywhere] +credential_process = aws_signing_helper credential-process \ + --certificate /path/to/certificate.pem \ + --private-key /path/to/private-key.pem \ + --trust-anchor-arn arn:aws:rolesanywhere:us-east-1:123456789012:trust-anchor/abc123 \ + --profile-arn arn:aws:rolesanywhere:us-east-1:123456789012:profile/def456 \ + --role-arn arn:aws:iam::123456789012:role/MyBedrockRole +``` + +**Usage**: Reference the profile in LiteLLM: + + + + +```python +from litellm import completion + +response = completion( + model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": "Hello!"}], + aws_profile_name="litellm-roles-anywhere", +) +``` + + + + +```yaml +model_list: + - model_name: bedrock-claude + litellm_params: + model: bedrock/anthropic.claude-3-sonnet-20240229-v1:0 + aws_profile_name: "litellm-roles-anywhere" +``` + + + + +See the [IAM Roles Anywhere Getting Started Guide](https://docs.aws.amazon.com/rolesanywhere/latest/userguide/getting-started.html) for trust anchor and profile setup. + Make the bedrock completion call diff --git a/docs/my-website/docs/providers/bedrock_agentcore.md b/docs/my-website/docs/providers/bedrock_agentcore.md index 43df7f82519..e3e352f7ab6 100644 --- a/docs/my-website/docs/providers/bedrock_agentcore.md +++ b/docs/my-website/docs/providers/bedrock_agentcore.md @@ -11,6 +11,12 @@ Call Bedrock AgentCore in the OpenAI Request/Response format. | Provider Route on LiteLLM | `bedrock/agentcore/{AGENT_RUNTIME_ARN}` | | Provider Doc | [AWS Bedrock AgentCore ↗](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agentcore_InvokeAgentRuntime.html) | +:::info + +This documentation is for **AgentCore Agents** (agent runtimes). If you want to use AgentCore MCP servers, add them as you would any other MCP server. See the [MCP documentation](https://docs.litellm.ai/docs/mcp) for details. + +::: + ## Quick Start ### Model Format to LiteLLM diff --git a/docs/my-website/docs/providers/zai.md b/docs/my-website/docs/providers/zai.md index 5055d0c1cdd..937ccd67680 100644 --- a/docs/my-website/docs/providers/zai.md +++ b/docs/my-website/docs/providers/zai.md @@ -19,7 +19,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[ {"role": "user", "content": "hello from litellm"} ], @@ -34,7 +34,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[ {"role": "user", "content": "hello from litellm"} ], @@ -51,7 +51,8 @@ We support ALL Z.AI GLM models, just set `zai/` as a prefix when sending complet | Model Name | Function Call | Notes | |------------|---------------|-------| -| glm-4.6 | `completion(model="zai/glm-4.6", messages)` | Latest flagship model, 200K context | +| glm-4.7 | `completion(model="zai/glm-4.7", messages)` | **Latest flagship**, 200K context, **Reasoning** | +| glm-4.6 | `completion(model="zai/glm-4.6", messages)` | 200K context | | glm-4.5 | `completion(model="zai/glm-4.5", messages)` | 128K context | | glm-4.5v | `completion(model="zai/glm-4.5v", messages)` | Vision model | | glm-4.5-x | `completion(model="zai/glm-4.5-x", messages)` | Premium tier | @@ -62,16 +63,17 @@ We support ALL Z.AI GLM models, just set `zai/` as a prefix when sending complet ## Model Pricing -| Model | Input ($/1M tokens) | Output ($/1M tokens) | Context Window | -|-------|---------------------|----------------------|----------------| -| glm-4.6 | $0.60 | $2.20 | 200K | -| glm-4.5 | $0.60 | $2.20 | 128K | -| glm-4.5v | $0.60 | $1.80 | 128K | -| glm-4.5-x | $2.20 | $8.90 | 128K | -| glm-4.5-air | $0.20 | $1.10 | 128K | -| glm-4.5-airx | $1.10 | $4.50 | 128K | -| glm-4-32b-0414-128k | $0.10 | $0.10 | 128K | -| glm-4.5-flash | **FREE** | **FREE** | 128K | +| Model | Input ($/1M tokens) | Output ($/1M tokens) | Cached Input ($/1M tokens) | Context Window | +|-------|---------------------|----------------------|---------------------------|----------------| +| glm-4.7 | $0.60 | $2.20 | $0.11 | 200K | +| glm-4.6 | $0.60 | $2.20 | - | 200K | +| glm-4.5 | $0.60 | $2.20 | - | 128K | +| glm-4.5v | $0.60 | $1.80 | - | 128K | +| glm-4.5-x | $2.20 | $8.90 | - | 128K | +| glm-4.5-air | $0.20 | $1.10 | - | 128K | +| glm-4.5-airx | $1.10 | $4.50 | - | 128K | +| glm-4-32b-0414-128k | $0.10 | $0.10 | - | 128K | +| glm-4.5-flash | **FREE** | **FREE** | - | 128K | ## Using with LiteLLM Proxy @@ -84,7 +86,7 @@ import os os.environ['ZAI_API_KEY'] = "" response = completion( - model="zai/glm-4.6", + model="zai/glm-4.7", messages=[{"role": "user", "content": "Hello, how are you?"}], ) @@ -98,9 +100,9 @@ print(response.choices[0].message.content) ```yaml model_list: - - model_name: glm-4.6 + - model_name: glm-4.7 litellm_params: - model: zai/glm-4.6 + model: zai/glm-4.7 api_key: os.environ/ZAI_API_KEY - model_name: glm-4.5-flash # Free tier litellm_params: @@ -121,7 +123,7 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ -H 'Content-Type: application/json' \ -H 'Authorization: Bearer sk-1234' \ -d '{ - "model": "glm-4.6", + "model": "glm-4.7", "messages": [ { "role": "user", diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 343cbd0e53f..87064d442ae 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -464,6 +464,9 @@ router_settings: | DATABASE_USER | Username for database connection | DATABASE_USERNAME | Alias for database user | DATABRICKS_API_BASE | Base URL for Databricks API +| DATABRICKS_CLIENT_ID | Client ID for Databricks OAuth M2M authentication (Service Principal application ID) +| DATABRICKS_CLIENT_SECRET | Client secret for Databricks OAuth M2M authentication +| DATABRICKS_USER_AGENT | Custom user agent string for Databricks API requests. Used for partner telemetry attribution | DAYS_IN_A_MONTH | Days in a month for calculation purposes. Default is 28 | DAYS_IN_A_WEEK | Days in a week for calculation purposes. Default is 7 | DAYS_IN_A_YEAR | Days in a year for calculation purposes. Default is 365 @@ -708,6 +711,7 @@ router_settings: | LITELLM_SSL_CIPHERS | SSL/TLS cipher configuration for faster handshakes. Controls cipher suite preferences for OpenSSL connections. | LITELLM_SECRET_AWS_KMS_LITELLM_LICENSE | AWS KMS encrypted license for LiteLLM | LITELLM_TOKEN | Access token for LiteLLM integration +| LITELLM_USER_AGENT | Custom user agent string for LiteLLM API requests. Used for partner telemetry attribution | LITELLM_PRINT_STANDARD_LOGGING_PAYLOAD | If true, prints the standard logging payload to the console - useful for debugging | LITELM_ENVIRONMENT | Environment for LiteLLM Instance. This is currently only logged to DeepEval to determine the environment for DeepEval integration. | LOGFIRE_TOKEN | Token for Logfire logging service diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index ba4ca190aa9..bc2f6a13362 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -116,7 +116,7 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ "role": "user", "content": "what llm are you" } - ], + ] } ' ``` diff --git a/docs/my-website/docs/proxy/token_auth.md b/docs/my-website/docs/proxy/token_auth.md index fe928a596cf..78cd144d56d 100644 --- a/docs/my-website/docs/proxy/token_auth.md +++ b/docs/my-website/docs/proxy/token_auth.md @@ -114,6 +114,189 @@ Set `JWT_PUBLIC_KEY_URL` in your environment to a comma-separated list of URLs f export JWT_PUBLIC_KEY_URL="https://demo.duendesoftware.com/.well-known/openid-configuration/jwks,https://accounts.google.com/.well-known/openid-configuration/jwks" ``` +### Kubernetes ServiceAccount Authentication + +Use Kubernetes ServiceAccount tokens to authenticate workloads running in your cluster. This is useful when you want pods to authenticate to LiteLLM using their native Kubernetes identity. + +#### Prerequisites + +1. Your Kubernetes cluster must have ServiceAccount token projection enabled (default in Kubernetes 1.20+) +2. Your cluster's OIDC issuer must be accessible (for EKS, GKE, AKS this is automatic) + +#### Step 1: Configure the OIDC Discovery URL + +Set `JWT_PUBLIC_KEY_URL` to your cluster's OIDC discovery endpoint: + + + + +```bash +# Get your EKS OIDC issuer URL +aws eks describe-cluster --name --query "cluster.identity.oidc.issuer" --output text + +# Set the JWKS URL (append /keys to the issuer URL) +export JWT_PUBLIC_KEY_URL="https://oidc.eks..amazonaws.com/id//keys" +``` + + + + +```bash +# GKE uses Google's OIDC provider +export JWT_PUBLIC_KEY_URL="https://container.googleapis.com/v1/projects//locations//clusters//jwks" +``` + + + + +```bash +# Get your AKS OIDC issuer URL +az aks show --name --resource-group --query "oidcIssuerProfile.issuerUrl" -o tsv + +# Set the JWKS URL +export JWT_PUBLIC_KEY_URL="/openid/v1/jwks" +``` + + + + +```bash +# For self-managed clusters, check your API server's --service-account-issuer flag +# The JWKS endpoint is typically at: +export JWT_PUBLIC_KEY_URL="https:///openid/v1/jwks" +``` + + + + +#### Step 2: Configure LiteLLM + +Configure LiteLLM to extract identity information from Kubernetes ServiceAccount tokens: + +```yaml +general_settings: + enable_jwt_auth: True + litellm_jwtauth: + # Use namespace as team identifier (resolves via team_alias in DB) + team_alias_jwt_field: "kubernetes\.io.namespace" +``` + +#### Step 3: Create ServiceAccount and Configure Pod + +Create a ServiceAccount with an associated secret and configure your pod to use the token: + +```yaml +apiVersion: v1 +kind: ServiceAccount +metadata: + name: my-llm-client + namespace: my-app +--- +apiVersion: v1 +kind: Secret +metadata: + name: my-llm-client-token + namespace: my-app + annotations: + kubernetes.io/service-account.name: my-llm-client +type: kubernetes.io/service-account-token +--- +apiVersion: v1 +kind: Pod +metadata: + name: llm-client-pod + namespace: my-app +spec: + serviceAccountName: my-llm-client + containers: + - name: app + image: my-app:latest + env: + - name: LITELLM_TOKEN + valueFrom: + secretKeyRef: + name: my-llm-client-token + key: token +``` + +Set the expected audience in LiteLLM: + +```bash +export JWT_AUDIENCE="https://kubernetes.default.svc" +``` + +#### Step 4: Create Team for Namespace + +Create a team in LiteLLM that matches the namespace (using `team_alias`): + +```bash +curl -X POST 'http://0.0.0.0:4000/team/new' \ +-H 'Authorization: Bearer ' \ +-H 'Content-Type: application/json' \ +-d '{ + "team_alias": "my-app", + "team_id": "my-app", + "models": ["gpt-4", "claude-sonnet-4-20250514"] +}' +``` + +#### Step 5: Use the Token + +From within the pod, the token is available in the `LITELLM_TOKEN` environment variable: + +```bash +# Make a request to LiteLLM using the env var +curl -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ +-H 'Content-Type: application/json' \ +-H "Authorization: Bearer $LITELLM_TOKEN" \ +-d '{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello!"}] +}' +``` + +#### Example: ServiceAccount Token Structure + +A Kubernetes ServiceAccount token looks like this: + +```json +{ + "aud": ["litellm-proxy"], + "exp": 1234567890, + "iat": 1234567890, + "iss": "https://oidc.eks.us-west-2.amazonaws.com/id/EXAMPLE", + "kubernetes.io": { + "namespace": "my-app", + "pod": { + "name": "llm-client-pod", + "uid": "pod-uid" + }, + "serviceaccount": { + "name": "my-llm-client", + "uid": "sa-uid" + } + }, + "nbf": 1234567890, + "sub": "system:serviceaccount:my-app:my-llm-client" +} +``` + +#### Advanced: Map Namespace to Team Using Name Resolution + +Use the `team_alias_jwt_field` to automatically resolve namespaces to teams: + +```yaml +general_settings: + enable_jwt_auth: True + litellm_jwtauth: + user_id_jwt_field: "sub" + # Map the namespace to team_alias in the database + team_alias_jwt_field: "kubernetes\.io.namespace" + user_id_upsert: true +``` + +This way, pods in namespace `production` automatically get associated with the team that has `team_alias: production`. + ### Set Accepted JWT Scope Names Change the string in JWT 'scopes', that litellm evaluates to see if a user has admin access. @@ -183,6 +366,62 @@ litellm_jwtauth: Now litellm will automatically update the spend for the user/team/org in the db for each call. +### Resolve by Name (Alias) Instead of ID + +Sometimes your JWT token contains human-readable names instead of database IDs. LiteLLM can resolve these names to IDs by looking them up in the database. + +**Use Case:** Your IDP provides team/org names in the JWT, but LiteLLM needs the actual database IDs for spend tracking and access control. + +```yaml +general_settings: + master_key: sk-1234 + enable_jwt_auth: True + litellm_jwtauth: + # Name-based fields (resolved via database lookup) + team_alias_jwt_field: "team_alias" # Resolves team by team_alias in DB + org_alias_jwt_field: "org_alias" # Resolves org by organization_alias in DB +``` + +**Expected JWT:** + +```json +{ + "sub": "user-123", + "team_alias": "engineering-team", + "org_alias": "acme-corp" +} +``` + +**How It Works:** + +1. LiteLLM extracts the name from the configured JWT field +2. Looks up the entity in the database by its alias field: + - Teams: `team_alias` column in `LiteLLM_TeamTable` + - Organizations: `organization_alias` column in `LiteLLM_OrganizationTable` +3. Uses the resolved ID for spend tracking and access control + +**Precedence:** ID fields always take precedence over name fields. If both `team_id_jwt_field` and `team_alias_jwt_field` are configured and both values exist in the JWT, the ID will be used. + +```yaml +# Example: ID takes precedence +litellm_jwtauth: + team_id_jwt_field: "team_id" # Used if present in JWT + team_alias_jwt_field: "team_alias" # Fallback if team_id not present +``` + +**Nested Fields:** Name fields also support dot notation for nested claims: + +```yaml +litellm_jwtauth: + team_alias_jwt_field: "organization.team.name" + org_alias_jwt_field: "company.name" +``` + +**Important Notes:** +- The entity (team/org) must already exist in the database with the matching alias +- Aliases should be unique - if multiple entities share the same alias, an error will be returned +- Name resolution adds a database lookup, so using IDs directly is slightly more performant + ### JWT Scopes Here's what scopes on JWT-Auth tokens look like diff --git a/docs/my-website/img/levo_logo.png b/docs/my-website/img/levo_logo.png new file mode 100644 index 00000000000..fdb72470b29 Binary files /dev/null and b/docs/my-website/img/levo_logo.png differ diff --git a/docs/my-website/img/levo_logo_dark.png b/docs/my-website/img/levo_logo_dark.png new file mode 100644 index 00000000000..70da632ee90 Binary files /dev/null and b/docs/my-website/img/levo_logo_dark.png differ diff --git a/docs/my-website/package-lock.json b/docs/my-website/package-lock.json index 8af06ec1a94..c5f15ebd5f8 100644 --- a/docs/my-website/package-lock.json +++ b/docs/my-website/package-lock.json @@ -8904,23 +8904,23 @@ "license": "ISC" }, "node_modules/body-parser": { - "version": "1.20.3", - "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.3.tgz", - "integrity": "sha512-7rAxByjUMqQ3/bHJy7D6OGXvx/MMc4IqBn/X0fcM1QUcAItpZrBEYhWGem+tzXH90c+G01ypMcYJBO9Y30203g==", + "version": "1.20.4", + "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.4.tgz", + "integrity": "sha512-ZTgYYLMOXY9qKU/57FAo8F+HA2dGX7bqGc71txDRC1rS4frdFI5R7NhluHxH6M0YItAP0sHB4uqAOcYKxO6uGA==", "license": "MIT", "dependencies": { - "bytes": "3.1.2", + "bytes": "~3.1.2", "content-type": "~1.0.5", "debug": "2.6.9", "depd": "2.0.0", - "destroy": "1.2.0", - "http-errors": "2.0.0", - "iconv-lite": "0.4.24", - "on-finished": "2.4.1", - "qs": "6.13.0", - "raw-body": "2.5.2", + "destroy": "~1.2.0", + "http-errors": "~2.0.1", + "iconv-lite": "~0.4.24", + "on-finished": "~2.4.1", + "qs": "~6.14.0", + "raw-body": "~2.5.3", "type-is": "~1.6.18", - "unpipe": "1.0.0" + "unpipe": "~1.0.0" }, "engines": { "node": ">= 0.8", @@ -8945,6 +8945,26 @@ "ms": "2.0.0" } }, + "node_modules/body-parser/node_modules/http-errors": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", + "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", + "license": "MIT", + "dependencies": { + "depd": "~2.0.0", + "inherits": "~2.0.4", + "setprototypeof": "~1.2.0", + "statuses": "~2.0.2", + "toidentifier": "~1.0.1" + }, + "engines": { + "node": ">= 0.8" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, "node_modules/body-parser/node_modules/iconv-lite": { "version": "0.4.24", "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz", @@ -8957,12 +8977,27 @@ "node": ">=0.10.0" } }, + "node_modules/body-parser/node_modules/inherits": { + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", + "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", + "license": "ISC" + }, "node_modules/body-parser/node_modules/ms": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz", "integrity": "sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A==", "license": "MIT" }, + "node_modules/body-parser/node_modules/statuses": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", + "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, "node_modules/bonjour-service": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/bonjour-service/-/bonjour-service-1.3.0.tgz", @@ -11873,39 +11908,39 @@ } }, "node_modules/express": { - "version": "4.21.2", - "resolved": "https://registry.npmjs.org/express/-/express-4.21.2.tgz", - "integrity": "sha512-28HqgMZAmih1Czt9ny7qr6ek2qddF4FclbMzwhCREB6OFfH+rXAnuNCwo1/wFvrtbgsQDb4kSbX9de9lFbrXnA==", + "version": "4.22.1", + "resolved": "https://registry.npmjs.org/express/-/express-4.22.1.tgz", + "integrity": "sha512-F2X8g9P1X7uCPZMA3MVf9wcTqlyNp7IhH5qPCI0izhaOIYXaW9L535tGA3qmjRzpH+bZczqq7hVKxTR4NWnu+g==", "license": "MIT", "dependencies": { "accepts": "~1.3.8", "array-flatten": "1.1.1", - "body-parser": "1.20.3", - "content-disposition": "0.5.4", + "body-parser": "~1.20.3", + "content-disposition": "~0.5.4", "content-type": "~1.0.4", - "cookie": "0.7.1", - "cookie-signature": "1.0.6", + "cookie": "~0.7.1", + "cookie-signature": "~1.0.6", "debug": "2.6.9", "depd": "2.0.0", "encodeurl": "~2.0.0", "escape-html": "~1.0.3", "etag": "~1.8.1", - "finalhandler": "1.3.1", - "fresh": "0.5.2", - "http-errors": "2.0.0", + "finalhandler": "~1.3.1", + "fresh": "~0.5.2", + "http-errors": "~2.0.0", "merge-descriptors": "1.0.3", "methods": "~1.1.2", - "on-finished": "2.4.1", + "on-finished": "~2.4.1", "parseurl": "~1.3.3", - "path-to-regexp": "0.1.12", + "path-to-regexp": "~0.1.12", "proxy-addr": "~2.0.7", - "qs": "6.13.0", + "qs": "~6.14.0", "range-parser": "~1.2.1", "safe-buffer": "5.2.1", - "send": "0.19.0", - "serve-static": "1.16.2", + "send": "~0.19.0", + "serve-static": "~1.16.2", "setprototypeof": "1.2.0", - "statuses": "2.0.1", + "statuses": "~2.0.1", "type-is": "~1.6.18", "utils-merge": "1.0.1", "vary": "~1.1.2" @@ -19281,12 +19316,12 @@ } }, "node_modules/qs": { - "version": "6.13.0", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.13.0.tgz", - "integrity": "sha512-+38qI9SOr8tfZ4QmJNplMUxqjbe7LKvvZgWdExBOmd+egZTtjLB67Gu0HRX3u/XOq7UU2Nx6nsjvS16Z9uwfpg==", + "version": "6.14.1", + "resolved": "https://registry.npmjs.org/qs/-/qs-6.14.1.tgz", + "integrity": "sha512-4EK3+xJl8Ts67nLYNwqw/dsFVnCf+qR7RgXSK9jEEm9unao3njwMDdmsdvoKBKHzxd7tCYz5e5M+SnMjdtXGQQ==", "license": "BSD-3-Clause", "dependencies": { - "side-channel": "^1.0.6" + "side-channel": "^1.1.0" }, "engines": { "node": ">=0.6" @@ -19362,15 +19397,15 @@ } }, "node_modules/raw-body": { - "version": "2.5.2", - "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.2.tgz", - "integrity": "sha512-8zGqypfENjCIqGhgXToC8aB2r7YrBX+AQAfIPs/Mlk+BtPTztOvTS01NRW/3Eh60J+a48lt8qsCzirQ6loCVfA==", + "version": "2.5.3", + "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.3.tgz", + "integrity": "sha512-s4VSOf6yN0rvbRZGxs8Om5CWj6seneMwK3oDb4lWDH0UPhWcxwOWw5+qk24bxq87szX1ydrwylIOp2uG1ojUpA==", "license": "MIT", "dependencies": { - "bytes": "3.1.2", - "http-errors": "2.0.0", - "iconv-lite": "0.4.24", - "unpipe": "1.0.0" + "bytes": "~3.1.2", + "http-errors": "~2.0.1", + "iconv-lite": "~0.4.24", + "unpipe": "~1.0.0" }, "engines": { "node": ">= 0.8" @@ -19385,6 +19420,26 @@ "node": ">= 0.8" } }, + "node_modules/raw-body/node_modules/http-errors": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", + "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", + "license": "MIT", + "dependencies": { + "depd": "~2.0.0", + "inherits": "~2.0.4", + "setprototypeof": "~1.2.0", + "statuses": "~2.0.2", + "toidentifier": "~1.0.1" + }, + "engines": { + "node": ">= 0.8" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, "node_modules/raw-body/node_modules/iconv-lite": { "version": "0.4.24", "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz", @@ -19397,6 +19452,21 @@ "node": ">=0.10.0" } }, + "node_modules/raw-body/node_modules/inherits": { + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", + "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", + "license": "ISC" + }, + "node_modules/raw-body/node_modules/statuses": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", + "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, "node_modules/rc": { "version": "1.2.8", "resolved": "https://registry.npmjs.org/rc/-/rc-1.2.8.tgz", diff --git a/docs/my-website/src/css/custom.css b/docs/my-website/src/css/custom.css index 2bc6a4cfdef..9fa4443afc9 100644 --- a/docs/my-website/src/css/custom.css +++ b/docs/my-website/src/css/custom.css @@ -28,3 +28,34 @@ --ifm-color-primary-lightest: #4fddbf; --docusaurus-highlighted-code-line-bg: rgba(0, 0, 0, 0.3); } + +/* Levo logo sizing and theme switching */ +.levo-logo-container { + position: relative; +} + +.levo-logo-container img, +.levo-logo-container picture, +.levo-logo-container .ideal-image { + max-width: 200px !important; + width: 200px !important; + height: auto !important; +} + +/* Show light logo by default, hide dark logo */ +.levo-logo-dark { + display: none !important; +} + +.levo-logo-light { + display: block !important; +} + +/* In dark mode, hide light logo and show dark logo */ +[data-theme='dark'] .levo-logo-light { + display: none !important; +} + +[data-theme='dark'] .levo-logo-dark { + display: block !important; +} diff --git a/docs/my-website/src/data/adopters/README.md b/docs/my-website/src/data/adopters/README.md new file mode 100644 index 00000000000..61a5215f802 --- /dev/null +++ b/docs/my-website/src/data/adopters/README.md @@ -0,0 +1,88 @@ +# LiteLLM Adopters + +This directory contains data for organizations that use LiteLLM in production. + +## Adding Your Organization + +We've made it super easy to add your organization! Just follow the steps below. + +### Quick Add (Recommended) + +**[Edit adopters.json on GitHub →](https://github.com/BerriAI/litellm/edit/main/docs/my-website/src/data/adopters/adopters.json)** + +This will open the GitHub editor in your browser where you can: + +1. Add your organization's entry to the JSON array +2. Commit your changes +3. GitHub will automatically create a pull request for you! + +No need to clone the repository or set up a development environment. + +### JSON Format + +Add your organization to the array in `adopters.json`: + +```json +{ + "name": "Your Organization Name", + "logoUrl": "https://yoursite.com/logo.svg", + "url": "https://yourcompany.com", + "description": "Brief description of how you use LiteLLM (shown on hover)" +} +``` + +### Fields + +- **`name`** (required): Your organization's display name +- **`logoUrl`** (required): URL to your logo - can be either: + - External URL: `https://yoursite.com/logo.svg` (easiest!) + - Local path: `/img/adopters/your-logo.svg` (requires uploading logo file) +- **`url`** (optional): Your organization's website (makes the logo clickable) +- **`description`** (optional): Brief description shown when users hover over your logo + +### Logo Options + +#### Option 1: External URL (Easiest) + +Simply provide a direct link to your logo hosted anywhere: + +```json +"logoUrl": "https://yourcompany.com/assets/logo.svg" +``` + +#### Option 2: Local Logo (Better Performance) + +If you prefer to host the logo locally: + +1. Add your logo to `docs/my-website/static/img/adopters/your-company.svg` +2. Reference it as: `"logoUrl": "/img/adopters/your-company.svg"` + +**Logo Specifications:** + +- **Format**: SVG preferred (PNG also acceptable) +- **Dimensions**: 240x160px or similar 3:2 ratio recommended +- **Background**: Transparent or white background works best + +### Example + +```json +{ + "name": "Acme Corporation", + "logoUrl": "https://acme.com/logo.svg", + "url": "https://acme.com", + "description": "Using LiteLLM to route requests across 50+ LLM providers" +} +``` + +### Display Order + +Adopters are displayed alphabetically by organization name, so your position will be determined automatically. + +### Need Help? + +If you have questions about adding your organization: + +- Ask in [GitHub Discussions](https://github.com/BerriAI/litellm/discussions) +- Join our [Discord community](https://discord.com/invite/wuPM9dRgDw) + +Thank you for supporting LiteLLM! 🚅 diff --git a/docs/my-website/src/data/adopters/adopters.json b/docs/my-website/src/data/adopters/adopters.json new file mode 100644 index 00000000000..52319c149e2 --- /dev/null +++ b/docs/my-website/src/data/adopters/adopters.json @@ -0,0 +1,8 @@ +[ + { + "name": "Your Logo Here", + "logoUrl": "/img/adopters/placeholder-company.svg", + "description": "Add your organization to show support for LiteLLM", + "url": "https://github.com/BerriAI/litellm/edit/main/docs/my-website/src/data/adopters/adopters.json" + } +] diff --git a/docs/my-website/src/data/adopters/index.js b/docs/my-website/src/data/adopters/index.js new file mode 100644 index 00000000000..b1a242dcc33 --- /dev/null +++ b/docs/my-website/src/data/adopters/index.js @@ -0,0 +1,23 @@ +import adoptersData from './adopters.json'; + +/** + * @typedef {Object} Adopter + * @property {string} name - The organization's display name + * @property {string} logoUrl - URL to the organization's logo + * @property {string} [url] - The organization's website URL + * @property {string} [description] - Brief description shown on hover + */ + +/** + * List of organizations using LiteLLM + * @type {Adopter[]} + */ +export const adopters = adoptersData; + +/** + * Adopters sorted alphabetically by name + * @type {Adopter[]} + */ +export const sortedAdopters = [...adopters].sort((a, b) => + a.name.localeCompare(b.name) +); diff --git a/docs/my-website/static/img/adopters/placeholder-company.svg b/docs/my-website/static/img/adopters/placeholder-company.svg new file mode 100644 index 00000000000..937dffc6eaf --- /dev/null +++ b/docs/my-website/static/img/adopters/placeholder-company.svg @@ -0,0 +1,8 @@ + + + + + + Add Your Logo + Click to contribute + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql new file mode 100644 index 00000000000..8eebb797e2c --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260102131258_add_metadata_urls_to_mcp_servers/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN "authorization_url" TEXT, +ADD COLUMN "registration_url" TEXT, +ADD COLUMN "token_url" TEXT; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index aac0b5b35de..ea47b6ed03b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + authorization_url String? + token_url String? + registration_url String? } // Generate Tokens for Proxy diff --git a/litellm/__init__.py b/litellm/__init__.py index a32d2d3ef90..4f8c57ef442 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -136,6 +136,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "gitlab", "cloudzero", "posthog", + "levo", ] cold_storage_custom_logger: Optional[_custom_logger_compatible_callbacks_literal] = None logged_real_time_event_types: Optional[Union[List[str], Literal["*"]]] = None @@ -553,6 +554,8 @@ docker_model_runner_models: Set = set() amazon_nova_models: Set = set() stability_models: Set = set() github_copilot_models: Set = set() +minimax_models: Set = set() +aws_polly_models: Set = set() def is_bedrock_pricing_only_model(key: str) -> bool: @@ -801,6 +804,10 @@ def add_known_models(): stability_models.add(key) elif value.get("litellm_provider") == "github_copilot": github_copilot_models.add(key) + elif value.get("litellm_provider") == "minimax": + minimax_models.add(key) + elif value.get("litellm_provider") == "aws_polly": + aws_polly_models.add(key) add_known_models() @@ -1005,6 +1012,8 @@ models_by_provider: dict = { "amazon_nova": amazon_nova_models, "stability": stability_models, "github_copilot": github_copilot_models, + "minimax": minimax_models, + "aws_polly": aws_polly_models, } # mapping for those models which have larger equivalents @@ -1049,8 +1058,8 @@ openai_image_generation_models = ["dall-e-2", "dall-e-3"] openai_video_generation_models = ["sora-2"] # timeout is lazy-loaded via __getattr__ -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.litellm_core_utils.core_helpers import remove_index_from_tool_calls +# get_llm_provider is lazy-loaded via __getattr__ +# remove_index_from_tool_calls is lazy-loaded via __getattr__ # Import KeyManagementSettings here (before utils import) because _key_management_settings # is accessed during import time in secret_managers/main.py (via dd_tracing -> datadog -> _service_logger -> utils) @@ -1499,6 +1508,7 @@ if TYPE_CHECKING: get_first_chars_messages: Callable[..., str] get_provider_fields: Callable[..., List] get_valid_models: Callable[..., list] + remove_index_from_tool_calls: Callable[..., None] # Response types - truly lazy loaded only (not in main.py or elsewhere) ModelResponseListIterator: Type[Any] @@ -1650,6 +1660,17 @@ def __getattr__(name: str) -> Any: LoggingCallbackManager = __getattr__("LoggingCallbackManager") _globals["logging_callback_manager"] = LoggingCallbackManager() return _globals["logging_callback_manager"] + + # Lazy load _service_logger module + if name == "_service_logger": + from ._lazy_imports import _get_litellm_globals + _globals = _get_litellm_globals() + # Check if already cached + if "_service_logger" not in _globals: + # Import the module lazily + import litellm._service_logger + _globals["_service_logger"] = litellm._service_logger + return _globals["_service_logger"] raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index c1b3e1df976..5ea9dd9b321 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -34,6 +34,7 @@ from ._lazy_imports_registry import ( DOTPROMPT_NAMES, LLM_CONFIG_NAMES, TYPES_NAMES, + LLM_PROVIDER_LOGIC_NAMES, # Import maps _UTILS_IMPORT_MAP, _COST_CALCULATOR_IMPORT_MAP, @@ -45,6 +46,7 @@ from ._lazy_imports_registry import ( _DOTPROMPT_IMPORT_MAP, _TYPES_IMPORT_MAP, _LLM_CONFIGS_IMPORT_MAP, + _LLM_PROVIDER_LOGIC_IMPORT_MAP, ) @@ -181,6 +183,8 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], Any]]: _LAZY_IMPORT_REGISTRY[name] = _lazy_import_llm_configs for name in TYPES_NAMES: _LAZY_IMPORT_REGISTRY[name] = _lazy_import_types + for name in LLM_PROVIDER_LOGIC_NAMES: + _LAZY_IMPORT_REGISTRY[name] = _lazy_import_llm_provider_logic return _LAZY_IMPORT_REGISTRY @@ -297,6 +301,11 @@ def _lazy_import_litellm_logging(name: str) -> Any: """Handler for litellm_logging module (Logging, modify_integration)""" return _generic_lazy_import(name, _LITELLM_LOGGING_IMPORT_MAP, "Litellm logging") + +def _lazy_import_llm_provider_logic(name: str) -> Any: + """Handler for LLM provider logic functions (get_llm_provider, etc.)""" + return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic") + # ============================================================================ # SPECIAL HANDLERS # ============================================================================ diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index e2f80a14391..93fa8b39af2 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -32,6 +32,7 @@ UTILS_NAMES = ( "ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse", "TranscriptionResponse", "TextCompletionResponse", "get_provider_fields", "ModelResponseListIterator", "get_valid_models", "timeout", + "get_llm_provider", "remove_index_from_tool_calls", ) # Token counter names that support lazy loading via _lazy_import_token_counter @@ -287,6 +288,12 @@ TYPES_NAMES = ( # is accessed during import time in secret_managers/main.py ) +# LLM provider logic names that support lazy loading via _lazy_import_llm_provider_logic +LLM_PROVIDER_LOGIC_NAMES = ( + "get_llm_provider", + "remove_index_from_tool_calls", +) + # Import maps for registry pattern - reduces repetition _UTILS_IMPORT_MAP = { "exception_type": (".utils", "exception_type"), @@ -330,6 +337,8 @@ _UTILS_IMPORT_MAP = { "ModelResponseListIterator": (".utils", "ModelResponseListIterator"), "get_valid_models": (".utils", "get_valid_models"), "timeout": (".timeout", "timeout"), + "get_llm_provider": ("litellm.litellm_core_utils.get_llm_provider_logic", "get_llm_provider"), + "remove_index_from_tool_calls": ("litellm.litellm_core_utils.core_helpers", "remove_index_from_tool_calls"), } _COST_CALCULATOR_IMPORT_MAP = { @@ -386,6 +395,11 @@ _TYPES_IMPORT_MAP = { "LoggingCallbackManager": ("litellm.litellm_core_utils.logging_callback_manager", "LoggingCallbackManager"), } +_LLM_PROVIDER_LOGIC_IMPORT_MAP = { + "get_llm_provider": ("litellm.litellm_core_utils.get_llm_provider_logic", "get_llm_provider"), + "remove_index_from_tool_calls": ("litellm.litellm_core_utils.core_helpers", "remove_index_from_tool_calls"), +} + _LLM_CONFIGS_IMPORT_MAP = { "AmazonConverseConfig": (".llms.bedrock.chat.converse_transformation", "AmazonConverseConfig"), "OpenAILikeChatConfig": (".llms.openai_like.chat.handler", "OpenAILikeChatConfig"), @@ -587,6 +601,7 @@ __all__ = [ "DOTPROMPT_NAMES", "LLM_CONFIG_NAMES", "TYPES_NAMES", + "LLM_PROVIDER_LOGIC_NAMES", # Import maps "_UTILS_IMPORT_MAP", "_COST_CALCULATOR_IMPORT_MAP", @@ -598,5 +613,6 @@ __all__ = [ "_DOTPROMPT_IMPORT_MAP", "_TYPES_IMPORT_MAP", "_LLM_CONFIGS_IMPORT_MAP", + "_LLM_PROVIDER_LOGIC_IMPORT_MAP", ] diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 9d3f990b1aa..58a52666d38 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -8,8 +8,10 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, + ChatCompletionImageObject, ChatCompletionRequest, ChatCompletionSystemMessage, + ChatCompletionTextObject, ChatCompletionToolCallFunctionChunk, ChatCompletionToolChoiceValues, ChatCompletionToolMessage, @@ -385,13 +387,36 @@ class GoogleGenAIAdapter: if role == "user": # Handle user messages with potential function responses - combined_text = "" + content_parts: List[ + Union[ChatCompletionTextObject, ChatCompletionImageObject] + ] = [] tool_messages: List[ChatCompletionToolMessage] = [] for part in parts: if isinstance(part, dict): if "text" in part: - combined_text += part["text"] + content_parts.append( + cast( + ChatCompletionTextObject, + {"type": "text", "text": part["text"]}, + ) + ) + elif "inline_data" in part: + # Handle Base64 image data + inline_data = part["inline_data"] + mime_type = inline_data.get("mime_type", "image/jpeg") + data = inline_data.get("data", "") + content_parts.append( + cast( + ChatCompletionImageObject, + { + "type": "image_url", + "image_url": { + "url": f"data:{mime_type};base64,{data}" + }, + }, + ) + ) elif "functionResponse" in part: # Transform function response to tool message func_response = part["functionResponse"] @@ -402,13 +427,33 @@ class GoogleGenAIAdapter: ) tool_messages.append(tool_message) elif isinstance(part, str): - combined_text += part + content_parts.append( + cast( + ChatCompletionTextObject, {"type": "text", "text": part} + ) + ) - # Add user message if there's text content - if combined_text: - messages.append( - ChatCompletionUserMessage(role="user", content=combined_text) - ) + # Add user message if there's content + if content_parts: + # If only one text part, use simple string format for backward compatibility + if ( + len(content_parts) == 1 + and isinstance(content_parts[0], dict) + and content_parts[0].get("type") == "text" + ): + text_part = cast(ChatCompletionTextObject, content_parts[0]) + messages.append( + ChatCompletionUserMessage( + role="user", content=text_part["text"] + ) + ) + else: + # Use multimodal format (array of content parts) + messages.append( + ChatCompletionUserMessage( + role="user", content=content_parts + ) + ) # Add tool messages messages.extend(tool_messages) @@ -468,7 +513,6 @@ class GoogleGenAIAdapter: Dict in Google GenAI generate_content response format """ - # Extract the main response content choice = response.choices[0] if response.choices else None if not choice: diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index fe0ce208ee6..6a76b57e7f7 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -243,14 +243,14 @@ class CustomGuardrail(CustomLogger): def _is_valid_response_type(self, result: Any) -> bool: """ Check if result is a valid LLMResponseTypes instance. - + Safely handles TypedDict types which don't support isinstance checks. For non-LiteLLM responses (like passthrough httpx.Response), returns True to allow them through. """ if result is None: return False - + try: # Try isinstance check on valid types that support it response_types = get_args(LLMResponseTypes) @@ -506,6 +506,7 @@ class CustomGuardrail(CustomLogger): duration: Optional[float] = None, masked_entity_count: Optional[Dict[str, int]] = None, guardrail_provider: Optional[str] = None, + event_type: Optional[GuardrailEventHooks] = None, ) -> None: """ Builds `StandardLoggingGuardrailInformation` and adds it to the request metadata so it can be used for logging to DataDog, Langfuse, etc. @@ -514,14 +515,19 @@ class CustomGuardrail(CustomLogger): guardrail_json_response = str(guardrail_json_response) from litellm.types.utils import GuardrailMode + # Use event_type if provided, otherwise fall back to self.event_hook + guardrail_mode: Union[GuardrailEventHooks, GuardrailMode, List[GuardrailEventHooks]] + if event_type is not None: + guardrail_mode = event_type + elif isinstance(self.event_hook, Mode): + guardrail_mode = GuardrailMode(**dict(self.event_hook.model_dump())) # type: ignore[typeddict-item] + else: + guardrail_mode = self.event_hook # type: ignore[assignment] + slg = StandardLoggingGuardrailInformation( guardrail_name=self.guardrail_name, guardrail_provider=guardrail_provider, - guardrail_mode=( - GuardrailMode(**self.event_hook.model_dump()) # type: ignore - if isinstance(self.event_hook, Mode) - else self.event_hook - ), + guardrail_mode=guardrail_mode, guardrail_response=guardrail_json_response, guardrail_status=guardrail_status, start_time=start_time, @@ -589,6 +595,7 @@ class CustomGuardrail(CustomLogger): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Add StandardLoggingGuardrailInformation to the request data @@ -605,6 +612,7 @@ class CustomGuardrail(CustomLogger): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) return response @@ -615,6 +623,7 @@ class CustomGuardrail(CustomLogger): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Add StandardLoggingGuardrailInformation to the request data @@ -628,6 +637,7 @@ class CustomGuardrail(CustomLogger): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) raise e @@ -712,16 +722,32 @@ def log_guardrail_information(func): Logs for: - pre_call - during_call - - TODO: log post_call. This is more involved since the logs are sent to DD, s3 before the guardrail is even run + - post_call """ import asyncio import functools + def _infer_event_type_from_function_name( + func_name: str, + ) -> Optional[GuardrailEventHooks]: + """Infer the actual event type from the function name""" + if func_name == "async_pre_call_hook": + return GuardrailEventHooks.pre_call + elif func_name == "async_moderation_hook": + return GuardrailEventHooks.during_call + elif func_name in ( + "async_post_call_success_hook", + "async_post_call_streaming_hook", + ): + return GuardrailEventHooks.post_call + return None + @functools.wraps(func) async def async_wrapper(*args, **kwargs): start_time = datetime.now() # Move start_time inside the wrapper self: CustomGuardrail = args[0] request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} + event_type = _infer_event_type_from_function_name(func.__name__) try: response = await func(*args, **kwargs) return self._process_response( @@ -730,6 +756,7 @@ def log_guardrail_information(func): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) except Exception as e: return self._process_error( @@ -738,6 +765,7 @@ def log_guardrail_information(func): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) @functools.wraps(func) @@ -745,18 +773,21 @@ def log_guardrail_information(func): start_time = datetime.now() # Move start_time inside the wrapper self: CustomGuardrail = args[0] request_data: dict = kwargs.get("data") or kwargs.get("request_data") or {} + event_type = _infer_event_type_from_function_name(func.__name__) try: response = func(*args, **kwargs) return self._process_response( response=response, request_data=request_data, duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) except Exception as e: return self._process_error( e=e, request_data=request_data, duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) @functools.wraps(func) diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index 1c8a5b883da..1c62ce9fcc3 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -25,6 +25,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.types.utils import StandardLoggingPayload API_EVENT_TYPES = Literal["llm_api_success", "llm_api_failure"] +LOG_FORMAT_TYPES = Literal["json_array", "ndjson", "single"] def load_compatible_callbacks() -> Dict: @@ -101,6 +102,7 @@ class GenericAPILogger(CustomBatchLogger): headers: Optional[dict] = None, event_types: Optional[List[API_EVENT_TYPES]] = None, callback_name: Optional[str] = None, + log_format: Optional[LOG_FORMAT_TYPES] = None, **kwargs, ): """ @@ -111,6 +113,7 @@ class GenericAPILogger(CustomBatchLogger): headers: Optional[dict] = None, event_types: Optional[List[API_EVENT_TYPES]] = None, callback_name: Optional[str] = None - If provided, loads config from generic_api_compatible_callbacks.json + log_format: Optional[LOG_FORMAT_TYPES] = None - Format for log output: "json_array" (default), "ndjson", or "single" """ ######################################################### # Check if callback_name is provided and load config @@ -135,6 +138,9 @@ class GenericAPILogger(CustomBatchLogger): if event_types is None and "event_types" in callback_config: event_types = callback_config["event_types"] + + if log_format is None and "log_format" in callback_config: + log_format = callback_config["log_format"] else: verbose_logger.warning( f"callback_name '{callback_name}' not found in generic_api_compatible_callbacks.json" @@ -156,8 +162,16 @@ class GenericAPILogger(CustomBatchLogger): self.endpoint: str = endpoint self.event_types: Optional[List[API_EVENT_TYPES]] = event_types self.callback_name: Optional[str] = callback_name + + # Validate and store log_format + if log_format is not None and log_format not in ["json_array", "ndjson", "single"]: + raise ValueError( + f"Invalid log_format: {log_format}. Must be one of: 'json_array', 'ndjson', 'single'" + ) + self.log_format: LOG_FORMAT_TYPES = log_format or "json_array" + verbose_logger.debug( - f"in init GenericAPILogger, callback_name: {self.callback_name}, endpoint {self.endpoint}, headers {self.headers}, event_types: {self.event_types}" + f"in init GenericAPILogger, callback_name: {self.callback_name}, endpoint {self.endpoint}, headers {self.headers}, event_types: {self.event_types}, log_format: {self.log_format}" ) ######################################################### @@ -289,25 +303,65 @@ class GenericAPILogger(CustomBatchLogger): async def async_send_batch(self): """ Sends the batch of messages to Generic API Endpoint + + Supports three formats: + - json_array: Sends all logs as a JSON array (default) + - ndjson: Sends logs as newline-delimited JSON + - single: Sends each log as individual HTTP request in parallel """ try: if not self.log_queue: return verbose_logger.debug( - f"Generic API Logger - about to flush {len(self.log_queue)} events" + f"Generic API Logger - about to flush {len(self.log_queue)} events in '{self.log_format}' format" ) - # make POST request to Generic API Endpoint - response = await self.async_httpx_client.post( - url=self.endpoint, - headers=self.headers, - data=safe_dumps(self.log_queue), - ) + if self.log_format == "single": + # Send each log as individual HTTP request in parallel + tasks = [] + for log_entry in self.log_queue: + task = self.async_httpx_client.post( + url=self.endpoint, + headers=self.headers, + data=safe_dumps(log_entry), + ) + tasks.append(task) - verbose_logger.debug( - f"Generic API Logger - sent batch to {self.endpoint}, status code {response.status_code}" - ) + # Execute all requests in parallel + responses = await asyncio.gather(*tasks, return_exceptions=True) + + # Log results + for idx, result in enumerate(responses): + if isinstance(result, Exception): + verbose_logger.exception( + f"Generic API Logger - Error sending log {idx}: {result}" + ) + else: + # result is a Response object + verbose_logger.debug( + f"Generic API Logger - sent log {idx}, status: {result.status_code}" # type: ignore + ) + else: + # Format the payload based on log_format + if self.log_format == "json_array": + data = safe_dumps(self.log_queue) + elif self.log_format == "ndjson": + data = "\n".join(safe_dumps(log) for log in self.log_queue) + else: + raise ValueError(f"Unknown log_format: {self.log_format}") + + # Make POST request + response = await self.async_httpx_client.post( + url=self.endpoint, + headers=self.headers, + data=data, + ) + + verbose_logger.debug( + f"Generic API Logger - sent batch to {self.endpoint}, " + f"status: {response.status_code}, format: {self.log_format}" + ) except Exception as e: verbose_logger.exception( diff --git a/litellm/integrations/generic_api/generic_api_compatible_callbacks.json b/litellm/integrations/generic_api/generic_api_compatible_callbacks.json index 6c8e5fd1b2a..12dc4ae643c 100644 --- a/litellm/integrations/generic_api/generic_api_compatible_callbacks.json +++ b/litellm/integrations/generic_api/generic_api_compatible_callbacks.json @@ -22,6 +22,7 @@ "headers": { "Content-Type": "application/json" }, - "environment_variables": ["SUMOLOGIC_WEBHOOK_URL"] + "environment_variables": ["SUMOLOGIC_WEBHOOK_URL"], + "log_format": "ndjson" } } \ No newline at end of file diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 10347bc7c67..f0f355b4895 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -3,14 +3,27 @@ import os import traceback from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + Union, + cast, +) from packaging.version import Version import litellm from litellm._logging import verbose_logger from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS -from litellm.litellm_core_utils.core_helpers import safe_deep_copy +from litellm.litellm_core_utils.core_helpers import ( + safe_deep_copy, + reconstruct_model_name, +) from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info from litellm.llms.custom_httpx.http_handler import _get_httpx_client from litellm.secret_managers.main import str_to_bool @@ -437,12 +450,17 @@ class LangFuseLogger: ) ) + custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) + model_name = reconstruct_model_name( + kwargs.get("model", ""), custom_llm_provider, metadata + ) + trace.generation( CreateGeneration( name=metadata.get("generation_name", "litellm-completion"), startTime=start_time, endTime=end_time, - model=kwargs["model"], + model=model_name, modelParameters=optional_params, prompt=input, completion=output, @@ -543,7 +561,9 @@ class LangFuseLogger: # as we want to fall back to litellm_call_id instead for better traceability. # Note: Users can still explicitly set a UUID trace_id via metadata["trace_id"] (highest priority) if trace_id is None and standard_logging_object is not None: - standard_trace_id = cast(Optional[str], standard_logging_object.get("trace_id")) + standard_trace_id = cast( + Optional[str], standard_logging_object.get("trace_id") + ) # Only use standard_logging_object.trace_id if it's not a UUID # UUIDs are 36 characters with hyphens in format: xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx # We check for this specific pattern to avoid rejecting valid trace_ids that happen to have hyphens @@ -575,7 +595,9 @@ class LangFuseLogger: mask_output = clean_metadata.pop("mask_output", False) # Look for masking function in the dedicated location first (set by scrub_sensitive_keys_in_metadata) # Fall back to metadata for backwards compatibility - masking_function = litellm_params.get("_langfuse_masking_function") or clean_metadata.pop("langfuse_masking_function", None) + masking_function = litellm_params.get( + "_langfuse_masking_function" + ) or clean_metadata.pop("langfuse_masking_function", None) # Apply custom masking function if provided if masking_function is not None and callable(masking_function): @@ -776,12 +798,17 @@ class LangFuseLogger: if system_fingerprint is not None: optional_params["system_fingerprint"] = system_fingerprint + custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) + model_name = reconstruct_model_name( + kwargs.get("model", ""), custom_llm_provider, metadata + ) + generation_params = { "name": generation_name, "id": clean_metadata.pop("generation_id", generation_id), "start_time": start_time, "end_time": end_time, - "model": kwargs["model"], + "model": model_name, "model_parameters": optional_params, "input": input if not mask_input else "redacted-by-litellm", "output": output if not mask_output else "redacted-by-litellm", @@ -918,7 +945,9 @@ class LangFuseLogger: return Version(self.langfuse_sdk_version) >= Version("2.7.3") @staticmethod - def _apply_masking_function(data: Any, masking_function: Callable[[Any], Any]) -> Any: + def _apply_masking_function( + data: Any, masking_function: Callable[[Any], Any] + ) -> Any: """ Apply a masking function to data, handling different data types. diff --git a/litellm/integrations/levo/README.md b/litellm/integrations/levo/README.md new file mode 100644 index 00000000000..cb18b1dbfb0 --- /dev/null +++ b/litellm/integrations/levo/README.md @@ -0,0 +1,125 @@ +# Levo AI Integration + +This integration enables sending LLM observability data to Levo AI using OpenTelemetry (OTLP) protocol. + +## Overview + +The Levo integration extends LiteLLM's OpenTelemetry support to automatically send traces to Levo's collector endpoint with proper authentication and routing headers. + +## Features + +- **Automatic OTLP Export**: Sends OpenTelemetry traces to Levo collector +- **Levo-Specific Headers**: Automatically includes `x-levo-organization-id` and `x-levo-workspace-id` for routing +- **Simple Configuration**: Just use `callbacks: ["levo"]` in your LiteLLM config +- **Environment-Based Setup**: Configure via environment variables + +## Quick Start + +### 1. Install Dependencies + +```bash +pip install opentelemetry-api opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-exporter-otlp-proto-grpc +``` + +### 2. Configure LiteLLM + +Add to your `litellm_config.yaml`: + +```yaml +litellm_settings: + callbacks: ["levo"] +``` + +### 3. Set Environment Variables + +```bash +export LEVOAI_API_KEY="" +export LEVOAI_ORG_ID="" +export LEVOAI_WORKSPACE_ID="" +export LEVOAI_COLLECTOR_URL="" +``` + +### 4. Start LiteLLM + +```bash +litellm --config config.yaml +``` + +All LLM requests will now automatically be sent to Levo! + +## Configuration + +### Required Environment Variables + +| Variable | Description | +|----------|-------------| +| `LEVOAI_API_KEY` | Your Levo API key for authentication | +| `LEVOAI_ORG_ID` | Your Levo organization ID for routing | +| `LEVOAI_WORKSPACE_ID` | Your Levo workspace ID for routing | +| `LEVOAI_COLLECTOR_URL` | Full collector endpoint URL from Levo support | + +### Optional Environment Variables + +| Variable | Description | Default | +|----------|-------------|---------| +| `LEVOAI_ENV_NAME` | Environment name for tagging traces | `None` | + +**Important**: The `LEVOAI_COLLECTOR_URL` is used exactly as provided. No path manipulation is performed. + +## How It Works + +1. **LevoLogger** extends LiteLLM's `OpenTelemetry` class +2. **Configuration** is read from environment variables via `get_levo_config()` +3. **OTLP Headers** are automatically set: + - `Authorization: Bearer {LEVOAI_API_KEY}` + - `x-levo-organization-id: {LEVOAI_ORG_ID}` + - `x-levo-workspace-id: {LEVOAI_WORKSPACE_ID}` +4. **Traces** are sent to the collector endpoint in OTLP format + +## Code Structure + +``` +litellm/integrations/levo/ +├── __init__.py # Exports LevoLogger +├── levo.py # LevoLogger implementation +└── README.md # This file +``` + +### Key Classes + +- **LevoLogger**: Extends `OpenTelemetry`, handles Levo-specific configuration +- **LevoConfig**: Pydantic model for Levo configuration (defined in `levo.py`) + +## Testing + +See the test files in `tests/test_litellm/integrations/levo/`: +- `test_levo.py`: Unit tests for configuration +- `test_levo_integration.py`: Integration tests for callback registration + +## Error Handling + +The integration validates all required environment variables at initialization: +- Missing `LEVOAI_API_KEY`: Raises `ValueError` with clear message +- Missing `LEVOAI_ORG_ID`: Raises `ValueError` with clear message +- Missing `LEVOAI_WORKSPACE_ID`: Raises `ValueError` with clear message +- Missing `LEVOAI_COLLECTOR_URL`: Raises `ValueError` with clear message + +## Integration with LiteLLM + +The Levo callback is registered in: +- `litellm/litellm_core_utils/custom_logger_registry.py`: Maps `"levo"` to `LevoLogger` +- `litellm/litellm_core_utils/litellm_logging.py`: Instantiates `LevoLogger` when `callbacks: ["levo"]` is used +- `litellm/__init__.py`: Added to `_custom_logger_compatible_callbacks_literal` + +## Documentation + +For detailed documentation, see: +- [LiteLLM Levo Integration Docs](../../../../docs/my-website/docs/observability/levo_integration.md) +- [Levo Documentation](https://docs.levo.ai) + +## Support + +For issues or questions: +- LiteLLM Issues: https://github.com/BerriAI/litellm/issues +- Levo Support: support@levo.ai + diff --git a/litellm/integrations/levo/__init__.py b/litellm/integrations/levo/__init__.py new file mode 100644 index 00000000000..7f4f84437d4 --- /dev/null +++ b/litellm/integrations/levo/__init__.py @@ -0,0 +1,3 @@ +from litellm.integrations.levo.levo import LevoLogger + +__all__ = ["LevoLogger"] diff --git a/litellm/integrations/levo/levo.py b/litellm/integrations/levo/levo.py new file mode 100644 index 00000000000..562f2fd9068 --- /dev/null +++ b/litellm/integrations/levo/levo.py @@ -0,0 +1,117 @@ +import os +from typing import TYPE_CHECKING, Any, Optional, Union + +from litellm.integrations.opentelemetry import OpenTelemetry + +if TYPE_CHECKING: + from opentelemetry.trace import Span as _Span + + from litellm.integrations.opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig + from litellm.types.integrations.arize import Protocol as _Protocol + + Protocol = _Protocol + OpenTelemetryConfig = _OpenTelemetryConfig + Span = Union[_Span, Any] +else: + Protocol = Any + OpenTelemetryConfig = Any + Span = Any + + +class LevoConfig: + """Configuration for Levo OTLP integration.""" + + def __init__( + self, + otlp_auth_headers: Optional[str], + protocol: Protocol, + endpoint: str, + ): + self.otlp_auth_headers = otlp_auth_headers + self.protocol = protocol + self.endpoint = endpoint + + +class LevoLogger(OpenTelemetry): + """Levo Logger that extends OpenTelemetry for OTLP integration.""" + + @staticmethod + def get_levo_config() -> LevoConfig: + """ + Retrieves the Levo configuration based on environment variables. + + Returns: + LevoConfig: Configuration object containing Levo OTLP settings. + + Raises: + ValueError: If required environment variables are missing. + """ + # Required environment variables + api_key = os.environ.get("LEVOAI_API_KEY", None) + org_id = os.environ.get("LEVOAI_ORG_ID", None) + workspace_id = os.environ.get("LEVOAI_WORKSPACE_ID", None) + collector_url = os.environ.get("LEVOAI_COLLECTOR_URL", None) + + # Validate required env vars + if not api_key: + raise ValueError( + "LEVOAI_API_KEY environment variable is required for Levo integration." + ) + if not org_id: + raise ValueError( + "LEVOAI_ORG_ID environment variable is required for Levo integration." + ) + if not workspace_id: + raise ValueError( + "LEVOAI_WORKSPACE_ID environment variable is required for Levo integration." + ) + if not collector_url: + raise ValueError( + "LEVOAI_COLLECTOR_URL environment variable is required for Levo integration. " + "Please contact Levo support to get your collector URL." + ) + + # Use collector URL exactly as provided by the user + endpoint = collector_url + protocol: Protocol = "otlp_http" + + # Build OTLP headers string + # Format: Authorization=Bearer {api_key},x-levo-organization-id={org_id},x-levo-workspace-id={workspace_id} + headers_parts = [f"Authorization=Bearer {api_key}"] + headers_parts.append(f"x-levo-organization-id={org_id}") + headers_parts.append(f"x-levo-workspace-id={workspace_id}") + + otlp_auth_headers = ",".join(headers_parts) + + return LevoConfig( + otlp_auth_headers=otlp_auth_headers, + protocol=protocol, + endpoint=endpoint, + ) + + async def async_health_check(self): + """ + Health check for Levo integration. + + Returns: + dict: Health status with status and message/error_message keys. + """ + try: + config = self.get_levo_config() + + if not config.otlp_auth_headers: + return { + "status": "unhealthy", + "error_message": "LEVOAI_API_KEY environment variable not set", + } + + return { + "status": "healthy", + "message": "Levo credentials are configured properly", + } + except ValueError as e: + return { + "status": "unhealthy", + "error_message": str(e), + } + diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 93dce578fe1..12e60bc25bb 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -48,6 +48,7 @@ else: LITELLM_TRACER_NAME = os.getenv("OTEL_TRACER_NAME", "litellm") LITELLM_METER_NAME = os.getenv("LITELLM_METER_NAME", "litellm") LITELLM_LOGGER_NAME = os.getenv("LITELLM_LOGGER_NAME", "litellm") +LITELLM_PROXY_REQUEST_SPAN_NAME = "Received Proxy Server Request" # Remove the hardcoded LITELLM_RESOURCE dictionary - we'll create it properly later RAW_REQUEST_SPAN_NAME = "raw_gen_ai_request" LITELLM_REQUEST_SPAN_NAME = "litellm_request" @@ -233,14 +234,16 @@ class OpenTelemetry(CustomLogger): trace.set_tracer_provider(tracer_provider) else: # Tracer provider explicitly provided (e.g., for testing) + # Do NOT call set_tracer_provider - the caller is responsible for managing global state + # If they want it to be global, they've already set it before passing it to us verbose_logger.debug( "OpenTelemetry: Using provided TracerProvider: %s", type(tracer_provider).__name__, ) - trace.set_tracer_provider(tracer_provider) - # grab our tracer - self.tracer = trace.get_tracer(LITELLM_TRACER_NAME) + # Grab our tracer from the TracerProvider (not from global context) + # This ensures we use the provided TracerProvider (e.g., for testing) + self.tracer = tracer_provider.get_tracer(LITELLM_TRACER_NAME) self.span_kind = SpanKind def _init_metrics(self, meter_provider): @@ -527,6 +530,7 @@ class OpenTelemetry(CustomLogger): # 3. Guardrail span self._create_guardrail_span(kwargs=kwargs, context=ctx) + return response ######################################################### @@ -557,9 +561,9 @@ class OpenTelemetry(CustomLogger): def _get_dynamic_otel_headers_from_kwargs(self, kwargs) -> Optional[dict]: """Extract dynamic headers from kwargs if available.""" - standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = ( - kwargs.get("standard_callback_dynamic_params") - ) + standard_callback_dynamic_params: Optional[ + StandardCallbackDynamicParams + ] = kwargs.get("standard_callback_dynamic_params") if not standard_callback_dynamic_params: return None @@ -607,18 +611,35 @@ class OpenTelemetry(CustomLogger): ) ctx, parent_span = self._get_span_context(kwargs) - if get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN"): - primary_span_parent = None - else: - primary_span_parent = parent_span - - # 1. Primary span - span = self._start_primary_span( - kwargs, response_obj, start_time, end_time, ctx, primary_span_parent + # Decide whether to create a primary span + # Always create if no parent span exists (backward compatibility) + # OR if USE_OTEL_LITELLM_REQUEST_SPAN is explicitly enabled + should_create_primary_span = parent_span is None or get_secret_bool( + "USE_OTEL_LITELLM_REQUEST_SPAN" ) - # 2. Raw‐request sub-span (if enabled) - self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) + if should_create_primary_span: + # Create a new litellm_request span + span = self._start_primary_span( + kwargs, response_obj, start_time, end_time, ctx + ) + # Raw-request sub-span (if enabled) - child of litellm_request span + self._maybe_log_raw_request( + kwargs, response_obj, start_time, end_time, span + ) + else: + # Do not create primary span (keep hierarchy shallow when parent exists) + from opentelemetry.trace import Status, StatusCode + + span = None + # Only set attributes if the span is still recording (not closed) + # Note: parent_span is guaranteed to be not None here + parent_span.set_status(Status(StatusCode.OK)) + self.set_attributes(parent_span, kwargs, response_obj) + # Raw-request as direct child of parent_span + self._maybe_log_raw_request( + kwargs, response_obj, start_time, end_time, parent_span + ) # 3. Guardrail span self._create_guardrail_span(kwargs=kwargs, context=ctx) @@ -628,12 +649,18 @@ class OpenTelemetry(CustomLogger): # 5. Semantic logs. if self.config.enable_events: - self._emit_semantic_logs(kwargs, response_obj, span) + log_span = span if span is not None else parent_span + if log_span is not None: + self._emit_semantic_logs(kwargs, response_obj, log_span) - # 6. End parent span (only if it wasn't reused as the primary span) - # If parent_span was reused as the primary span, it was already ended in _start_primary_span - if parent_span is not None and parent_span is not span: - parent_span.end(end_time=self._to_ns(datetime.now())) + # 6. Do NOT end parent span - it should be managed by its creator + # External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM + # However, proxy-created spans should be closed here + if ( + parent_span is not None + and parent_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME + ): + parent_span.end(end_time=self._to_ns(end_time)) def _start_primary_span( self, @@ -642,16 +669,19 @@ class OpenTelemetry(CustomLogger): start_time, end_time, context, - parent_span: Optional[Span] = None, ): from opentelemetry.trace import Status, StatusCode otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs) - span = parent_span or otel_tracer.start_span( + + # Always create a new span + # The parent relationship is preserved through the context parameter + span = otel_tracer.start_span( name=self._get_span_name(kwargs), start_time=self._to_ns(start_time), context=context, ) + span.set_status(Status(StatusCode.OK)) self.set_attributes(span, kwargs, response_obj) span.end(end_time=self._to_ns(end_time)) @@ -764,10 +794,10 @@ class OpenTelemetry(CustomLogger): return float(val) # isinstance(val, str) - parse datetime string (with or without microseconds) try: - return datetime.strptime(val, '%Y-%m-%d %H:%M:%S.%f').timestamp() + return datetime.strptime(val, "%Y-%m-%d %H:%M:%S.%f").timestamp() except ValueError: try: - return datetime.strptime(val, '%Y-%m-%d %H:%M:%S').timestamp() + return datetime.strptime(val, "%Y-%m-%d %H:%M:%S").timestamp() except ValueError: return None @@ -775,23 +805,23 @@ class OpenTelemetry(CustomLogger): """Record Time to First Token (TTFT) metric for streaming requests.""" optional_params = kwargs.get("optional_params", {}) is_streaming = optional_params.get("stream", False) - + if not (self._time_to_first_token_histogram and is_streaming): return - + # Use api_call_start_time for precision (matches Prometheus implementation) # This excludes LiteLLM overhead and measures pure LLM API latency api_call_start_time = kwargs.get("api_call_start_time", None) completion_start_time = kwargs.get("completion_start_time", None) - + if api_call_start_time is not None and completion_start_time is not None: # Convert to timestamps if needed (handles datetime, float, and string) api_call_start_ts = self._to_timestamp(api_call_start_time) completion_start_ts = self._to_timestamp(completion_start_time) - + if api_call_start_ts is None or completion_start_ts is None: return # Skip recording if conversion failed - + time_to_first_token_seconds = completion_start_ts - api_call_start_ts self._time_to_first_token_histogram.record( time_to_first_token_seconds, attributes=common_attrs @@ -806,38 +836,40 @@ class OpenTelemetry(CustomLogger): common_attrs: dict, ): """Record Time Per Output Token (TPOT) metric. - + Calculated as: generation_time / completion_tokens - For streaming: uses end_time - completion_start_time (time to generate all tokens after first) - For non-streaming: uses end_time - api_call_start_time (total generation time) """ if not self._time_per_output_token_histogram: return - + # Get completion tokens from response_obj completion_tokens = None if response_obj and (usage := response_obj.get("usage")): completion_tokens = usage.get("completion_tokens") - + if completion_tokens is None or completion_tokens <= 0: return - + # Calculate generation time completion_start_time = kwargs.get("completion_start_time", None) api_call_start_time = kwargs.get("api_call_start_time", None) - + # Convert end_time to timestamp (handles datetime, float, and string) end_time_ts = self._to_timestamp(end_time) if end_time_ts is None: # Fallback to duration_s if conversion failed generation_time_seconds = duration_s if generation_time_seconds > 0: - time_per_output_token_seconds = generation_time_seconds / completion_tokens + time_per_output_token_seconds = ( + generation_time_seconds / completion_tokens + ) self._time_per_output_token_histogram.record( time_per_output_token_seconds, attributes=common_attrs ) return - + if completion_start_time is not None: # Streaming: use completion_start_time (when first token arrived) # This measures time to generate all tokens after the first one @@ -858,7 +890,7 @@ class OpenTelemetry(CustomLogger): else: # Fallback: use duration_s (already calculated as (end_time - start_time).total_seconds()) generation_time_seconds = duration_s - + if generation_time_seconds > 0: time_per_output_token_seconds = generation_time_seconds / completion_tokens self._time_per_output_token_histogram.record( @@ -872,37 +904,37 @@ class OpenTelemetry(CustomLogger): common_attrs: dict, ): """Record Total Generation Time (response duration) metric. - + Measures pure LLM API generation time: end_time - api_call_start_time This excludes LiteLLM overhead and measures only the LLM provider's response time. Works for both streaming and non-streaming requests. - + Mirrors Prometheus's litellm_llm_api_latency_metric. Uses kwargs.get("end_time") with fallback to parameter for consistency with Prometheus. """ if not self._response_duration_histogram: return - + api_call_start_time = kwargs.get("api_call_start_time", None) if api_call_start_time is None: return - + # Use end_time from kwargs if available (matches Prometheus), otherwise use parameter # For streaming: end_time is when the stream completes (final chunk received) # For non-streaming: end_time is when the response is received _end_time = kwargs.get("end_time") or end_time if _end_time is None: _end_time = datetime.now() - + # Convert to timestamps if needed (handles datetime, float, and string) api_call_start_ts = self._to_timestamp(api_call_start_time) end_time_ts = self._to_timestamp(_end_time) - + if api_call_start_ts is None or end_time_ts is None: return # Skip recording if conversion failed - + response_duration_seconds = end_time_ts - api_call_start_ts - + if response_duration_seconds > 0: self._response_duration_histogram.record( response_duration_seconds, attributes=common_attrs @@ -1065,26 +1097,49 @@ class OpenTelemetry(CustomLogger): ) _parent_context, parent_otel_span = self._get_span_context(kwargs) - # Span 1: Requst sent to litellm SDK - otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs) - span = otel_tracer.start_span( - name=self._get_span_name(kwargs), - start_time=self._to_ns(start_time), - context=_parent_context, + # Decide whether to create a primary span + # Always create if no parent span exists (backward compatibility) + # OR if USE_OTEL_LITELLM_REQUEST_SPAN is explicitly enabled + should_create_primary_span = parent_otel_span is None or get_secret_bool( + "USE_OTEL_LITELLM_REQUEST_SPAN" ) - span.set_status(Status(StatusCode.ERROR)) - self.set_attributes(span, kwargs, response_obj) - # Record exception information using OTEL standard method - self._record_exception_on_span(span=span, kwargs=kwargs) + if should_create_primary_span: + # Span 1: Request sent to litellm SDK + otel_tracer: Tracer = self.get_tracer_to_use_for_request(kwargs) + span = otel_tracer.start_span( + name=self._get_span_name(kwargs), + start_time=self._to_ns(start_time), + context=_parent_context, + ) + span.set_status(Status(StatusCode.ERROR)) + self.set_attributes(span, kwargs, response_obj) - span.end(end_time=self._to_ns(end_time)) + # Record exception information using OTEL standard method + self._record_exception_on_span(span=span, kwargs=kwargs) + + span.end(end_time=self._to_ns(end_time)) + else: + # When parent span exists and USE_OTEL_LITELLM_REQUEST_SPAN=false, + # record error on parent span (keeps hierarchy shallow) + # Only set attributes if the span is still recording (not closed) + # Note: parent_otel_span is guaranteed to be not None here + if parent_otel_span.is_recording(): + parent_otel_span.set_status(Status(StatusCode.ERROR)) + self.set_attributes(parent_otel_span, kwargs, response_obj) + self._record_exception_on_span(span=parent_otel_span, kwargs=kwargs) # Create span for guardrail information self._create_guardrail_span(kwargs=kwargs, context=_parent_context) - if parent_otel_span is not None: - parent_otel_span.end(end_time=self._to_ns(datetime.now())) + # Do NOT end parent span - it should be managed by its creator + # External spans (from Langfuse, user code, HTTP headers, global context) must not be closed by LiteLLM + # However, proxy-created spans should be closed here + if ( + parent_otel_span is not None + and parent_otel_span.name == LITELLM_PROXY_REQUEST_SPAN_NAME + ): + parent_otel_span.end(end_time=self._to_ns(end_time)) def _record_exception_on_span(self, span: Span, kwargs: dict): """ @@ -1263,7 +1318,9 @@ class OpenTelemetry(CustomLogger): ) return elif self.callback_name == "weave_otel": - from litellm.integrations.weave.weave_otel import set_weave_otel_attributes + from litellm.integrations.weave.weave_otel import ( + set_weave_otel_attributes, + ) set_weave_otel_attributes(span, kwargs, response_obj) return @@ -1994,9 +2051,9 @@ class OpenTelemetry(CustomLogger): """ Create a span for the received proxy server request. """ - + return self.tracer.start_span( - name="Received Proxy Server Request", + name=LITELLM_PROXY_REQUEST_SPAN_NAME, start_time=self._to_ns(start_time), context=self.get_traceparent_from_header(headers=headers), kind=self.span_kind.SERVER, diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 20f1357a1c8..c01f7481277 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -214,7 +214,7 @@ class PrometheusLogger(CustomLogger): # Remaining Rate Limit for model self.litellm_remaining_requests_metric = self._gauge_factory( - "litellm_remaining_requests", + "litellm_remaining_requests_metric", "LLM Deployment Analytics - remaining requests for model, returned from LLM API Provider", labelnames=self.get_labels_for_metric( "litellm_remaining_requests_metric" @@ -222,7 +222,7 @@ class PrometheusLogger(CustomLogger): ) self.litellm_remaining_tokens_metric = self._gauge_factory( - "litellm_remaining_tokens", + "litellm_remaining_tokens_metric", "remaining tokens for model, returned from LLM API Provider", labelnames=self.get_labels_for_metric( "litellm_remaining_tokens_metric" diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 9378ca71f54..dadb36f3fd7 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -38,18 +38,18 @@ def safe_divide_seconds( def safe_divide( - numerator: Union[int, float], - denominator: Union[int, float], - default: Union[int, float] = 0 + numerator: Union[int, float], + denominator: Union[int, float], + default: Union[int, float] = 0, ) -> Union[int, float]: """ Safely divide two numbers, returning a default value if denominator is zero. - + Args: numerator: The number to divide denominator: The number to divide by default: Value to return if denominator is zero (defaults to 0) - + Returns: The result of numerator/denominator, or default if denominator is zero """ @@ -153,7 +153,8 @@ def get_metadata_variable_name_from_kwargs( - LiteLLM is now moving to using `litellm_metadata` for our metadata """ return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - + + def get_litellm_metadata_from_kwargs(kwargs: dict): """ Helper to get litellm metadata from all litellm request kwargs @@ -176,6 +177,25 @@ def get_litellm_metadata_from_kwargs(kwargs: dict): return {} +def reconstruct_model_name( + model_name: str, + custom_llm_provider: Optional[str], + metadata: dict, +) -> str: + """Reconstruct full model name with provider prefix for logging.""" + # Check if deployment model name from router metadata is available (has original prefix) + deployment_model_name = metadata.get("deployment") + if deployment_model_name and "/" in deployment_model_name: + # Use the deployment model name which preserves the original provider prefix + return deployment_model_name + elif custom_llm_provider and model_name and "/" not in model_name: + # Only add prefix for Bedrock (not for direct Anthropic API) + # This ensures Bedrock models get the prefix while direct Anthropic models don't + if custom_llm_provider == "bedrock": + return f"{custom_llm_provider}/{model_name}" + return model_name + + # Helper functions used for OTEL logging def _get_parent_otel_span_from_kwargs( kwargs: Optional[dict] = None, @@ -246,8 +266,8 @@ def safe_deep_copy(data): Safe Deep Copy The LiteLLM request may contain objects that cannot be pickled/deep-copied - (e.g., tracing spans, locks, clients). - + (e.g., tracing spans, locks, clients). + This helper deep-copies each top-level key independently; on failure keeps original ref """ @@ -306,23 +326,23 @@ def safe_deep_copy(data): def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: """ Recursively filter out Exception objects and callable objects from dicts/lists. - + This is a defensive utility to prevent deepcopy failures when exception objects are accidentally stored in parameter dictionaries (e.g., optional_params). Also filters callable objects (functions) to prevent JSON serialization errors. Exceptions and callables should not be stored in params - this function removes them. - + Args: data: The data structure to filter (dict, list, or any other type) max_depth: Maximum recursion depth to prevent infinite loops - + Returns: Filtered data structure with Exception and callable objects removed, or None if the entire input was an Exception or callable """ if max_depth <= 0: return data - + # Skip exception objects if isinstance(data, Exception): return None @@ -333,7 +353,7 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: obj_type_name = type(data).__name__ if obj_type_name in ["Logging", "LiteLLMLoggingObj"]: return None - + if isinstance(data, dict): result: dict[str, Any] = {} for k, v in data.items(): @@ -352,7 +372,9 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: result_list: list[Any] = [] for item in data: # Skip exception and callable items - if isinstance(item, Exception) or (callable(item) and not isinstance(item, type)): + if isinstance(item, Exception) or ( + callable(item) and not isinstance(item, type) + ): continue try: filtered = filter_exceptions_from_params(item, max_depth - 1) @@ -366,37 +388,35 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: return data -def filter_internal_params(data: dict, additional_internal_params: Optional[set] = None) -> dict: +def filter_internal_params( + data: dict, additional_internal_params: Optional[set] = None +) -> dict: """ Filter out LiteLLM internal parameters that shouldn't be sent to provider APIs. - + This removes internal/MCP-related parameters that are used by LiteLLM internally but should not be included in API requests to providers. - + Args: data: Dictionary of parameters to filter additional_internal_params: Optional set of additional internal parameter names to filter - + Returns: Filtered dictionary with internal parameters removed """ if not isinstance(data, dict): return data - + # Known internal parameters that should never be sent to provider APIs internal_params = { "skip_mcp_handler", "mcp_handler_context", "_skip_mcp_handler", } - + # Add any additional internal params if provided if additional_internal_params: internal_params.update(additional_internal_params) - + # Filter out internal parameters - return { - k: v - for k, v in data.items() - if k not in internal_params - } \ No newline at end of file + return {k: v for k, v in data.items() if k not in internal_params} diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index fa2ff42e1df..47cbcb8aec9 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -76,6 +76,7 @@ class CustomLoggerRegistry: "arize_phoenix": OpenTelemetry, "langtrace": OpenTelemetry, "weave_otel": OpenTelemetry, + "levo": OpenTelemetry, "mlflow": MlflowLogger, "langfuse": LangfusePromptManagement, "otel": OpenTelemetry, diff --git a/litellm/litellm_core_utils/dot_notation_indexing.py b/litellm/litellm_core_utils/dot_notation_indexing.py index 6e293a4cb77..1e835004e94 100644 --- a/litellm/litellm_core_utils/dot_notation_indexing.py +++ b/litellm/litellm_core_utils/dot_notation_indexing.py @@ -9,6 +9,7 @@ Custom implementation with zero external dependencies. Supported syntax: - "field" - top-level field - "parent.child" - nested field +- "parent\\.with\\.dots.child" - keys containing dots (escape with backslash) - "array[*]" - all array elements (wildcard) - "array[0]" - specific array element (index) - "array[*].field" - field in all array elements @@ -47,6 +48,9 @@ def get_nested_value( 'value' >>> get_nested_value(data, "a.b.d", "default") 'default' + >>> data = {"kubernetes.io": {"namespace": "default"}} + >>> get_nested_value(data, "kubernetes\\.io.namespace") + 'default' """ if not key_path: return default @@ -58,8 +62,11 @@ def get_nested_value( else key_path ) - # Split the key path into parts - parts = key_path.split(".") + # Split the key path into parts, respecting escaped dots (\.) + # Use a temporary placeholder, split on unescaped dots, then restore + placeholder = "\x00" + parts = key_path.replace("\\.", placeholder).split(".") + parts = [p.replace(placeholder, ".") for p in parts] # Traverse through the dictionary current: Any = data diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 164e2a73e65..b753e9fa8b5 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -229,10 +229,10 @@ def get_llm_provider( # noqa: PLR0915 elif endpoint == "https://api.ai21.com/studio/v1": custom_llm_provider = "ai21_chat" dynamic_api_key = get_secret_str("AI21_API_KEY") - elif endpoint == "https://codestral.mistral.ai/v1": + elif endpoint == "codestral.mistral.ai/v1/chat/completions": custom_llm_provider = "codestral" dynamic_api_key = get_secret_str("CODESTRAL_API_KEY") - elif endpoint == "https://codestral.mistral.ai/v1": + elif endpoint == "codestral.mistral.ai/v1/fim/completions": custom_llm_provider = "text-completion-codestral" dynamic_api_key = get_secret_str("CODESTRAL_API_KEY") elif endpoint == "app.empower.dev/api/v1": diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index c0090d1c3e7..cd324935562 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -59,6 +59,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.deepeval.deepeval import DeepEvalLogger from litellm.integrations.mlflow import MlflowLogger from litellm.integrations.sqs import SQSLogger +from litellm.litellm_core_utils.core_helpers import reconstruct_model_name from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( StandardBuiltInToolCostTracking, @@ -332,9 +333,9 @@ class Logging(LiteLLMLoggingBaseClass): self.litellm_trace_id: str = litellm_trace_id or str(uuid.uuid4()) self.function_id = function_id self.streaming_chunks: List[Any] = [] # for generating complete stream response - self.sync_streaming_chunks: List[Any] = ( - [] - ) # for generating complete stream response + self.sync_streaming_chunks: List[ + Any + ] = [] # for generating complete stream response self.log_raw_request_response = log_raw_request_response # Initialize dynamic callbacks @@ -719,9 +720,9 @@ class Logging(LiteLLMLoggingBaseClass): prompt_spec=prompt_spec, dynamic_callback_params=dynamic_callback_params, ): - self.model_call_details["prompt_integration"] = ( - logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = logger.__class__.__name__ return logger except Exception: # If check fails, continue to next logger @@ -789,9 +790,9 @@ class Logging(LiteLLMLoggingBaseClass): if anthropic_cache_control_logger := AnthropicCacheControlHook.get_custom_logger_for_anthropic_cache_control_hook( non_default_params ): - self.model_call_details["prompt_integration"] = ( - anthropic_cache_control_logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = anthropic_cache_control_logger.__class__.__name__ return anthropic_cache_control_logger ######################################################### @@ -803,9 +804,9 @@ class Logging(LiteLLMLoggingBaseClass): internal_usage_cache=None, llm_router=None, ) - self.model_call_details["prompt_integration"] = ( - vector_store_custom_logger.__class__.__name__ - ) + self.model_call_details[ + "prompt_integration" + ] = vector_store_custom_logger.__class__.__name__ # Add to global callbacks so post-call hooks are invoked if ( vector_store_custom_logger @@ -865,9 +866,9 @@ class Logging(LiteLLMLoggingBaseClass): model ): # if model name was changes pre-call, overwrite the initial model call name with the new one self.model_call_details["model"] = model - self.model_call_details["litellm_params"]["api_base"] = ( - self._get_masked_api_base(additional_args.get("api_base", "")) - ) + self.model_call_details["litellm_params"][ + "api_base" + ] = self._get_masked_api_base(additional_args.get("api_base", "")) def pre_call(self, input, api_key, model=None, additional_args={}): # noqa: PLR0915 # Log the exact input to the LLM API @@ -896,10 +897,10 @@ class Logging(LiteLLMLoggingBaseClass): try: # [Non-blocking Extra Debug Information in metadata] if turn_off_message_logging is True: - _metadata["raw_request"] = ( - "redacted by litellm. \ + _metadata[ + "raw_request" + ] = "redacted by litellm. \ 'litellm.turn_off_message_logging=True'" - ) else: curl_command = self._get_request_curl_command( api_base=additional_args.get("api_base", ""), @@ -910,34 +911,34 @@ class Logging(LiteLLMLoggingBaseClass): _metadata["raw_request"] = str(curl_command) # split up, so it's easier to parse in the UI - self.model_call_details["raw_request_typed_dict"] = ( - RawRequestTypedDict( - raw_request_api_base=str( - additional_args.get("api_base") or "" - ), - raw_request_body=self._get_raw_request_body( - additional_args.get("complete_input_dict", {}) - ), - # NOTE: setting ignore_sensitive_headers to True will cause - # the Authorization header to be leaked when calls to the health - # endpoint are made and fail. - raw_request_headers=self._get_masked_headers( - additional_args.get("headers", {}) or {}, - ), - error=None, - ) + self.model_call_details[ + "raw_request_typed_dict" + ] = RawRequestTypedDict( + raw_request_api_base=str( + additional_args.get("api_base") or "" + ), + raw_request_body=self._get_raw_request_body( + additional_args.get("complete_input_dict", {}) + ), + # NOTE: setting ignore_sensitive_headers to True will cause + # the Authorization header to be leaked when calls to the health + # endpoint are made and fail. + raw_request_headers=self._get_masked_headers( + additional_args.get("headers", {}) or {}, + ), + error=None, ) except Exception as e: - self.model_call_details["raw_request_typed_dict"] = ( - RawRequestTypedDict( - error=str(e), - ) + self.model_call_details[ + "raw_request_typed_dict" + ] = RawRequestTypedDict( + error=str(e), ) - _metadata["raw_request"] = ( - "Unable to Log \ + _metadata[ + "raw_request" + ] = "Unable to Log \ raw request: {}".format( - str(e) - ) + str(e) ) if getattr(self, "logger_fn", None) and callable(self.logger_fn): try: @@ -1238,13 +1239,13 @@ class Logging(LiteLLMLoggingBaseClass): for callback in callbacks: try: if isinstance(callback, CustomLogger): - response: Optional[MCPPostCallResponseObject] = ( - await callback.async_post_mcp_tool_call_hook( - kwargs=kwargs, - response_obj=post_mcp_tool_call_response_obj, - start_time=start_time, - end_time=end_time, - ) + response: Optional[ + MCPPostCallResponseObject + ] = await callback.async_post_mcp_tool_call_hook( + kwargs=kwargs, + response_obj=post_mcp_tool_call_response_obj, + start_time=start_time, + end_time=end_time, ) ###################################################################### # if any of the callbacks modify the response, use the modified response @@ -1423,9 +1424,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details["response_cost_failure_debug_information"] = ( - debug_info - ) + self.model_call_details[ + "response_cost_failure_debug_information" + ] = debug_info return None try: @@ -1451,9 +1452,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( f"response_cost_failure_debug_information: {debug_info}" ) - self.model_call_details["response_cost_failure_debug_information"] = ( - debug_info - ) + self.model_call_details[ + "response_cost_failure_debug_information" + ] = debug_info return None @@ -1603,16 +1604,16 @@ class Logging(LiteLLMLoggingBaseClass): result=logging_result ) - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=logging_result, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=logging_result, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) def _transform_usage_objects(self, result): @@ -1667,9 +1668,9 @@ class Logging(LiteLLMLoggingBaseClass): end_time = datetime.datetime.now() if self.completion_start_time is None: self.completion_start_time = end_time - self.model_call_details["completion_start_time"] = ( - self.completion_start_time - ) + self.model_call_details[ + "completion_start_time" + ] = self.completion_start_time self.model_call_details["log_event_type"] = "successful_api_call" self.model_call_details["end_time"] = end_time @@ -1706,21 +1707,21 @@ class Logging(LiteLLMLoggingBaseClass): end_time=end_time, ) elif isinstance(result, dict) or isinstance(result, list): - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=result, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=result, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) elif standard_logging_object is not None: - self.model_call_details["standard_logging_object"] = ( - standard_logging_object - ) + self.model_call_details[ + "standard_logging_object" + ] = standard_logging_object else: self.model_call_details["response_cost"] = None @@ -1870,23 +1871,23 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.debug( "Logging Details LiteLLM-Success Call streaming complete" ) - self.model_call_details["complete_streaming_response"] = ( - complete_streaming_response - ) - self.model_call_details["response_cost"] = ( - self._response_cost_calculator(result=complete_streaming_response) - ) + self.model_call_details[ + "complete_streaming_response" + ] = complete_streaming_response + self.model_call_details[ + "response_cost" + ] = self._response_cost_calculator(result=complete_streaming_response) ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=complete_streaming_response, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=complete_streaming_response, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) callbacks = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_success_callbacks, @@ -2214,10 +2215,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = ( - self.model_call_details.get( - "complete_streaming_response", {} - ) + self.model_call_details[ + "complete_response" + ] = self.model_call_details.get( + "complete_streaming_response", {} ) result = self.model_call_details["complete_response"] openMeterLogger.log_success_event( @@ -2256,10 +2257,10 @@ class Logging(LiteLLMLoggingBaseClass): ) else: if self.stream and complete_streaming_response: - self.model_call_details["complete_response"] = ( - self.model_call_details.get( - "complete_streaming_response", {} - ) + self.model_call_details[ + "complete_response" + ] = self.model_call_details.get( + "complete_streaming_response", {} ) result = self.model_call_details["complete_response"] @@ -2402,9 +2403,9 @@ class Logging(LiteLLMLoggingBaseClass): if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") - self.model_call_details["async_complete_streaming_response"] = ( - complete_streaming_response - ) + self.model_call_details[ + "async_complete_streaming_response" + ] = complete_streaming_response try: if self.model_call_details.get("cache_hit", False) is True: @@ -2415,10 +2416,10 @@ class Logging(LiteLLMLoggingBaseClass): model_call_details=self.model_call_details ) # base_model defaults to None if not set on model_info - self.model_call_details["response_cost"] = ( - self._response_cost_calculator( - result=complete_streaming_response - ) + self.model_call_details[ + "response_cost" + ] = self._response_cost_calculator( + result=complete_streaming_response ) verbose_logger.debug( @@ -2431,16 +2432,16 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = None ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj=complete_streaming_response, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="success", - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj=complete_streaming_response, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="success", + standard_built_in_tools_params=self.standard_built_in_tools_params, ) callbacks = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_async_success_callbacks, @@ -2676,18 +2677,18 @@ class Logging(LiteLLMLoggingBaseClass): ## STANDARDIZED LOGGING PAYLOAD - self.model_call_details["standard_logging_object"] = ( - get_standard_logging_object_payload( - kwargs=self.model_call_details, - init_response_obj={}, - start_time=start_time, - end_time=end_time, - logging_obj=self, - status="failure", - error_str=str(exception), - original_exception=exception, - standard_built_in_tools_params=self.standard_built_in_tools_params, - ) + self.model_call_details[ + "standard_logging_object" + ] = get_standard_logging_object_payload( + kwargs=self.model_call_details, + init_response_obj={}, + start_time=start_time, + end_time=end_time, + logging_obj=self, + status="failure", + error_str=str(exception), + original_exception=exception, + standard_built_in_tools_params=self.standard_built_in_tools_params, ) return start_time, end_time @@ -3301,7 +3302,9 @@ class Logging(LiteLLMLoggingBaseClass): # Deep copy result and add usage result_copy = result.model_copy(deep=True) - result_copy.usage = usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) + result_copy.usage = ( + usage.model_dump() if hasattr(usage, "model_dump") else dict(usage) + ) return result_copy @@ -3629,9 +3632,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 endpoint=arize_config.endpoint, ) - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}" for callback in _in_memory_loggers: if ( isinstance(callback, ArizeLogger) @@ -3642,7 +3645,6 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _in_memory_loggers.append(_arize_otel_logger) return _arize_otel_logger # type: ignore elif logging_integration == "arize_phoenix": - from litellm.integrations.opentelemetry import ( OpenTelemetry, OpenTelemetryConfig, @@ -3658,13 +3660,13 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"{existing_attrs},openinference.project.name={arize_phoenix_config.project_name}" else: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"openinference.project.name={arize_phoenix_config.project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"openinference.project.name={arize_phoenix_config.project_name}" # Set Phoenix project name from environment variable phoenix_project_name = os.environ.get("PHOENIX_PROJECT_NAME", None) @@ -3672,19 +3674,19 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 existing_attrs = os.environ.get("OTEL_RESOURCE_ATTRIBUTES", "") # Add openinference.project.name attribute if existing_attrs: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"{existing_attrs},openinference.project.name={phoenix_project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"{existing_attrs},openinference.project.name={phoenix_project_name}" else: - os.environ["OTEL_RESOURCE_ATTRIBUTES"] = ( - f"openinference.project.name={phoenix_project_name}" - ) + os.environ[ + "OTEL_RESOURCE_ATTRIBUTES" + ] = f"openinference.project.name={phoenix_project_name}" # auth can be disabled on local deployments of arize phoenix if arize_phoenix_config.otlp_auth_headers is not None: - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - arize_phoenix_config.otlp_auth_headers - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = arize_phoenix_config.otlp_auth_headers for callback in _in_memory_loggers: if ( @@ -3697,6 +3699,31 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 ) _in_memory_loggers.append(_arize_phoenix_otel_logger) return _arize_phoenix_otel_logger # type: ignore + elif logging_integration == "levo": + from litellm.integrations.levo.levo import LevoLogger + from litellm.integrations.opentelemetry import ( + OpenTelemetry, + OpenTelemetryConfig, + ) + + levo_config = LevoLogger.get_levo_config() + otel_config = OpenTelemetryConfig( + exporter=levo_config.protocol, + endpoint=levo_config.endpoint, + headers=levo_config.otlp_auth_headers, + ) + + # Check if LevoLogger instance already exists + for callback in _in_memory_loggers: + if ( + isinstance(callback, LevoLogger) + and callback.callback_name == "levo" + ): + return callback # type: ignore + + _levo_otel_logger = LevoLogger(config=otel_config, callback_name="levo") + _in_memory_loggers.append(_levo_otel_logger) + return _levo_otel_logger # type: ignore elif logging_integration == "otel": from litellm.integrations.opentelemetry import OpenTelemetry @@ -3816,9 +3843,9 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 exporter="otlp_http", endpoint="https://langtrace.ai/api/trace", ) - os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( - f"api_key={os.getenv('LANGTRACE_API_KEY')}" - ) + os.environ[ + "OTEL_EXPORTER_OTLP_TRACES_HEADERS" + ] = f"api_key={os.getenv('LANGTRACE_API_KEY')}" for callback in _in_memory_loggers: if ( isinstance(callback, OpenTelemetry) @@ -4589,10 +4616,10 @@ class StandardLoggingPayloadSetup: for key in StandardLoggingHiddenParams.__annotations__.keys(): if key in hidden_params: if key == "additional_headers": - clean_hidden_params["additional_headers"] = ( - StandardLoggingPayloadSetup.get_additional_headers( - hidden_params[key] - ) + clean_hidden_params[ + "additional_headers" + ] = StandardLoggingPayloadSetup.get_additional_headers( + hidden_params[key] ) else: clean_hidden_params[key] = hidden_params[key] # type: ignore @@ -4898,25 +4925,6 @@ def _extract_response_obj_and_hidden_params( return response_obj, hidden_params -def _reconstruct_model_name( - model_name: str, - custom_llm_provider: Optional[str], - metadata: dict, -) -> str: - """Reconstruct full model name with provider prefix for logging.""" - # Check if deployment model name from router metadata is available (has original prefix) - deployment_model_name = metadata.get("deployment") - if deployment_model_name and "/" in deployment_model_name: - # Use the deployment model name which preserves the original provider prefix - return deployment_model_name - elif custom_llm_provider and model_name and "/" not in model_name: - # Only add prefix for Bedrock (not for direct Anthropic API) - # This ensures Bedrock models get the prefix while direct Anthropic models don't - if custom_llm_provider == "bedrock": - return f"{custom_llm_provider}/{model_name}" - return model_name - - def get_standard_logging_object_payload( kwargs: Optional[dict], init_response_obj: Union[Any, BaseModel, dict], @@ -5049,7 +5057,7 @@ def get_standard_logging_object_payload( # This ensures Bedrock models like "us.anthropic.claude-3-5-sonnet-20240620-v1:0" # are logged as "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0" custom_llm_provider = cast(Optional[str], kwargs.get("custom_llm_provider")) - model_name = _reconstruct_model_name( + model_name = reconstruct_model_name( kwargs.get("model", "") or "", custom_llm_provider, metadata ) @@ -5205,9 +5213,9 @@ def scrub_sensitive_keys_in_metadata(litellm_params: Optional[dict]): ): for k, v in metadata["user_api_key_metadata"].items(): if k == "logging": # prevent logging user logging keys - cleaned_user_api_key_metadata[k] = ( - "scrubbed_by_litellm_for_sensitive_keys" - ) + cleaned_user_api_key_metadata[ + k + ] = "scrubbed_by_litellm_for_sensitive_keys" else: cleaned_user_api_key_metadata[k] = v diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 232d9bfc5d1..cbc0763382c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -161,6 +161,15 @@ def _get_token_base_cost( prompt_base_cost = cast(float, _get_cost_per_unit(model_info, input_cost_key)) completion_base_cost = cast(float, _get_cost_per_unit(model_info, output_cost_key)) + + # For image generation models that don't have output_cost_per_token, + # use output_cost_per_image_token as the base cost (all output tokens are image tokens) + if completion_base_cost == 0.0 or completion_base_cost is None: + output_image_cost = _get_cost_per_unit( + model_info, "output_cost_per_image_token", None + ) + if output_image_cost is not None: + completion_base_cost = cast(float, output_image_cost) cache_creation_cost = cast( float, _get_cost_per_unit(model_info, cache_creation_cost_key) ) @@ -342,6 +351,7 @@ class PromptTokensDetailsResult(TypedDict): cache_creation_token_details: Optional[CacheCreationTokenDetails] text_tokens: int audio_tokens: int + image_tokens: int character_count: int image_count: int video_length_seconds: int @@ -374,6 +384,10 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: cast(Optional[int], getattr(usage.prompt_tokens_details, "audio_tokens", 0)) or 0 ) + image_tokens = ( + cast(Optional[int], getattr(usage.prompt_tokens_details, "image_tokens", 0)) + or 0 + ) character_count = ( cast( Optional[int], @@ -398,6 +412,7 @@ def _parse_prompt_tokens_details(usage: Usage) -> PromptTokensDetailsResult: cache_creation_token_details=cache_creation_token_details, text_tokens=text_tokens, audio_tokens=audio_tokens, + image_tokens=image_tokens, character_count=character_count, image_count=image_count, video_length_seconds=video_length_seconds, @@ -470,6 +485,11 @@ def _calculate_input_cost( model_info, "input_cost_per_audio_token", prompt_tokens_details["audio_tokens"] ) + ### IMAGE TOKEN COST (for gpt-image-1 and similar models) + prompt_cost += calculate_cost_component( + model_info, "input_cost_per_image_token", prompt_tokens_details["image_tokens"] + ) + ### CACHE WRITING COST - Now uses tiered pricing prompt_cost += calculate_cache_writing_cost( cache_creation_tokens=prompt_tokens_details["cache_creation_tokens"], @@ -533,6 +553,7 @@ def generic_cost_per_token( cache_creation_token_details=None, text_tokens=usage.prompt_tokens, audio_tokens=0, + image_tokens=0, character_count=0, image_count=0, video_length_seconds=0, @@ -583,12 +604,22 @@ def generic_cost_per_token( reasoning_tokens = completion_tokens_details["reasoning_tokens"] image_tokens = completion_tokens_details["image_tokens"] - # Only assume all tokens are text if there's NO breakdown at all - # If image_tokens, audio_tokens, or reasoning_tokens exist, respect text_tokens=0 + # Handle text_tokens calculation: + # 1. If text_tokens is explicitly provided and > 0, use it + # 2. If there's a breakdown (reasoning/audio/image tokens), calculate text_tokens as the remainder + # 3. If no breakdown at all, assume all completion_tokens are text_tokens has_token_breakdown = image_tokens > 0 or audio_tokens > 0 or reasoning_tokens > 0 - if text_tokens == 0 and not has_token_breakdown: - text_tokens = usage.completion_tokens - is_text_tokens_total = True + if text_tokens == 0: + if has_token_breakdown: + # Calculate text tokens as remainder when we have a breakdown + # This handles cases like OpenAI's reasoning models where text_tokens isn't provided + text_tokens = max( + 0, usage.completion_tokens - reasoning_tokens - audio_tokens - image_tokens + ) + else: + # No breakdown at all, all tokens are text tokens + text_tokens = usage.completion_tokens + is_text_tokens_total = True ## TEXT COST completion_cost = float(text_tokens) * completion_base_cost @@ -782,6 +813,50 @@ class CostCalculatorUtils: model=model, image_response=completion_response, ) + elif custom_llm_provider == litellm.LlmProviders.OPENAI.value: + # Check if this is a gpt-image model (token-based pricing) + model_lower = model.lower() + if "gpt-image-1" in model_lower: + from litellm.llms.openai.image_generation.cost_calculator import ( + cost_calculator as openai_gpt_image_cost_calculator, + ) + + return openai_gpt_image_cost_calculator( + model=model, + image_response=completion_response, + custom_llm_provider=custom_llm_provider, + ) + # Fall through to default for DALL-E models + return default_image_cost_calculator( + model=model, + quality=quality, + custom_llm_provider=custom_llm_provider, + n=n, + size=size, + optional_params=optional_params, + ) + elif custom_llm_provider == litellm.LlmProviders.AZURE.value: + # Check if this is a gpt-image model (token-based pricing) + model_lower = model.lower() + if "gpt-image-1" in model_lower: + from litellm.llms.openai.image_generation.cost_calculator import ( + cost_calculator as openai_gpt_image_cost_calculator, + ) + + return openai_gpt_image_cost_calculator( + model=model, + image_response=completion_response, + custom_llm_provider=custom_llm_provider, + ) + # Fall through to default for DALL-E models + return default_image_cost_calculator( + model=model, + quality=quality, + custom_llm_provider=custom_llm_provider, + n=n, + size=size, + optional_params=optional_params, + ) else: return default_image_cost_calculator( model=model, diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index 59d2a8a8dd0..bbe28e3ec2c 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -445,25 +445,43 @@ def convert_to_model_response_object( # noqa: PLR0915 hidden_params["additional_headers"] = additional_headers ### CHECK IF ERROR IN RESPONSE ### - openrouter returns these in the dictionary + # Some OpenAI-compatible providers (e.g., Apertis) return empty error objects + # even on success. Only raise if the error contains meaningful data. if ( response_object is not None and "error" in response_object and response_object["error"] is not None ): - error_args = {"status_code": 422, "message": "Error in response object"} - if isinstance(response_object["error"], dict): - if "code" in response_object["error"]: - error_args["status_code"] = response_object["error"]["code"] - if "message" in response_object["error"]: - if isinstance(response_object["error"]["message"], dict): - message_str = json.dumps(response_object["error"]["message"]) - else: - message_str = str(response_object["error"]["message"]) - error_args["message"] = message_str - raised_exception = Exception() - setattr(raised_exception, "status_code", error_args["status_code"]) - setattr(raised_exception, "message", error_args["message"]) - raise raised_exception + error_obj = response_object["error"] + has_meaningful_error = False + + if isinstance(error_obj, dict): + # Check if error dict has non-empty message or non-null code + error_message = error_obj.get("message", "") + error_code = error_obj.get("code") + has_meaningful_error = bool(error_message) or error_code is not None + elif isinstance(error_obj, str): + # String error is meaningful if non-empty + has_meaningful_error = bool(error_obj) + else: + # Any other truthy value is considered meaningful + has_meaningful_error = True + + if has_meaningful_error: + error_args = {"status_code": 422, "message": "Error in response object"} + if isinstance(error_obj, dict): + if "code" in error_obj: + error_args["status_code"] = error_obj["code"] + if "message" in error_obj: + if isinstance(error_obj["message"], dict): + message_str = json.dumps(error_obj["message"]) + else: + message_str = str(error_obj["message"]) + error_args["message"] = message_str + raised_exception = Exception() + setattr(raised_exception, "status_code", error_args["status_code"]) + setattr(raised_exception, "message", error_args["message"]) + raise raised_exception try: if response_type == "completion" and ( diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index b78484816da..4f76a5bad03 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -166,6 +166,7 @@ class LoggingCallbackManager: endpoint = callback_config.get("endpoint") headers = callback_config.get("headers") event_types = callback_config.get("event_types") + log_format = callback_config.get("log_format") if endpoint is None or headers is None: verbose_logger.warning( @@ -180,6 +181,7 @@ class LoggingCallbackManager: and cached_logger.endpoint == endpoint and cached_logger.headers == headers and cached_logger.event_types == event_types + and cached_logger.log_format == log_format ): return cached_logger @@ -187,6 +189,7 @@ class LoggingCallbackManager: endpoint=endpoint, headers=headers, event_types=event_types, + log_format=log_format, ) _generic_api_logger_cache[callback] = new_logger return new_logger diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 20b0bc92fb7..13a83956edd 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -51,6 +51,7 @@ class LoggingWorker: self._worker_task: Optional[asyncio.Task] = None self._running_tasks: set[asyncio.Task] = set() self._sem: Optional[asyncio.Semaphore] = None + self._bound_loop: Optional[asyncio.AbstractEventLoop] = None self._last_aggressive_clear_time: float = 0.0 self._aggressive_clear_in_progress: bool = False @@ -58,9 +59,27 @@ class LoggingWorker: atexit.register(self._flush_on_exit) def _ensure_queue(self) -> None: - """Initialize the queue if it doesn't exist.""" + """Initialize the queue if it doesn't exist or if event loop has changed.""" + try: + current_loop = asyncio.get_running_loop() + except RuntimeError: + # No running loop, can't initialize + return + + # Check if we need to reinitialize due to event loop change + if self._queue is not None and self._bound_loop is not current_loop: + verbose_logger.debug( + "LoggingWorker: Event loop changed, reinitializing queue and worker" + ) + # Clear old state - these are bound to the old loop + self._queue = None + self._sem = None + self._worker_task = None + self._running_tasks.clear() + if self._queue is None: self._queue = asyncio.Queue(maxsize=self.max_queue_size) + self._bound_loop = current_loop def start(self) -> None: """Start the logging worker. Idempotent - safe to call multiple times.""" @@ -126,7 +145,7 @@ class LoggingWorker: # Capture the current context when enqueueing task = LoggingTask(coroutine=coroutine, context=contextvars.copy_context()) - + try: self._queue.put_nowait(task) except asyncio.QueueFull: @@ -141,15 +160,15 @@ class LoggingWorker: """ if self._aggressive_clear_in_progress: return False - + try: loop = asyncio.get_running_loop() current_time = loop.time() time_since_last_clear = current_time - self._last_aggressive_clear_time - + if time_since_last_clear < LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS: return False - + return True except RuntimeError: # No event loop running, drop the task @@ -158,8 +177,8 @@ class LoggingWorker: def _mark_aggressive_clear_started(self) -> None: """ Mark that an aggressive clear operation has started. - - Note: This should only be called after _should_start_aggressive_clear() + + Note: This should only be called after _should_start_aggressive_clear() returns True, which guarantees an event loop exists. """ loop = asyncio.get_running_loop() @@ -171,7 +190,7 @@ class LoggingWorker: Handle queue full condition by either starting an aggressive clear or scheduling a delayed retry. """ - + if self._should_start_aggressive_clear(): self._mark_aggressive_clear_started() # Schedule clearing as async task so enqueue returns immediately (non-blocking) @@ -191,7 +210,8 @@ class LoggingWorker: time_since_last_clear = current_time - self._last_aggressive_clear_time remaining_cooldown = max( 0.0, - LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS - time_since_last_clear + LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS + - time_since_last_clear, ) # Add a small buffer (10% of cooldown or 50ms, whichever is larger) to ensure # cooldown has expired and aggressive clear has completed @@ -212,7 +232,7 @@ class LoggingWorker: # Check that we have a running event loop (will raise RuntimeError if not) asyncio.get_running_loop() delay = self._calculate_retry_delay() - + # Schedule the retry as a background task asyncio.create_task(self._retry_enqueue_task(task, delay)) except RuntimeError: @@ -225,11 +245,11 @@ class LoggingWorker: This is called as a background task from _schedule_delayed_enqueue_retry. """ await asyncio.sleep(delay) - + # Try to enqueue the task directly, preserving its original context if self._queue is None: return - + try: self._queue.put_nowait(task) except asyncio.QueueFull: @@ -243,15 +263,17 @@ class LoggingWorker: """ if self._queue is None: return [] - + # Calculate items based on percentage of queue size - items_to_extract = (self.max_queue_size * LOGGING_WORKER_CLEAR_PERCENTAGE) // 100 + items_to_extract = ( + self.max_queue_size * LOGGING_WORKER_CLEAR_PERCENTAGE + ) // 100 # Use actual queue size to avoid unnecessary iterations actual_size = self._queue.qsize() if actual_size == 0: return [] items_to_extract = min(items_to_extract, actual_size) - + # Extract tasks from queue (using list comprehension would require wrapping in try/except) extracted_tasks = [] for _ in range(items_to_extract): @@ -259,10 +281,12 @@ class LoggingWorker: extracted_tasks.append(self._queue.get_nowait()) except asyncio.QueueEmpty: break - + return extracted_tasks - async def _aggressively_clear_queue_async(self, new_task: Optional[LoggingTask] = None) -> None: + async def _aggressively_clear_queue_async( + self, new_task: Optional[LoggingTask] = None + ) -> None: """ Aggressively clear the queue by extracting and processing items. This is called when the queue is full to prevent dropping logs. @@ -271,18 +295,20 @@ class LoggingWorker: try: if self._queue is None: return - + extracted_tasks = self._extract_tasks_from_queue() - + # Add new task to extracted tasks to process directly if new_task is not None: extracted_tasks.append(new_task) - + # Process extracted tasks directly if extracted_tasks: await self._process_extracted_tasks(extracted_tasks) except Exception as e: - verbose_logger.exception(f"LoggingWorker error during aggressive clear: {e}") + verbose_logger.exception( + f"LoggingWorker error during aggressive clear: {e}" + ) finally: # Always reset the flag even if an error occurs self._aggressive_clear_in_progress = False @@ -291,7 +317,7 @@ class LoggingWorker: """Process a single task and mark it done.""" if self._queue is None: return - + try: await asyncio.wait_for( task["context"].run(asyncio.create_task, task["coroutine"]), @@ -310,7 +336,7 @@ class LoggingWorker: """ if not tasks or self._queue is None: return - + # Process all tasks concurrently for maximum speed await asyncio.gather(*[self._process_single_task(task) for task in tasks]) @@ -361,10 +387,7 @@ class LoggingWorker: for _ in range(MAX_ITERATIONS_TO_CLEAR_QUEUE): # Check if we've exceeded the maximum time - if ( - asyncio.get_event_loop().time() - start_time - >= MAX_TIME_TO_CLEAR_QUEUE - ): + if asyncio.get_event_loop().time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE: verbose_logger.warning( f"clear_queue exceeded max_time of {MAX_TIME_TO_CLEAR_QUEUE}s, stopping early" ) @@ -381,6 +404,9 @@ class LoggingWorker: except Exception: # Suppress errors during cleanup pass + finally: + # Clear reference to prevent memory leaks + task = None self._queue.task_done() # If you're using join() elsewhere except asyncio.QueueEmpty: break @@ -410,7 +436,7 @@ class LoggingWorker: This ensures callbacks queued by async completions are processed even when the script exits before the worker loop can handle them. - + Note: All logging in this method is wrapped to handle cases where logging handlers are closed during shutdown. """ @@ -423,7 +449,9 @@ class LoggingWorker: return queue_size = self._queue.qsize() - self._safe_log("info", f"[LoggingWorker] atexit: Flushing {queue_size} remaining events...") + self._safe_log( + "info", f"[LoggingWorker] atexit: Flushing {queue_size} remaining events..." + ) # Create a new event loop since the original is closed loop = asyncio.new_event_loop() @@ -438,7 +466,7 @@ class LoggingWorker: if loop.time() - start_time >= MAX_TIME_TO_CLEAR_QUEUE: self._safe_log( "warning", - f"[LoggingWorker] atexit: Reached time limit ({MAX_TIME_TO_CLEAR_QUEUE}s), stopping flush" + f"[LoggingWorker] atexit: Reached time limit ({MAX_TIME_TO_CLEAR_QUEUE}s), stopping flush", ) break @@ -456,8 +484,14 @@ class LoggingWorker: except Exception: # Silent failure to not break user's program pass + finally: + # Clear reference to prevent memory leaks + task = None - self._safe_log("info", f"[LoggingWorker] atexit: Successfully flushed {processed} events!") + self._safe_log( + "info", + f"[LoggingWorker] atexit: Successfully flushed {processed} events!", + ) finally: loop.close() diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ca2a092dbc8..b100b9b516b 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1087,9 +1087,35 @@ def _parse_content_for_reasoning( return None, message_text +def _extract_base64_data(image_url: str) -> str: + """ + Extract pure base64 data from an image URL. + + If the URL is a data URL (e.g., "data:image/png;base64,iVBOR..."), + extract and return only the base64 data portion. + Otherwise, return the original URL unchanged. + + This is needed for providers like Ollama that expect pure base64 data + rather than full data URLs. + + Args: + image_url: The image URL or data URL to process + + Returns: + The base64 data if it's a data URL, otherwise the original URL + """ + if image_url.startswith("data:") and ";base64," in image_url: + return image_url.split(";base64,", 1)[1] + return image_url + + def extract_images_from_message(message: AllMessageValues) -> List[str]: """ - Extract images from a message + Extract images from a message. + + For data URLs (e.g., "data:image/png;base64,iVBOR..."), only the base64 + data portion is extracted. This is required for providers like Ollama + that expect pure base64 data rather than full data URLs. """ images = [] message_content = message.get("content") @@ -1098,7 +1124,7 @@ def extract_images_from_message(message: AllMessageValues) -> List[str]: image_url = m.get("image_url") if image_url: if isinstance(image_url, str): - images.append(image_url) + images.append(_extract_base64_data(image_url)) elif isinstance(image_url, dict) and "url" in image_url: - images.append(image_url["url"]) + images.append(_extract_base64_data(image_url["url"])) return images diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 6cc6c229f56..12570a02de7 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -930,7 +930,8 @@ def create_anthropic_image_param( # Check if the image URL is an HTTP/HTTPS URL if image_url.startswith("http://") or image_url.startswith("https://"): - # For Bedrock invoke, always convert URLs to base64 (Bedrock invoke doesn't support URLs) + # For Bedrock invoke and Vertex AI Anthropic, always convert URLs to base64 + # as these providers don't support URL sources for images if is_bedrock_invoke or image_url.startswith("http://"): base64_url = convert_url_to_base64(url=image_url) image_chunk = convert_to_anthropic_image_obj( @@ -1496,9 +1497,10 @@ def convert_to_gemini_tool_call_result( content_type = content.get("type", "") if content_type == "text": content_str += content.get("text", "") - elif content_type == "input_image": - # Extract image for inline_data (for Computer Use screenshots) - image_url = content.get("image_url", "") + elif content_type in ("input_image", "image_url"): + # Extract image for inline_data (for Computer Use screenshots and tool results) + image_url_data = content.get("image_url", "") + image_url = image_url_data.get("url", "") if isinstance(image_url_data, dict) else image_url_data if image_url: # Convert image to base64 blob format for Gemini @@ -2022,9 +2024,12 @@ def anthropic_messages_pt( # noqa: PLR0915 "format": image_url_value.get("format"), } # Bedrock invoke models have format: invoke/... + # Vertex AI Anthropic also doesn't support URL sources for images is_bedrock_invoke = model.lower().startswith("invoke/") + is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False + force_base64 = is_bedrock_invoke or is_vertex_ai _anthropic_content_element = create_anthropic_image_param( - image_url_input, format=format, is_bedrock_invoke=is_bedrock_invoke + image_url_input, format=format, is_bedrock_invoke=force_base64 ) _content_element = add_cache_control_to_content( anthropic_content_element=_anthropic_content_element, diff --git a/litellm/llms/__init__.py b/litellm/llms/__init__.py index 15c035ceec8..c73f0b22b4b 100644 --- a/litellm/llms/__init__.py +++ b/litellm/llms/__init__.py @@ -45,6 +45,7 @@ def get_cost_for_web_search_request( return 0.0 elif custom_llm_provider == "xai": from .xai.cost_calculator import cost_per_web_search_request + return cost_per_web_search_request(usage=usage, model_info=model_info) else: return None @@ -110,6 +111,21 @@ def discover_guardrail_translation_mappings() -> ( verbose_logger.error(f"Error processing {module_path}: {e}") continue + try: + from litellm.proxy._experimental.mcp_server.guardrail_translation import ( + guardrail_translation_mappings as mcp_guardrail_translation_mappings, + ) + + discovered_mappings.update(mcp_guardrail_translation_mappings) + verbose_logger.debug( + "Loaded MCP guardrail translation mappings: %s", + list(mcp_guardrail_translation_mappings.keys()), + ) + except ImportError: + verbose_logger.debug( + "MCP guardrail translation mappings not available; skipping" + ) + verbose_logger.debug( f"Discovered {len(discovered_mappings)} guardrail translation mappings: {list(discovered_mappings.keys())}" ) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 6bdc17f7979..c71edcdc2d1 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -54,7 +54,10 @@ from litellm.types.utils import ( CompletionTokensDetailsWrapper, ) from litellm.types.utils import Message as LitellmMessage -from litellm.types.utils import PromptTokensDetailsWrapper, ServerToolUse +from litellm.types.utils import ( + PromptTokensDetailsWrapper, + ServerToolUse, +) from litellm.utils import ( ModelResponse, Usage, @@ -204,9 +207,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) # Relevant issue: https://github.com/BerriAI/litellm/issues/7755 def get_cache_control_headers(self) -> dict: + # Anthropic no longer requires the prompt-caching beta header + # Prompt caching now works automatically when cache_control is used in messages + # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching return { "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", } def _map_tool_choice( @@ -1034,7 +1039,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): anthropic_messages = anthropic_messages_pt( model=model, messages=messages, - llm_provider="anthropic", + llm_provider=self.custom_llm_provider or "anthropic", ) except Exception as e: raise AnthropicError( diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 098694f15ae..fcbe9823ed4 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -12,7 +12,11 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.types.llms.anthropic import AllAnthropicToolsValues, AnthropicMcpServerTool, ANTHROPIC_HOSTED_TOOLS +from litellm.types.llms.anthropic import ( + ANTHROPIC_HOSTED_TOOLS, + AllAnthropicToolsValues, + AnthropicMcpServerTool, +) from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import TokenCountResponse @@ -273,8 +277,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): beta_header = self.get_computer_tool_beta_header(computer_tool_used) betas.append(beta_header) - if prompt_caching_set: - betas.append("prompt-caching-2024-07-31") + # Anthropic no longer requires the prompt-caching beta header + # Prompt caching now works automatically when cache_control is used in messages + # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching if file_id_used: betas.append("files-api-2025-04-14") @@ -305,8 +310,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): container_with_skills_used: bool = False, ) -> dict: betas = set() - if prompt_caching_set: - betas.add("prompt-caching-2024-07-31") + # Anthropic no longer requires the prompt-caching beta header + # Prompt caching now works automatically when cache_control is used in messages + # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching if computer_tool_used: beta_header = self.get_computer_tool_beta_header(computer_tool_used) betas.add(beta_header) diff --git a/litellm/llms/azure_ai/anthropic/messages_transformation.py b/litellm/llms/azure_ai/anthropic/messages_transformation.py index 73dc84167ab..55818cc07d6 100644 --- a/litellm/llms/azure_ai/anthropic/messages_transformation.py +++ b/litellm/llms/azure_ai/anthropic/messages_transformation.py @@ -48,7 +48,12 @@ class AzureAnthropicMessagesConfig(AnthropicMessagesConfig): headers = BaseAzureLLM._base_validate_azure_environment( headers=headers, litellm_params=litellm_params_obj ) - + + # Azure Anthropic uses x-api-key header (not api-key) + # Convert api-key to x-api-key if present + if "api-key" in headers and "x-api-key" not in headers: + headers["x-api-key"] = headers.pop("api-key") + # Set anthropic-version header if "anthropic-version" not in headers: headers["anthropic-version"] = "2023-06-01" diff --git a/litellm/llms/gemini/google_genai/transformation.py b/litellm/llms/gemini/google_genai/transformation.py index d8692bb6a3a..d3b55b13ad4 100644 --- a/litellm/llms/gemini/google_genai/transformation.py +++ b/litellm/llms/gemini/google_genai/transformation.py @@ -312,7 +312,9 @@ class GoogleGenAIConfig(BaseGoogleGenAIGenerateContentConfig, VertexLLM): ) request_dict = cast(dict, typed_generate_content_request) - + + if system_instruction is not None: + request_dict["systemInstruction"] = system_instruction return request_dict def transform_generate_content_response( diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index 9e6497e66ab..71956158f52 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -15,7 +15,7 @@ def _prepare_ollama_embedding_payload( ) -> Dict[str, Any]: data: Dict[str, Any] = {"model": model, "input": prompts} - special_optional_params = ["truncate", "options", "keep_alive"] + special_optional_params = ["truncate", "options", "keep_alive","dimensions"] for k, v in optional_params.items(): if k in special_optional_params: diff --git a/litellm/llms/openai/image_generation/cost_calculator.py b/litellm/llms/openai/image_generation/cost_calculator.py new file mode 100644 index 00000000000..35caaf6e9b1 --- /dev/null +++ b/litellm/llms/openai/image_generation/cost_calculator.py @@ -0,0 +1,63 @@ +""" +Cost calculator for OpenAI image generation models (gpt-image-1, gpt-image-1-mini) + +These models use token-based pricing instead of pixel-based pricing like DALL-E. +""" + +from typing import Optional + +from litellm import verbose_logger +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.utils import ImageResponse + + +def cost_calculator( + model: str, + image_response: ImageResponse, + custom_llm_provider: Optional[str] = None, +) -> float: + """ + Calculate cost for OpenAI gpt-image-1 and gpt-image-1-mini models. + + Uses the same usage format as Responses API, so we reuse the helper + to transform to chat completion format and use generic_cost_per_token. + + Args: + model: The model name (e.g., "gpt-image-1", "gpt-image-1-mini") + image_response: The ImageResponse containing usage data + custom_llm_provider: Optional provider name + + Returns: + float: Total cost in USD + """ + usage = getattr(image_response, "usage", None) + + if usage is None: + verbose_logger.debug( + f"No usage data available for {model}, cannot calculate token-based cost" + ) + return 0.0 + + # Transform ImageUsage to Usage using the existing helper + # ImageUsage has the same format as ResponseAPIUsage + chat_usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + usage + ) + + # Use generic_cost_per_token for cost calculation + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=chat_usage, + custom_llm_provider=custom_llm_provider or "openai", + ) + + total_cost = prompt_cost + completion_cost + + verbose_logger.debug( + f"OpenAI gpt-image cost calculation for {model}: " + f"prompt_cost=${prompt_cost:.6f}, completion_cost=${completion_cost:.6f}, " + f"total=${total_cost:.6f}" + ) + + return total_cost diff --git a/litellm/llms/sap/chat/transformation.py b/litellm/llms/sap/chat/transformation.py index 01ceb72c0de..e13abca59f8 100755 --- a/litellm/llms/sap/chat/transformation.py +++ b/litellm/llms/sap/chat/transformation.py @@ -203,9 +203,13 @@ class GenAIHubOrchestrationConfig(OpenAIGPTConfig): headers: dict, ) -> dict: supported_params = self.get_supported_openai_params(model) + # Include extra params that passed validation (e.g., thinking_config for Gemini models via allowed_openai_params) + extra_params = [k for k in optional_params if k not in supported_params and k not in {"tools", "model_version"}] + supported_params = supported_params + extra_params model_params = { k: v for k, v in optional_params.items() if k in supported_params } + model_version = optional_params.pop("model_version", "latest") template = [] for message in messages: diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index baa825bfcca..2bbdfa17cde 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -383,7 +383,18 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 and isinstance(_message_content, str) ): assistant_text = _message_content - assistant_content.append(PartType(text=assistant_text)) # type: ignore + # Check if message has thought_signatures in provider_specific_fields + provider_specific_fields = assistant_msg.get("provider_specific_fields") + thought_signatures = None + if provider_specific_fields and isinstance(provider_specific_fields, dict): + thought_signatures = provider_specific_fields.get("thought_signatures") + + # If we have thought signatures, add them to the part + if thought_signatures and isinstance(thought_signatures, list) and len(thought_signatures) > 0: + # Use the first signature for the text part (Gemini expects one signature per part) + assistant_content.append(PartType(text=assistant_text, thoughtSignature=thought_signatures[0])) # type: ignore + else: + assistant_content.append(PartType(text=assistant_text)) # type: ignore ## HANDLE ASSISTANT FUNCTION CALL if ( diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index a5cc3dca8c1..ba1788a217f 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -552,24 +552,46 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "Invalid tool={}. Use `litellm.set_verbose` or `litellm --detailed_debug` to see raw request." ) - # Only include function_declarations if there are actual functions - _tools = Tools() +# Build list of Tool objects - each Tool should contain exactly one type + # per Vertex AI API spec: "A Tool object should contain exactly one type of Tool" + _tools_list: List[Tools] = [] + + # Function declarations can be grouped together in one Tool if gtool_func_declarations: - _tools["function_declarations"] = gtool_func_declarations + func_tool = Tools() + func_tool["function_declarations"] = gtool_func_declarations + _tools_list.append(func_tool) + + # Each special tool type must be in its own Tool object if googleSearch is not None: - _tools[VertexToolName.GOOGLE_SEARCH.value] = googleSearch + search_tool = Tools() + search_tool[VertexToolName.GOOGLE_SEARCH.value] = googleSearch + _tools_list.append(search_tool) if googleSearchRetrieval is not None: - _tools[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval + retrieval_tool = Tools() + retrieval_tool[VertexToolName.GOOGLE_SEARCH_RETRIEVAL.value] = googleSearchRetrieval + _tools_list.append(retrieval_tool) if enterpriseWebSearch is not None: - _tools[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch + enterprise_tool = Tools() + enterprise_tool[VertexToolName.ENTERPRISE_WEB_SEARCH.value] = enterpriseWebSearch + _tools_list.append(enterprise_tool) if code_execution is not None: - _tools[VertexToolName.CODE_EXECUTION.value] = code_execution + code_tool = Tools() + code_tool[VertexToolName.CODE_EXECUTION.value] = code_execution + _tools_list.append(code_tool) if urlContext is not None: - _tools[VertexToolName.URL_CONTEXT.value] = urlContext + url_tool = Tools() + url_tool[VertexToolName.URL_CONTEXT.value] = urlContext + _tools_list.append(url_tool) if googleMaps is not None: - _tools[VertexToolName.GOOGLE_MAPS.value] = googleMaps + maps_tool = Tools() + maps_tool[VertexToolName.GOOGLE_MAPS.value] = googleMaps + _tools_list.append(maps_tool) if computerUse is not None: - _tools[VertexToolName.COMPUTER_USE.value] = computerUse + computer_tool = Tools() + computer_tool[VertexToolName.COMPUTER_USE.value] = computerUse + _tools_list.append(computer_tool) + # Add retrieval config to toolConfig if googleMaps has location data if google_maps_retrieval_config is not None: @@ -579,7 +601,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "retrievalConfig" ] = google_maps_retrieval_config - return [_tools] + return _tools_list def _map_response_schema(self, value: dict) -> dict: old_schema = deepcopy(value) @@ -1210,6 +1232,25 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): thinking_blocks.append(block) return thinking_blocks + def _extract_thought_signatures_from_parts( + self, parts: List[HttpxPartType] + ) -> Optional[List[str]]: + """Extract thoughtSignature values from parts. + + Per Google's docs, thoughtSignature is returned for multi-turn context preservation + and can appear on parts even without thought: true (e.g., regular text responses, + function calls). This method extracts all thoughtSignature values from parts. + + Returns: + List of thoughtSignature strings if any are found, None otherwise + """ + signatures: List[str] = [] + for part in parts: + signature = part.get("thoughtSignature") + if signature is not None: + signatures.append(signature) + return signatures if signatures else None + def _extract_image_response_from_parts( self, parts: List[HttpxPartType] ) -> Optional[List[ImageURLListItem]]: @@ -1620,6 +1661,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): from litellm.types.utils import Delta, StreamingChoices annotations = chat_completion_message.get("annotations") # type: ignore + provider_specific_fields = chat_completion_message.get("provider_specific_fields") # type: ignore # create a streaming choice object choice = StreamingChoices( finish_reason=VertexGeminiConfig._check_finish_reason( @@ -1633,6 +1675,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): images=image_response, function_call=functions, annotations=annotations, # type: ignore + provider_specific_fields=provider_specific_fields, ), logprobs=chat_completion_logprobs, enhancements=None, @@ -1811,6 +1854,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) ) + # Extract thoughtSignatures from parts (can exist without thought: true) + thought_signatures = ( + VertexGeminiConfig()._extract_thought_signatures_from_parts( + parts=candidate["content"]["parts"] + ) + ) + if audio_response is not None: cast(Dict[str, Any], chat_completion_message)[ "audio" @@ -1876,6 +1926,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): reasoning_content = "\n".join(reasoning_content_parts) chat_completion_message["reasoning_content"] = reasoning_content + # Store thoughtSignatures in provider_specific_fields + if thought_signatures is not None: + if "provider_specific_fields" not in chat_completion_message: + chat_completion_message["provider_specific_fields"] = {} + chat_completion_message["provider_specific_fields"]["thought_signatures"] = thought_signatures # type: ignore + if isinstance(model_response, ModelResponseStream): choice = VertexGeminiConfig._create_streaming_choice( chat_completion_message=chat_completion_message, diff --git a/litellm/llms/zai/chat/transformation.py b/litellm/llms/zai/chat/transformation.py index 47b314d4e0d..4380256f0a4 100644 --- a/litellm/llms/zai/chat/transformation.py +++ b/litellm/llms/zai/chat/transformation.py @@ -20,7 +20,7 @@ class ZAIChatConfig(OpenAIGPTConfig): return api_base, dynamic_api_key def get_supported_openai_params(self, model: str) -> list: - return [ + base_params = [ "max_tokens", "stream", "stream_options", @@ -31,3 +31,12 @@ class ZAIChatConfig(OpenAIGPTConfig): "tool_choice", ] + import litellm + + try: + if litellm.supports_reasoning(model=model, custom_llm_provider=self.custom_llm_provider): + base_params.append("thinking") + except Exception: + pass + + return base_params diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 513a4a554e0..d32adf54b5e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -249,6 +249,30 @@ "/v1/images/generations" ] }, + "aiml/google/imagen-4.0-ultra-generate-001": { + "litellm_provider": "aiml", + "metadata": { + "notes": "Imagen 4.0 Ultra Generate API - Photorealistic image generation with precise text rendering" + }, + "mode": "image_generation", + "output_cost_per_image": 0.063, + "source": "https://docs.aimlapi.com/api-references/image-models/google/imagen-4-ultra-generate", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, + "aiml/google/nano-banana-pro": { + "litellm_provider": "aiml", + "metadata": { + "notes": "Gemini 3 Pro Image (Nano Banana Pro) - Advanced text-to-image generation with reasoning and 4K resolution support" + }, + "mode": "image_generation", + "output_cost_per_image": 0.1575, + "source": "https://docs.aimlapi.com/api-references/image-models/google/gemini-3-pro-image-preview", + "supported_endpoints": [ + "/v1/images/generations" + ] + }, "amazon.nova-canvas-v1:0": { "litellm_provider": "bedrock", "max_input_tokens": 2600, @@ -3508,6 +3532,40 @@ "supports_service_tier": true, "supports_vision": true }, + "azure/gpt-5.2-chat": { + "cache_read_input_token_cost": 1.75e-07, + "cache_read_input_token_cost_priority": 3.5e-07, + "input_cost_per_token": 1.75e-06, + "input_cost_per_token_priority": 3.5e-06, + "litellm_provider": "azure", + "max_input_tokens": 128000, + "max_output_tokens": 16384, + "max_tokens": 16384, + "mode": "chat", + "output_cost_per_token": 1.4e-05, + "output_cost_per_token_priority": 2.8e-05, + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/responses" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text" + ], + "supports_function_calling": true, + "supports_native_streaming": true, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "azure/gpt-5.2-chat-2025-12-11": { "cache_read_input_token_cost": 1.75e-07, "cache_read_input_token_cost_priority": 3.5e-07, @@ -3605,12 +3663,16 @@ "supports_web_search": true }, "azure/gpt-image-1": { - "input_cost_per_pixel": 4.0054321e-08, + "cache_read_input_image_token_cost": 2.5e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 1e-05, + "input_cost_per_token": 5e-06, "litellm_provider": "azure", "mode": "image_generation", - "output_cost_per_pixel": 0.0, + "output_cost_per_image_token": 4e-05, "supported_endpoints": [ - "/v1/images/generations" + "/v1/images/generations", + "/v1/images/edits" ] }, "azure/hd/1024-x-1024/dall-e-3": { @@ -3713,12 +3775,16 @@ ] }, "azure/gpt-image-1-mini": { - "input_cost_per_pixel": 8.0566406e-09, + "cache_read_input_image_token_cost": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_image_token": 2.5e-06, + "input_cost_per_token": 2e-06, "litellm_provider": "azure", "mode": "image_generation", - "output_cost_per_pixel": 0.0, + "output_cost_per_image_token": 8e-06, "supported_endpoints": [ - "/v1/images/generations" + "/v1/images/generations", + "/v1/images/edits" ] }, "azure/gpt-image-1.5": { @@ -10885,13 +10951,13 @@ "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/deepseek-v3p2": { - "input_cost_per_token": 1.2e-06, + "input_cost_per_token": 5.6e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 163840, "max_output_tokens": 163840, "max_tokens": 163840, "mode": "chat", - "output_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.68e-06, "source": "https://fireworks.ai/models/fireworks/deepseek-v3p2", "supports_function_calling": true, "supports_reasoning": true, @@ -16922,6 +16988,336 @@ "supports_vision": true, "supports_pdf_input": true }, + "low/1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.034, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.05, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.05, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.133, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.20, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.20, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.034, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.05, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.05, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.133, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.20, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.20, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "gpt-5": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, @@ -17683,16 +18079,16 @@ "supports_vision": true }, "gpt-image-1": { - "input_cost_per_image": 0.042, - "input_cost_per_pixel": 4.0054321e-08, - "input_cost_per_token": 0.000005, - "input_cost_per_image_token": 0.00001, + "cache_read_input_image_token_cost": 2.5e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 1e-05, + "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", - "output_cost_per_pixel": 0.0, - "output_cost_per_token": 0.00004, + "output_cost_per_image_token": 4e-05, "supported_endpoints": [ - "/v1/images/generations" + "/v1/images/generations", + "/v1/images/edits" ] }, "gpt-image-1-mini": { @@ -18117,6 +18513,18 @@ "supports_response_schema": false, "supports_tool_choice": true }, + "groq/gemma-7b-it": { + "input_cost_per_token": 5e-08, + "litellm_provider": "groq", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 8e-08, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_tool_choice": true + }, "groq/meta-llama/llama-guard-4-12b": { "input_cost_per_token": 2e-07, "litellm_provider": "groq", @@ -29241,6 +29649,20 @@ "supports_vision": true, "supports_web_search": true }, + "zai/glm-4.7": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.6": { "input_cost_per_token": 6e-07, "output_cost_per_token": 2.2e-06, diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e0fd610e678 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/__init__.py @@ -0,0 +1,16 @@ +"""Guardrail translation mapping for MCP tool calls.""" + +from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) +from litellm.types.utils import CallTypes + +# This mapping lives alongside the MCP server implementation because MCP +# integrations are managed by the proxy subsystem, not litellm.llms providers. +# Unified guardrails import this module explicitly to register the handler. + +guardrail_translation_mappings = { + CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, +} + +__all__ = ["guardrail_translation_mappings", "MCPGuardrailTranslationHandler"] diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py new file mode 100644 index 00000000000..8d6d236b884 --- /dev/null +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -0,0 +1,89 @@ +""" +MCP Guardrail Handler for Unified Guardrails. + +This handler works with the synthetic "messages" payload generated by +`ProxyLogging._convert_mcp_to_llm_format`, which always produces a single user +message whose `content` string encodes the MCP tool name and arguments. The +handler simply feeds that text through the configured guardrail and writes the +result back onto the message. +""" + +from typing import TYPE_CHECKING, Any, Dict, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from mcp.types import CallToolResult + + +class MCPGuardrailTranslationHandler(BaseTranslation): + """Guardrail translation handler for MCP tool calls.""" + + async def process_input_messages( + self, + data: Dict[str, Any], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + ) -> Dict[str, Any]: + messages = data.get("messages") + if not isinstance(messages, list) or not messages: + verbose_proxy_logger.debug("MCP Guardrail: No messages to process") + return data + + first_message = messages[0] + content: Optional[str] = None + if isinstance(first_message, dict): + content = first_message.get("content") + else: + content = getattr(first_message, "content", None) + + if not isinstance(content, str): + verbose_proxy_logger.debug( + "MCP Guardrail: Message content missing or not a string", + ) + return data + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=GenericGuardrailAPIInputs(texts=[content]), + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + guardrailed_texts = ( + guardrailed_inputs.get("texts", []) if guardrailed_inputs else [] + ) + + if guardrailed_texts: + new_content = guardrailed_texts[0] + if isinstance(first_message, dict): + first_message["content"] = new_content + else: + setattr(first_message, "content", new_content) + + verbose_proxy_logger.debug( + "MCP Guardrail: Updated content for tool %s", + data.get("mcp_tool_name"), + ) + else: + verbose_proxy_logger.debug( + "MCP Guardrail: Guardrail returned no text updates for tool %s", + data.get("mcp_tool_name"), + ) + + return data + + async def process_output_response( + self, + response: "CallToolResult", + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + user_api_key_dict: Optional[Any] = None, + ) -> Any: + # Not implemented: MCP guardrail translation never calls this path today. + verbose_proxy_logger.debug( + "MCP Guardrail: Output processing not implemented for MCP tools", + ) + return response diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2260649e8b2..d193fc27fb9 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -11,7 +11,7 @@ import datetime import hashlib import json import re -from typing import Any, Dict, List, Optional, Set, Tuple, Union, cast +from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast from urllib.parse import urlparse from fastapi import HTTPException @@ -30,6 +30,7 @@ from pydantic import AnyUrl import litellm from litellm._logging import verbose_logger +from litellm.types.utils import CallTypes from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.experimental_mcp_client.client import MCPClient from litellm.llms.custom_httpx.http_handler import get_async_httpx_client @@ -536,9 +537,12 @@ class MCPServerManager: client_secret=client_secret_value or getattr(mcp_server, "client_secret", None), scopes=resolved_scopes, - authorization_url=getattr(mcp_oauth_metadata, "authorization_url", None), - token_url=getattr(mcp_oauth_metadata, "token_url", None), - registration_url=getattr(mcp_oauth_metadata, "registration_url", None), + authorization_url=mcp_server.authorization_url + or getattr(mcp_oauth_metadata, "authorization_url", None), + token_url=mcp_server.token_url + or getattr(mcp_oauth_metadata, "token_url", None), + registration_url=mcp_server.registration_url + or getattr(mcp_oauth_metadata, "registration_url", None), command=getattr(mcp_server, "command", None), args=getattr(mcp_server, "args", None) or [], env=env_dict, @@ -548,7 +552,7 @@ class MCPServerManager: ) return new_server - async def add_update_server(self, mcp_server: LiteLLM_MCPServerTable): + async def add_server(self, mcp_server: LiteLLM_MCPServerTable): try: if mcp_server.server_id not in self.registry: new_server = await self.build_mcp_server_from_table(mcp_server) @@ -559,6 +563,17 @@ class MCPServerManager: verbose_logger.debug(f"Failed to add MCP server: {str(e)}") raise e + async def update_server(self, mcp_server: LiteLLM_MCPServerTable): + try: + if mcp_server.server_id in self.registry: + new_server = await self.build_mcp_server_from_table(mcp_server) + self.registry[mcp_server.server_id] = new_server + verbose_logger.debug(f"Updated MCP Server: {new_server.name}") + + except Exception as e: + verbose_logger.debug(f"Failed to udpate MCP server: {str(e)}") + raise e + def get_all_mcp_server_ids(self) -> Set[str]: """ Get all MCP server IDs @@ -1662,11 +1677,11 @@ class MCPServerManager: ) try: - # Use standard pre_call_hook with call_type="mcp_call" + # Use standard pre_call_hook modified_data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_auth, # type: ignore data=synthetic_llm_data, - call_type="mcp_call", # type: ignore + call_type=CallTypes.call_mcp_tool.value, ) if modified_data: # Convert response back to MCP format and apply modifications @@ -1723,7 +1738,7 @@ class MCPServerManager: proxy_logging_obj.during_call_hook( user_api_key_dict=user_api_key_auth, data=synthetic_llm_data, - call_type="mcp_call", # type: ignore + call_type=CallTypes.call_mcp_tool.value, ) ) @@ -1879,7 +1894,7 @@ class MCPServerManager: ######################################################### # Pre MCP Tool Call Hook # Allow validation and modification of tool calls before execution - # Using standard pre_call_hook with call_type="mcp_call" + # Using standard pre_call_hook ######################################################### if proxy_logging_obj: await self.pre_call_tool_check( @@ -2040,7 +2055,7 @@ class MCPServerManager: verbose_logger.debug( f"Adding server to registry: {server.server_id} ({server.server_name})" ) - await self.add_update_server(server) + await self.add_server(server) verbose_logger.debug( f"Registry now contains {len(self.get_registry())} servers" @@ -2127,7 +2142,7 @@ class MCPServerManager: async def health_check_server( self, server_id: str, mcp_auth_header: Optional[str] = None - ) -> Dict[str, Any]: + ) -> LiteLLM_MCPServerTable: """ Perform a health check on a specific MCP server. @@ -2138,206 +2153,186 @@ class MCPServerManager: Returns: Dict containing health check results """ - import time from datetime import datetime server = self.get_mcp_server_by_id(server_id) if not server: - return { - "server_id": server_id, - "server_name": None, - "status": "unknown", - "error": "Server not found", - "last_health_check": datetime.now().isoformat(), - "response_time_ms": None, - } - - start_time = time.time() - try: - # Try to get tools from the server as a health check - tools = await self._get_tools_from_server(server, mcp_auth_header) - response_time = (time.time() - start_time) * 1000 - - return { - "server_id": server_id, - "server_name": server.name, - "status": "healthy", - "tools_count": len(tools), - "last_health_check": datetime.now().isoformat(), - "response_time_ms": round(response_time, 2), - "error": None, - } - except Exception as e: - response_time = (time.time() - start_time) * 1000 - error_message = str(e) - - return { - "server_id": server_id, - "server_name": server.name, - "status": "unhealthy", - "last_health_check": datetime.now().isoformat(), - "response_time_ms": round(response_time, 2), - "error": error_message, - } - - async def health_check_all_servers( - self, mcp_auth_header: Optional[str] = None - ) -> Dict[str, Any]: - """ - Perform health checks on all MCP servers. - - Args: - mcp_auth_header: Optional authentication header for the MCP servers - - Returns: - Dict containing health check results for all servers - """ - all_servers = self.get_registry() - results = {} - - for server_id, server in all_servers.items(): - results[server_id] = await self.health_check_server( - server_id, mcp_auth_header + verbose_logger.warning(f"MCP Server {server_id} not found") + return LiteLLM_MCPServerTable( + server_id=server_id, + server_name=None, + transport=MCPTransport.http, # Default transport for not found servers + status="unknown", + health_check_error="Server not found", + last_health_check=datetime.now(), ) - return results + status: Literal["healthy", "unhealthy", "unknown"] = "unknown" + health_check_error = None - async def health_check_allowed_servers( - self, - user_api_key_auth: Optional[UserAPIKeyAuth] = None, - mcp_auth_header: Optional[str] = None, - ) -> Dict[str, Any]: - """ - Perform health checks on all MCP servers that the user has access to. + # Check if we should skip health check based on auth configuration + should_skip_health_check = False - Args: - user_api_key_auth: User authentication info for access control - mcp_auth_header: Optional authentication header for the MCP servers + # Skip if auth_type is oauth2 + if server.auth_type == MCPAuth.oauth2: + should_skip_health_check = True + # Skip if auth_type is not none and authentication_token is missing + elif ( + server.auth_type + and server.auth_type != MCPAuth.none + and not server.authentication_token + ): + should_skip_health_check = True - Returns: - Dict containing health check results for accessible servers - """ - # Get allowed servers for the user - allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) + if not should_skip_health_check: + extra_headers = {} + if server.static_headers: + extra_headers.update(server.static_headers) - # Perform health checks on allowed servers - results = {} - for server_id in allowed_server_ids: - results[server_id] = await self.health_check_server( - server_id, mcp_auth_header + client = self._create_mcp_client( + server=server, + mcp_auth_header=None, + extra_headers=extra_headers, + stdio_env=None, ) - return results + try: + + async def _noop(session): + return "ok" + + # Add timeout wrapper to prevent hanging + await asyncio.wait_for(client.run_with_session(_noop), timeout=10.0) + status = "healthy" + except asyncio.TimeoutError: + health_check_error = "Health check timed out after 10 seconds" + status = "unhealthy" + except Exception as e: + health_check_error = str(e) + status = "unhealthy" + + return LiteLLM_MCPServerTable( + server_id=server.server_id, + server_name=server.server_name, + alias=server.alias, + description=( + server.mcp_info.get("description") if server.mcp_info else None + ), + url=server.url, + transport=server.transport, + auth_type=server.auth_type, + created_at=datetime.now(), + updated_at=datetime.now(), + teams=[], + mcp_access_groups=server.access_groups or [], + allowed_tools=server.allowed_tools or [], + extra_headers=server.extra_headers or [], + mcp_info=server.mcp_info, + static_headers=server.static_headers, + status=status, + last_health_check=datetime.now(), + health_check_error=health_check_error, + command=getattr(server, "command", None), + args=getattr(server, "args", None) or [], + env=getattr(server, "env", None) or {}, + authorization_url=server.authorization_url, + token_url=server.token_url, + registration_url=server.registration_url, + ) async def get_all_mcp_servers_with_health_and_teams( self, user_api_key_auth: Optional[UserAPIKeyAuth] = None, - include_health: bool = True, + server_ids: Optional[List[str]] = None, ) -> List[LiteLLM_MCPServerTable]: """ Get all MCP servers that the user has access to, with health status and team information. Args: user_api_key_auth: User authentication info for access control - include_health: Whether to include health check information + server_ids: Optional list of server IDs to filter. If provided, only these servers + will be checked (subject to access control). If None, all accessible servers are checked. Returns: List of MCP server objects with health and team data """ - from litellm.proxy._experimental.mcp_server.db import ( - get_all_mcp_servers, - get_mcp_servers, - ) - from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view - from litellm.proxy.proxy_server import prisma_client # Get allowed server IDs allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) - # Get servers from database + # Filter by requested server_ids if provided + if server_ids: + # Only check servers that are both requested AND accessible + target_server_ids = [sid for sid in server_ids if sid in allowed_server_ids] + else: + # Check all accessible servers + target_server_ids = allowed_server_ids + + # Run health checks concurrently + tasks = [self.health_check_server(server_id) for server_id in target_server_ids] + results = await asyncio.gather(*tasks) + + # Filter out None results (servers that were not found) + list_mcp_servers = [server for server in results if server is not None] + + return list_mcp_servers + + async def get_all_allowed_mcp_servers( + self, + user_api_key_auth: Optional[UserAPIKeyAuth] = None, + ) -> List[LiteLLM_MCPServerTable]: + """ + Get all MCP servers that the user has access to. + + Args: + user_api_key_auth: User authentication info for access control + + Returns: + List of MCP server objects without health status + """ + from datetime import datetime + + # Get allowed server IDs + allowed_server_ids = await self.get_allowed_mcp_servers(user_api_key_auth) + list_mcp_servers: List[LiteLLM_MCPServerTable] = [] - if prisma_client is not None: - list_mcp_servers = await get_mcp_servers(prisma_client, allowed_server_ids) - # If admin, also get all servers from database - if user_api_key_auth and _user_has_admin_view(user_api_key_auth): - all_mcp_servers = await get_all_mcp_servers(prisma_client) - for server in all_mcp_servers: - if server.server_id not in allowed_server_ids: - list_mcp_servers.append(server) + for server_id in allowed_server_ids: + server = self.get_mcp_server_by_id(server_id) + if not server: + verbose_logger.warning(f"MCP Server {server_id} not found in registry") + continue - # Add config.yaml servers - for _server_id, _server_config in self.config_mcp_servers.items(): - if _server_id in allowed_server_ids: - list_mcp_servers.append( - LiteLLM_MCPServerTable( - **{ - **_server_config.model_dump(), - "created_at": datetime.datetime.now(), - "updated_at": datetime.datetime.now(), - "description": ( - _server_config.mcp_info.get("description") - if _server_config.mcp_info - else None - ), - "allowed_tools": _server_config.allowed_tools or [], - "mcp_info": _server_config.mcp_info, - "mcp_access_groups": _server_config.access_groups or [], - "extra_headers": _server_config.extra_headers or [], - "command": getattr(_server_config, "command", None), - "args": getattr(_server_config, "args", None) or [], - "env": getattr(_server_config, "env", None) or {}, - } - ) - ) - - # Get team information for non-admin users - server_to_teams_map: Dict[str, List[Dict[str, str]]] = {} - if ( - user_api_key_auth - and not _user_has_admin_view(user_api_key_auth) - and prisma_client is not None - ): - teams = await prisma_client.db.litellm_teamtable.find_many( - include={"object_permission": True} + # Build LiteLLM_MCPServerTable without health check + mcp_server_table = LiteLLM_MCPServerTable( + server_id=server.server_id, + server_name=server.server_name, + alias=server.alias, + description=( + server.mcp_info.get("description") if server.mcp_info else None + ), + url=server.url, + transport=server.transport, + auth_type=server.auth_type, + created_at=datetime.now(), + updated_at=datetime.now(), + teams=[], + mcp_access_groups=server.access_groups or [], + allowed_tools=server.allowed_tools or [], + extra_headers=server.extra_headers or [], + mcp_info=server.mcp_info, + static_headers=server.static_headers, + status=None, # No health check performed + last_health_check=None, # No health check performed + health_check_error=None, + command=getattr(server, "command", None), + args=getattr(server, "args", None) or [], + env=getattr(server, "env", None) or {}, + authorization_url=server.authorization_url, + token_url=server.token_url, + registration_url=server.registration_url, ) - - user_teams = [] - for team in teams: - if team.members_with_roles: - for member in team.members_with_roles: - if ( - "user_id" in member - and member["user_id"] is not None - and member["user_id"] == user_api_key_auth.user_id - ): - user_teams.append(team) - - # Create a mapping of server_id to teams that have access to it - for team in user_teams: - if team.object_permission and team.object_permission.mcp_servers: - for server_id in team.object_permission.mcp_servers: - if server_id not in server_to_teams_map: - server_to_teams_map[server_id] = [] - server_to_teams_map[server_id].append( - { - "team_id": team.team_id, - "team_alias": team.team_alias, - "organization_id": team.organization_id, - } - ) - - ## mark invalid servers w/ reason for being invalid - valid_server_ids = self.get_all_mcp_server_ids() - for server in list_mcp_servers: - if server.server_id not in valid_server_ids: - server.status = "unhealthy" - ## try adding server to registry to get error - try: - await self.add_update_server(server) - except Exception as e: - server.health_check_error = str(e) - server.health_check_error = "Server is not in in memory registry yet. This could be a temporary sync issue." + list_mcp_servers.append(mcp_server_table) return list_mcp_servers diff --git a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py index e4969df131e..8cdb381cf39 100644 --- a/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py +++ b/litellm/proxy/_experimental/mcp_server/openapi_to_mcp_generator.py @@ -3,7 +3,9 @@ This module is used to generate MCP tools from OpenAPI specs. """ import json +from pathlib import PurePosixPath from typing import Any, Dict, Optional +from urllib.parse import quote import httpx @@ -17,6 +19,29 @@ BASE_URL = "" HEADERS: Dict[str, str] = {} +def _sanitize_path_parameter_value(param_value: Any, param_name: str) -> str: + """Ensure path params cannot introduce directory traversal.""" + if param_value is None: + return "" + + value_str = str(param_value) + if value_str == "": + return "" + + normalized_value = value_str.replace("\\", "/") + if "/" in normalized_value: + raise ValueError( + f"Path parameter '{param_name}' must not contain path separators" + ) + + if any(part in {".", ".."} for part in PurePosixPath(normalized_value).parts): + raise ValueError( + f"Path parameter '{param_name}' cannot include '.' or '..' segments" + ) + + return quote(value_str, safe="") + + def load_openapi_spec(filepath: str) -> Dict[str, Any]: """Load OpenAPI specification from JSON file.""" with open(filepath, "r") as f: @@ -144,12 +169,18 @@ def create_tool_function( url = base_url + path # Replace path parameters using original names from OpenAPI spec + # Apply path traversal validation and URL encoding for param_name in path_params: param_value = kwargs.get(param_name, "") if param_value: + try: + # Sanitize and encode path parameter to prevent traversal attacks + safe_value = _sanitize_path_parameter_value(param_value, param_name) + except ValueError as exc: + return "Invalid path parameter: " + str(exc) # Replace {param_name} or {{param_name}} in URL - url = url.replace("{" + param_name + "}", str(param_value)) - url = url.replace("{{" + param_name + "}}", str(param_value)) + url = url.replace("{" + param_name + "}", safe_value) + url = url.replace("{{" + param_name + "}}", safe_value) # Build query params using original parameter names params: Dict[str, Any] = {} diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 891b52db7af..4c947b99ba3 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -1,5 +1,4 @@ import importlib -import traceback from typing import Dict, List, Optional, Union from fastapi import APIRouter, Depends, Query, Request @@ -347,17 +346,16 @@ if MCP_AVAILABLE: except Exception as e: verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True) - stack_trace = traceback.format_exc() return { "status": "error", - "message": f"An internal error has occurred: {str(e)}", - "stack_trace": stack_trace, + "message": "An internal error has occurred while testing the MCP server.", } - @router.post("/test/connection") + @router.post("/test/connection", dependencies=[Depends(user_api_key_auth)]) async def test_connection( request: Request, new_mcp_server_request: NewMCPServerRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Test if we can connect to the provided MCP server before adding it diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 570a6bb9f3b..7abfe7a96bc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -412,7 +412,6 @@ class LiteLLMRoutes(enum.Enum): agent_routes = [ "/v1/agents", "/agents", - "/a2a/{agent_id}", "/a2a/{agent_id}/message/send", "/a2a/{agent_id}/message/stream", @@ -830,9 +829,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} - model_max_budget: Optional[dict] = ( - {} - ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} + model_max_budget: Optional[ + dict + ] = {} # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} model_config = ConfigDict(protected_namespaces=()) model_rpm_limit: Optional[dict] = None @@ -1035,6 +1034,9 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None @model_validator(mode="before") @classmethod @@ -1092,6 +1094,9 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None @model_validator(mode="before") @classmethod @@ -1141,6 +1146,9 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): command: Optional[str] = None args: List[str] = Field(default_factory=list) env: Dict[str, str] = Field(default_factory=dict) + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None class MakeMCPServersPublicRequest(LiteLLMPydanticObjectBase): @@ -1160,6 +1168,9 @@ class NewSkillRequest(LiteLLMPydanticObjectBase): file_name: Optional[str] = None # Original filename file_type: Optional[str] = None # MIME type (e.g., "application/zip") metadata: Optional[Dict[str, Any]] = None + authorization_url: Optional[str] = None + token_url: Optional[str] = None + registration_url: Optional[str] = None class UpdateSkillRequest(LiteLLMPydanticObjectBase): @@ -1347,12 +1358,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model @model_validator(mode="before") @classmethod @@ -1374,12 +1385,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model class DeleteCustomerRequest(LiteLLMPydanticObjectBase): @@ -1464,15 +1475,15 @@ class NewTeamRequest(TeamBase): ] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating tpm model_tpm_limit: Optional[Dict[str, int]] = None - team_member_budget: Optional[float] = ( - None # allow user to set a budget for all team members - ) - team_member_rpm_limit: Optional[int] = ( - None # allow user to set RPM limit for all team members - ) - team_member_tpm_limit: Optional[int] = ( - None # allow user to set TPM limit for all team members - ) + team_member_budget: Optional[ + float + ] = None # allow user to set a budget for all team members + team_member_rpm_limit: Optional[ + int + ] = None # allow user to set RPM limit for all team members + team_member_tpm_limit: Optional[ + int + ] = None # allow user to set TPM limit for all team members team_member_key_duration: Optional[str] = None # e.g. "1d", "1w", "1m" allowed_vector_store_indexes: Optional[List[AllowedVectorStoreIndexItem]] = None @@ -1558,9 +1569,9 @@ class BlockKeyRequest(LiteLLMPydanticObjectBase): class AddTeamCallback(LiteLLMPydanticObjectBase): callback_name: str - callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = ( - "success_and_failure" - ) + callback_type: Optional[ + Literal["success", "failure", "success_and_failure"] + ] = "success_and_failure" callback_vars: Dict[str, str] @model_validator(mode="before") @@ -1788,9 +1799,10 @@ class DynamoDBArgs(LiteLLMPydanticObjectBase): class PassThroughGuardrailSettings(LiteLLMPydanticObjectBase): """ Settings for a specific guardrail on a passthrough endpoint. - + Allows field-level targeting for guardrail execution. """ + request_fields: Optional[List[str]] = Field( default=None, description="JSONPath expressions for input field targeting (pre_call). Examples: 'query', 'documents[*].text', 'messages[*].content'. If not specified, guardrail runs on entire request payload.", @@ -1871,9 +1883,9 @@ class ConfigList(LiteLLMPydanticObjectBase): stored_in_db: Optional[bool] field_default_value: Any premium_field: bool = False - nested_fields: Optional[List[FieldDetail]] = ( - None # For nested dictionary or Pydantic fields - ) + nested_fields: Optional[ + List[FieldDetail] + ] = None # For nested dictionary or Pydantic fields class UserHeaderMapping(LiteLLMPydanticObjectBase): @@ -2152,6 +2164,7 @@ class UserAPIKeyAuth( user_rpm_limit: Optional[int] = None user_email: Optional[str] = None request_route: Optional[str] = None + user: Optional[Any] = None # Expanded user object when expand=user is used model_config = ConfigDict(arbitrary_types_allowed=True) @@ -2258,9 +2271,9 @@ class LiteLLM_OrganizationMembershipTable(LiteLLMPydanticObjectBase): budget_id: Optional[str] = None created_at: datetime updated_at: datetime - user: Optional[Any] = ( - None # You might want to replace 'Any' with a more specific type if available - ) + user: Optional[ + Any + ] = None # You might want to replace 'Any' with a more specific type if available litellm_budget_table: Optional[LiteLLM_BudgetTable] = None model_config = ConfigDict(protected_namespaces=()) @@ -2702,7 +2715,7 @@ class AllCallbacks(LiteLLMPydanticObjectBase): "TRACELOOP_API_KEY", ], ui_callback_name="Traceloop", - ) + ) class SpendLogsMetadata(TypedDict): @@ -2736,9 +2749,7 @@ class SpendLogsMetadata(TypedDict): cold_storage_object_key: Optional[ str ] # S3/GCS object key for cold storage retrieval - litellm_overhead_time_ms: Optional[ - float - ] # LiteLLM overhead time in milliseconds + litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds cost_breakdown: Optional[ CostBreakdown ] # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.) @@ -3222,9 +3233,9 @@ class TeamModelDeleteRequest(BaseModel): # Organization Member Requests class OrganizationMemberAddRequest(OrgMemberAddRequest): organization_id: str - max_budget_in_organization: Optional[float] = ( - None # Users max budget within the organization - ) + max_budget_in_organization: Optional[ + float + ] = None # Users max budget within the organization class OrganizationMemberDeleteRequest(MemberDeleteRequest): @@ -3439,9 +3450,9 @@ class ProviderBudgetResponse(LiteLLMPydanticObjectBase): Maps provider names to their budget configs. """ - providers: Dict[str, ProviderBudgetResponseObject] = ( - {} - ) # Dictionary mapping provider names to their budget configurations + providers: Dict[ + str, ProviderBudgetResponseObject + ] = {} # Dictionary mapping provider names to their budget configurations class ProxyStateVariables(TypedDict): @@ -3560,8 +3571,16 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): default=None, description="If no team_id given, default permissions/spend-tracking to this team.s", ) + team_alias_jwt_field: Optional[str] = Field( + default=None, + description="The field in the JWT token that stores the team name/alias. Will be resolved to team_id via database lookup.", + ) org_id_jwt_field: Optional[str] = None + org_alias_jwt_field: Optional[str] = Field( + default=None, + description="The field in the JWT token that stores the organization name/alias. Will be resolved to org_id via database lookup.", + ) user_id_jwt_field: Optional[str] = None user_email_jwt_field: Optional[str] = None user_allowed_email_domain: Optional[str] = None @@ -3576,9 +3595,9 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): enforce_rbac: bool = False roles_jwt_field: Optional[str] = None # v2 on role mappings role_mappings: Optional[List[RoleMapping]] = None - object_id_jwt_field: Optional[str] = ( - None # can be either user / team, inferred from the role mapping - ) + object_id_jwt_field: Optional[ + str + ] = None # can be either user / team, inferred from the role mapping scope_mappings: Optional[List[ScopeMapping]] = None enforce_scope_based_access: bool = False enforce_team_based_model_access: bool = False @@ -3729,13 +3748,16 @@ class DailyOrganizationSpendTransaction(BaseDailySpendTransaction): class DailyUserSpendTransaction(BaseDailySpendTransaction): user_id: str + class DailyEndUserSpendTransaction(BaseDailySpendTransaction): end_user_id: str + class DailyTagSpendTransaction(BaseDailySpendTransaction): request_id: Optional[str] tag: str + class DailyAgentSpendTransaction(BaseDailySpendTransaction): agent_id: str @@ -3767,8 +3789,8 @@ class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): flat_model_file_ids: List[str] created_by: Optional[str] updated_by: Optional[str] - storage_backend: Optional[str] = None - storage_url: Optional[str] = None + storage_backend: Optional[str] = None + storage_url: Optional[str] = None class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): @@ -3799,4 +3821,4 @@ class LiteLLM_ManagedVectorStoresTable(LiteLLMPydanticObjectBase): class ResponseLiteLLM_ManagedVectorStore(TypedDict, total=False): - vector_store: LiteLLM_ManagedVectorStoresTable + vector_store: LiteLLM_ManagedVectorStoresTable \ No newline at end of file diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e2e90abeb1b..de4973ecc69 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1372,6 +1372,195 @@ async def get_team_object( ) +@log_db_metrics +async def get_team_object_by_alias( + team_alias: str, + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + parent_otel_span: Optional["Span"] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, +) -> LiteLLM_TeamTableCachedObj: + """ + Look up a team by its team_alias (name) in the database. + + Args: + team_alias: The team name/alias to look up + prisma_client: Database client + user_api_key_cache: Cache for storing results + parent_otel_span: Optional OpenTelemetry span + proxy_logging_obj: Optional proxy logging object + + Returns: + LiteLLM_TeamTableCachedObj: The team object if found + + Raises: + HTTPException: If team doesn't exist or multiple teams have the same alias + """ + if prisma_client is None: + raise Exception( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + + # Check cache first (keyed by alias) + cache_key = "team_alias:{}".format(team_alias) + + cached_team_obj = await _get_team_object_from_cache( + key=cache_key, + proxy_logging_obj=proxy_logging_obj, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + ) + + if cached_team_obj is not None: + return cached_team_obj + + # Query database by team_alias + try: + teams = await prisma_client.db.litellm_teamtable.find_many( + where={"team_alias": team_alias} + ) + + if not teams: + raise HTTPException( + status_code=404, + detail={ + "error": f"Team with alias '{team_alias}' doesn't exist in db. Create team via `/team/new` call." + }, + ) + + if len(teams) > 1: + raise HTTPException( + status_code=400, + detail={ + "error": f"Multiple teams found with alias '{team_alias}'. Please use team_id_jwt_field instead or ensure team aliases are unique." + }, + ) + + team = teams[0] + team_obj = LiteLLM_TeamTableCachedObj(**team.model_dump()) + + # Cache the result by both alias and team_id + await user_api_key_cache.async_set_cache( + key=cache_key, + value=team_obj, + ttl=DEFAULT_IN_MEMORY_TTL, + ) + # Also cache by team_id for consistency + team_id_cache_key = "team_id:{}".format(team_obj.team_id) + await user_api_key_cache.async_set_cache( + key=team_id_cache_key, + value=team_obj, + ttl=DEFAULT_IN_MEMORY_TTL, + ) + + return team_obj + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception( + "Error looking up team by alias: %s", team_alias + ) + raise HTTPException( + status_code=500, + detail={ + "error": f"Error looking up team by alias '{team_alias}': {str(e)}" + }, + ) + + +@log_db_metrics +async def get_org_object_by_alias( + org_alias: str, + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + parent_otel_span: Optional["Span"] = None, + proxy_logging_obj: Optional[ProxyLogging] = None, +) -> Optional[LiteLLM_OrganizationTable]: + """ + Look up an organization by its organization_alias in the database. + + Args: + org_alias: The organization name/alias to look up + prisma_client: Database client + user_api_key_cache: Cache for storing results + parent_otel_span: Optional OpenTelemetry span + proxy_logging_obj: Optional proxy logging object + + Returns: + LiteLLM_OrganizationTable if found, None otherwise + + Raises: + HTTPException: If organization not found or multiple orgs have the same alias + """ + if prisma_client is None: + raise Exception( + "No DB Connected. See - https://docs.litellm.ai/docs/proxy/virtual_keys" + ) + + # Check cache first (keyed by alias) + cache_key = "org_alias:{}".format(org_alias) + cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key) + if cached_org_obj is not None: + if isinstance(cached_org_obj, dict): + return LiteLLM_OrganizationTable(**cached_org_obj) + elif isinstance(cached_org_obj, LiteLLM_OrganizationTable): + return cached_org_obj + + # Query database by organization_alias + try: + orgs = await prisma_client.db.litellm_organizationtable.find_many( + where={"organization_alias": org_alias} + ) + + if not orgs: + raise HTTPException( + status_code=404, + detail={ + "error": f"Organization with alias '{org_alias}' doesn't exist in db. Create organization via `/organization/new` call." + }, + ) + + if len(orgs) > 1: + raise HTTPException( + status_code=400, + detail={ + "error": f"Multiple organizations found with alias '{org_alias}'. Please use org_id_jwt_field instead or ensure organization aliases are unique." + }, + ) + + org = orgs[0] + org_obj = LiteLLM_OrganizationTable(**org.model_dump()) + + # Cache the result + await user_api_key_cache.async_set_cache( + key=cache_key, + value=org_obj.model_dump(), + ttl=DEFAULT_IN_MEMORY_TTL, + ) + # Also cache by org_id for consistency + await user_api_key_cache.async_set_cache( + key="org_id:{}".format(org_obj.organization_id), + value=org_obj.model_dump(), + ttl=DEFAULT_IN_MEMORY_TTL, + ) + + return org_obj + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception( + "Error looking up organization by alias: %s", org_alias + ) + raise HTTPException( + status_code=500, + detail={ + "error": f"Error looking up organization by alias '{org_alias}': {str(e)}" + }, + ) + + class ExperimentalUIJWTToken: @staticmethod def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 17ff0de9f7b..33667b5d8d9 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -48,10 +48,12 @@ from .auth_checks import ( get_actual_routes, get_end_user_object, get_org_object, + get_org_object_by_alias, get_role_based_models, get_role_based_routes, get_team_membership, get_team_object, + get_team_object_by_alias, get_user_object, ) @@ -194,10 +196,13 @@ class JWTHandler: def is_required_team_id(self) -> bool: """ Returns: - - True: if 'team_id_jwt_field' is set - - False: if not + - True: if 'team_id_jwt_field' or 'team_alias_jwt_field' is set + - False: if neither is set """ - if self.litellm_jwtauth.team_id_jwt_field is None: + if ( + self.litellm_jwtauth.team_id_jwt_field is None + and self.litellm_jwtauth.team_alias_jwt_field is None + ): return False return True @@ -240,6 +245,31 @@ class JWTHandler: team_id = default_value return team_id + def get_team_alias(self, token: dict, default_value: Optional[str]) -> Optional[str]: + """ + Extract team name/alias from JWT token using the configured team_alias_jwt_field. + + Args: + token: The decoded JWT token dictionary + default_value: Default value to return if field not found + + Returns: + The team alias from the token, or default_value if not found + """ + try: + if self.litellm_jwtauth.team_alias_jwt_field is not None: + team_alias = get_nested_value( + data=token, + key_path=self.litellm_jwtauth.team_alias_jwt_field, + default=default_value, + ) + return team_alias + else: + team_alias = None + except KeyError: + team_alias = default_value + return team_alias + def is_upsert_user_id(self, valid_user_email: Optional[bool] = None) -> bool: """ Returns: @@ -383,6 +413,31 @@ class JWTHandler: org_id = default_value return org_id + def get_org_alias(self, token: dict, default_value: Optional[str]) -> Optional[str]: + """ + Extract organization name/alias from JWT token using the configured org_alias_jwt_field. + + Args: + token: The decoded JWT token dictionary + default_value: Default value to return if field not found + + Returns: + The organization alias from the token, or default_value if not found + """ + try: + if self.litellm_jwtauth.org_alias_jwt_field is not None: + org_alias = get_nested_value( + data=token, + key_path=self.litellm_jwtauth.org_alias_jwt_field, + default=default_value, + ) + return org_alias + else: + org_alias = None + except KeyError: + org_alias = default_value + return org_alias + def get_scopes(self, token: dict) -> List[str]: try: if isinstance(token["scope"], str): @@ -813,18 +868,14 @@ class JWTAuthManager: parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, ) -> Tuple[Optional[str], Optional[LiteLLM_TeamTable]]: - """Find and validate specific team ID""" + """Find and validate specific team ID from team_id_jwt_field or team_alias_jwt_field""" individual_team_id = jwt_handler.get_team_id( token=jwt_valid_token, default_value=None ) - if not individual_team_id and jwt_handler.is_required_team_id() is True: - raise Exception( - f"No team id found in token. Checked team_id field '{jwt_handler.litellm_jwtauth.team_id_jwt_field}'" - ) - - ## VALIDATE TEAM OBJECT ### team_object: Optional[LiteLLM_TeamTable] = None + + # First try to get team by team_id if individual_team_id: team_object = await get_team_object( team_id=individual_team_id, @@ -834,6 +885,37 @@ class JWTAuthManager: proxy_logging_obj=proxy_logging_obj, team_id_upsert=jwt_handler.litellm_jwtauth.team_id_upsert, ) + return individual_team_id, team_object + + # If no team_id found, try to resolve via team_alias_jwt_field + team_alias = jwt_handler.get_team_alias( + token=jwt_valid_token, default_value=None + ) + if team_alias: + verbose_proxy_logger.info( + f"JWT Auth: Resolving team by alias: '{team_alias}'" + ) + team_object = await get_team_object_by_alias( + team_alias=team_alias, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if team_object: + individual_team_id = team_object.team_id + verbose_proxy_logger.info( + f"JWT Auth: Resolved team_alias='{team_alias}' to team_id='{individual_team_id}'" + ) + return individual_team_id, team_object + + # Check if team is required but not found + if jwt_handler.is_required_team_id() is True: + team_id_field = jwt_handler.litellm_jwtauth.team_id_jwt_field + team_alias_field = jwt_handler.litellm_jwtauth.team_alias_jwt_field + raise Exception( + f"No team found in token. Checked team_id field '{team_id_field}' and team_alias field '{team_alias_field}'" + ) return individual_team_id, team_object @@ -942,13 +1024,16 @@ class JWTAuthManager: parent_otel_span: Optional[Span], proxy_logging_obj: ProxyLogging, route: str, + org_alias: Optional[str] = None, ) -> Tuple[ Optional[LiteLLM_UserTable], Optional[LiteLLM_OrganizationTable], - Optional[LiteLLM_EndUserTable], + Optional[LiteLLM_EndUserTable], Optional[LiteLLM_TeamMembership], ]: - """Get user, org, and end user objects""" + """Get user, org, and end user objects. Also resolves org aliases to IDs if configured.""" + + # Get org object - first try by ID, then by alias org_object: Optional[LiteLLM_OrganizationTable] = None if org_id: org_object = ( @@ -962,6 +1047,21 @@ class JWTAuthManager: if org_id else None ) + elif org_alias: + verbose_proxy_logger.info( + f"JWT Auth: Resolving org by alias: '{org_alias}'" + ) + org_object = await get_org_object_by_alias( + org_alias=org_alias, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if org_object: + verbose_proxy_logger.info( + f"JWT Auth: Resolved org_alias='{org_alias}' to org_id='{org_object.organization_id}'" + ) user_object: Optional[LiteLLM_UserTable] = None if user_id: @@ -1304,6 +1404,8 @@ class JWTAuthManager: parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, ) + # Extract alias fields for resolution (if configured) + org_alias = jwt_handler.get_org_alias(token=jwt_valid_token, default_value=None) # Get other objects user_object, org_object, end_user_object, team_membership_object = ( @@ -1320,9 +1422,13 @@ class JWTAuthManager: parent_otel_span=parent_otel_span, proxy_logging_obj=proxy_logging_obj, route=route, + org_alias=org_alias, ) ) + # Derive org_id from org_object if resolved by alias + resolved_org_id = org_object.organization_id if org_object else org_id + await JWTAuthManager.sync_user_role_and_teams( jwt_handler=jwt_handler, jwt_valid_token=jwt_valid_token, @@ -1345,10 +1451,9 @@ class JWTAuthManager: ) # check if user is proxy admin - if user_object and user_object.user_role == LitellmUserRoles.PROXY_ADMIN: - is_proxy_admin = True - else: - is_proxy_admin = False + is_proxy_admin = bool( + user_object and user_object.user_role == LitellmUserRoles.PROXY_ADMIN + ) return JWTAuthBuilderResult( is_proxy_admin=is_proxy_admin, @@ -1356,7 +1461,7 @@ class JWTAuthManager: team_object=team_object, user_id=user_id, user_object=user_object, - org_id=org_id, + org_id=resolved_org_id, # Use resolved org_id (from alias lookup if applicable) org_object=org_object, end_user_id=end_user_id, end_user_object=end_user_object, diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index fb9757ca647..8cc33ce6cdd 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -34,6 +34,59 @@ from litellm.secret_managers.main import get_secret_bool from litellm.types.proxy.ui_sso import ReturnedUITokenObject +async def expire_previous_ui_session_tokens( + user_id: str, prisma_client: Optional[PrismaClient] +) -> None: + """ + Expire (block) all other valid UI session tokens for a user. + + This prevents accumulation of multiple valid UI session tokens that + are supposed to be short-lived test keys. Only affects keys with + team_id = "litellm-dashboard" and that haven't expired yet. + + Args: + user_id: The user ID whose previous UI session tokens should be expired + prisma_client: Database client for performing the update + """ + if prisma_client is None: + return + + try: + from datetime import datetime, timezone + + current_time = datetime.now(timezone.utc) + + # Find all unblocked AND non-expired UI session tokens for this user + ui_session_tokens = await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": user_id, + "team_id": "litellm-dashboard", + "OR": [ + {"blocked": None}, # Tokens that have never been blocked (null) + {"blocked": False}, # Tokens explicitly set to not blocked + ], + "expires": {"gt": current_time}, # Only get tokens that haven't expired + } + ) + + if not ui_session_tokens: + return + + # Block all the found tokens + tokens_to_block = [token.token for token in ui_session_tokens if token.token] + + if tokens_to_block: + await prisma_client.db.litellm_verificationtoken.update_many( + where={"token": {"in": tokens_to_block}}, + data={"blocked": True} + ) + + except Exception: + # Silently fail - don't block login if cleanup fails + # This is a best-effort operation + pass + + def get_ui_credentials(master_key: Optional[str]) -> tuple[str, str]: """ Get UI username and password from environment variables or master key. @@ -85,7 +138,7 @@ class LoginResult: self.login_method = login_method -async def authenticate_user( +async def authenticate_user( # noqa: PLR0915 username: str, password: str, master_key: Optional[str], @@ -174,6 +227,10 @@ async def authenticate_user( ) if os.getenv("DATABASE_URL") is not None: + # Expire any previous UI session tokens for this user + await expire_previous_ui_session_tokens( + user_id=key_user_id, prisma_client=prisma_client + ) response = await generate_key_helper_fn( request_type="key", **{ @@ -260,6 +317,11 @@ async def authenticate_user( hash_password, _password ): if os.getenv("DATABASE_URL") is not None: + # Expire any previous UI session tokens for this user + await expire_previous_ui_session_tokens( + user_id=user_id, prisma_client=prisma_client + ) + response = await generate_key_helper_fn( request_type="key", **{ # type: ignore diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 495d4db304c..9b53d9a3a80 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -551,6 +551,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 valid_token = UserAPIKeyAuth( api_key=None, team_id=team_id, + team_alias=( + team_object.team_alias if team_object is not None else None + ), team_tpm_limit=( team_object.tpm_limit if team_object is not None else None ), diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 62c997659bd..59032189438 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -449,6 +449,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) + event_type = ( + GuardrailEventHooks.pre_call + if source == "INPUT" + else GuardrailEventHooks.post_call + ) + try: httpx_response = await self.async_handler.post( url=prepared_request.url, @@ -469,6 +475,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) # Re-raise the exception to maintain existing behavior raise @@ -486,6 +493,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) ######################################################### if httpx_response.status_code == 200: @@ -605,10 +613,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): """ Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions. - If `self.mask_request_content` or `self.mask_response_content` is set to `True`, + If `self.mask_request_content` or `self.mask_response_content` is set to `True`, then use the output from the guardrail to mask the request or response content. - - However, even with masking enabled, content with action="BLOCKED" should still + + However, even with masking enabled, content with action="BLOCKED" should still raise an exception, only content with action="ANONYMIZED" should be masked. """ @@ -731,9 +739,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[ + Union[BedrockGuardrailResponse, str] + ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data @@ -803,9 +811,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ######################################################### ########## 1. Make the Bedrock API request ########## ######################################################### - bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = ( - None - ) + bedrock_guardrail_response: Optional[ + Union[BedrockGuardrailResponse, str] + ] = None try: bedrock_guardrail_response = await self.make_bedrock_api_request( source="INPUT", messages=filtered_messages, request_data=data @@ -1296,11 +1304,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, ) - if bedrock_response.get("action") == "BLOCKED": - raise Exception( - f"Content blocked by Bedrock guardrail: {bedrock_response.get('reason', 'Unknown reason')}" - ) - # Apply any masking that was applied by the guardrail output_list = bedrock_response.get("output") diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index 6915286a2d7..59381149809 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -97,6 +97,7 @@ class DynamoAIGuardrails(CustomGuardrail): async def _call_dynamoai_guardrails( self, messages: List[Dict[str, Any]], + event_type: GuardrailEventHooks, text_type: str = "input", request_data: Optional[dict] = None, ) -> DynamoAIResponse: @@ -157,6 +158,7 @@ class DynamoAIGuardrails(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json @@ -177,6 +179,7 @@ class DynamoAIGuardrails(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -332,6 +335,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=_messages, text_type="input", request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -380,6 +384,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=_messages, text_type="input", request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -460,6 +465,7 @@ class DynamoAIGuardrails(CustomGuardrail): messages=dynamoai_messages, text_type="output", request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 55fa17c21e7..2fc05213640 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -108,6 +108,7 @@ class IBMGuardrailDetector(CustomGuardrail): async def _call_detector_server( self, contents: List[str], + event_type: GuardrailEventHooks, request_data: Optional[dict] = None, ) -> List[List[IBMDetectorDetection]]: """ @@ -142,7 +143,6 @@ class IBMGuardrailDetector(CustomGuardrail): ) try: - response = await self.async_handler.post( url=self.api_url, json=payload, @@ -172,6 +172,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json @@ -192,6 +193,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -199,6 +201,7 @@ class IBMGuardrailDetector(CustomGuardrail): async def _call_orchestrator( self, content: str, + event_type: GuardrailEventHooks, request_data: Optional[dict] = None, ) -> List[IBMDetectorDetection]: """ @@ -258,6 +261,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) return response_json.get("detections", []) @@ -278,6 +282,7 @@ class IBMGuardrailDetector(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) raise @@ -472,6 +477,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -500,6 +506,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.debug( @@ -557,6 +564,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -585,6 +593,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.debug( @@ -673,6 +682,7 @@ class IBMGuardrailDetector(CustomGuardrail): result = await self._call_detector_server( contents=contents_to_check, request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( @@ -702,6 +712,7 @@ class IBMGuardrailDetector(CustomGuardrail): orchestrator_result = await self._call_orchestrator( content=content, request_data=data, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py index 6d4ed089818..953275acf14 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py +++ b/litellm/proxy/guardrails/guardrail_hooks/javelin/javelin.py @@ -83,6 +83,7 @@ class JavelinGuardrail(CustomGuardrail): async def call_javelin_guard( self, request: JavelinGuardRequest, + event_type: GuardrailEventHooks, ) -> JavelinGuardResponse: """ Call the Javelin guard API. @@ -158,6 +159,7 @@ class JavelinGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), + event_type=event_type, ) async def async_pre_call_hook( @@ -208,7 +210,9 @@ class JavelinGuardrail(CustomGuardrail): config=self.config if self.config else {}, ) - javelin_response = await self.call_javelin_guard(request=javelin_guard_request) + javelin_response = await self.call_javelin_guard( + request=javelin_guard_request, event_type=GuardrailEventHooks.pre_call + ) assessments = javelin_response.get("assessments", []) reject_prompt = "" diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 6d98866eadf..732331349e0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -70,6 +70,7 @@ class LakeraAIGuardrail(CustomGuardrail): self, messages: List[AllMessageValues], request_data: Dict, + event_type: GuardrailEventHooks, ) -> Tuple[LakeraAIResponse, Dict]: """ Call the Lakera AI v2 guard API. @@ -128,6 +129,7 @@ class LakeraAIGuardrail(CustomGuardrail): end_time=datetime.now().timestamp(), duration=(datetime.now() - start_time).total_seconds(), masked_entity_count=masked_entity_count, + event_type=event_type, ) def _mask_pii_in_messages( @@ -214,6 +216,7 @@ class LakeraAIGuardrail(CustomGuardrail): lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( messages=new_messages, request_data=data, + event_type=GuardrailEventHooks.pre_call, ) ######################################################### @@ -279,6 +282,7 @@ class LakeraAIGuardrail(CustomGuardrail): lakera_guardrail_response, masked_entity_count = await self.call_v2_guard( messages=new_messages, request_data=data, + event_type=GuardrailEventHooks.during_call, ) ######################################################### diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 51136c29eca..a12eb2486d2 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -295,7 +295,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): filters = ( list(filter_results.values()) if isinstance(filter_results, dict) - else filter_results if isinstance(filter_results, list) else [] + else filter_results + if isinstance(filter_results, list) + else [] ) # Prefer sanitized text from deidentifyResult if present @@ -327,6 +329,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): start_time: Optional[float] = None, end_time: Optional[float] = None, duration: Optional[float] = None, + event_type: Optional[GuardrailEventHooks] = None, ): """ Override to store only the Model Armor API response, not the entire data dict. @@ -351,6 +354,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): duration=duration, start_time=start_time, end_time=end_time, + event_type=event_type, ) return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py index a0ea90ccf21..1794751a08c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma.py @@ -119,9 +119,7 @@ class NomaGuardrail(CustomGuardrail): self.api_base = api_base or os.environ.get( "NOMA_API_BASE", NomaGuardrail._DEFAULT_API_BASE ) - self.application_id = application_id or os.environ.get( - "NOMA_APPLICATION_ID" - ) + self.application_id = application_id or os.environ.get("NOMA_APPLICATION_ID") self.default_application_id = "litellm" if monitor_mode is None: @@ -163,6 +161,7 @@ class NomaGuardrail(CustomGuardrail): self, request_data: dict, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Optional[str]: """Shared logic for processing user message checks""" start_time = datetime.now() @@ -213,6 +212,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if self.monitor_mode: @@ -242,6 +242,7 @@ class NomaGuardrail(CustomGuardrail): request_data: dict, response: LLMResponse, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Optional[str]: """Shared logic for processing LLM response checks""" @@ -293,6 +294,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if self.monitor_mode: @@ -578,7 +580,6 @@ class NomaGuardrail(CustomGuardrail): data: dict, call_type: CallTypesLiteral, ) -> Optional[Union[Exception, str, dict]]: - verbose_proxy_logger.debug("Running Noma pre-call hook") if ( @@ -602,7 +603,9 @@ class NomaGuardrail(CustomGuardrail): return data try: - return await self._check_user_message(data, user_api_key_dict) + return await self._check_user_message( + data, user_api_key_dict, GuardrailEventHooks.pre_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_user_message_check with "blocked" status raise @@ -619,6 +622,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.pre_call, ) verbose_proxy_logger.error(f"Noma pre-call hook failed: {str(e)}") @@ -650,7 +654,9 @@ class NomaGuardrail(CustomGuardrail): return data try: - return await self._check_user_message(data, user_api_key_dict) + return await self._check_user_message( + data, user_api_key_dict, GuardrailEventHooks.during_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_user_message_check with "blocked" status raise @@ -667,6 +673,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.during_call, ) verbose_proxy_logger.error(f"Noma moderation hook failed: {str(e)}") @@ -700,7 +707,9 @@ class NomaGuardrail(CustomGuardrail): return response try: - return await self._check_llm_response(data, response, user_api_key_dict) + return await self._check_llm_response( + data, response, user_api_key_dict, GuardrailEventHooks.post_call + ) except NomaBlockedMessage: # Blocked requests were already logged in _process_llm_response_check with "blocked" status raise @@ -717,6 +726,7 @@ class NomaGuardrail(CustomGuardrail): start_time=start_time.timestamp(), end_time=start_time.timestamp(), duration=0.0, + event_type=GuardrailEventHooks.post_call, ) verbose_proxy_logger.error(f"Noma post-call hook failed: {str(e)}") @@ -728,9 +738,12 @@ class NomaGuardrail(CustomGuardrail): self, request_data: dict, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Union[Exception, str, dict, None]: """Check user message for policy violations""" - user_message = await self._process_user_message_check(request_data, user_auth) + user_message = await self._process_user_message_check( + request_data, user_auth, event_type + ) if not user_message: return request_data @@ -741,10 +754,11 @@ class NomaGuardrail(CustomGuardrail): request_data: dict, response: LLMResponse, user_auth: UserAPIKeyAuth, + event_type: Optional[GuardrailEventHooks] = None, ) -> Any: """Check LLM response for policy violations""" content = await self._process_llm_response_check( - request_data, response, user_auth + request_data, response, user_auth, event_type ) if not content: return response @@ -858,7 +872,10 @@ class NomaGuardrail(CustomGuardrail): if isinstance(assembled_model_response, ModelResponse): try: processed_response = await self._check_llm_response( - request_data, assembled_model_response, user_api_key_dict + request_data, + assembled_model_response, + user_api_key_dict, + GuardrailEventHooks.post_call, ) except NomaBlockedMessage: raise diff --git a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py index 88145ae9e47..02e481acddd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py +++ b/litellm/proxy/guardrails/guardrail_hooks/panw_prisma_airs/panw_prisma_airs.py @@ -24,6 +24,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypesLiteral, ModelResponse if TYPE_CHECKING: @@ -523,6 +524,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): scan_result: Dict[str, Any], data: Dict[str, Any], start_time: datetime, + event_type: GuardrailEventHooks, is_response: bool = False, ) -> Optional[Dict[str, Any]]: """Handle API errors with fail-open/fail-closed logic.""" @@ -542,6 +544,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=duration, + event_type=event_type, ) if scan_result.get("_always_block"): @@ -735,7 +738,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): return self._handle_api_error_with_logging( - scan_result, data, start_time, is_response=False + scan_result, + data, + start_time, + is_response=False, + event_type=GuardrailEventHooks.pre_call, ) end_time = datetime.now() @@ -749,6 +756,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=GuardrailEventHooks.pre_call, ) action = scan_result.get("action", "block") @@ -872,7 +880,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): self._handle_api_error_with_logging( - scan_result, data, start_time, is_response=True + scan_result, + data, + start_time, + is_response=True, + event_type=GuardrailEventHooks.post_call, ) return response @@ -887,6 +899,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=GuardrailEventHooks.post_call, ) action = scan_result.get("action", "block") @@ -1066,7 +1079,11 @@ class PanwPrismaAirsHandler(CustomGuardrail): if scan_result.get("_is_transient") or scan_result.get("_always_block"): self._handle_api_error_with_logging( - scan_result, request_data, start_time, is_response=True + scan_result, + request_data, + start_time, + is_response=True, + event_type=EventHooks.post_call, ) for chunk in all_chunks: yield chunk @@ -1083,6 +1100,7 @@ class PanwPrismaAirsHandler(CustomGuardrail): start_time=start_time.timestamp(), end_time=end_time.timestamp(), duration=(end_time - start_time).total_seconds(), + event_type=EventHooks.post_call, ) # Add guardrail to applied guardrails header for observability diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index a1bbf36ac0c..f66341fde5c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -28,7 +28,6 @@ class UnifiedLLMGuardrails(CustomLogger): self, **kwargs, ): - # store kwargs as optional_params self.optional_params = kwargs @@ -63,6 +62,9 @@ class UnifiedLLMGuardrails(CustomLogger): return data event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call + if call_type == CallTypes.call_mcp_tool.value: + event_type = GuardrailEventHooks.pre_mcp_call + if ( guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True @@ -114,6 +116,9 @@ class UnifiedLLMGuardrails(CustomLogger): return data event_type: GuardrailEventHooks = GuardrailEventHooks.during_call + if call_type == CallTypes.call_mcp_tool.value: + event_type = GuardrailEventHooks.during_mcp_call + if ( guardrail_to_apply.should_run_guardrail(data=data, event_type=event_type) is not True @@ -128,7 +133,10 @@ class UnifiedLLMGuardrails(CustomLogger): endpoint_guardrail_translation_mappings = ( load_guardrail_translation_mappings() ) - if call_type is not None and CallTypes(call_type) not in endpoint_guardrail_translation_mappings: + if ( + call_type is not None + and CallTypes(call_type) not in endpoint_guardrail_translation_mappings + ): return data endpoint_translation = endpoint_guardrail_translation_mappings[ @@ -180,8 +188,8 @@ class UnifiedLLMGuardrails(CustomLogger): call_type: Optional[CallTypesLiteral] = None if user_api_key_dict.request_route is not None: call_types = get_call_types_for_route(user_api_key_dict.request_route) - if call_types is not None and len(call_types) > 0: # type: ignore - call_type = call_types[0] # type: ignore + if call_types is not None and len(call_types) > 0: # type: ignore + call_type = call_types[0] # type: ignore if call_type is None: call_type = _infer_call_type(call_type=None, completion_response=response) # type: ignore @@ -330,7 +338,6 @@ class UnifiedLLMGuardrails(CustomLogger): # Process chunk based on sampling rate if chunk_counter % sampling_rate == 0: - verbose_proxy_logger.debug( "Processing streaming chunk %s (sampling_rate=%s) with guardrail %s", chunk_counter, diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 2d86f74a41c..e43da32565a 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -55,6 +55,18 @@ async def new_budget( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) + # Validate budget values are not negative + if budget_obj.max_budget is not None and budget_obj.max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"} + ) + if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"} + ) + # if no budget_reset_at date is set, but a budget_duration is given, then set budget_reset_at initially to the first completed duration interval in future if budget_obj.budget_reset_at is None and budget_obj.budget_duration is not None: budget_obj.budget_reset_at = datetime.utcnow() + timedelta( @@ -107,6 +119,18 @@ async def update_budget( if budget_obj.budget_id is None: raise HTTPException(status_code=400, detail={"error": "budget_id is required"}) + # Validate budget values are not negative + if budget_obj.max_budget is not None and budget_obj.max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {budget_obj.max_budget}"} + ) + if budget_obj.soft_budget is not None and budget_obj.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"soft_budget cannot be negative. Received: {budget_obj.soft_budget}"} + ) + response = await prisma_client.db.litellm_budgettable.update( where={"budget_id": budget_obj.budget_id}, data={ diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index cc2ac908149..4d72bc86257 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1069,6 +1069,18 @@ async def generate_key_fn( verbose_proxy_logger.debug("entered /key/generate") + # Validate budget values are not negative + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + ) + if data.soft_budget is not None and data.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + ) + if user_custom_key_generate is not None: if asyncio.iscoroutinefunction(user_custom_key_generate): result = await user_custom_key_generate(data) # type: ignore @@ -1502,6 +1514,13 @@ async def update_key_fn( ) try: + # Validate budget values are not negative + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + ) + data_json: dict = data.model_dump(exclude_unset=True, exclude_none=True) key = data_json.pop("key") @@ -3020,10 +3039,14 @@ async def list_keys( description="Column to sort by (e.g. 'user_id', 'created_at', 'spend')", ), sort_order: str = Query(default="desc", description="Sort order ('asc' or 'desc')"), + expand: Optional[List[str]] = Query(None, description="Expand related objects (e.g. 'user')"), ) -> KeyListResponseObject: """ List all keys for a given user / team / organization. + Parameters: + expand: Optional[List[str]] - Expand related objects (e.g. 'user' to include user information) + Returns: { "keys": List[str] or List[UserAPIKeyAuth], @@ -3031,6 +3054,9 @@ async def list_keys( "current_page": int, "total_pages": int, } + + When expand includes "user", each key object will include a "user" field with the associated user object. + Note: When expand=user is specified, full key objects are returned regardless of the return_full_object parameter. """ try: from litellm.proxy.proxy_server import prisma_client @@ -3080,6 +3106,7 @@ async def list_keys( include_created_by_keys=include_created_by_keys, sort_by=sort_by, sort_order=sort_order, + expand=expand, ) verbose_proxy_logger.debug("Successfully prepared response") @@ -3215,45 +3242,17 @@ def _validate_sort_params( return order_by -async def _list_key_helper( - prisma_client: PrismaClient, - page: int, - size: int, +def _build_key_filter_conditions( user_id: Optional[str], team_id: Optional[str], organization_id: Optional[str], key_alias: Optional[str], key_hash: Optional[str], - exclude_team_id: Optional[str] = None, - return_full_object: bool = False, - admin_team_ids: Optional[ - List[str] - ] = None, # New parameter for teams where user is admin - include_created_by_keys: bool = False, - sort_by: Optional[str] = None, - sort_order: str = "desc", -) -> KeyListResponseObject: - """ - Helper function to list keys - Args: - page: int - size: int - user_id: Optional[str] - team_id: Optional[str] - key_alias: Optional[str] - exclude_team_id: Optional[str] # exclude a specific team_id - return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token - admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin - - Returns: - KeyListResponseObject - { - "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types - "total_count": int, - "current_page": int, - "total_pages": int, - } - """ + exclude_team_id: Optional[str], + admin_team_ids: Optional[List[str]], + include_created_by_keys: bool, +) -> Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]]: + """Build filter conditions for key listing.""" # Prepare filter conditions where: Dict[str, Union[str, Dict[str, Any], List[Dict[str, Any]]]] = {} where.update(_get_condition_to_filter_out_ui_session_tokens()) @@ -3294,6 +3293,59 @@ async def _list_key_helper( where.update(or_conditions[0]) verbose_proxy_logger.debug(f"Filter conditions: {where}") + return where + + +async def _list_key_helper( + prisma_client: PrismaClient, + page: int, + size: int, + user_id: Optional[str], + team_id: Optional[str], + organization_id: Optional[str], + key_alias: Optional[str], + key_hash: Optional[str], + exclude_team_id: Optional[str] = None, + return_full_object: bool = False, + admin_team_ids: Optional[ + List[str] + ] = None, # New parameter for teams where user is admin + include_created_by_keys: bool = False, + sort_by: Optional[str] = None, + sort_order: str = "desc", + expand: Optional[List[str]] = None, +) -> KeyListResponseObject: + """ + Helper function to list keys + Args: + page: int + size: int + user_id: Optional[str] + team_id: Optional[str] + key_alias: Optional[str] + exclude_team_id: Optional[str] # exclude a specific team_id + return_full_object: bool # when true, will return UserAPIKeyAuth objects instead of just the token + admin_team_ids: Optional[List[str]] # list of team IDs where the user is an admin + + Returns: + KeyListResponseObject + { + "keys": List[str] or List[UserAPIKeyAuth], # Updated to reflect possible return types + "total_count": int, + "current_page": int, + "total_pages": int, + } + """ + where = _build_key_filter_conditions( + user_id=user_id, + team_id=team_id, + organization_id=organization_id, + key_alias=key_alias, + key_hash=key_hash, + exclude_team_id=exclude_team_id, + admin_team_ids=admin_team_ids, + include_created_by_keys=include_created_by_keys, + ) # Calculate skip for pagination skip = (page - 1) * size @@ -3334,13 +3386,28 @@ async def _list_key_helper( # Calculate total pages total_pages = -(-total_count // size) # Ceiling division + # Fetch user information if expand includes "user" + user_map = {} + if expand and "user" in expand: + user_ids = [key.user_id for key in keys if key.user_id] + if user_ids: + users = await prisma_client.db.litellm_usertable.find_many( + where={"user_id": {"in": list(set(user_ids))}} # Remove duplicates + ) + user_map = {user.user_id: user for user in users} + # Prepare response key_list: List[Union[str, UserAPIKeyAuth]] = [] for key in keys: key_dict = key.dict() # Attach object_permission if object_permission_id is set key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) - if return_full_object is True: + + # Include user information if expand includes "user" + if expand and "user" in expand and key.user_id and key.user_id in user_map: + key_dict["user"] = user_map[key.user_id].dict() + + if return_full_object is True or (expand and "user" in expand): key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object else: _token = key_dict.get("token") diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index f0eddcc8683..9111f53a517 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -16,7 +16,7 @@ Endpoints here: import importlib from dataclasses import dataclass from datetime import datetime, timedelta -from typing import Any, Dict, Iterable, List, Optional +from typing import Any, Dict, Iterable, List, Literal, Optional from fastapi import ( APIRouter, @@ -24,6 +24,7 @@ from fastapi import ( Form, Header, HTTPException, + Query, Request, Response, status, @@ -208,6 +209,9 @@ if MCP_AVAILABLE: command=payload.command, args=payload.args, env=payload.env, + authorization_url=payload.authorization_url, + token_url=payload.token_url, + registration_url=payload.registration_url, ) def get_prisma_client_or_throw(message: str): @@ -296,117 +300,6 @@ if MCP_AVAILABLE: access_groups_list = sorted(list(access_groups)) return {"access_groups": access_groups_list} - @router.get( - "/server/{server_id}/health", - description="Perform health check on a specific MCP server", - dependencies=[Depends(user_api_key_auth)], - ) - async def health_check_mcp_server( - server_id: str, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - ): - """ - Perform a health check on the MCP server specified by the `server_id` - Parameters: - - server_id: str - Required. The unique identifier of the mcp server to health check. - ``` - curl --location 'http://localhost:4000/v1/mcp/server/{server_id}/health' \ - --header 'Authorization: Bearer your_api_key_here' - ``` - """ - # Check if server exists and user has access - prisma_client = get_prisma_client_or_throw( - "Database not connected. Connect a database to your proxy" - ) - - # check to see if server exists for all users - mcp_server = await get_mcp_server(prisma_client, server_id) - if mcp_server is None: - raise HTTPException( - status_code=status.HTTP_404_NOT_FOUND, - detail={"error": f"MCP Server with id {server_id} not found"}, - ) - - # Implement authz restriction from requested user - if not _user_has_admin_view(user_api_key_dict): - # Perform authz check to filter the mcp servers user has access to - mcp_server_records = await get_all_mcp_servers_for_user( - prisma_client, user_api_key_dict - ) - exists = does_mcp_server_exist(mcp_server_records, server_id) - - if not exists: - raise HTTPException( - status_code=status.HTTP_403_FORBIDDEN, - detail={ - "error": f"User does not have permission to access mcp server with id {server_id}. You can only access mcp servers that you have access to." - }, - ) - - # Perform health check using server manager - try: - health_result = await global_mcp_server_manager.health_check_server( - server_id - ) - return health_result - except Exception as e: - verbose_proxy_logger.exception( - f"Error performing health check on MCP server {server_id}: {str(e)}" - ) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error performing health check: {str(e)}"}, - ) - - @router.get( - "/server/health", - description="Perform health check on all accessible MCP servers", - dependencies=[Depends(user_api_key_auth)], - ) - async def health_check_all_mcp_servers( - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), - ): - """ - Perform health checks on all MCP servers accessible to the user - ``` - curl --location 'http://localhost:4000/v1/mcp/server/health' \ - --header 'Authorization: Bearer your_api_key_here' - ``` - """ - # Use server manager to get health checks for allowed servers - try: - all_health_results = ( - await global_mcp_server_manager.health_check_allowed_servers( - user_api_key_auth=user_api_key_dict - ) - ) - - return { - "total_servers": len(all_health_results), - "healthy_count": len( - [r for r in all_health_results.values() if r["status"] == "healthy"] - ), - "unhealthy_count": len( - [ - r - for r in all_health_results.values() - if r["status"] == "unhealthy" - ] - ), - "unknown_count": len( - [r for r in all_health_results.values() if r["status"] == "unknown"] - ), - "servers": all_health_results, - } - except Exception as e: - verbose_proxy_logger.exception( - f"Error performing health checks on MCP servers: {str(e)}" - ) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail={"error": f"Error performing health checks: {str(e)}"}, - ) - ## FastAPI Routes @router.get( "/server", @@ -429,7 +322,7 @@ if MCP_AVAILABLE: aggregated_servers: Dict[str, LiteLLM_MCPServerTable] = {} for auth_context in auth_contexts: - servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams( + servers = await global_mcp_server_manager.get_all_allowed_mcp_servers( user_api_key_auth=auth_context ) for server in servers: @@ -447,6 +340,56 @@ if MCP_AVAILABLE: server.mcp_info["is_public"] = True return redacted_mcp_servers + @router.get( + "/server/health", + description="Health check for MCP servers", + dependencies=[Depends(user_api_key_auth)], + ) + async def health_check_servers( + server_ids: Optional[List[str]] = Query( + None, + description="Server IDs to check. If not provided, checks all accessible servers.", + ), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + ): + """ + Perform health checks on one or more MCP servers. + + Parameters: + - server_ids: Optional list of server IDs. If not provided, checks all accessible servers. + + Returns: + - Health check results for requested servers + + ``` + # Check all accessible servers + curl --location 'http://localhost:4000/v1/mcp/server/health' \ + --header 'Authorization: Bearer your_api_key_here' + + # Check specific servers + curl --location 'http://localhost:4000/v1/mcp/server/health?server_ids=server-1&server_ids=server-2' \ + --header 'Authorization: Bearer your_api_key_here' + ``` + """ + auth_contexts = await build_effective_auth_contexts(user_api_key_dict) + + server_status_map: Dict[ + str, Optional[Literal["healthy", "unhealthy", "unknown"]] + ] = {} + for auth_context in auth_contexts: + servers = await global_mcp_server_manager.get_all_mcp_servers_with_health_and_teams( + user_api_key_auth=auth_context, + server_ids=server_ids, + ) + for server in servers: + if server.server_id not in server_status_map: + server_status_map[server.server_id] = server.status + + return [ + {"server_id": server_id, "status": status} + for server_id, status in server_status_map.items() + ] + @router.get( "/server/{server_id}", description="Returns the mcp server info", @@ -484,15 +427,11 @@ if MCP_AVAILABLE: server_id ) # Update the server object with health check results - mcp_server.status = health_result.get("status", "unknown") - mcp_server.last_health_check = ( - datetime.fromisoformat( - health_result.get("last_health_check", datetime.now().isoformat()) - ) - if health_result.get("last_health_check") - else None + mcp_server.status = ( + health_result.status if health_result.status else "unknown" ) - mcp_server.health_check_error = health_result.get("error") + mcp_server.last_health_check = health_result.last_health_check + mcp_server.health_check_error = health_result.health_check_error except Exception as e: verbose_proxy_logger.debug( f"Error performing health check on server {server_id}: {e}" @@ -512,7 +451,7 @@ if MCP_AVAILABLE: exists = does_mcp_server_exist(mcp_server_records, server_id) if exists: - await global_mcp_server_manager.add_update_server(mcp_server) + await global_mcp_server_manager.add_server(mcp_server) return _redact_mcp_credentials(mcp_server) else: raise HTTPException( @@ -586,7 +525,7 @@ if MCP_AVAILABLE: payload, touched_by=user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, ) - await global_mcp_server_manager.add_update_server(new_mcp_server) + await global_mcp_server_manager.add_server(new_mcp_server) # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() @@ -867,7 +806,7 @@ if MCP_AVAILABLE: "error": f"MCP Server not found, passed server_id={payload.server_id}" }, ) - await global_mcp_server_manager.add_update_server(mcp_server_record_updated) + await global_mcp_server_manager.update_server(mcp_server_record_updated) # Ensure registry is up to date by reloading from database await global_mcp_server_manager.reload_servers_from_database() diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 99b37c765a8..a088f46d67d 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -24,8 +24,11 @@ from litellm.proxy.management_endpoints.budget_management_endpoints import ( new_budget, update_budget, ) -from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field -from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.proxy.management_endpoints.common_utils import ( + _set_object_metadata_field, + _user_has_admin_view, +) from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) @@ -34,11 +37,10 @@ from litellm.proxy.management_helpers.utils import ( management_endpoint_wrapper, ) from litellm.proxy.utils import PrismaClient -from litellm.utils import _update_dictionary from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) -from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity +from litellm.utils import _update_dictionary router = APIRouter() @@ -168,6 +170,18 @@ async def new_organization( status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value} ) + # Validate budget values are not negative + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + ) + if data.soft_budget is not None and data.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + ) + user_object_correct_type: Optional[LiteLLM_UserTable] = None if user_api_key_dict.user_id is not None: @@ -414,6 +428,18 @@ async def update_organization( # Create validated data model data = LiteLLM_OrganizationTableUpdate(**raw_data_with_flat_budget_fields) + # Validate budget values are not negative + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + ) + if data.soft_budget is not None and data.soft_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"soft_budget cannot be negative. Received: {data.soft_budget}"} + ) + if data.updated_by is None: data.updated_by = user_api_key_dict.user_id diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c6fab9a73f0..76c607f5c49 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -732,6 +732,18 @@ async def new_team( # noqa: PLR0915 if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) + # Validate budget values are not negative + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + ) + if data.team_member_budget is not None and data.team_member_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"} + ) + # Check if license is over limit total_teams = await prisma_client.db.litellm_teamtable.count() if total_teams and _license_check.is_team_count_over_limit( @@ -1169,7 +1181,7 @@ def validate_team_org_change( "/team/update", tags=["team management"], dependencies=[Depends(user_api_key_auth)] ) @management_endpoint_wrapper -async def update_team( +async def update_team( # noqa: PLR0915 data: UpdateTeamRequest, http_request: Request, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), @@ -1254,6 +1266,18 @@ async def update_team( raise HTTPException(status_code=400, detail={"error": "No team id passed in"}) verbose_proxy_logger.debug("/team/update - %s", data) + # Validate budget values are not negative + if data.max_budget is not None and data.max_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"max_budget cannot be negative. Received: {data.max_budget}"} + ) + if data.team_member_budget is not None and data.team_member_budget < 0: + raise HTTPException( + status_code=400, + detail={"error": f"team_member_budget cannot be negative. Received: {data.team_member_budget}"} + ) + existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( where={"team_id": data.team_id} ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f56c0c2b07a..06525e39133 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4626,7 +4626,7 @@ class ProxyStartupEvent: verbose_proxy_logger.info("Batch cost check job scheduled successfully") except Exception as e: - verbose_proxy_logger.error(f"Failed to setup batch cost checking: {e}") + verbose_proxy_logger.debug(f"Failed to setup batch cost checking: {e}") verbose_proxy_logger.debug( "Checking batch cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..." ) @@ -4657,7 +4657,7 @@ class ProxyStartupEvent: verbose_proxy_logger.info("Responses cost check job scheduled successfully") except Exception as e: - verbose_proxy_logger.error(f"Failed to setup responses cost checking: {e}") + verbose_proxy_logger.debug(f"Failed to setup responses cost checking: {e}") verbose_proxy_logger.debug( "Checking responses cost for LiteLLM Managed Files is an Enterprise Feature. Skipping..." ) @@ -7322,13 +7322,9 @@ async def model_info_v2( """ global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router - if llm_router is None: - raise HTTPException( - status_code=500, - detail={ - "error": f"No model list passed, models router={llm_router}. You can add a model through the config.yaml or on the LiteLLM Admin UI." - }, - ) + # Return empty data array when no models are configured (graceful handling for fresh installs) + if llm_router is None or not llm_router.model_list: + return {"data": []} if prisma_client is None: raise HTTPException( @@ -8228,14 +8224,9 @@ async def model_group_info( """ global llm_model_list, general_settings, user_config_file_path, proxy_config, llm_router - if llm_model_list is None: - raise HTTPException( - status_code=500, detail={"error": "LLM Model List not loaded in"} - ) - if llm_router is None: - raise HTTPException( - status_code=500, detail={"error": "LLM Router is not loaded in"} - ) + # Return empty data array when no models are configured (graceful handling for fresh installs) + if llm_model_list is None or llm_router is None or not llm_model_list: + return {"data": []} from litellm.proxy.utils import get_available_models_for_user diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 9916bdf6923..ec1c4619527 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -1680,6 +1680,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "MINIMAX", + "provider_display_name": "MiniMax", + "litellm_provider": "minimax", + "credential_fields": [ + { + "key": "api_key", + "label": "API Key", + "placeholder": "your-minimax-api-key", + "tooltip": "MiniMax API Key from https://platform.minimaxi.com/", + "required": true, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "api_base", + "label": "API Base URL", + "placeholder": "https://api.minimax.io/v1", + "tooltip": "International: https://api.minimax.io/v1, China: https://api.minimaxi.com/v1", + "required": false, + "field_type": "text", + "options": null, + "default_value": "https://api.minimax.io/v1" + } + ], + "default_model_placeholder": "minimax/MiniMax-M2" + }, { "provider": "MOONSHOT", "provider_display_name": "Moonshot", @@ -2865,7 +2893,7 @@ "key": "api_base", "label": "API Base", "placeholder": null, - "tooltip": null, + "tooltip": "Base URL of your WatsonX instance", "required": false, "field_type": "text", "options": null, @@ -2875,14 +2903,54 @@ "key": "api_key", "label": "API Key", "placeholder": null, - "tooltip": null, + "tooltip": "IBM Cloud API key. Required if not using Token or Zen API Key", "required": false, "field_type": "password", "options": null, "default_value": null + }, + { + "key": "token", + "label": "IAM Token", + "placeholder": null, + "tooltip": "Pre-generated IAM bearer token. Use instead of API Key if you manage tokens externally", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "zen_api_key", + "label": "Zen API Key", + "placeholder": null, + "tooltip": "Zen API Key for Cloud Pak for Data deployments. Use instead of API Key for on-premises", + "required": false, + "field_type": "password", + "options": null, + "default_value": null + }, + { + "key": "project_id", + "label": "Project ID", + "placeholder": null, + "tooltip": "Optional: Your Watsonx.ai Project ID", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "space_id", + "label": "Deployment Space ID", + "placeholder": null, + "tooltip": "Optional: Watsonx.ai Deployment Space ID", + "required": false, + "field_type": "text", + "options": null, + "default_value": null } ], - "default_model_placeholder": "gpt-3.5-turbo" + "default_model_placeholder": "watsonx/ibm/granite-3-3-8b-instruct" }, { "provider": "WATSONX_TEXT", diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index aac0b5b35de..ea47b6ed03b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + authorization_url String? + token_url String? + registration_url String? } // Generate Tokens for Proxy diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 1861c44b699..1c457d7bf4c 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,7 +11,10 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB, REDACTED_BY_LITELM_STRING -from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs +from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, + reconstruct_model_name, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.proxy.utils import PrismaClient, hash_token @@ -100,9 +103,9 @@ def _get_spend_logs_metadata( clean_metadata["applied_guardrails"] = applied_guardrails clean_metadata["batch_models"] = batch_models clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata - clean_metadata["vector_store_request_metadata"] = ( - _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) - ) + clean_metadata[ + "vector_store_request_metadata" + ] = _get_vector_store_request_for_spend_logs_payload(vector_store_request_metadata) clean_metadata["guardrail_information"] = guardrail_information clean_metadata["usage_object"] = usage_object clean_metadata["model_map_information"] = model_map_information @@ -393,6 +396,9 @@ def get_logging_payload( # noqa: PLR0915 # Extract agent_id for A2A requests (set directly on model_call_details) agent_id: Optional[str] = kwargs.get("agent_id") + custom_llm_provider = kwargs.get("custom_llm_provider") + raw_model = cast(str, kwargs.get("model") or "") + model_name = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) try: payload: SpendLogsPayload = SpendLogsPayload( @@ -403,7 +409,7 @@ def get_logging_payload( # noqa: PLR0915 startTime=_ensure_datetime_utc(start_time), endTime=_ensure_datetime_utc(end_time), completionStartTime=_ensure_datetime_utc(completion_start_time), - model=kwargs.get("model", "") or "", + model=model_name, user=metadata.get("user_api_key_user_id", "") or "", team_id=metadata.get("user_api_key_team_id", "") or "", organization_id=metadata.get("user_api_key_org_id") or "", @@ -449,7 +455,7 @@ def get_logging_payload( # noqa: PLR0915 # Explicitly clear large intermediate objects to reduce memory pressure del response_obj_dict, usage, clean_metadata, additional_usage_values - + return payload except Exception as e: verbose_proxy_logger.exception( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index d595db4a2e0..d1a78534dae 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -151,25 +151,25 @@ def _get_email_logger_class(): """ Determine which email logger class to use based on environment variables. Priority: SendGrid > Resend > SMTP > BaseEmailLogger (fallback) - + Returns: The email logger class to use, or None if BaseEmailLogger is not available """ if BaseEmailLogger is None: return None - + # Check for SendGrid API key if SendGridEmailLogger is not None and os.getenv("SENDGRID_API_KEY"): return SendGridEmailLogger - + # Check for Resend API key if ResendEmailLogger is not None and os.getenv("RESEND_API_KEY"): return ResendEmailLogger - + # Check for SMTP configuration if SMTPEmailLogger is not None and os.getenv("SMTP_HOST"): return SMTPEmailLogger - + # Fallback to BaseEmailLogger (though it won't actually send emails) return BaseEmailLogger @@ -452,7 +452,6 @@ class ProxyLogging: litellm.logging_callback_manager.add_litellm_callback(self.service_logging_obj) # type: ignore for callback in litellm.callbacks: if isinstance(callback, str): - callback = litellm.litellm_core_utils.litellm_logging._init_custom_logger_compatible_class( # type: ignore cast(_custom_logger_compatible_callbacks_literal, callback), internal_usage_cache=self.internal_usage_cache.dual_cache, @@ -965,7 +964,7 @@ class ProxyLogging: # Determine the event type based on call type event_type = GuardrailEventHooks.pre_call - if call_type == "mcp_call": + if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.pre_mcp_call # Check if the guardrail should run for this request @@ -1038,7 +1037,6 @@ class ProxyLogging: data.pop("prompt_id", None) if custom_logger and prompt_spec is not None: - ( model, messages, @@ -1261,7 +1259,7 @@ class ProxyLogging: from litellm.types.guardrails import GuardrailEventHooks event_type = GuardrailEventHooks.during_call - if call_type == "mcp_call": + if call_type == CallTypes.call_mcp_tool.value: event_type = GuardrailEventHooks.during_mcp_call if ( @@ -1270,7 +1268,7 @@ class ProxyLogging: ): continue # Convert user_api_key_dict to proper format for async_moderation_hook - if call_type == "mcp_call": + if call_type == CallTypes.call_mcp_tool.value: user_api_key_auth_dict = self._convert_user_api_key_auth_to_dict( user_api_key_dict ) @@ -1288,7 +1286,6 @@ class ProxyLogging: call_type=call_type, ) else: - guardrail_task = callback.async_moderation_hook( data=data, user_api_key_dict=user_api_key_auth_dict, # type: ignore @@ -1337,7 +1334,7 @@ class ProxyLogging: if self.alerting is None: # do nothing if alerting is not switched on return - + if "slack" in self.alerting: await self.slack_alerting_instance.budget_alerts( type=type, @@ -1548,7 +1545,10 @@ class ProxyLogging: traceback_str=traceback_str, ) # If callback returned an HTTPException, use it (first one wins) - if isinstance(hook_result, HTTPException) and transformed_exception is None: + if ( + isinstance(hook_result, HTTPException) + and transformed_exception is None + ): transformed_exception = hook_result except HTTPException as e: # If callback raised an HTTPException, use it (first one wins) @@ -1849,7 +1849,6 @@ class ProxyLogging: current_response = response for callback in litellm.callbacks: - _callback: Optional[CustomLogger] = None if isinstance(callback, str): _callback = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( @@ -3568,11 +3567,13 @@ class ProxyUpdateSpend: ) # Atomically read and remove logs to process (protected by lock) async with prisma_client._spend_log_transactions_lock: - logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] + logs_to_process = prisma_client.spend_log_transactions[ + :MAX_LOGS_PER_INTERVAL + ] # Remove the logs we're about to process - prisma_client.spend_log_transactions = ( - prisma_client.spend_log_transactions[len(logs_to_process):] - ) + prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ + len(logs_to_process) : + ] start_time = time.time() try: for i in range(n_retry_times + 1): @@ -3675,9 +3676,7 @@ async def update_spend( # noqa: PLR0915 # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - verbose_proxy_logger.debug( - "Spend Logs transactions: {}".format(queue_size) - ) + verbose_proxy_logger.debug("Spend Logs transactions: {}".format(queue_size)) # Process spend log transactions when called directly. # This keeps backwards compatibility with the old behavior. @@ -3699,19 +3698,19 @@ async def update_spend_logs_job( ): """ Job to process spend_log_transactions queue. - + This job is triggered based on queue size rather than time. Processes spend log transactions when the queue reaches a threshold. """ n_retry_times = 3 - + # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - + if queue_size == 0: return - + await ProxyUpdateSpend.update_spend_logs( n_retry_times=n_retry_times, prisma_client=prisma_client, @@ -3728,7 +3727,7 @@ async def _monitor_spend_logs_queue( """ Background task that monitors the spend_log_transactions queue size and triggers processing when the threshold is reached. - + Args: prisma_client: Prisma client instance db_writer_client: Optional HTTP handler for external spend logs endpoint @@ -3738,23 +3737,23 @@ async def _monitor_spend_logs_queue( SPEND_LOG_QUEUE_POLL_INTERVAL, SPEND_LOG_QUEUE_SIZE_THRESHOLD, ) - + threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL max_backoff = 30.0 # Maximum backoff interval in seconds backoff_multiplier = 1.5 # Exponential backoff multiplier current_interval = base_interval - + verbose_proxy_logger.info( f"Starting spend logs queue monitor (threshold: {threshold}, poll_interval: {base_interval}s)" ) - + while True: try: # Check queue size with lock protection async with prisma_client._spend_log_transactions_lock: queue_size = len(prisma_client.spend_log_transactions) - + if queue_size > 0: if queue_size >= threshold: verbose_proxy_logger.debug( @@ -3767,8 +3766,10 @@ async def _monitor_spend_logs_queue( f"Spend logs queue size ({queue_size}) below threshold ({threshold}), processing with backoff" ) # Exponential backoff when below threshold but still processing - current_interval = min(current_interval * backoff_multiplier, max_backoff) - + current_interval = min( + current_interval * backoff_multiplier, max_backoff + ) + await update_spend_logs_job( prisma_client=prisma_client, db_writer_client=db_writer_client, @@ -3776,8 +3777,10 @@ async def _monitor_spend_logs_queue( ) else: # Exponential backoff when no logs to process - current_interval = min(current_interval * backoff_multiplier, max_backoff) - + current_interval = min( + current_interval * backoff_multiplier, max_backoff + ) + await asyncio.sleep(current_interval) except Exception as e: verbose_proxy_logger.error( @@ -3788,7 +3791,6 @@ async def _monitor_spend_logs_queue( await asyncio.sleep(current_interval) - def _raise_failed_update_spend_exception( e: Exception, start_time: float, proxy_logging_obj: ProxyLogging ): diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index fe686598141..e26a2477b1b 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -3,7 +3,7 @@ from typing import Any, Optional, cast import litellm -from litellm import get_llm_provider +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 2eea28f6cc1..9cdcd3894e0 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -205,7 +205,15 @@ class LiteLLM_Proxy_MCP_Handler: else: tool_name = getattr(mcp_tool, "name", None) - if tool_name and tool_name in allowed_tool_names: + if not tool_name: + continue + + if tool_name in allowed_tool_names: + filtered_tools.append(mcp_tool) + continue + + unprefixed_name, _ = split_server_prefix_from_name(tool_name) + if unprefixed_name in allowed_tool_names: filtered_tools.append(mcp_tool) return filtered_tools diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index ad99609e905..a92b5d25a37 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -26,7 +26,7 @@ from litellm.types.llms.openai import ( from litellm.types.responses.main import DecodedResponseId from litellm.types.utils import ( CompletionTokensDetailsWrapper, - PromptTokensDetails, + PromptTokensDetailsWrapper, SpecialEnums, Usage, ) @@ -431,7 +431,12 @@ class ResponseAPILoggingUtils: def _transform_response_api_usage_to_chat_usage( usage_input: Optional[Union[dict, ResponseAPIUsage]], ) -> Usage: - """Tranforms the ResponseAPIUsage object to a Usage object""" + """ + Transforms ResponseAPIUsage or ImageUsage to a Usage object. + + Both have the same spec with input_tokens, output_tokens, and + input_tokens_details (text_tokens, image_tokens). + """ if usage_input is None: return Usage( prompt_tokens=0, @@ -445,18 +450,19 @@ class ResponseAPILoggingUtils: ) prompt_tokens: int = response_api_usage.input_tokens or 0 completion_tokens: int = response_api_usage.output_tokens or 0 - prompt_tokens_details: Optional[PromptTokensDetails] = None + prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None if response_api_usage.input_tokens_details: - prompt_tokens_details = PromptTokensDetails( - cached_tokens=response_api_usage.input_tokens_details.cached_tokens, - audio_tokens=response_api_usage.input_tokens_details.audio_tokens, + prompt_tokens_details = PromptTokensDetailsWrapper( + cached_tokens=getattr(response_api_usage.input_tokens_details, "cached_tokens", None), + audio_tokens=getattr(response_api_usage.input_tokens_details, "audio_tokens", None), + text_tokens=getattr(response_api_usage.input_tokens_details, "text_tokens", None), + image_tokens=getattr(response_api_usage.input_tokens_details, "image_tokens", None), ) completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None - if response_api_usage.output_tokens_details: + output_tokens_details = getattr(response_api_usage, "output_tokens_details", None) + if output_tokens_details: completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=getattr( - response_api_usage.output_tokens_details, "reasoning_tokens", None - ) + reasoning_tokens=getattr(output_tokens_details, "reasoning_tokens", None) ) chat_usage = Usage( diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index c6804b1ad4c..69d6ab9b6e2 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -7,7 +7,7 @@ import re from re import Match from typing import Dict, List, Optional, Tuple -from litellm import get_llm_provider +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm._logging import verbose_router_logger diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index ceeae958a80..c2912558cab 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1197,6 +1197,39 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): # Define private attributes using PrivateAttr _hidden_params: dict = PrivateAttr(default_factory=dict) + @property + def output_text(self) -> str: + """ + Convenience property that aggregates all `output_text` items from the `output` list. + + If no `output_text` content blocks exist, then an empty string is returned. + + This matches the OpenAI SDK's Response.output_text behavior. + """ + texts: List[str] = [] + for output_item in self.output: + # Handle both dict and object access patterns + if isinstance(output_item, dict): + item_type = output_item.get("type") + content = output_item.get("content", []) + else: + item_type = getattr(output_item, "type", None) + content = getattr(output_item, "content", []) + + if item_type == "message": + for content_item in content: + if isinstance(content_item, dict): + content_type = content_item.get("type") + text = content_item.get("text", "") + else: + content_type = getattr(content_item, "type", None) + text = getattr(content_item, "text", "") or "" + + if content_type == "output_text": + texts.append(text) + + return "".join(texts) + class ResponsesAPIStreamEvents(str, Enum): """ diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 144e503acdf..3eec67d9d26 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -142,6 +142,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): ] # only for vertex ai models input_cost_per_query: Optional[float] # only for rerank models input_cost_per_image: Optional[float] # only for vertex ai models + input_cost_per_image_token: Optional[float] # for gpt-image-1 and similar models input_cost_per_audio_per_second: Optional[float] # only for vertex ai models input_cost_per_video_per_second: Optional[float] # only for vertex ai models input_cost_per_second: Optional[float] # for OpenAI Speech models @@ -1300,7 +1301,7 @@ class CacheCreationTokenDetails(BaseModel): class PromptTokensDetailsWrapper( PromptTokensDetails -): # wrapper for older openai versions +): # extends with image generation fields (text_tokens, image_tokens) text_tokens: Optional[int] = None """Text tokens sent to the model.""" diff --git a/litellm/utils.py b/litellm/utils.py index 102df5d595e..df0b2317123 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -34,7 +34,6 @@ from inspect import iscoroutine from io import StringIO from os.path import abspath, dirname, join -import aiohttp import dotenv import httpx import openai @@ -48,21 +47,16 @@ from tiktoken import Encoding from tokenizers import Tokenizer import litellm -import litellm._service_logger # for storing API inputs, outputs, and metadata + import litellm.litellm_core_utils -import litellm.litellm_core_utils.audio_utils.utils +# audio_utils.utils is lazy-loaded - only imported when needed for transcription calls import litellm.litellm_core_utils.json_validation_rule -import litellm.llms -import litellm.llms.gemini from litellm._lazy_imports import ( _get_default_encoding, _get_modified_max_tokens, _get_token_counter_new, ) from litellm._uuid import uuid -from litellm.caching._internal_lru_cache import lru_cache_wrapper -from litellm.caching.caching import DualCache -from litellm.caching.caching_handler import CachingHandlerResponse, LLMCachingHandler from litellm.constants import ( DEFAULT_CHAT_COMPLETION_PARAM_VALUES, DEFAULT_EMBEDDING_PARAM_VALUES, @@ -77,90 +71,84 @@ from litellm.constants import ( OPENAI_EMBEDDING_PARAMS, TOOL_CHOICE_OBJECT_TOKEN_COUNT, ) -from litellm.integrations.custom_guardrail import CustomGuardrail -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.vector_store_integrations.base_vector_store import ( - BaseVectorStore, -) -# Import cached imports utilities -from litellm.litellm_core_utils.cached_imports import ( - get_coroutine_checker, - get_litellm_logging_class, - get_set_callbacks, -) -from litellm.litellm_core_utils.core_helpers import ( - get_litellm_metadata_from_kwargs, - map_finish_reason, - process_response_headers, -) -from litellm.litellm_core_utils.credential_accessor import CredentialAccessor -from litellm.litellm_core_utils.dot_notation_indexing import ( - delete_nested_value, - is_nested_path, -) -from litellm.litellm_core_utils.exception_mapping_utils import ( - _get_response_headers, - exception_type, - get_error_message, -) -from litellm.litellm_core_utils.get_litellm_params import ( - _get_base_model_from_litellm_call_metadata, - get_litellm_params, -) -from litellm.litellm_core_utils.get_llm_provider_logic import ( - _is_non_openai_azure_model, - get_llm_provider, -) -from litellm.litellm_core_utils.get_supported_openai_params import ( - get_supported_openai_params, -) -from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe -from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( - LiteLLMResponseObjectHandler, - _handle_invalid_parallel_tool_calls, - convert_to_model_response_object, - convert_to_streaming_response, - convert_to_streaming_response_async, -) -from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base -from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( - get_formatted_prompt, -) -from litellm.litellm_core_utils.llm_response_utils.get_headers import ( - get_response_headers, -) -from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( - ResponseMetadata, -) -from litellm.litellm_core_utils.prompt_templates.common_utils import ( - _parse_content_for_reasoning, -) -from litellm.litellm_core_utils.redact_messages import ( - LiteLLMLoggingObject, - redact_message_input_output_from_logging, -) -from litellm.litellm_core_utils.rules import Rules -from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.llms.base_llm.google_genai.transformation import ( - BaseGoogleGenAIGenerateContentConfig, -) -from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig -from litellm.llms.base_llm.search.transformation import BaseSearchConfig -from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig -from litellm.llms.bedrock.common_utils import BedrockModelInfo -from litellm.llms.cohere.common_utils import CohereModelInfo -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from litellm.llms.mistral.ocr.transformation import MistralOCRConfig -from litellm.router_utils.get_retry_from_policy import ( - get_num_retries_from_retry_policy, - reset_retry_policy, -) -from litellm.secret_managers.main import get_secret -from litellm.types.llms.anthropic import ( - ANTHROPIC_API_ONLY_HEADERS, - AnthropicThinkingParam, -) + + +_CachingHandlerResponse = None +_LLMCachingHandler = None +_CustomGuardrail = None +_CustomLogger = None + + +def _get_cached_custom_logger(): + """ + Get cached CustomLogger class. + Lazy imports on first call to avoid loading custom_logger at import time. + Subsequent calls use cached class for better performance. + """ + global _CustomLogger + if _CustomLogger is None: + from litellm.integrations.custom_logger import CustomLogger + _CustomLogger = CustomLogger + return _CustomLogger + + +def _get_cached_custom_guardrail(): + """ + Get cached CustomGuardrail class. + Lazy imports on first call to avoid loading custom_guardrail at import time. + Subsequent calls use cached class for better performance. + """ + global _CustomGuardrail + if _CustomGuardrail is None: + from litellm.integrations.custom_guardrail import CustomGuardrail + _CustomGuardrail = CustomGuardrail + return _CustomGuardrail + + +def _get_cached_caching_handler_response(): + """ + Get cached CachingHandlerResponse class. + Lazy imports on first call to avoid loading caching_handler at import time. + Subsequent calls use cached class for better performance. + """ + global _CachingHandlerResponse + if _CachingHandlerResponse is None: + from litellm.caching.caching_handler import CachingHandlerResponse + _CachingHandlerResponse = CachingHandlerResponse + return _CachingHandlerResponse + + +def _get_cached_llm_caching_handler(): + """ + Get cached LLMCachingHandler class. + Lazy imports on first call to avoid loading caching_handler at import time. + Subsequent calls use cached class for better performance. + """ + global _LLMCachingHandler + if _LLMCachingHandler is None: + from litellm.caching.caching_handler import LLMCachingHandler + _LLMCachingHandler = LLMCachingHandler + return _LLMCachingHandler + + +# Cached lazy import for audio_utils.utils +# Module-level cache to avoid repeated imports while preserving memory benefits +_audio_utils_module = None + + +def _get_cached_audio_utils(): + """ + Get cached audio_utils.utils module. + Lazy imports on first call to avoid loading audio_utils.utils at import time. + Subsequent calls use cached module for better performance. + """ + global _audio_utils_module + if _audio_utils_module is None: + import litellm.litellm_core_utils.audio_utils.utils + _audio_utils_module = litellm.litellm_core_utils.audio_utils.utils + return _audio_utils_module + from litellm.types.llms.openai import ( AllMessageValues, AllPromptValues, @@ -171,7 +159,6 @@ from litellm.types.llms.openai import ( OpenAITextCompletionUserMessage, OpenAIWebSearchOptions, ) -from litellm.types.rerank import RerankResponse from litellm.types.utils import FileTypes # type: ignore from litellm.types.utils import ( OPENAI_RESPONSE_HEADERS, @@ -256,16 +243,7 @@ from typing import ( from openai import OpenAIError as OriginalError -from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( - update_response_metadata, -) -from litellm.litellm_core_utils.thread_pool_executor import executor -from litellm.llms.base_llm.anthropic_messages.transformation import ( - BaseAnthropicMessagesConfig, -) -from litellm.llms.base_llm.audio_transcription.transformation import ( - BaseAudioTranscriptionConfig, -) +# These are lazy loaded via __getattr__ from litellm.llms.base_llm.base_utils import ( BaseLLMModelInfo, type_to_response_format_param, @@ -274,31 +252,126 @@ from litellm.llms.base_llm.base_utils import ( if TYPE_CHECKING: # Heavy types that are only needed for type checking; avoid importing # their modules at runtime during `litellm` import. + from litellm.caching.caching_handler import CachingHandlerResponse, LLMCachingHandler + from litellm.integrations.custom_logger import CustomLogger from litellm.llms.base_llm.files.transformation import BaseFilesConfig from litellm.proxy._types import AllowedModelRegion + # Type stubs for lazy-loaded functions to help mypy understand their types + # These imports allow mypy to understand the types when these are accessed via __getattr__ + from litellm.litellm_core_utils.exception_mapping_utils import exception_type + from litellm.litellm_core_utils.get_llm_provider_logic import ( + _is_non_openai_azure_model, + get_llm_provider, + ) + from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params, + ) + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + LiteLLMResponseObjectHandler, + _handle_invalid_parallel_tool_calls, + convert_to_model_response_object, + convert_to_streaming_response, + convert_to_streaming_response_async, + ) + from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + ResponseMetadata, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _parse_content_for_reasoning, + ) + from litellm.litellm_core_utils.redact_messages import ( + LiteLLMLoggingObject, + redact_message_input_output_from_logging, + ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig, + ) + from litellm.llms.base_llm.ocr.transformation import BaseOCRConfig + from litellm.llms.base_llm.search.transformation import BaseSearchConfig + from litellm.llms.base_llm.text_to_speech.transformation import BaseTextToSpeechConfig + from litellm.llms.bedrock.common_utils import BedrockModelInfo + from litellm.llms.cohere.common_utils import CohereModelInfo + from litellm.llms.mistral.ocr.transformation import MistralOCRConfig + # Type stubs for lazy-loaded functions and classes + from litellm.litellm_core_utils.cached_imports import ( + get_coroutine_checker, + get_litellm_logging_class, + get_set_callbacks, + ) + from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs, + map_finish_reason, + process_response_headers, + ) + from litellm.litellm_core_utils.dot_notation_indexing import ( + delete_nested_value, + is_nested_path, + ) + from litellm.litellm_core_utils.get_litellm_params import ( + _get_base_model_from_litellm_call_metadata, + get_litellm_params, + ) + from litellm.litellm_core_utils.llm_request_utils import _ensure_extra_body_is_safe + from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( + get_formatted_prompt, + ) + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers, + ) + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata, + ) + from litellm.litellm_core_utils.rules import Rules + from litellm.litellm_core_utils.thread_pool_executor import executor + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig, + ) + from litellm.llms.base_llm.audio_transcription.transformation import ( + BaseAudioTranscriptionConfig, + ) + from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler + from litellm.router_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy, + reset_retry_policy, + ) + from litellm.secret_managers.main import get_secret + # Type stubs for lazy-loaded config classes and types + from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig + from litellm.llms.base_llm.containers.transformation import BaseContainerConfig + from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig + from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig + from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, + ) + from litellm.llms.base_llm.image_variations.transformation import ( + BaseImageVariationConfig, + ) + from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig + from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig + from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig + from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig + from litellm.llms.base_llm.vector_store_files.transformation import ( + BaseVectorStoreFilesConfig, + ) + from litellm.llms.base_llm.videos.transformation import BaseVideoConfig + from litellm.types.llms.anthropic import ( + ANTHROPIC_API_ONLY_HEADERS, + AnthropicThinkingParam, + ) + from litellm.types.rerank import RerankResponse + from litellm.types.llms.openai import ( + ChatCompletionDeltaToolCallChunk, + ChatCompletionToolCallChunk, + ChatCompletionToolCallFunctionChunk, + ) + from litellm.types.router import LiteLLM_Params -from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.llms.base_llm.completion.transformation import BaseTextCompletionConfig -from litellm.llms.base_llm.containers.transformation import BaseContainerConfig -from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig -from litellm.llms.base_llm.image_edit.transformation import BaseImageEditConfig -from litellm.llms.base_llm.image_generation.transformation import ( - BaseImageGenerationConfig, -) -from litellm.llms.base_llm.image_variations.transformation import ( - BaseImageVariationConfig, -) -from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig -from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.base_llm.skills.transformation import BaseSkillsAPIConfig -from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig -from litellm.llms.base_llm.vector_store_files.transformation import ( - BaseVectorStoreFilesConfig, -) -from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from ._logging import _is_debugging_on, verbose_logger from .caching.caching import ( @@ -326,12 +399,6 @@ from .exceptions import ( UnprocessableEntityError, UnsupportedParamsError, ) -from .types.llms.openai import ( - ChatCompletionDeltaToolCallChunk, - ChatCompletionToolCallChunk, - ChatCompletionToolCallFunctionChunk, -) -from .types.router import LiteLLM_Params if TYPE_CHECKING: from litellm import MockException @@ -487,7 +554,7 @@ def _add_custom_logger_callback_to_specific_event( def _custom_logger_class_exists_in_success_callbacks( - callback_class: CustomLogger, + callback_class: "CustomLogger", ) -> bool: """ Returns True if an instance of the custom logger exists in litellm.success_callback or litellm._async_success_callback @@ -503,7 +570,7 @@ def _custom_logger_class_exists_in_success_callbacks( def _custom_logger_class_exists_in_failure_callbacks( - callback_class: CustomLogger, + callback_class: "CustomLogger", ) -> bool: """ Returns True if an instance of the custom logger exists in litellm.failure_callback or litellm._async_failure_callback @@ -536,6 +603,7 @@ def get_applied_guardrails(kwargs: Dict[str, Any]) -> List[str]: request_guardrails = get_request_guardrails(kwargs) applied_guardrails = [] + CustomGuardrail = _get_cached_custom_guardrail() for callback in litellm.callbacks: if callback is not None and isinstance(callback, CustomGuardrail): if callback.guardrail_name is not None: @@ -547,10 +615,22 @@ def get_applied_guardrails(kwargs: Dict[str, Any]) -> List[str]: return applied_guardrails +def _get_utils_globals() -> dict: + """ + Get the globals dictionary of the utils module. + + This is where we cache imported attributes so we don't import them twice. + """ + return sys.modules[__name__].__dict__ + + def load_credentials_from_list(kwargs: dict): """ Updates kwargs with the credentials if credential_name in kwarg """ + # Access CredentialAccessor via module to trigger lazy loading if needed + CredentialAccessor = getattr(sys.modules[__name__], 'CredentialAccessor') + credential_name = kwargs.get("litellm_credential_name") if credential_name and litellm.credential_list: credential_accessor = CredentialAccessor.get_credential_values(credential_name) @@ -560,7 +640,7 @@ def load_credentials_from_list(kwargs: dict): def get_dynamic_callbacks( - dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]], + dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]], ) -> List: returned_callbacks = litellm.callbacks.copy() if dynamic_callbacks: @@ -694,7 +774,7 @@ def function_setup( # noqa: PLR0915 function_id: Optional[str] = kwargs["id"] if "id" in kwargs else None ## DYNAMIC CALLBACKS ## - dynamic_callbacks: Optional[List[Union[str, Callable, CustomLogger]]] = ( + dynamic_callbacks: Optional[List[Union[str, Callable, "CustomLogger"]]] = ( kwargs.pop("callbacks", None) ) all_callbacks = get_dynamic_callbacks(dynamic_callbacks=dynamic_callbacks) @@ -739,6 +819,7 @@ def function_setup( # noqa: PLR0915 + litellm.failure_callback ) ) + get_set_callbacks = getattr(sys.modules[__name__], 'get_set_callbacks') get_set_callbacks()(callback_list=callback_list, function_id=function_id) ## ASYNC CALLBACKS if len(litellm.input_callback) > 0: @@ -795,16 +876,16 @@ def function_setup( # noqa: PLR0915 litellm.failure_callback.pop(index) ### DYNAMIC CALLBACKS ### dynamic_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] + List[Union[str, Callable, "CustomLogger"]] ] = None dynamic_async_success_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] + List[Union[str, Callable, "CustomLogger"]] ] = None dynamic_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] + List[Union[str, Callable, "CustomLogger"]] ] = None dynamic_async_failure_callbacks: Optional[ - List[Union[str, Callable, CustomLogger]] + List[Union[str, Callable, "CustomLogger"]] ] = None if kwargs.get("success_callback", None) is not None and isinstance( kwargs["success_callback"], list @@ -866,6 +947,7 @@ def function_setup( # noqa: PLR0915 elif kwargs.get("messages", None): messages = kwargs["messages"] ### PRE-CALL RULES ### + Rules = getattr(sys.modules[__name__], 'Rules') if ( Rules.has_pre_call_rules() and isinstance(messages, list) @@ -965,7 +1047,9 @@ def function_setup( # noqa: PLR0915 or call_type == CallTypes.transcription.value ): _file_obj: FileTypes = args[1] if len(args) > 1 else kwargs["file"] - file_checksum = litellm.litellm_core_utils.audio_utils.utils.get_audio_file_content_hash( + # Lazy import audio_utils.utils only when needed for transcription calls + audio_utils = _get_cached_audio_utils() + file_checksum = audio_utils.get_audio_file_content_hash( file_obj=_file_obj ) if "metadata" in kwargs: @@ -996,6 +1080,7 @@ def function_setup( # noqa: PLR0915 call_type=call_type, ): stream = True + get_litellm_logging_class = getattr(sys.modules[__name__], 'get_litellm_logging_class') logging_obj = get_litellm_logging_class()( # Victim for object pool model=model, # type: ignore messages=messages, @@ -1079,6 +1164,8 @@ def _get_wrapper_num_retries( if num_retries is None: num_retries = litellm.num_retries if kwargs.get("retry_policy", None): + get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') + reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') retry_policy_num_retries = get_num_retries_from_retry_policy( exception=exception, retry_policy=kwargs.get("retry_policy"), @@ -1122,6 +1209,7 @@ async def async_pre_call_deployment_hook(kwargs: Dict[str, Any], call_type: str) modified_kwargs = kwargs.copy() + CustomLogger = _get_cached_custom_logger() for callback in litellm.callbacks: if isinstance(callback, CustomLogger): result = await callback.async_pre_call_deployment_hook( @@ -1144,6 +1232,7 @@ async def async_post_call_success_deployment_hook( except ValueError: typed_call_type = None # unknown call type + CustomLogger = _get_cached_custom_logger() for callback in litellm.callbacks: if isinstance(callback, CustomLogger): result = await callback.async_post_call_success_deployment_hook( @@ -1262,6 +1351,7 @@ def post_call_processing( def client(original_function): # noqa: PLR0915 + Rules = getattr(sys.modules[__name__], 'Rules') rules_obj = Rules() @wraps(original_function) @@ -1323,7 +1413,8 @@ def client(original_function): # noqa: PLR0915 ## LOAD CREDENTIALS load_credentials_from_list(kwargs) kwargs["litellm_logging_obj"] = logging_obj - _llm_caching_handler: LLMCachingHandler = LLMCachingHandler( + LLMCachingHandler = _get_cached_llm_caching_handler() + _llm_caching_handler: "LLMCachingHandler" = LLMCachingHandler( original_function=original_function, request_kwargs=kwargs, start_time=start_time, @@ -1378,7 +1469,7 @@ def client(original_function): # noqa: PLR0915 ): # allow users to control returning cached responses from the completion function # checking cache verbose_logger.debug("INSIDE CHECKING SYNC CACHE") - caching_handler_response: CachingHandlerResponse = ( + caching_handler_response: "CachingHandlerResponse" = ( _llm_caching_handler._sync_get_cache( model=model or "", original_function=original_function, @@ -1446,6 +1537,7 @@ def client(original_function): # noqa: PLR0915 ) else: # RETURN RESULT + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1489,6 +1581,7 @@ def client(original_function): # noqa: PLR0915 # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation ctx = contextvars.copy_context() + executor = getattr(sys.modules[__name__], 'executor') executor.submit( ctx.run, logging_obj.success_handler, @@ -1497,6 +1590,7 @@ def client(original_function): # noqa: PLR0915 end_time, ) # RETURN RESULT + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1513,6 +1607,8 @@ def client(original_function): # noqa: PLR0915 kwargs.get("num_retries", None) or litellm.num_retries or None ) if kwargs.get("retry_policy", None): + get_num_retries_from_retry_policy = getattr(sys.modules[__name__], 'get_num_retries_from_retry_policy') + reset_retry_policy = getattr(sys.modules[__name__], 'reset_retry_policy') num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), @@ -1569,7 +1665,8 @@ def client(original_function): # noqa: PLR0915 logging_obj: Optional[LiteLLMLoggingObject] = kwargs.get( "litellm_logging_obj", None ) - _llm_caching_handler: LLMCachingHandler = LLMCachingHandler( + LLMCachingHandler = _get_cached_llm_caching_handler() + _llm_caching_handler: "LLMCachingHandler" = LLMCachingHandler( original_function=original_function, request_kwargs=kwargs, start_time=start_time, @@ -1608,7 +1705,7 @@ def client(original_function): # noqa: PLR0915 print_verbose( f"ASYNC kwargs[caching]: {kwargs.get('caching', False)}; litellm.cache: {litellm.cache}; kwargs.get('cache'): {kwargs.get('cache', None)}" ) - _caching_handler_response: Optional[CachingHandlerResponse] = ( + _caching_handler_response: "Optional[CachingHandlerResponse]" = ( await _llm_caching_handler._async_get_cache( model=model or "", original_function=original_function, @@ -1683,6 +1780,7 @@ def client(original_function): # noqa: PLR0915 chunks, messages=kwargs.get("messages", None) ) else: + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1747,6 +1845,7 @@ def client(original_function): # noqa: PLR0915 end_time=end_time, ) + update_response_metadata = getattr(sys.modules[__name__], 'update_response_metadata') update_response_metadata( result=result, logging_obj=logging_obj, @@ -1822,6 +1921,7 @@ def client(original_function): # noqa: PLR0915 setattr(e, "timeout", timeout) raise e + get_coroutine_checker = getattr(sys.modules[__name__], 'get_coroutine_checker') is_coroutine = get_coroutine_checker().is_async_callable(original_function) # Return the appropriate wrapper based on the original function type @@ -2182,6 +2282,7 @@ def supports_response_schema( """ ## GET LLM PROVIDER ## try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider ) @@ -2879,6 +2980,9 @@ def get_optional_params_embeddings( # noqa: PLR0915 additional_drop_params: Optional[List[str]] = None, **kwargs, ): + # Lazy load get_supported_openai_params + get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') + # retrieve all parameters passed to the function passed_params = locals() custom_llm_provider = passed_params.pop("custom_llm_provider", None) @@ -3148,6 +3252,21 @@ def get_optional_params_embeddings( # noqa: PLR0915 drop_params=drop_params if drop_params is not None else False, ) + elif custom_llm_provider == "ollama": + if 'dimensions' in non_default_params: + optional_params['dimensions']=non_default_params.pop('dimensions') + if len(non_default_params.keys()) > 0: + if ( + litellm.drop_params is True or drop_params is True + ): # drop the unsupported non-default values + keys = list(non_default_params.keys()) + for k in keys: + non_default_params.pop(k, None) + else: + raise UnsupportedParamsError( + status_code=500, + message=f"Setting {non_default_params} is not supported by {custom_llm_provider}. To drop it from the call, set `litellm.drop_params = True`.", + ) elif ( custom_llm_provider != "openai" and custom_llm_provider != "azure" @@ -3666,6 +3785,7 @@ def get_optional_params( # noqa: PLR0915 message=f"{custom_llm_provider} does not support parameters: {list(unsupported_params.keys())}, for model={model}. To drop these, set `litellm.drop_params=True` or for proxy:\n\n`litellm_settings:\n drop_params: true`\n. \n If you want to use these params dynamically send allowed_openai_params={list(unsupported_params.keys())} in your request.", ) + get_supported_openai_params = getattr(sys.modules[__name__], 'get_supported_openai_params') supported_params = get_supported_openai_params( model=model, custom_llm_provider=custom_llm_provider ) @@ -3920,6 +4040,7 @@ def get_optional_params( # noqa: PLR0915 ), ) elif custom_llm_provider == "bedrock": + BedrockModelInfo = getattr(sys.modules[__name__], 'BedrockModelInfo') bedrock_route = BedrockModelInfo.get_bedrock_route(model) bedrock_base_model = BedrockModelInfo.get_base_model(model) if bedrock_route == "converse" or bedrock_route == "converse_like": @@ -4344,6 +4465,8 @@ def get_optional_params( # noqa: PLR0915 # Apply nested drops from additional_drop_params if additional_drop_params: + is_nested_path = getattr(sys.modules[__name__], 'is_nested_path') + delete_nested_value = getattr(sys.modules[__name__], 'delete_nested_value') nested_paths = [p for p in additional_drop_params if is_nested_path(p)] for path in nested_paths: optional_params = delete_nested_value(optional_params, path) @@ -4393,6 +4516,7 @@ def add_provider_specific_params_to_optional_params( else: processed_extra_body = initial_extra_body + _ensure_extra_body_is_safe = getattr(sys.modules[__name__], '_ensure_extra_body_is_safe') optional_params["extra_body"] = _ensure_extra_body_is_safe( extra_body=processed_extra_body ) @@ -4803,6 +4927,7 @@ def get_max_tokens(model: str) -> Optional[int]: return litellm.model_cost[model]["max_output_tokens"] elif "max_tokens" in litellm.model_cost[model]: return litellm.model_cost[model]["max_tokens"] + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model, custom_llm_provider, _, _ = get_llm_provider(model=model) if custom_llm_provider == "huggingface": max_tokens = _get_max_position_embeddings(model_name=model) @@ -4923,6 +5048,7 @@ def _get_potential_model_names( if custom_llm_provider is None: # Get custom_llm_provider try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') split_model, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: split_model = model @@ -5216,6 +5342,9 @@ def _get_model_info_helper( # noqa: PLR0915 input_cost_per_audio_token=_model_info.get( "input_cost_per_audio_token", None ), + input_cost_per_image_token=_model_info.get( + "input_cost_per_image_token", None + ), input_cost_per_token_batches=_model_info.get( "input_cost_per_token_batches" ), @@ -5642,6 +5771,7 @@ def validate_environment( # noqa: PLR0915 } ## EXTRACT LLM PROVIDER - if model name provided try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') _, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: custom_llm_provider = None @@ -6204,6 +6334,7 @@ def register_prompt_template( complete_model = model potential_models = [complete_model] try: + get_llm_provider = getattr(sys.modules[__name__], 'get_llm_provider') model = get_llm_provider(model=model)[0] potential_models.append(model) except Exception: @@ -6289,6 +6420,7 @@ class TextCompletionStreamWrapper: except StopIteration: raise StopIteration except Exception as e: + exception_type = getattr(sys.modules[__name__], 'exception_type') raise exception_type( model=self.model, custom_llm_provider=self.custom_llm_provider or "", @@ -6910,14 +7042,14 @@ def _get_base_model_from_metadata(model_call_details=None): return _base_model metadata = litellm_params.get("metadata", {}) - base_model_from_metadata = _get_base_model_from_litellm_call_metadata( - metadata=metadata - ) + _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') + base_model_from_metadata = _get_base_model_from_litellm_call_metadata(metadata=metadata) if base_model_from_metadata is not None: return base_model_from_metadata # Also check litellm_metadata (used by Responses API and other generic API calls) litellm_metadata = litellm_params.get("litellm_metadata", {}) + _get_base_model_from_litellm_call_metadata = getattr(sys.modules[__name__], '_get_base_model_from_litellm_call_metadata') return _get_base_model_from_litellm_call_metadata(metadata=litellm_metadata) return None @@ -7329,6 +7461,7 @@ class ProviderConfigManager: litellm.LlmProviders.COHERE_CHAT == provider or litellm.LlmProviders.COHERE == provider ): + CohereModelInfo = getattr(sys.modules[__name__], 'CohereModelInfo') route = CohereModelInfo.get_cohere_route(model) if route == "v2": return litellm.CohereV2ChatConfig() @@ -8179,6 +8312,7 @@ class ProviderConfigManager: return get_vertex_ai_ocr_config(model=model) + MistralOCRConfig = getattr(sys.modules[__name__], 'MistralOCRConfig') PROVIDER_TO_CONFIG_MAP = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, } @@ -8319,6 +8453,7 @@ def get_end_user_id_for_cost_tracking( service_type: "litellm_logging" or "prometheus" - used to allow prometheus only disable cost tracking. """ + get_litellm_metadata_from_kwargs = getattr(sys.modules[__name__], 'get_litellm_metadata_from_kwargs') _metadata = cast( dict, get_litellm_metadata_from_kwargs(dict(litellm_params=litellm_params)) ) @@ -8610,16 +8745,672 @@ def should_run_mock_completion( return False -# Re-export encoding from main.py for backward compatibility -# This allows tests to import: from litellm.utils import encoding -# We use a lazy import to avoid loading main.py at utils.py import time -def __getattr__(name: str) -> Any: +def __getattr__(name: str) -> Any: # noqa: PLR0915 """Lazy import handler for utils module""" + _globals = _get_utils_globals() + + # Lazy load encoding from main.py to avoid heavy tiktoken import if name == "encoding": - # Cache it in the module's __dict__ for subsequent accesses - import sys - - from litellm.main import encoding as _encoding - sys.modules[__name__].__dict__["encoding"] = _encoding - return _encoding + # Check if already cached + if "encoding" not in _globals: + from litellm.main import encoding as _encoding + _globals["encoding"] = _encoding + return _globals["encoding"] + + # Lazy load BaseVectorStore to avoid loading it at module import time + if name == "BaseVectorStore": + # Check if already cached + if "BaseVectorStore" not in _globals: + from litellm.integrations.vector_store_integrations.base_vector_store import ( + BaseVectorStore as _BaseVectorStore, + ) + _globals["BaseVectorStore"] = _BaseVectorStore + return _globals["BaseVectorStore"] + + # Lazy load CredentialAccessor to avoid loading it at module import time + if name == "CredentialAccessor": + # Check if already cached + if "CredentialAccessor" not in _globals: + from litellm.litellm_core_utils.credential_accessor import ( + CredentialAccessor as _CredentialAccessor, + ) + _globals["CredentialAccessor"] = _CredentialAccessor + return _globals["CredentialAccessor"] + + # Lazy load exception_mapping_utils functions to avoid loading at module import time + if name == "exception_type": + # Check if already cached + if "exception_type" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + exception_type as _exception_type, + ) + _globals["exception_type"] = _exception_type + return _globals["exception_type"] + + if name == "get_error_message": + # Check if already cached + if "get_error_message" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + get_error_message as _get_error_message, + ) + _globals["get_error_message"] = _get_error_message + return _globals["get_error_message"] + + if name == "_get_response_headers": + # Check if already cached + if "_get_response_headers" not in _globals: + from litellm.litellm_core_utils.exception_mapping_utils import ( + _get_response_headers as __get_response_headers, + ) + _globals["_get_response_headers"] = __get_response_headers + return _globals["_get_response_headers"] + + # Lazy load get_llm_provider_logic functions to avoid loading at module import time + if name == "get_llm_provider": + # Check if already cached + if "get_llm_provider" not in _globals: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider as _get_llm_provider, + ) + _globals["get_llm_provider"] = _get_llm_provider + return _globals["get_llm_provider"] + + if name == "_is_non_openai_azure_model": + # Check if already cached + if "_is_non_openai_azure_model" not in _globals: + from litellm.litellm_core_utils.get_llm_provider_logic import ( + _is_non_openai_azure_model as __is_non_openai_azure_model, + ) + _globals["_is_non_openai_azure_model"] = __is_non_openai_azure_model + return _globals["_is_non_openai_azure_model"] + + # Lazy load get_supported_openai_params to avoid loading at module import time + if name == "get_supported_openai_params": + # Check if already cached + if "get_supported_openai_params" not in _globals: + from litellm.litellm_core_utils.get_supported_openai_params import ( + get_supported_openai_params as _get_supported_openai_params, + ) + _globals["get_supported_openai_params"] = _get_supported_openai_params + return _globals["get_supported_openai_params"] + + # Lazy load convert_dict_to_response functions to avoid loading at module import time + if name == "LiteLLMResponseObjectHandler": + # Check if already cached + if "LiteLLMResponseObjectHandler" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + LiteLLMResponseObjectHandler as _LiteLLMResponseObjectHandler, + ) + _globals["LiteLLMResponseObjectHandler"] = _LiteLLMResponseObjectHandler + return _globals["LiteLLMResponseObjectHandler"] + + if name == "_handle_invalid_parallel_tool_calls": + # Check if already cached + if "_handle_invalid_parallel_tool_calls" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + _handle_invalid_parallel_tool_calls as __handle_invalid_parallel_tool_calls, + ) + _globals["_handle_invalid_parallel_tool_calls"] = __handle_invalid_parallel_tool_calls + return _globals["_handle_invalid_parallel_tool_calls"] + + if name == "convert_to_model_response_object": + # Check if already cached + if "convert_to_model_response_object" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_model_response_object as _convert_to_model_response_object, + ) + _globals["convert_to_model_response_object"] = _convert_to_model_response_object + return _globals["convert_to_model_response_object"] + + if name == "convert_to_streaming_response": + # Check if already cached + if "convert_to_streaming_response" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response as _convert_to_streaming_response, + ) + _globals["convert_to_streaming_response"] = _convert_to_streaming_response + return _globals["convert_to_streaming_response"] + + if name == "convert_to_streaming_response_async": + # Check if already cached + if "convert_to_streaming_response_async" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_streaming_response_async as _convert_to_streaming_response_async, + ) + _globals["convert_to_streaming_response_async"] = _convert_to_streaming_response_async + return _globals["convert_to_streaming_response_async"] + + # Lazy load get_api_base to avoid loading at module import time + if name == "get_api_base": + # Check if already cached + if "get_api_base" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_api_base import ( + get_api_base as _get_api_base, + ) + _globals["get_api_base"] = _get_api_base + return _globals["get_api_base"] + + # Lazy load ResponseMetadata to avoid loading at module import time + if name == "ResponseMetadata": + # Check if already cached + if "ResponseMetadata" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + ResponseMetadata as _ResponseMetadata, + ) + _globals["ResponseMetadata"] = _ResponseMetadata + return _globals["ResponseMetadata"] + + # Lazy load _parse_content_for_reasoning to avoid loading at module import time + if name == "_parse_content_for_reasoning": + # Check if already cached + if "_parse_content_for_reasoning" not in _globals: + from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _parse_content_for_reasoning as __parse_content_for_reasoning, + ) + _globals["_parse_content_for_reasoning"] = __parse_content_for_reasoning + return _globals["_parse_content_for_reasoning"] + + # Lazy load redact_messages to avoid loading at module import time + if name == "LiteLLMLoggingObject": + # Check if already cached + if "LiteLLMLoggingObject" not in _globals: + from litellm.litellm_core_utils.redact_messages import ( + LiteLLMLoggingObject as _LiteLLMLoggingObject, + ) + _globals["LiteLLMLoggingObject"] = _LiteLLMLoggingObject + return _globals["LiteLLMLoggingObject"] + + if name == "redact_message_input_output_from_logging": + # Check if already cached + if "redact_message_input_output_from_logging" not in _globals: + from litellm.litellm_core_utils.redact_messages import ( + redact_message_input_output_from_logging as _redact_message_input_output_from_logging, + ) + _globals["redact_message_input_output_from_logging"] = _redact_message_input_output_from_logging + return _globals["redact_message_input_output_from_logging"] + + # Lazy load CustomStreamWrapper to avoid loading at module import time + if name == "CustomStreamWrapper": + # Check if already cached + if "CustomStreamWrapper" not in _globals: + from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper as _CustomStreamWrapper, + ) + _globals["CustomStreamWrapper"] = _CustomStreamWrapper + return _globals["CustomStreamWrapper"] + + # Lazy load BaseGoogleGenAIGenerateContentConfig to avoid loading at module import time + if name == "BaseGoogleGenAIGenerateContentConfig": + # Check if already cached + if "BaseGoogleGenAIGenerateContentConfig" not in _globals: + from litellm.llms.base_llm.google_genai.transformation import ( + BaseGoogleGenAIGenerateContentConfig as _BaseGoogleGenAIGenerateContentConfig, + ) + _globals["BaseGoogleGenAIGenerateContentConfig"] = _BaseGoogleGenAIGenerateContentConfig + return _globals["BaseGoogleGenAIGenerateContentConfig"] + + # Lazy load BaseOCRConfig to avoid loading at module import time + if name == "BaseOCRConfig": + # Check if already cached + if "BaseOCRConfig" not in _globals: + from litellm.llms.base_llm.ocr.transformation import ( + BaseOCRConfig as _BaseOCRConfig, + ) + _globals["BaseOCRConfig"] = _BaseOCRConfig + return _globals["BaseOCRConfig"] + + # Lazy load BaseSearchConfig to avoid loading at module import time + if name == "BaseSearchConfig": + # Check if already cached + if "BaseSearchConfig" not in _globals: + from litellm.llms.base_llm.search.transformation import ( + BaseSearchConfig as _BaseSearchConfig, + ) + _globals["BaseSearchConfig"] = _BaseSearchConfig + return _globals["BaseSearchConfig"] + + # Lazy load BaseTextToSpeechConfig to avoid loading at module import time + if name == "BaseTextToSpeechConfig": + # Check if already cached + if "BaseTextToSpeechConfig" not in _globals: + from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig as _BaseTextToSpeechConfig, + ) + _globals["BaseTextToSpeechConfig"] = _BaseTextToSpeechConfig + return _globals["BaseTextToSpeechConfig"] + + # Lazy load BedrockModelInfo to avoid loading at module import time + if name == "BedrockModelInfo": + # Check if already cached + if "BedrockModelInfo" not in _globals: + from litellm.llms.bedrock.common_utils import ( + BedrockModelInfo as _BedrockModelInfo, + ) + _globals["BedrockModelInfo"] = _BedrockModelInfo + return _globals["BedrockModelInfo"] + + # Lazy load CohereModelInfo to avoid loading at module import time + if name == "CohereModelInfo": + # Check if already cached + if "CohereModelInfo" not in _globals: + from litellm.llms.cohere.common_utils import ( + CohereModelInfo as _CohereModelInfo, + ) + _globals["CohereModelInfo"] = _CohereModelInfo + return _globals["CohereModelInfo"] + + # Lazy load MistralOCRConfig to avoid loading at module import time + if name == "MistralOCRConfig": + # Check if already cached + if "MistralOCRConfig" not in _globals: + from litellm.llms.mistral.ocr.transformation import ( + MistralOCRConfig as _MistralOCRConfig, + ) + _globals["MistralOCRConfig"] = _MistralOCRConfig + return _globals["MistralOCRConfig"] + + # Lazy load Rules to avoid loading at module import time + if name == "Rules": + # Check if already cached + if "Rules" not in _globals: + from litellm.litellm_core_utils.rules import Rules as _Rules + _globals["Rules"] = _Rules + return _globals["Rules"] + + # Lazy load AsyncHTTPHandler and HTTPHandler to avoid loading at module import time + if name == "AsyncHTTPHandler": + # Check if already cached + if "AsyncHTTPHandler" not in _globals: + from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler as _AsyncHTTPHandler, + ) + _globals["AsyncHTTPHandler"] = _AsyncHTTPHandler + return _globals["AsyncHTTPHandler"] + + if name == "HTTPHandler": + # Check if already cached + if "HTTPHandler" not in _globals: + from litellm.llms.custom_httpx.http_handler import ( + HTTPHandler as _HTTPHandler, + ) + _globals["HTTPHandler"] = _HTTPHandler + return _globals["HTTPHandler"] + + # Lazy load get_num_retries_from_retry_policy and reset_retry_policy to avoid loading at module import time + if name == "get_num_retries_from_retry_policy": + # Check if already cached + if "get_num_retries_from_retry_policy" not in _globals: + from litellm.router_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, + ) + _globals["get_num_retries_from_retry_policy"] = _get_num_retries_from_retry_policy + return _globals["get_num_retries_from_retry_policy"] + + if name == "reset_retry_policy": + # Check if already cached + if "reset_retry_policy" not in _globals: + from litellm.router_utils.get_retry_from_policy import ( + reset_retry_policy as _reset_retry_policy, + ) + _globals["reset_retry_policy"] = _reset_retry_policy + return _globals["reset_retry_policy"] + + # Lazy load get_secret to avoid loading at module import time + if name == "get_secret": + # Check if already cached + if "get_secret" not in _globals: + from litellm.secret_managers.main import get_secret as _get_secret + _globals["get_secret"] = _get_secret + return _globals["get_secret"] + + # Lazy load cached_imports functions to avoid loading at module import time + if name == "get_coroutine_checker": + # Check if already cached + if "get_coroutine_checker" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_coroutine_checker as _get_coroutine_checker, + ) + _globals["get_coroutine_checker"] = _get_coroutine_checker + return _globals["get_coroutine_checker"] + + if name == "get_litellm_logging_class": + # Check if already cached + if "get_litellm_logging_class" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_litellm_logging_class as _get_litellm_logging_class, + ) + _globals["get_litellm_logging_class"] = _get_litellm_logging_class + return _globals["get_litellm_logging_class"] + + if name == "get_set_callbacks": + # Check if already cached + if "get_set_callbacks" not in _globals: + from litellm.litellm_core_utils.cached_imports import ( + get_set_callbacks as _get_set_callbacks, + ) + _globals["get_set_callbacks"] = _get_set_callbacks + return _globals["get_set_callbacks"] + + # Lazy load core_helpers functions to avoid loading at module import time + if name == "get_litellm_metadata_from_kwargs": + # Check if already cached + if "get_litellm_metadata_from_kwargs" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + get_litellm_metadata_from_kwargs as _get_litellm_metadata_from_kwargs, + ) + _globals["get_litellm_metadata_from_kwargs"] = _get_litellm_metadata_from_kwargs + return _globals["get_litellm_metadata_from_kwargs"] + + if name == "map_finish_reason": + # Check if already cached + if "map_finish_reason" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + map_finish_reason as _map_finish_reason, + ) + _globals["map_finish_reason"] = _map_finish_reason + return _globals["map_finish_reason"] + + if name == "process_response_headers": + # Check if already cached + if "process_response_headers" not in _globals: + from litellm.litellm_core_utils.core_helpers import ( + process_response_headers as _process_response_headers, + ) + _globals["process_response_headers"] = _process_response_headers + return _globals["process_response_headers"] + + # Lazy load dot_notation_indexing functions to avoid loading at module import time + if name == "delete_nested_value": + # Check if already cached + if "delete_nested_value" not in _globals: + from litellm.litellm_core_utils.dot_notation_indexing import ( + delete_nested_value as _delete_nested_value, + ) + _globals["delete_nested_value"] = _delete_nested_value + return _globals["delete_nested_value"] + + if name == "is_nested_path": + # Check if already cached + if "is_nested_path" not in _globals: + from litellm.litellm_core_utils.dot_notation_indexing import ( + is_nested_path as _is_nested_path, + ) + _globals["is_nested_path"] = _is_nested_path + return _globals["is_nested_path"] + + # Lazy load get_litellm_params functions to avoid loading at module import time + if name == "_get_base_model_from_litellm_call_metadata": + # Check if already cached + if "_get_base_model_from_litellm_call_metadata" not in _globals: + from litellm.litellm_core_utils.get_litellm_params import ( + _get_base_model_from_litellm_call_metadata as __get_base_model_from_litellm_call_metadata, + ) + _globals["_get_base_model_from_litellm_call_metadata"] = __get_base_model_from_litellm_call_metadata + return _globals["_get_base_model_from_litellm_call_metadata"] + + if name == "get_litellm_params": + # Check if already cached + if "get_litellm_params" not in _globals: + from litellm.litellm_core_utils.get_litellm_params import ( + get_litellm_params as _get_litellm_params, + ) + _globals["get_litellm_params"] = _get_litellm_params + return _globals["get_litellm_params"] + + # Lazy load _ensure_extra_body_is_safe to avoid loading at module import time + if name == "_ensure_extra_body_is_safe": + # Check if already cached + if "_ensure_extra_body_is_safe" not in _globals: + from litellm.litellm_core_utils.llm_request_utils import ( + _ensure_extra_body_is_safe as __ensure_extra_body_is_safe, + ) + _globals["_ensure_extra_body_is_safe"] = __ensure_extra_body_is_safe + return _globals["_ensure_extra_body_is_safe"] + + # Lazy load get_formatted_prompt to avoid loading at module import time + if name == "get_formatted_prompt": + # Check if already cached + if "get_formatted_prompt" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_formatted_prompt import ( + get_formatted_prompt as _get_formatted_prompt, + ) + _globals["get_formatted_prompt"] = _get_formatted_prompt + return _globals["get_formatted_prompt"] + + # Lazy load get_response_headers to avoid loading at module import time + if name == "get_response_headers": + # Check if already cached + if "get_response_headers" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.get_headers import ( + get_response_headers as _get_response_headers, + ) + _globals["get_response_headers"] = _get_response_headers + return _globals["get_response_headers"] + + # Lazy load update_response_metadata to avoid loading at module import time + if name == "update_response_metadata": + # Check if already cached + if "update_response_metadata" not in _globals: + from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( + update_response_metadata as _update_response_metadata, + ) + _globals["update_response_metadata"] = _update_response_metadata + return _globals["update_response_metadata"] + + # Lazy load executor to avoid loading at module import time + if name == "executor": + # Check if already cached + if "executor" not in _globals: + from litellm.litellm_core_utils.thread_pool_executor import ( + executor as _executor, + ) + _globals["executor"] = _executor + return _globals["executor"] + + # Lazy load BaseAnthropicMessagesConfig to avoid loading at module import time + if name == "BaseAnthropicMessagesConfig": + # Check if already cached + if "BaseAnthropicMessagesConfig" not in _globals: + from litellm.llms.base_llm.anthropic_messages.transformation import ( + BaseAnthropicMessagesConfig as _BaseAnthropicMessagesConfig, + ) + _globals["BaseAnthropicMessagesConfig"] = _BaseAnthropicMessagesConfig + return _globals["BaseAnthropicMessagesConfig"] + + # Lazy load BaseAudioTranscriptionConfig to avoid loading at module import time + if name == "BaseAudioTranscriptionConfig": + # Check if already cached + if "BaseAudioTranscriptionConfig" not in _globals: + from litellm.llms.base_llm.audio_transcription.transformation import ( + BaseAudioTranscriptionConfig as _BaseAudioTranscriptionConfig, + ) + _globals["BaseAudioTranscriptionConfig"] = _BaseAudioTranscriptionConfig + return _globals["BaseAudioTranscriptionConfig"] + + # Lazy load BaseBatchesConfig to avoid loading at module import time + if name == "BaseBatchesConfig": + # Check if already cached + if "BaseBatchesConfig" not in _globals: + from litellm.llms.base_llm.batches.transformation import ( + BaseBatchesConfig as _BaseBatchesConfig, + ) + _globals["BaseBatchesConfig"] = _BaseBatchesConfig + return _globals["BaseBatchesConfig"] + + # Lazy load BaseContainerConfig to avoid loading at module import time + if name == "BaseContainerConfig": + # Check if already cached + if "BaseContainerConfig" not in _globals: + from litellm.llms.base_llm.containers.transformation import ( + BaseContainerConfig as _BaseContainerConfig, + ) + _globals["BaseContainerConfig"] = _BaseContainerConfig + return _globals["BaseContainerConfig"] + + # Lazy load BaseEmbeddingConfig to avoid loading at module import time + if name == "BaseEmbeddingConfig": + # Check if already cached + if "BaseEmbeddingConfig" not in _globals: + from litellm.llms.base_llm.embedding.transformation import ( + BaseEmbeddingConfig as _BaseEmbeddingConfig, + ) + _globals["BaseEmbeddingConfig"] = _BaseEmbeddingConfig + return _globals["BaseEmbeddingConfig"] + + # Lazy load BaseImageEditConfig to avoid loading at module import time + if name == "BaseImageEditConfig": + # Check if already cached + if "BaseImageEditConfig" not in _globals: + from litellm.llms.base_llm.image_edit.transformation import ( + BaseImageEditConfig as _BaseImageEditConfig, + ) + _globals["BaseImageEditConfig"] = _BaseImageEditConfig + return _globals["BaseImageEditConfig"] + + # Lazy load BaseImageGenerationConfig to avoid loading at module import time + if name == "BaseImageGenerationConfig": + # Check if already cached + if "BaseImageGenerationConfig" not in _globals: + from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig as _BaseImageGenerationConfig, + ) + _globals["BaseImageGenerationConfig"] = _BaseImageGenerationConfig + return _globals["BaseImageGenerationConfig"] + + # Lazy load BaseImageVariationConfig to avoid loading at module import time + if name == "BaseImageVariationConfig": + # Check if already cached + if "BaseImageVariationConfig" not in _globals: + from litellm.llms.base_llm.image_variations.transformation import ( + BaseImageVariationConfig as _BaseImageVariationConfig, + ) + _globals["BaseImageVariationConfig"] = _BaseImageVariationConfig + return _globals["BaseImageVariationConfig"] + + # Lazy load BasePassthroughConfig to avoid loading at module import time + if name == "BasePassthroughConfig": + # Check if already cached + if "BasePassthroughConfig" not in _globals: + from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig as _BasePassthroughConfig, + ) + _globals["BasePassthroughConfig"] = _BasePassthroughConfig + return _globals["BasePassthroughConfig"] + + # Lazy load BaseRealtimeConfig to avoid loading at module import time + if name == "BaseRealtimeConfig": + # Check if already cached + if "BaseRealtimeConfig" not in _globals: + from litellm.llms.base_llm.realtime.transformation import ( + BaseRealtimeConfig as _BaseRealtimeConfig, + ) + _globals["BaseRealtimeConfig"] = _BaseRealtimeConfig + return _globals["BaseRealtimeConfig"] + + # Lazy load BaseRerankConfig to avoid loading at module import time + if name == "BaseRerankConfig": + # Check if already cached + if "BaseRerankConfig" not in _globals: + from litellm.llms.base_llm.rerank.transformation import ( + BaseRerankConfig as _BaseRerankConfig, + ) + _globals["BaseRerankConfig"] = _BaseRerankConfig + return _globals["BaseRerankConfig"] + + # Lazy load BaseVectorStoreConfig to avoid loading at module import time + if name == "BaseVectorStoreConfig": + # Check if already cached + if "BaseVectorStoreConfig" not in _globals: + from litellm.llms.base_llm.vector_store.transformation import ( + BaseVectorStoreConfig as _BaseVectorStoreConfig, + ) + _globals["BaseVectorStoreConfig"] = _BaseVectorStoreConfig + return _globals["BaseVectorStoreConfig"] + + # Lazy load BaseVectorStoreFilesConfig to avoid loading at module import time + if name == "BaseVectorStoreFilesConfig": + # Check if already cached + if "BaseVectorStoreFilesConfig" not in _globals: + from litellm.llms.base_llm.vector_store_files.transformation import ( + BaseVectorStoreFilesConfig as _BaseVectorStoreFilesConfig, + ) + _globals["BaseVectorStoreFilesConfig"] = _BaseVectorStoreFilesConfig + return _globals["BaseVectorStoreFilesConfig"] + + # Lazy load BaseVideoConfig to avoid loading at module import time + if name == "BaseVideoConfig": + # Check if already cached + if "BaseVideoConfig" not in _globals: + from litellm.llms.base_llm.videos.transformation import ( + BaseVideoConfig as _BaseVideoConfig, + ) + _globals["BaseVideoConfig"] = _BaseVideoConfig + return _globals["BaseVideoConfig"] + + # Lazy load ANTHROPIC_API_ONLY_HEADERS to avoid loading at module import time + if name == "ANTHROPIC_API_ONLY_HEADERS": + # Check if already cached + if "ANTHROPIC_API_ONLY_HEADERS" not in _globals: + from litellm.types.llms.anthropic import ( + ANTHROPIC_API_ONLY_HEADERS as _ANTHROPIC_API_ONLY_HEADERS, + ) + _globals["ANTHROPIC_API_ONLY_HEADERS"] = _ANTHROPIC_API_ONLY_HEADERS + return _globals["ANTHROPIC_API_ONLY_HEADERS"] + + # Lazy load AnthropicThinkingParam to avoid loading at module import time + if name == "AnthropicThinkingParam": + # Check if already cached + if "AnthropicThinkingParam" not in _globals: + from litellm.types.llms.anthropic import ( + AnthropicThinkingParam as _AnthropicThinkingParam, + ) + _globals["AnthropicThinkingParam"] = _AnthropicThinkingParam + return _globals["AnthropicThinkingParam"] + + # Lazy load RerankResponse to avoid loading at module import time + if name == "RerankResponse": + # Check if already cached + if "RerankResponse" not in _globals: + from litellm.types.rerank import RerankResponse as _RerankResponse + _globals["RerankResponse"] = _RerankResponse + return _globals["RerankResponse"] + + # Lazy load ChatCompletionDeltaToolCallChunk to avoid loading at module import time + if name == "ChatCompletionDeltaToolCallChunk": + # Check if already cached + if "ChatCompletionDeltaToolCallChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionDeltaToolCallChunk as _ChatCompletionDeltaToolCallChunk, + ) + _globals["ChatCompletionDeltaToolCallChunk"] = _ChatCompletionDeltaToolCallChunk + return _globals["ChatCompletionDeltaToolCallChunk"] + + # Lazy load ChatCompletionToolCallChunk to avoid loading at module import time + if name == "ChatCompletionToolCallChunk": + # Check if already cached + if "ChatCompletionToolCallChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionToolCallChunk as _ChatCompletionToolCallChunk, + ) + _globals["ChatCompletionToolCallChunk"] = _ChatCompletionToolCallChunk + return _globals["ChatCompletionToolCallChunk"] + + # Lazy load ChatCompletionToolCallFunctionChunk to avoid loading at module import time + if name == "ChatCompletionToolCallFunctionChunk": + # Check if already cached + if "ChatCompletionToolCallFunctionChunk" not in _globals: + from litellm.types.llms.openai import ( + ChatCompletionToolCallFunctionChunk as _ChatCompletionToolCallFunctionChunk, + ) + _globals["ChatCompletionToolCallFunctionChunk"] = _ChatCompletionToolCallFunctionChunk + return _globals["ChatCompletionToolCallFunctionChunk"] + + # Lazy load LiteLLM_Params to avoid loading at module import time + if name == "LiteLLM_Params": + # Check if already cached + if "LiteLLM_Params" not in _globals: + from litellm.types.router import LiteLLM_Params as _LiteLLM_Params + _globals["LiteLLM_Params"] = _LiteLLM_Params + return _globals["LiteLLM_Params"] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4651107c5b8..81b4469f24c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3663,12 +3663,16 @@ "supports_web_search": true }, "azure/gpt-image-1": { - "input_cost_per_pixel": 4.0054321e-08, + "cache_read_input_image_token_cost": 2.5e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 1e-05, + "input_cost_per_token": 5e-06, "litellm_provider": "azure", "mode": "image_generation", - "output_cost_per_pixel": 0.0, + "output_cost_per_image_token": 4e-05, "supported_endpoints": [ - "/v1/images/generations" + "/v1/images/generations", + "/v1/images/edits" ] }, "azure/hd/1024-x-1024/dall-e-3": { @@ -3771,12 +3775,16 @@ ] }, "azure/gpt-image-1-mini": { - "input_cost_per_pixel": 8.0566406e-09, + "cache_read_input_image_token_cost": 2.5e-07, + "cache_read_input_token_cost": 2e-07, + "input_cost_per_image_token": 2.5e-06, + "input_cost_per_token": 2e-06, "litellm_provider": "azure", "mode": "image_generation", - "output_cost_per_pixel": 0.0, + "output_cost_per_image_token": 8e-06, "supported_endpoints": [ - "/v1/images/generations" + "/v1/images/generations", + "/v1/images/edits" ] }, "azure/gpt-image-1.5": { @@ -10943,13 +10951,13 @@ "supports_tool_choice": true }, "fireworks_ai/accounts/fireworks/models/deepseek-v3p2": { - "input_cost_per_token": 1.2e-06, + "input_cost_per_token": 5.6e-07, "litellm_provider": "fireworks_ai", "max_input_tokens": 163840, "max_output_tokens": 163840, "max_tokens": 163840, "mode": "chat", - "output_cost_per_token": 1.2e-06, + "output_cost_per_token": 1.68e-06, "source": "https://fireworks.ai/models/fireworks/deepseek-v3p2", "supports_function_calling": true, "supports_reasoning": true, @@ -11632,6 +11640,7 @@ "supports_tool_choice": true }, "gemini-1.5-flash": { + "deprecation_date": "2025-09-29", "input_cost_per_audio_per_second": 2e-06, "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, "input_cost_per_character": 1.875e-08, @@ -11736,6 +11745,7 @@ "supports_vision": true }, "gemini-1.5-flash-exp-0827": { + "deprecation_date": "2025-09-29", "input_cost_per_audio_per_second": 2e-06, "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, "input_cost_per_character": 1.875e-08, @@ -11770,6 +11780,7 @@ "supports_vision": true }, "gemini-1.5-flash-preview-0514": { + "deprecation_date": "2025-09-29", "input_cost_per_audio_per_second": 2e-06, "input_cost_per_audio_per_second_above_128k_tokens": 4e-06, "input_cost_per_character": 1.875e-08, @@ -11803,6 +11814,7 @@ "supports_vision": true }, "gemini-1.5-pro": { + "deprecation_date": "2025-09-29", "input_cost_per_audio_per_second": 3.125e-05, "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, "input_cost_per_character": 3.125e-07, @@ -11890,6 +11902,7 @@ "supports_vision": true }, "gemini-1.5-pro-preview-0215": { + "deprecation_date": "2025-09-29", "input_cost_per_audio_per_second": 3.125e-05, "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, "input_cost_per_character": 3.125e-07, @@ -11917,6 +11930,7 @@ "supports_tool_choice": true }, "gemini-1.5-pro-preview-0409": { + "deprecation_date": "2025-09-29", "input_cost_per_audio_per_second": 3.125e-05, "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, "input_cost_per_character": 3.125e-07, @@ -11943,6 +11957,7 @@ "supports_tool_choice": true }, "gemini-1.5-pro-preview-0514": { + "deprecation_date": "2025-09-29", "input_cost_per_audio_per_second": 3.125e-05, "input_cost_per_audio_per_second_above_128k_tokens": 6.25e-05, "input_cost_per_character": 3.125e-07, @@ -12214,6 +12229,7 @@ "tpm": 250000 }, "gemini-2.0-flash-preview-image-generation": { + "deprecation_date": "2025-11-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, @@ -12252,6 +12268,7 @@ "supports_web_search": true }, "gemini-2.0-flash-thinking-exp": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 0.0, "input_cost_per_audio_per_second": 0, "input_cost_per_audio_per_second_above_128k_tokens": 0, @@ -12300,6 +12317,7 @@ "supports_web_search": true }, "gemini-2.0-flash-thinking-exp-01-21": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 0.0, "input_cost_per_audio_per_second": 0, "input_cost_per_audio_per_second_above_128k_tokens": 0, @@ -12486,6 +12504,7 @@ "tpm": 8000000 }, "gemini-2.5-flash-image-preview": { + "deprecation_date": "2026-01-15", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -12796,6 +12815,7 @@ "tpm": 8000000 }, "gemini-2.5-flash-lite-preview-06-17": { + "deprecation_date": "2025-11-18", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 1e-07, @@ -12885,6 +12905,7 @@ "supports_web_search": true }, "gemini-2.5-flash-preview-05-20": { + "deprecation_date": "2025-11-18", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -13156,6 +13177,7 @@ "supports_web_search": true }, "gemini-2.5-pro-preview-03-25": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 3.125e-07, "input_cost_per_audio_token": 1.25e-06, "input_cost_per_token": 1.25e-06, @@ -13201,6 +13223,7 @@ "supports_web_search": true }, "gemini-2.5-pro-preview-05-06": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 3.125e-07, "input_cost_per_audio_token": 1.25e-06, "input_cost_per_token": 1.25e-06, @@ -13416,6 +13439,7 @@ "tpm": 10000000 }, "gemini/gemini-1.5-flash": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 7.5e-08, "input_cost_per_token_above_128k_tokens": 1.5e-07, "litellm_provider": "gemini", @@ -13499,6 +13523,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-flash-8b": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 0, "input_cost_per_token_above_128k_tokens": 0, "litellm_provider": "gemini", @@ -13525,6 +13550,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-flash-8b-exp-0827": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 0, "input_cost_per_token_above_128k_tokens": 0, "litellm_provider": "gemini", @@ -13550,6 +13576,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-flash-8b-exp-0924": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 0, "input_cost_per_token_above_128k_tokens": 0, "litellm_provider": "gemini", @@ -13576,6 +13603,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-flash-exp-0827": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 0, "input_cost_per_token_above_128k_tokens": 0, "litellm_provider": "gemini", @@ -13601,6 +13629,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-flash-latest": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 7.5e-08, "input_cost_per_token_above_128k_tokens": 1.5e-07, "litellm_provider": "gemini", @@ -13627,6 +13656,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-pro": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 3.5e-06, "input_cost_per_token_above_128k_tokens": 7e-06, "litellm_provider": "gemini", @@ -13688,6 +13718,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-pro-exp-0801": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 3.5e-06, "input_cost_per_token_above_128k_tokens": 7e-06, "litellm_provider": "gemini", @@ -13707,6 +13738,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-pro-exp-0827": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 0, "input_cost_per_token_above_128k_tokens": 0, "litellm_provider": "gemini", @@ -13726,6 +13758,7 @@ "tpm": 4000000 }, "gemini/gemini-1.5-pro-latest": { + "deprecation_date": "2025-09-29", "input_cost_per_token": 3.5e-06, "input_cost_per_token_above_128k_tokens": 7e-06, "litellm_provider": "gemini", @@ -13908,6 +13941,7 @@ "tpm": 4000000 }, "gemini/gemini-2.0-flash-lite-preview-02-05": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 1.875e-08, "input_cost_per_audio_token": 7.5e-08, "input_cost_per_token": 7.5e-08, @@ -13945,6 +13979,7 @@ "tpm": 10000000 }, "gemini/gemini-2.0-flash-live-001": { + "deprecation_date": "2025-12-09", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 2.1e-06, "input_cost_per_image": 2.1e-06, @@ -13993,6 +14028,7 @@ "tpm": 250000 }, "gemini/gemini-2.0-flash-preview-image-generation": { + "deprecation_date": "2025-11-14", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1e-07, @@ -14032,6 +14068,7 @@ "tpm": 10000000 }, "gemini/gemini-2.0-flash-thinking-exp": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 0.0, "input_cost_per_audio_per_second": 0, "input_cost_per_audio_per_second_above_128k_tokens": 0, @@ -14081,6 +14118,7 @@ "tpm": 4000000 }, "gemini/gemini-2.0-flash-thinking-exp-01-21": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 0.0, "input_cost_per_audio_per_second": 0, "input_cost_per_audio_per_second_above_128k_tokens": 0, @@ -14269,6 +14307,7 @@ "tpm": 8000000 }, "gemini/gemini-2.5-flash-image-preview": { + "deprecation_date": "2026-01-15", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -14589,6 +14628,7 @@ "tpm": 250000 }, "gemini/gemini-2.5-flash-lite-preview-06-17": { + "deprecation_date": "2025-11-18", "cache_read_input_token_cost": 2.5e-08, "input_cost_per_audio_token": 5e-07, "input_cost_per_token": 1e-07, @@ -14680,6 +14720,7 @@ "tpm": 250000 }, "gemini/gemini-2.5-flash-preview-05-20": { + "deprecation_date": "2025-11-18", "cache_read_input_token_cost": 7.5e-08, "input_cost_per_audio_token": 1e-06, "input_cost_per_token": 3e-07, @@ -15026,6 +15067,7 @@ "tpm": 250000 }, "gemini/gemini-2.5-pro-preview-03-25": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 3.125e-07, "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1.25e-06, @@ -15066,6 +15108,7 @@ "tpm": 10000000 }, "gemini/gemini-2.5-pro-preview-05-06": { + "deprecation_date": "2025-12-02", "cache_read_input_token_cost": 3.125e-07, "input_cost_per_audio_token": 7e-07, "input_cost_per_token": 1.25e-06, @@ -15341,6 +15384,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "gemini/imagen-3.0-generate-002": { + "deprecation_date": "2025-11-10", "litellm_provider": "gemini", "mode": "image_generation", "output_cost_per_image": 0.04, @@ -15407,6 +15451,7 @@ ] }, "gemini/veo-3.0-fast-generate-preview": { + "deprecation_date": "2025-11-12", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -15421,6 +15466,7 @@ ] }, "gemini/veo-3.0-generate-preview": { + "deprecation_date": "2025-11-12", "litellm_provider": "gemini", "max_input_tokens": 1024, "max_tokens": 1024, @@ -16980,6 +17026,336 @@ "supports_vision": true, "supports_pdf_input": true }, + "low/1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.034, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.05, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.05, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.133, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.20, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.20, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1024-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1024-x-1536/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1536-x-1024/gpt-image-1.5": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "low/1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.034, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.05, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "medium/1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.05, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.133, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.20, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "high/1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.20, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "standard/1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1024-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.009, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1024-x-1536/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, + "1536-x-1024/gpt-image-1.5-2025-12-16": { + "input_cost_per_image": 0.013, + "litellm_provider": "openai", + "mode": "image_generation", + "supported_endpoints": [ + "/v1/images/generations", + "/v1/images/edits" + ], + "supports_vision": true, + "supports_pdf_input": true + }, "gpt-5": { "cache_read_input_token_cost": 1.25e-07, "cache_read_input_token_cost_flex": 6.25e-08, @@ -17741,16 +18117,16 @@ "supports_vision": true }, "gpt-image-1": { - "input_cost_per_image": 0.042, - "input_cost_per_pixel": 4.0054321e-08, - "input_cost_per_token": 0.000005, - "input_cost_per_image_token": 0.00001, + "cache_read_input_image_token_cost": 2.5e-06, + "cache_read_input_token_cost": 1.25e-06, + "input_cost_per_image_token": 1e-05, + "input_cost_per_token": 5e-06, "litellm_provider": "openai", "mode": "image_generation", - "output_cost_per_pixel": 0.0, - "output_cost_per_token": 0.00004, + "output_cost_per_image_token": 4e-05, "supported_endpoints": [ - "/v1/images/generations" + "/v1/images/generations", + "/v1/images/edits" ] }, "gpt-image-1-mini": { @@ -18175,6 +18551,18 @@ "supports_response_schema": false, "supports_tool_choice": true }, + "groq/gemma-7b-it": { + "input_cost_per_token": 5e-08, + "litellm_provider": "groq", + "max_input_tokens": 8192, + "max_output_tokens": 8192, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 8e-08, + "supports_function_calling": true, + "supports_response_schema": false, + "supports_tool_choice": true + }, "groq/meta-llama/llama-guard-4-12b": { "input_cost_per_token": 2e-07, "litellm_provider": "groq", @@ -24776,6 +25164,7 @@ "source": "https://docs.mistral.ai/capabilities/code_generation/" }, "text-embedding-004": { + "deprecation_date": "2026-01-14", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -27546,6 +27935,7 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" }, "vertex_ai/imagen-3.0-generate-002": { + "deprecation_date": "2025-11-10", "litellm_provider": "vertex_ai-image-models", "mode": "image_generation", "output_cost_per_image": 0.04, @@ -28056,6 +28446,7 @@ ] }, "vertex_ai/veo-3.0-fast-generate-preview": { + "deprecation_date": "2025-11-12", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -28070,6 +28461,7 @@ ] }, "vertex_ai/veo-3.0-generate-preview": { + "deprecation_date": "2025-11-12", "litellm_provider": "vertex_ai-video-models", "max_input_tokens": 1024, "max_tokens": 1024, @@ -29299,6 +29691,20 @@ "supports_vision": true, "supports_web_search": true }, + "zai/glm-4.7": { + "cache_creation_input_token_cost": 0, + "cache_read_input_token_cost": 1.1e-07, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.2e-06, + "litellm_provider": "zai", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "mode": "chat", + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true, + "source": "https://docs.z.ai/guides/overview/pricing" + }, "zai/glm-4.6": { "input_cost_per_token": 6e-07, "output_cost_per_token": 2.2e-06, diff --git a/poetry.lock b/poetry.lock index ee97c00594c..a0a0f8540e5 100644 --- a/poetry.lock +++ b/poetry.lock @@ -2273,75 +2273,6 @@ googleapis-common-protos = {version = ">=1.56.0,<2.0.0", extras = ["grpc"]} grpcio = ">=1.44.0,<2.0.0" protobuf = ">=3.20.2,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<7.0.0" -[[package]] -name = "grpcio" -version = "1.67.1" -description = "HTTP/2-based RPC framework" -optional = false -python-versions = ">=3.8" -groups = ["main", "dev", "proxy-dev"] -markers = "python_version < \"3.14\"" -files = [ - {file = "grpcio-1.67.1-cp310-cp310-linux_armv7l.whl", hash = "sha256:8b0341d66a57f8a3119b77ab32207072be60c9bf79760fa609c5609f2deb1f3f"}, - {file = "grpcio-1.67.1-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:f5a27dddefe0e2357d3e617b9079b4bfdc91341a91565111a21ed6ebbc51b22d"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:43112046864317498a33bdc4797ae6a268c36345a910de9b9c17159d8346602f"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c9b929f13677b10f63124c1a410994a401cdd85214ad83ab67cc077fc7e480f0"}, - {file = "grpcio-1.67.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e7d1797a8a3845437d327145959a2c0c47c05947c9eef5ff1a4c80e499dcc6fa"}, - {file = "grpcio-1.67.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:0489063974d1452436139501bf6b180f63d4977223ee87488fe36858c5725292"}, - {file = "grpcio-1.67.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:9fd042de4a82e3e7aca44008ee2fb5da01b3e5adb316348c21980f7f58adc311"}, - {file = "grpcio-1.67.1-cp310-cp310-win32.whl", hash = "sha256:638354e698fd0c6c76b04540a850bf1db27b4d2515a19fcd5cf645c48d3eb1ed"}, - {file = "grpcio-1.67.1-cp310-cp310-win_amd64.whl", hash = "sha256:608d87d1bdabf9e2868b12338cd38a79969eaf920c89d698ead08f48de9c0f9e"}, - {file = "grpcio-1.67.1-cp311-cp311-linux_armv7l.whl", hash = "sha256:7818c0454027ae3384235a65210bbf5464bd715450e30a3d40385453a85a70cb"}, - {file = "grpcio-1.67.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:ea33986b70f83844cd00814cee4451055cd8cab36f00ac64a31f5bb09b31919e"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:c7a01337407dd89005527623a4a72c5c8e2894d22bead0895306b23c6695698f"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:80b866f73224b0634f4312a4674c1be21b2b4afa73cb20953cbbb73a6b36c3cc"}, - {file = "grpcio-1.67.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f9fff78ba10d4250bfc07a01bd6254a6d87dc67f9627adece85c0b2ed754fa96"}, - {file = "grpcio-1.67.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:8a23cbcc5bb11ea7dc6163078be36c065db68d915c24f5faa4f872c573bb400f"}, - {file = "grpcio-1.67.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:1a65b503d008f066e994f34f456e0647e5ceb34cfcec5ad180b1b44020ad4970"}, - {file = "grpcio-1.67.1-cp311-cp311-win32.whl", hash = "sha256:e29ca27bec8e163dca0c98084040edec3bc49afd10f18b412f483cc68c712744"}, - {file = "grpcio-1.67.1-cp311-cp311-win_amd64.whl", hash = "sha256:786a5b18544622bfb1e25cc08402bd44ea83edfb04b93798d85dca4d1a0b5be5"}, - {file = "grpcio-1.67.1-cp312-cp312-linux_armv7l.whl", hash = "sha256:267d1745894200e4c604958da5f856da6293f063327cb049a51fe67348e4f953"}, - {file = "grpcio-1.67.1-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:85f69fdc1d28ce7cff8de3f9c67db2b0ca9ba4449644488c1e0303c146135ddb"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:f26b0b547eb8d00e195274cdfc63ce64c8fc2d3e2d00b12bf468ece41a0423a0"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4422581cdc628f77302270ff839a44f4c24fdc57887dc2a45b7e53d8fc2376af"}, - {file = "grpcio-1.67.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1d7616d2ded471231c701489190379e0c311ee0a6c756f3c03e6a62b95a7146e"}, - {file = "grpcio-1.67.1-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:8a00efecde9d6fcc3ab00c13f816313c040a28450e5e25739c24f432fc6d3c75"}, - {file = "grpcio-1.67.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:699e964923b70f3101393710793289e42845791ea07565654ada0969522d0a38"}, - {file = "grpcio-1.67.1-cp312-cp312-win32.whl", hash = "sha256:4e7b904484a634a0fff132958dabdb10d63e0927398273917da3ee103e8d1f78"}, - {file = "grpcio-1.67.1-cp312-cp312-win_amd64.whl", hash = "sha256:5721e66a594a6c4204458004852719b38f3d5522082be9061d6510b455c90afc"}, - {file = "grpcio-1.67.1-cp313-cp313-linux_armv7l.whl", hash = "sha256:aa0162e56fd10a5547fac8774c4899fc3e18c1aa4a4759d0ce2cd00d3696ea6b"}, - {file = "grpcio-1.67.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:beee96c8c0b1a75d556fe57b92b58b4347c77a65781ee2ac749d550f2a365dc1"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_aarch64.whl", hash = "sha256:a93deda571a1bf94ec1f6fcda2872dad3ae538700d94dc283c672a3b508ba3af"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:0e6f255980afef598a9e64a24efce87b625e3e3c80a45162d111a461a9f92955"}, - {file = "grpcio-1.67.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e838cad2176ebd5d4a8bb03955138d6589ce9e2ce5d51c3ada34396dbd2dba8"}, - {file = "grpcio-1.67.1-cp313-cp313-musllinux_1_1_i686.whl", hash = "sha256:a6703916c43b1d468d0756c8077b12017a9fcb6a1ef13faf49e67d20d7ebda62"}, - {file = "grpcio-1.67.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:917e8d8994eed1d86b907ba2a61b9f0aef27a2155bca6cbb322430fc7135b7bb"}, - {file = "grpcio-1.67.1-cp313-cp313-win32.whl", hash = "sha256:e279330bef1744040db8fc432becc8a727b84f456ab62b744d3fdb83f327e121"}, - {file = "grpcio-1.67.1-cp313-cp313-win_amd64.whl", hash = "sha256:fa0c739ad8b1996bd24823950e3cb5152ae91fca1c09cc791190bf1627ffefba"}, - {file = "grpcio-1.67.1-cp38-cp38-linux_armv7l.whl", hash = "sha256:178f5db771c4f9a9facb2ab37a434c46cb9be1a75e820f187ee3d1e7805c4f65"}, - {file = "grpcio-1.67.1-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0f3e49c738396e93b7ba9016e153eb09e0778e776df6090c1b8c91877cc1c426"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:24e8a26dbfc5274d7474c27759b54486b8de23c709d76695237515bc8b5baeab"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3b6c16489326d79ead41689c4b84bc40d522c9a7617219f4ad94bc7f448c5085"}, - {file = "grpcio-1.67.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:60e6a4dcf5af7bbc36fd9f81c9f372e8ae580870a9e4b6eafe948cd334b81cf3"}, - {file = "grpcio-1.67.1-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:95b5f2b857856ed78d72da93cd7d09b6db8ef30102e5e7fe0961fe4d9f7d48e8"}, - {file = "grpcio-1.67.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:b49359977c6ec9f5d0573ea4e0071ad278ef905aa74e420acc73fd28ce39e9ce"}, - {file = "grpcio-1.67.1-cp38-cp38-win32.whl", hash = "sha256:f5b76ff64aaac53fede0cc93abf57894ab2a7362986ba22243d06218b93efe46"}, - {file = "grpcio-1.67.1-cp38-cp38-win_amd64.whl", hash = "sha256:804c6457c3cd3ec04fe6006c739579b8d35c86ae3298ffca8de57b493524b771"}, - {file = "grpcio-1.67.1-cp39-cp39-linux_armv7l.whl", hash = "sha256:a25bdea92b13ff4d7790962190bf6bf5c4639876e01c0f3dda70fc2769616335"}, - {file = "grpcio-1.67.1-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:cdc491ae35a13535fd9196acb5afe1af37c8237df2e54427be3eecda3653127e"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:85f862069b86a305497e74d0dc43c02de3d1d184fc2c180993aa8aa86fbd19b8"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ec74ef02010186185de82cc594058a3ccd8d86821842bbac9873fd4a2cf8be8d"}, - {file = "grpcio-1.67.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:01f616a964e540638af5130469451cf580ba8c7329f45ca998ab66e0c7dcdb04"}, - {file = "grpcio-1.67.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:299b3d8c4f790c6bcca485f9963b4846dd92cf6f1b65d3697145d005c80f9fe8"}, - {file = "grpcio-1.67.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:60336bff760fbb47d7e86165408126f1dded184448e9a4c892189eb7c9d3f90f"}, - {file = "grpcio-1.67.1-cp39-cp39-win32.whl", hash = "sha256:5ed601c4c6008429e3d247ddb367fe8c7259c355757448d7c1ef7bd4a6739e8e"}, - {file = "grpcio-1.67.1-cp39-cp39-win_amd64.whl", hash = "sha256:5db70d32d6703b89912af16d6d45d78406374a8b8ef0d28140351dd0ec610e98"}, - {file = "grpcio-1.67.1.tar.gz", hash = "sha256:3dc2ed4cabea4dc14d5e708c2b426205956077cc5de419b4d4079315017e9732"}, -] - -[package.extras] -protobuf = ["grpcio-tools (>=1.67.1)"] - [[package]] name = "grpcio" version = "1.76.0" @@ -2349,7 +2280,6 @@ description = "HTTP/2-based RPC framework" optional = false python-versions = ">=3.9" groups = ["main", "dev", "proxy-dev"] -markers = "python_version >= \"3.14\"" files = [ {file = "grpcio-1.76.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:65a20de41e85648e00305c1bb09a3598f840422e522277641145a32d42dcefcc"}, {file = "grpcio-1.76.0-cp310-cp310-macosx_11_0_universal2.whl", hash = "sha256:40ad3afe81676fd9ec6d9d406eda00933f218038433980aa19d401490e46ecde"}, @@ -8051,4 +7981,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "b010d9da7f5a765670932b78d720aae4fcb819daba050683ee125b4367972419" +content-hash = "7eed2b2c25173a275ac83c55fd901b9b84663b1d7daa54f0e78b30bf1c8f0e3e" diff --git a/pyproject.toml b/pyproject.toml index f929fb94cb0..3b09119a748 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -72,7 +72,7 @@ soundfile = {version = "^0.12.1", optional = true} # - 1.68.0-1.68.1 has reconnect bug (https://github.com/grpc/grpc/issues/38290) # - 1.75.0+ has Python 3.14 wheels and bug fix grpcio = [ - {version = ">=1.62.3,<1.68.0", python = "<3.14"}, + {version = ">=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0", python = "<3.14"}, {version = ">=1.75.0", python = ">=3.14"}, ] diff --git a/requirements.txt b/requirements.txt index 3bc968c8cb8..06a7c17336c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -41,7 +41,7 @@ opentelemetry-api==1.25.0 opentelemetry-sdk==1.25.0 opentelemetry-exporter-otlp==1.25.0 # grpcio: 1.68.0-1.68.1 has reconnect bug (#38290), 1.75+ has Python 3.14 wheels + fix -grpcio>=1.62.3,<1.68.0; python_version < "3.14" +grpcio>=1.62.3,!=1.68.*,!=1.69.*,!=1.70.*,!=1.71.0,!=1.71.1,!=1.72.0,!=1.72.1,!=1.73.0; python_version < "3.14" grpcio>=1.75.0; python_version >= "3.14" sentry_sdk==2.21.0 # for sentry error handling detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests diff --git a/schema.prisma b/schema.prisma index aac0b5b35de..ea47b6ed03b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -208,6 +208,9 @@ model LiteLLM_MCPServerTable { command String? args String[] @default([]) env Json? @default("{}") + authorization_url String? + token_url String? + registration_url String? } // Generate Tokens for Proxy diff --git a/tests/code_coverage_tests/memory_test.py b/tests/code_coverage_tests/memory_test.py new file mode 100644 index 00000000000..1ce93191992 --- /dev/null +++ b/tests/code_coverage_tests/memory_test.py @@ -0,0 +1,616 @@ +""" +Memory Violation Detection Test + +Detects bad memory patterns in the LiteLLM codebase that can lead to memory leaks or OOMs. + +The detector uses a modular pattern-based system. To add detection for new memory patterns: + +1. Create a Pattern subclass implementing get_pattern_name(), visit_assign(), and check_cleanup() + - You can extend the Pattern class with additional methods as needed for your detection logic +2. Add the pattern to MemoryViolationDetector.DEFAULT_PATTERNS + +Currently detects: +- queue.get() / queue.get_nowait() operations where variables aren't set to None +- Class-level data structures that have add operations during runtime without size limits: + * Built-in: list, dict, set + * Collections: deque, defaultdict, Counter, OrderedDict, ChainMap + * Queues: queue.Queue, asyncio.Queue (if unbounded, i.e., no maxsize parameter) + * Heap operations: heapq.heappush(), heapq.heapreplace(), heapq.heappushpop() on class-level lists +""" + +import ast +import os +from abc import ABC, abstractmethod +from typing import List, Dict, Any, Optional, Sequence + + +class Pattern(ABC): + """Base class for memory violation detection patterns""" + + @abstractmethod + def get_pattern_name(self) -> str: + """Return unique identifier for this violation type""" + pass + + @abstractmethod + def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]: + """Detect memory-sensitive operations in assignment. Returns list of {line, var_name, call} dicts.""" + pass + + @abstractmethod + def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt], + context: Dict[str, Any]) -> List[Dict[str, Any]]: + """Verify variables are set to None. Returns list of violation dicts.""" + pass + + +class QueueGetPattern(Pattern): + """Detects queue.get()/get_nowait() operations that aren't cleared""" + + def get_pattern_name(self) -> str: + return "queue_reference_not_cleared" + + def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]: + """Detect queue.get() or queue.get_nowait() calls where object name contains 'queue'""" + operations = [] + + if isinstance(node.value, ast.Call): + func = node.value.func + if isinstance(func, ast.Attribute) and func.attr in ("get", "get_nowait"): + obj_name = context["get_attr_string"](func.value) + if "queue" in obj_name.lower() and node.targets and isinstance(node.targets[0], ast.Name): + operations.append({ + "line": node.lineno, + "var_name": node.targets[0].id, + "call": context["get_call_string"](node.value), + }) + + return operations + + def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt], + context: Dict[str, Any]) -> List[Dict[str, Any]]: + """Flag queue variables that aren't set to None""" + violations = [] + is_var_set_to_none = context["is_var_set_to_none"] + current_function = context["current_function"] + file_path = context["file_path"] + + queue_vars = {op["var_name"]: op["line"] for op in operations} + + for var_name, line_num in queue_vars.items(): + if not is_var_set_to_none(var_name, function_body): + violations.append({ + "line": line_num, + "type": self.get_pattern_name(), + "var_name": var_name, + "function": current_function, + "file_path": file_path, + "message": ( + f"Queue variable '{var_name}' in function " + f"'{current_function}' is not set to None after use. " + f"If the runtime is overwhelmed, this can cause OOM (Out of Memory) errors." + ), + }) + + return violations + + +class UnboundedDataStructurePattern(Pattern): + """Detects class-level data structures (lists, dicts, sets) that can grow unbounded""" + + def get_pattern_name(self) -> str: + return "unbounded_data_structure" + + def visit_assign(self, node: ast.Assign, context: Dict[str, Any]) -> List[Dict[str, Any]]: + """Detect list/dict/set creations that are at class level""" + operations = [] + + # Check if this is a data structure creation + is_data_structure = False + structure_type = None + + if isinstance(node.value, (ast.List, ast.Dict, ast.Set)): + is_data_structure = True + if isinstance(node.value, ast.List): + structure_type = "list" + elif isinstance(node.value, ast.Dict): + structure_type = "dict" + elif isinstance(node.value, ast.Set): + structure_type = "set" + elif isinstance(node.value, ast.Call): + # Check for list(), dict(), set() calls + func = node.value.func + if isinstance(func, ast.Name): + if func.id in ("list", "dict", "set"): + is_data_structure = True + structure_type = func.id + elif isinstance(func, ast.Attribute): + # Handle cases like collections.defaultdict(list), collections.deque(), etc. + obj_name = context["get_attr_string"](func.value) + attr_name = func.attr + + # Check for collections module data structures + if "collections" in obj_name.lower() or "collections" in str(func.value): + if attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"): + # For deque, we track it and let size checks determine if it's bounded + # (deque with maxlen parameter is bounded, but we detect that via size checks) + is_data_structure = True + structure_type = attr_name + elif attr_name in ("list", "dict", "set"): + # collections.defaultdict(list) pattern + is_data_structure = True + structure_type = "defaultdict" if "defaultdict" in obj_name.lower() else attr_name + # Check for queue.Queue, asyncio.Queue (if unbounded) + elif "queue" in obj_name.lower() or "asyncio" in obj_name.lower(): + if attr_name == "Queue": + # Check if maxsize is set (bounded queue) + has_maxsize = False + for keyword in node.value.keywords: + if keyword.arg == "maxsize": + has_maxsize = True + break + if not has_maxsize: + is_data_structure = True + structure_type = "queue" + # Direct attribute access like deque(), Counter(), etc. + elif attr_name in ("deque", "defaultdict", "Counter", "OrderedDict", "ChainMap"): + is_data_structure = True + structure_type = attr_name + + if is_data_structure and node.targets and isinstance(node.targets[0], ast.Name): + scope = context.get("current_scope", "function") + # Only track if it's at class level (not module level) + if scope == "class": + operations.append({ + "line": node.lineno, + "var_name": node.targets[0].id, + "structure_type": structure_type, + "scope": scope, + "call": context["get_call_string"](node.value) if isinstance(node.value, ast.Call) else f"{structure_type}()", + }) + + return operations + + def check_cleanup(self, operations: List[Dict[str, Any]], function_body: List[ast.stmt], + context: Dict[str, Any]) -> List[Dict[str, Any]]: + """Flag persistent data structures that have add operations without size limits""" + violations = [] + current_function = context["current_function"] + current_scope = context.get("current_scope", "function") + file_path = context["file_path"] + get_attr_string = context["get_attr_string"] + + # Skip if this is initialization code (module-level, class-level, or __init__ methods) + # Only flag operations in regular methods/functions that can be called during runtime + is_initialization = ( + current_scope in ("module", "class") or + current_function in ("__init__", "__new__", "__class_init__") or + current_function is None # Module-level code + ) + + if is_initialization: + return violations # Don't flag initialization code + + # Track which variables have add operations and size checks + var_add_operations = {} # var_name -> list of lines with add operations + var_size_checks = {} # var_name -> has size limit check + + # Build a set of variable names to check + tracked_vars = {op["var_name"]: op for op in operations} + + # Scan body for operations on these variables + for stmt in function_body: + for node in ast.walk(stmt): + # Check for method calls that add items + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + attr_name = node.func.attr + obj_name = get_attr_string(node.func.value) + + # Check if this is an add operation on one of our tracked variables + for var_name, op in tracked_vars.items(): + structure_type = op["structure_type"] + + # Match variable name (exact or as attribute) + if obj_name == var_name or obj_name.endswith(f".{var_name}") or obj_name.endswith(f"['{var_name}']"): + # Check for add operations + add_ops = { + "list": ["append", "extend", "insert"], + "dict": ["update", "setdefault"], + "set": ["add", "update"], + "deque": ["append", "appendleft", "extend", "extendleft", "insert"], + "defaultdict": ["update", "setdefault"], + "Counter": ["update"], + "OrderedDict": ["update", "setdefault"], + "ChainMap": ["new_child"], + "queue": ["put", "put_nowait"], + } + + if attr_name in add_ops.get(structure_type, []): + if var_name not in var_add_operations: + var_add_operations[var_name] = [] + var_add_operations[var_name].append(node.lineno) + + # Check for size limit checks (len() calls, maxsize/maxlen attributes) + if (attr_name in ("__len__",) or + "maxsize" in attr_name.lower() or + "max_size" in attr_name.lower() or + attr_name == "maxlen"): # For deque + var_size_checks[var_name] = True + + # Check for heapq operations on tracked lists (heapq.heappush, heapq.heappop) + if isinstance(node, ast.Call): + func = node.func + # Check for heapq.heappush(list_var, item) or heapq.heappop(list_var) + if isinstance(func, ast.Attribute): + func_obj = get_attr_string(func.value) + func_name = func.attr + # Check if it's a heapq operation + if func_obj == "heapq" and func_name in ("heappush", "heapreplace", "heappushpop"): + # First argument should be our tracked variable + if len(node.args) > 0: + arg_name = get_attr_string(node.args[0]) + for var_name, op in tracked_vars.items(): + if op["structure_type"] == "list" and ( + arg_name == var_name or arg_name.endswith(f".{var_name}") + ): + if var_name not in var_add_operations: + var_add_operations[var_name] = [] + var_add_operations[var_name].append(node.lineno) + + # Check for dict item assignment: dict[key] = value + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Subscript): + target_name = get_attr_string(target.value) + for var_name in tracked_vars: + if target_name == var_name or target_name.endswith(f".{var_name}"): + if var_name not in var_add_operations: + var_add_operations[var_name] = [] + var_add_operations[var_name].append(node.lineno) + + # Check for augmented assignment: list += [...] + if isinstance(node, ast.AugAssign): + target_name = get_attr_string(node.target) + for var_name in tracked_vars: + if target_name == var_name or target_name.endswith(f".{var_name}"): + if var_name not in var_add_operations: + var_add_operations[var_name] = [] + var_add_operations[var_name].append(node.lineno) + + # Check for size comparisons in conditionals + if isinstance(node, (ast.If, ast.While, ast.Assert)): + test = getattr(node, "test", None) + if test: + for comp_node in ast.walk(test): + if isinstance(comp_node, ast.Compare): + left_str = get_attr_string(comp_node.left) if hasattr(comp_node, "left") else "" + # Check for len() calls + if isinstance(comp_node.left, ast.Call): + call_func = comp_node.left.func + if isinstance(call_func, ast.Name) and call_func.id == "len": + if len(comp_node.left.args) > 0: + arg_name = get_attr_string(comp_node.left.args[0]) + for var_name in tracked_vars: + if arg_name == var_name or arg_name.endswith(f".{var_name}"): + # Check if comparing to a limit + for comparator in comp_node.comparators: + if isinstance(comparator, ast.Constant): + var_size_checks[var_name] = True + elif isinstance(comparator, ast.Name): + # Could be a constant like MAX_SIZE + if "max" in comparator.id.lower() or "limit" in comparator.id.lower(): + var_size_checks[var_name] = True + # Handle deprecated ast.Num for Python < 3.8 + try: + Num = getattr(ast, "Num", None) + if Num and isinstance(comparator, Num): + var_size_checks[var_name] = True + except (AttributeError, TypeError): + pass + # Check for direct variable comparisons + for var_name in tracked_vars: + if var_name in left_str: + for comparator in comp_node.comparators: + if isinstance(comparator, ast.Constant): + var_size_checks[var_name] = True + # Handle deprecated ast.Num for Python < 3.8 + try: + Num = getattr(ast, "Num", None) + if Num and isinstance(comparator, Num): + var_size_checks[var_name] = True + except (AttributeError, TypeError): + pass + + # Flag violations: persistent structures with add operations but no size checks + for op in operations: + var_name = op["var_name"] + structure_type = op["structure_type"] + + if var_name in var_add_operations and var_name not in var_size_checks: + violations.append({ + "line": op["line"], + "type": self.get_pattern_name(), + "var_name": var_name, + "function": current_function or "class-level", + "file_path": file_path, + "message": ( + f"Class-level {structure_type} '{var_name}' " + f"has add operations (lines {var_add_operations[var_name]}) but no size limit checks. " + f"This can lead to unbounded memory growth and OOM errors during runtime." + ), + }) + + return violations + + +class MemoryViolationDetector(ast.NodeVisitor): + """AST visitor that detects memory violations using registered patterns""" + + DEFAULT_PATTERNS: List[Pattern] = [QueueGetPattern(), UnboundedDataStructurePattern()] + + def __init__(self, file_path: str, patterns: Optional[Sequence[Pattern]] = None): + self.file_path = file_path + self.violations: List[Dict[str, Any]] = [] + self.current_function: Optional[str] = None + self.current_scope: str = "module" # Track current scope: module, class, function + self.patterns = self.DEFAULT_PATTERNS if patterns is None else patterns + self.ast_tree: Optional[ast.Module] = None # Store full AST for module-level checks + + self.pattern_operations: Dict[str, List[Dict[str, Any]]] = { + pattern.get_pattern_name(): [] for pattern in self.patterns + } + + # Track class-level operations separately (for checking in functions) + self.class_level_operations: Dict[str, List[Dict[str, Any]]] = { + pattern.get_pattern_name(): [] for pattern in self.patterns + } + + self._context = { + "get_call_string": self._get_call_string, + "get_attr_string": self._get_attr_string, + "is_var_set_to_none": self._is_var_set_to_none, + "current_function": None, + "current_scope": "module", + "file_path": file_path, + } + + def visit_ClassDef(self, node): + """Track class scope""" + old_scope = self.current_scope + self.current_scope = "class" + self._context["current_scope"] = "class" + + self.generic_visit(node) + + self.current_scope = old_scope + self._context["current_scope"] = old_scope + + def visit_FunctionDef(self, node): + """Track function scope and check cleanup after visiting""" + old_function = self.current_function + old_scope = self.current_scope + self.current_function = node.name + self.current_scope = "function" + self._context["current_function"] = node.name + self._context["current_scope"] = "function" + + for pattern_name in self.pattern_operations: + self.pattern_operations[pattern_name] = [] + + self.generic_visit(node) + self._check_function_cleanup(node) + + self.current_function = old_function + self.current_scope = old_scope + self._context["current_function"] = old_function + self._context["current_scope"] = old_scope + + def visit_AsyncFunctionDef(self, node): + """Track async function scope and check cleanup after visiting""" + old_function = self.current_function + old_scope = self.current_scope + self.current_function = node.name + self.current_scope = "function" + self._context["current_function"] = node.name + self._context["current_scope"] = "function" + + for pattern_name in self.pattern_operations: + self.pattern_operations[pattern_name] = [] + + self.generic_visit(node) + self._check_function_cleanup(node) + + self.current_function = old_function + self.current_scope = old_scope + self._context["current_function"] = old_function + self._context["current_scope"] = old_scope + + def visit_Assign(self, node): + """Detect memory-sensitive operations in assignments""" + for pattern in self.patterns: + operations = pattern.visit_assign(node, self._context) + # Track function-level operations + self.pattern_operations[pattern.get_pattern_name()].extend(operations) + # Track class-level operations separately (for checking in functions) + for op in operations: + if op.get("scope") == "class": + self.class_level_operations[pattern.get_pattern_name()].append(op) + + self.generic_visit(node) + + def _check_function_cleanup(self, node): + """Check cleanup for all detected operations""" + for pattern in self.patterns: + operations = self.pattern_operations[pattern.get_pattern_name()] + if operations: + violations = pattern.check_cleanup(operations, node.body, self._context) + self.violations.extend(violations) + + # For UnboundedDataStructurePattern, also check if this function modifies class-level structures + if isinstance(pattern, UnboundedDataStructurePattern): + class_ops = self.class_level_operations[pattern.get_pattern_name()] + if class_ops and self.current_function not in ("__init__", "__new__", "__class_init__", None): + # Check if this regular function modifies class-level structures + violations = pattern.check_cleanup(class_ops, node.body, self._context) + self.violations.extend(violations) + + def _check_module_level_cleanup(self): + """Check cleanup for module/class level operations""" + # Module-level operations are now checked when visiting functions + # This method is kept for potential future use but doesn't need to do anything + # since we only want to flag runtime modifications in functions, not initialization code + pass + + def _is_var_set_to_none(self, var_name: str, body: List[ast.stmt]) -> bool: + """Check if variable is set to None after its initial assignment""" + assignment_line = None + for stmt in body: + for node in ast.walk(stmt): + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == var_name: + assignment_line = node.lineno + break + if assignment_line: + break + if assignment_line: + break + + if not assignment_line: + return False + + for stmt in body: + for node in ast.walk(stmt): + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == var_name and node.lineno > assignment_line: + if isinstance(node.value, ast.Constant) and node.value.value is None: + return True + try: + NameConstant = getattr(ast, "NameConstant", None) + if NameConstant and isinstance(node.value, NameConstant): + if getattr(node.value, "value", None) is None: + return True + except (AttributeError, TypeError): + pass + return False + + def _get_call_string(self, node: ast.Call) -> str: + """Get string representation of function call""" + try: + if hasattr(ast, "unparse"): + return ast.unparse(node) + elif isinstance(node.func, ast.Attribute): + return f"{self._get_attr_string(node.func.value)}.{node.func.attr}()" + return str(node) + except Exception: + return str(node) + + def _get_attr_string(self, node: ast.AST) -> str: + """Get string representation of attribute access""" + if isinstance(node, ast.Name): + return node.id + elif isinstance(node, ast.Attribute): + return f"{self._get_attr_string(node.value)}.{node.attr}" + return str(node) + + +def check_file_for_memory_violations(file_path: str, patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]: + """Check a single file for memory violations""" + try: + with open(file_path, "r", encoding="utf-8") as f: + content = f.read() + + if "test" in file_path.lower() or "__pycache__" in file_path: + return [] + + tree = ast.parse(content, filename=file_path) + detector = MemoryViolationDetector(file_path, patterns) + detector.ast_tree = tree # Store AST for potential future use + detector.visit(tree) + # Class-level operations are checked when visiting functions + return detector.violations + except Exception as e: + print(f"Error parsing {file_path}: {e}") + return [] + + +def check_directory_for_memory_violations(directory_path: str, ignore_patterns: Optional[List[str]] = None, + patterns: Optional[Sequence[Pattern]] = None) -> List[Dict[str, Any]]: + """Recursively scan directory for memory violations""" + if ignore_patterns is None: + ignore_patterns = ["__pycache__", ".pyc", "site-packages", "venv", ".venv", "env", ".env", "node_modules", "tests"] + + all_violations = [] + for root, _dirs, files in os.walk(directory_path): + if any(pattern in root for pattern in ignore_patterns): + continue + for file in files: + if file.endswith(".py"): + violations = check_file_for_memory_violations(os.path.join(root, file), patterns) + all_violations.extend(violations) + return all_violations + + +def main(): + """Run memory violation detection on codebase""" + codebase_path = "./litellm" + + print("=" * 80) + print("MEMORY VIOLATION DETECTION TEST") + print("=" * 80) + print(f"Scanning: {codebase_path}") + print(f"Active patterns: {', '.join(p.get_pattern_name() for p in MemoryViolationDetector.DEFAULT_PATTERNS)}") + print() + + violations = check_directory_for_memory_violations(codebase_path) + + if violations: + by_type = {} + for v in violations: + vtype = v["type"] + if vtype not in by_type: + by_type[vtype] = [] + by_type[vtype].append(v) + + print("MEMORY VIOLATIONS FOUND:") + print("=" * 80) + + total = len(violations) + for vtype, vlist in by_type.items(): + print(f"\n{vtype.upper().replace('_', ' ')}: {len(vlist)} violation(s)") + print("-" * 80) + for v in vlist[:10]: + print(f" [VIOLATION] {v['file_path'] if 'file_path' in v else 'unknown'}:{v['line']}") + print(f" Function: {v['function']}") + print(f" Variable: {v['var_name']}") + print(f" {v['message']}") + print() + if len(vlist) > 10: + print(f" ... and {len(vlist) - 10} more violations of this type") + + print("=" * 80) + print(f"TOTAL VIOLATIONS: {total}") + print() + print("RECOMMENDATIONS:") + print(" 1. Set queue variables to None after use: obj = queue.get(); ...; obj = None") + print(" 2. Use bounded queues to prevent unbounded accumulation") + print(" 3. Process items faster than they're added, or drain queues periodically") + print(" 4. For class-level data structures (lists, dicts, sets) that are modified at runtime:") + print(" - Add size limit checks: if len(data) >= MAX_SIZE: ...") + print(" - Implement periodic cleanup or use bounded collections") + print(" - Consider using collections.deque with maxlen for lists") + print("=" * 80) + + first_v = violations[0] + raise Exception( + f"Found {total} memory violations! " + f"First violation: {first_v.get('file_path', 'unknown')}:{first_v['line']} - " + f"{first_v['message']}" + ) + else: + print("OK No memory violations found!") + + +if __name__ == "__main__": + main() diff --git a/tests/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index 8331738baba..20e6a381e5b 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -37,6 +37,7 @@ IGNORE_FUNCTIONS = [ "_split_text", # max depth set. "_delete_nested_value_custom", # max depth set (bounded by number of path segments). "filter_exceptions_from_params", # max depth set (default 20) to prevent infinite recursion. + "__getattr__", # lazy loading pattern in litellm/__init__.py with proper caching to prevent infinite recursion. ] diff --git a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py index 7e269f21451..c151150f634 100644 --- a/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py +++ b/tests/llm_translation/test_llm_response_utils/test_convert_dict_to_chat_completion.py @@ -903,3 +903,159 @@ def test_convert_to_model_response_object_with_thinking_content(): resp: ModelResponse = convert_to_model_response_object(**args) assert resp is not None assert resp.choices[0].message.reasoning_content is not None + + +def test_convert_to_model_response_object_with_empty_error_object(): + """ + Test that convert_to_model_response_object handles empty error objects gracefully. + + This is a regression test for issue #18407 where providers like Apertis return + empty error objects even on successful responses, causing spurious APIErrors. + + The error object structure: + { + "error": { + "message": "", + "type": "", + "param": "", + "code": null + } + } + """ + response_object = { + "model": "minimax-m2.1", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hey! I'm doing well, thanks for asking!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 49, + "completion_tokens": 87, + "total_tokens": 136, + }, + "error": { + "message": "", + "type": "", + "param": "", + "code": None, + }, + } + + # This should NOT raise an exception + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + assert result.model == "minimax-m2.1" + assert len(result.choices) == 1 + assert result.choices[0].message.content == "Hey! I'm doing well, thanks for asking!" + + +def test_convert_to_model_response_object_with_real_error(): + """ + Test that convert_to_model_response_object still raises for real errors. + + Ensures the empty error fix doesn't break legitimate error handling. + """ + response_object = { + "error": { + "message": "Rate limit exceeded", + "type": "rate_limit_error", + "param": None, + "code": 429, + }, + } + + with pytest.raises(Exception) as exc_info: + convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + # The exception should have the error message + assert hasattr(exc_info.value, "message") + assert "Rate limit exceeded" in str(exc_info.value.message) + + +def test_convert_to_model_response_object_with_empty_dict_error(): + """ + Test that convert_to_model_response_object handles completely empty error dict. + """ + response_object = { + "model": "test-model", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "Hello!", + }, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + "error": {}, # Completely empty error object + } + + # This should NOT raise an exception + result = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + + assert isinstance(result, ModelResponse) + assert result.choices[0].message.content == "Hello!" + + +def test_convert_to_model_response_object_with_error_code_only(): + """ + Test that errors with only a code (no message) are still treated as real errors. + """ + response_object = { + "error": { + "message": "", + "code": 500, + }, + } + + with pytest.raises(Exception): + convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object=response_object, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) diff --git a/tests/local_testing/test_anthropic_prompt_caching.py b/tests/local_testing/test_anthropic_prompt_caching.py index 0926bd17b70..c8589dd8844 100644 --- a/tests/local_testing/test_anthropic_prompt_caching.py +++ b/tests/local_testing/test_anthropic_prompt_caching.py @@ -104,7 +104,6 @@ async def test_litellm_anthropic_prompt_caching_tools(): ], extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) @@ -112,11 +111,12 @@ async def test_litellm_anthropic_prompt_caching_tools(): print("call args=", mock_post.call_args) expected_url = "https://api.anthropic.com/v1/messages" + # Note: anthropic-beta header for prompt-caching is no longer required + # Anthropic now supports prompt caching automatically when cache_control is used expected_headers = { "accept": "application/json", "content-type": "application/json", "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", "x-api-key": "mock_api_key", } @@ -285,7 +285,6 @@ async def test_anthropic_api_prompt_caching_basic(): max_tokens=10, extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) @@ -356,7 +355,6 @@ async def test_anthropic_api_prompt_caching_basic_with_cache_creation(): max_tokens=10, extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) @@ -645,7 +643,6 @@ async def test_litellm_anthropic_prompt_caching_system(): ], extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) @@ -657,7 +654,6 @@ async def test_litellm_anthropic_prompt_caching_system(): "accept": "application/json", "content-type": "application/json", "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", "x-api-key": "mock_api_key", } diff --git a/tests/local_testing/test_gemini_reasoning_content.py b/tests/local_testing/test_gemini_reasoning_content.py index 758616b4c94..f1f9c2ab512 100644 --- a/tests/local_testing/test_gemini_reasoning_content.py +++ b/tests/local_testing/test_gemini_reasoning_content.py @@ -1,4 +1,5 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig +from litellm.llms.vertex_ai.gemini.transformation import _gemini_convert_messages_with_history def test_thought_true_creates_thinking_block(): @@ -36,3 +37,108 @@ def test_thought_signature_without_thought_does_not_create_block(): config = VertexGeminiConfig() thinking_blocks = config._extract_thinking_blocks_from_parts(parts) assert thinking_blocks == [] + + +def test_extract_thought_signatures_from_regular_parts(): + """ + Test that thoughtSignatures are extracted from regular text parts (without thought=True). + This is the key feature for Gemini 3 multi-turn context preservation. + """ + parts = [{"text": "I am Gemini", "thoughtSignature": "sig-regular-123"}] + config = VertexGeminiConfig() + + # Should NOT create thinking block + thinking_blocks = config._extract_thinking_blocks_from_parts(parts) + assert thinking_blocks == [] + + # Should extract thought signature + signatures = config._extract_thought_signatures_from_parts(parts) + assert signatures is not None + assert len(signatures) == 1 + assert signatures[0] == "sig-regular-123" + + +def test_extract_multiple_thought_signatures(): + """ + Test extraction of multiple thoughtSignatures from different parts. + """ + parts = [ + {"text": "Part 1", "thoughtSignature": "sig-1"}, + {"text": "Part 2", "thoughtSignature": "sig-2"}, + {"text": "Part 3"} # No signature + ] + config = VertexGeminiConfig() + signatures = config._extract_thought_signatures_from_parts(parts) + + assert signatures is not None + assert len(signatures) == 2 + assert signatures[0] == "sig-1" + assert signatures[1] == "sig-2" + + +def test_round_trip_thought_signature_in_conversation(): + """ + Test that thoughtSignatures are properly round-tripped through conversation history. + This ensures multi-turn context preservation works correctly. + """ + messages = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": "Hi there", + "provider_specific_fields": { + "thought_signatures": ["sig-round-trip-abc"] + } + }, + {"role": "user", "content": "How are you?"} + ] + + gemini_contents = _gemini_convert_messages_with_history(messages) + + # Find the assistant (model) message + model_message = None + for content in gemini_contents: + if content.get("role") == "model": + model_message = content + break + + assert model_message is not None + assert len(model_message["parts"]) >= 1 + + # Check that the text part has the thoughtSignature + text_part = model_message["parts"][0] + assert text_part["text"] == "Hi there" + assert "thoughtSignature" in text_part + assert text_part["thoughtSignature"] == "sig-round-trip-abc" + + +def test_round_trip_without_thought_signature_still_works(): + """ + Test that messages without thoughtSignatures continue to work normally. + This ensures backward compatibility. + """ + messages = [ + {"role": "user", "content": "Hello"}, + { + "role": "assistant", + "content": "Hi there" + }, + {"role": "user", "content": "How are you?"} + ] + + gemini_contents = _gemini_convert_messages_with_history(messages) + + # Find the assistant (model) message + model_message = None + for content in gemini_contents: + if content.get("role") == "model": + model_message = content + break + + assert model_message is not None + assert len(model_message["parts"]) >= 1 + + # Check that the text part works without thoughtSignature + text_part = model_message["parts"][0] + assert text_part["text"] == "Hi there" + assert "thoughtSignature" not in text_part diff --git a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json index 26b712c1cf2..bd2f06b502e 100644 --- a/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json +++ b/tests/logging_callback_tests/langfuse_expected_request_body/completion_with_bedrock_call.json @@ -54,7 +54,7 @@ "id": "time-14-13-16-469836_chatcmpl-3803a9e9-aa68-4493-94d9-247f354830d6", "endTime": "2025-05-26T14:13:16.795438-07:00", "completionStartTime": "2025-05-26T14:13:16.795438-07:00", - "model": "anthropic.claude-3-5-sonnet-20240620-v1:0", + "model": "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0", "modelParameters": { "aws_region": "us-east-1" }, diff --git a/tests/logging_callback_tests/test_generic_api_callback.py b/tests/logging_callback_tests/test_generic_api_callback.py index 3ddf84f2939..c3e1171e96a 100644 --- a/tests/logging_callback_tests/test_generic_api_callback.py +++ b/tests/logging_callback_tests/test_generic_api_callback.py @@ -207,3 +207,253 @@ async def test_generic_api_callback_multiple_logs(): assert ( payload_item["response"]["choices"][0]["message"]["content"] == "hi" ), "Response should be hi" + + +@pytest.mark.asyncio +async def test_generic_api_callback_ndjson_format(): + """ + Test the GenericAPILogger callback with ndjson log format. + Validates that logs are sent as newline-delimited JSON. + """ + # Create a mock for the async_httpx_client's post method + mock_post = AsyncMock() + mock_post.return_value.status_code = 200 + mock_post.return_value.text = "OK" + + # Set up an endpoint for testing + test_endpoint = "https://example.com/api/logs" + test_headers = {"Authorization": "Bearer test_token"} + os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint + + # Initialize the GenericAPILogger with ndjson format + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + headers=test_headers, + flush_interval=1, + log_format="ndjson" # Set NDJSON format + ) + generic_logger.async_httpx_client.post = mock_post + litellm.callbacks = [generic_logger] + + # Make multiple completion calls to generate multiple logs + for i in range(3): + response = await litellm.acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": f"Hello, world! {i}"}], + mock_response="hi", + user="test_user", + ) + + # Wait for async flush + await asyncio.sleep(3) + + # Assert httpx post was called + mock_post.assert_called_once() + + # Get the actual request body from the mock + actual_url = mock_post.call_args[1]["url"] + assert actual_url == test_endpoint, f"Expected URL {test_endpoint}, got {actual_url}" + + # Get the data sent + ndjson_data = mock_post.call_args[1]["data"] + print("##########\n") + print("ndjson_data:", ndjson_data) + print("##########\n") + + # Validate it's NDJSON format (newline-delimited) + assert isinstance(ndjson_data, str), "Data should be a string for NDJSON" + + # Split by newlines and parse each line + lines = ndjson_data.strip().split("\n") + assert len(lines) == 3, f"Expected 3 lines of NDJSON, got {len(lines)}" + + # Validate each line is valid JSON + for i, line in enumerate(lines): + payload_item = json.loads(line) + payload_item = StandardLoggingPayload(**payload_item) + + # Basic assertions + assert payload_item["response_cost"] > 0, "Response cost should be greater than 0" + assert payload_item["model"] == "gpt-4o", "Model should be gpt-4o" + assert payload_item["model_parameters"]["user"] == "test_user", "User should be test_user" + + +@pytest.mark.asyncio +async def test_generic_api_callback_single_format(): + """ + Test the GenericAPILogger callback with single log format. + Validates that each log is sent as an individual request in parallel. + """ + # Create a mock for the async_httpx_client's post method + mock_post = AsyncMock() + mock_post.return_value.status_code = 200 + mock_post.return_value.text = "OK" + + # Set up an endpoint for testing + test_endpoint = "https://example.com/api/logs" + test_headers = {"Authorization": "Bearer test_token"} + os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint + + # Initialize the GenericAPILogger with single format + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + headers=test_headers, + flush_interval=1, # Quick flush to trigger batch send + log_format="single" # Set single format + ) + generic_logger.async_httpx_client.post = mock_post + litellm.callbacks = [generic_logger] + + # Make 3 completion calls + for i in range(3): + response = await litellm.acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": f"Hello, world! {i}"}], + mock_response="hi", + user="test_user", + ) + + # Wait for async flush + await asyncio.sleep(3) + + # Assert httpx post was called 3 times (once per log in batch) + assert mock_post.call_count == 3, f"Expected 3 calls, got {mock_post.call_count}" + + # Validate each call sent a single log object (not an array) + for call_idx in range(3): + call_args = mock_post.call_args_list[call_idx] + json_data = call_args[1]["data"] + + print(f"########## Call {call_idx} ##########") + print("json_data:", json_data) + + # Parse and validate - should be a single object, not an array + actual_request = json.loads(json_data) + assert isinstance(actual_request, dict), f"Call {call_idx}: Expected dict, got {type(actual_request)}" + + # Validate it's a valid StandardLoggingPayload + payload_item = StandardLoggingPayload(**actual_request) + assert payload_item["response_cost"] > 0, "Response cost should be greater than 0" + assert payload_item["model"] == "gpt-4o", "Model should be gpt-4o" + + +@pytest.mark.asyncio +async def test_generic_api_callback_json_array_format_explicit(): + """ + Test the GenericAPILogger callback with explicit json_array format. + Validates backward compatibility when explicitly set to json_array. + """ + # Create a mock for the async_httpx_client's post method + mock_post = AsyncMock() + mock_post.return_value.status_code = 200 + mock_post.return_value.text = "OK" + + # Set up an endpoint for testing + test_endpoint = "https://example.com/api/logs" + test_headers = {"Authorization": "Bearer test_token"} + os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint + + # Initialize the GenericAPILogger with explicit json_array format + generic_logger = GenericAPILogger( + endpoint=test_endpoint, + headers=test_headers, + flush_interval=1, + log_format="json_array" # Explicitly set json_array + ) + generic_logger.async_httpx_client.post = mock_post + litellm.callbacks = [generic_logger] + + # Make multiple completion calls + for i in range(5): + response = await litellm.acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": f"Hello, world! {i}"}], + mock_response="hi", + user="test_user", + ) + + # Wait for async flush + await asyncio.sleep(3) + + # Assert httpx post was called once (batched) + mock_post.assert_called_once() + + # Get the data and validate it's a JSON array + json_data = mock_post.call_args[1]["data"] + actual_request = json.loads(json_data) + + assert isinstance(actual_request, list), "Request body should be a list (JSON array)" + assert len(actual_request) == 5, f"Expected 5 items, got {len(actual_request)}" + + # Validate each item + for payload_item in actual_request: + payload_item = StandardLoggingPayload(**payload_item) + assert payload_item["response_cost"] > 0, "Response cost should be greater than 0" + assert payload_item["model"] == "gpt-4o", "Model should be gpt-4o" + + +@pytest.mark.asyncio +async def test_generic_api_callback_sumologic_uses_ndjson(): + """ + Test that the sumologic callback uses ndjson format by default + when loaded from generic_api_compatible_callbacks.json + """ + # Create a mock for the async_httpx_client's post method + mock_post = AsyncMock() + mock_post.return_value.status_code = 200 + mock_post.return_value.text = "OK" + + # Set environment variable for sumologic + os.environ["SUMOLOGIC_WEBHOOK_URL"] = "https://collectors.sumologic.com/receiver/v1/http/test123" + + # Initialize using callback_name (loads from JSON config) + generic_logger = GenericAPILogger( + callback_name="sumologic", + flush_interval=1 + ) + generic_logger.async_httpx_client.post = mock_post + litellm.callbacks = [generic_logger] + + # Verify the logger has ndjson format + assert generic_logger.log_format == "ndjson", "Sumologic should use ndjson format" + + # Make completion calls + for i in range(2): + await litellm.acompletion( + model="gpt-4o", + messages=[{"role": "user", "content": f"Test {i}"}], + mock_response="response", + user="test_user", + ) + + # Wait for async flush + await asyncio.sleep(3) + + # Assert httpx post was called + mock_post.assert_called_once() + + # Verify NDJSON format + ndjson_data = mock_post.call_args[1]["data"] + assert isinstance(ndjson_data, str), "Data should be a string for NDJSON" + + lines = ndjson_data.strip().split("\n") + assert len(lines) == 2, f"Expected 2 lines of NDJSON, got {len(lines)}" + + # Each line should be valid JSON + for line in lines: + json.loads(line) # Will raise if invalid JSON + + +@pytest.mark.asyncio +async def test_generic_api_callback_invalid_log_format(): + """ + Test that invalid log_format values raise a ValueError + """ + test_endpoint = "https://example.com/api/logs" + os.environ["GENERIC_LOGGER_ENDPOINT"] = test_endpoint + + with pytest.raises(ValueError, match="Invalid log_format"): + GenericAPILogger( + endpoint=test_endpoint, + log_format="invalid_format" # type: ignore # Intentionally invalid for testing + ) diff --git a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py index c8ceded4cf2..3d0682d9033 100644 --- a/tests/logging_callback_tests/test_opentelemetry_unit_tests.py +++ b/tests/logging_callback_tests/test_opentelemetry_unit_tests.py @@ -40,11 +40,15 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest): @pytest.mark.asyncio async def test_opentelemetry_integration(self): """ - Unit test to confirm the parent otel span is ended. + Unit test to confirm external parent otel spans are NOT ended by LiteLLM. + + External spans (passed via metadata) should be managed by their creators, + not by LiteLLM. This prevents premature closure of spans from Langfuse, + user code, or other external observability tools. """ # Reset all callbacks to ensure clean state litellm.logging_callback_manager._reset_all_callbacks() - + parent_otel_span = MagicMock() litellm.callbacks = ["otel"] @@ -57,8 +61,9 @@ class TestOpentelemetryUnitTests(BaseLoggingCallbackTest): await asyncio.sleep(1) - # Verify span was ended (may be called multiple times due to callback architecture) - parent_otel_span.end.assert_called() + # Verify external span was NOT ended by LiteLLM + # External spans should only be closed by their creators + parent_otel_span.end.assert_not_called() def test_init_tracing_respects_existing_tracer_provider(self): """ diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 9242dfc75f4..8fb0e80cc39 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -1057,7 +1057,7 @@ async def test_mcp_server_manager_config_integration_with_database(): ) # Test the add_update_server method (this tests our fix) - await test_manager.add_update_server(db_server) + await test_manager.add_server(db_server) # Verify the server was added with correct access_groups registry = test_manager.get_registry() @@ -1381,7 +1381,7 @@ async def test_add_update_server_with_alias(): mock_mcp_server.token_url = None # Add server to manager - await test_manager.add_update_server(mock_mcp_server) + await test_manager.add_server(mock_mcp_server) # Verify server was added with correct name (should use alias) assert "test-server-123" in test_manager.registry @@ -1421,7 +1421,7 @@ async def test_add_update_server_without_alias(): mock_mcp_server.token_url = None # Add server to manager - await test_manager.add_update_server(mock_mcp_server) + await test_manager.add_server(mock_mcp_server) # Verify server was added with correct name (should use server_name) assert "test-server-123" in test_manager.registry @@ -1461,7 +1461,7 @@ async def test_add_update_server_fallback_to_server_id(): mock_mcp_server.token_url = None # Add server to manager - await test_manager.add_update_server(mock_mcp_server) + await test_manager.add_server(mock_mcp_server) # Verify server was added with correct name (should use server_id) assert "test-server-123" in test_manager.registry diff --git a/tests/old_proxy_tests/tests/test_anthropic_context_caching.py b/tests/old_proxy_tests/tests/test_anthropic_context_caching.py index 7a153295f35..6b37873df4e 100644 --- a/tests/old_proxy_tests/tests/test_anthropic_context_caching.py +++ b/tests/old_proxy_tests/tests/test_anthropic_context_caching.py @@ -30,7 +30,6 @@ response = client.chat.completions.create( ], extra_headers={ "anthropic-version": "2023-06-01", - "anthropic-beta": "prompt-caching-2024-07-31", }, ) diff --git a/tests/store_model_in_db_tests/test_mcp_servers.py b/tests/store_model_in_db_tests/test_mcp_servers.py index 8197ce2693b..a369cce83c0 100644 --- a/tests/store_model_in_db_tests/test_mcp_servers.py +++ b/tests/store_model_in_db_tests/test_mcp_servers.py @@ -139,7 +139,7 @@ async def test_create_mcp_server_direct(): mock_get_prisma.return_value = mock_prisma # Mock server manager - mock_manager.add_update_server = mock.AsyncMock() + mock_manager.add_server = mock.AsyncMock() mock_manager.reload_servers_from_database = mock.AsyncMock() # Set up test data @@ -195,7 +195,7 @@ async def test_create_mcp_server_direct(): # Verify mocks were called mock_get_server.assert_called_once_with(mock_prisma, server_id) mock_create.assert_called_once() - mock_manager.add_update_server.assert_called_once_with(expected_response) + mock_manager.add_server.assert_called_once_with(expected_response) @pytest.mark.asyncio @@ -379,7 +379,7 @@ async def test_edit_mcp_server_redacts_credentials(): mock_prisma = mock.Mock() mock_get_prisma.return_value = mock_prisma - mock_manager.add_update_server = mock.AsyncMock() + mock_manager.update_server = mock.AsyncMock() mock_manager.reload_servers_from_database = mock.AsyncMock() server_id = str(uuid.uuid4()) @@ -417,7 +417,7 @@ async def test_edit_mcp_server_redacts_credentials(): mock_validate.assert_called_once() mock_update.assert_awaited_once() - mock_manager.add_update_server.assert_called_once_with(updated_server) + mock_manager.update_server.assert_called_once_with(updated_server) mock_manager.reload_servers_from_database.assert_awaited_once() def test_validate_mcp_server_name_direct(): """ diff --git a/tests/test_litellm/google_genai/test_google_genai_adapter.py b/tests/test_litellm/google_genai/test_google_genai_adapter.py index e8882a1acb3..135881ad209 100644 --- a/tests/test_litellm/google_genai/test_google_genai_adapter.py +++ b/tests/test_litellm/google_genai/test_google_genai_adapter.py @@ -1197,6 +1197,131 @@ async def test_agenerate_content_x_goog_api_key_header(): # Verify other expected headers assert headers.get("Content-Type") == "application/json", f"Expected Content-Type application/json, got {headers.get('Content-Type')}" - + print(f"✓ Test passed: x-goog-api-key header correctly set to {api_key_value}") print(f"✓ All headers: {list(headers.keys())}") + + +def test_inline_data_base64_image_transformation(): + """Test transformation of Gemini inline_data (Base64 images) to OpenAI format""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with Base64 image + model = "gpt-4-vision-preview" + contents = { + "role": "user", + "parts": [ + {"text": "What's in this image?"}, + { + "inline_data": { + "mime_type": "image/jpeg", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + } + } + ] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is an array (multimodal format) + content = completion_request["messages"][0]["content"] + assert isinstance(content, list), "Content should be a list for multimodal messages" + assert len(content) == 2, "Should have 2 content parts (text + image)" + + # Verify text part + text_part = content[0] + assert text_part["type"] == "text" + assert text_part["text"] == "What's in this image?" + + # Verify image part + image_part = content[1] + assert image_part["type"] == "image_url" + assert "image_url" in image_part + assert "url" in image_part["image_url"] + assert image_part["image_url"]["url"].startswith("data:image/jpeg;base64,") + assert "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" in image_part["image_url"]["url"] + + +def test_inline_data_image_only_transformation(): + """Test transformation of Gemini inline_data with only image (no text)""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with only Base64 image (no text) + model = "gpt-4-vision-preview" + contents = { + "role": "user", + "parts": [ + { + "inline_data": { + "mime_type": "image/png", + "data": "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg==" + } + } + ] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is an array (multimodal format) + content = completion_request["messages"][0]["content"] + assert isinstance(content, list), "Content should be a list for multimodal messages" + assert len(content) == 1, "Should have 1 content part (image only)" + + # Verify image part + image_part = content[0] + assert image_part["type"] == "image_url" + assert "image_url" in image_part + assert "url" in image_part["image_url"] + assert image_part["image_url"]["url"].startswith("data:image/png;base64,") + + +def test_inline_data_backward_compatibility_text_only(): + """Test that pure text messages still use simple string format (backward compatibility)""" + from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter + + adapter = GoogleGenAIAdapter() + + # Test input with only text (no images) + model = "gpt-3.5-turbo" + contents = { + "role": "user", + "parts": [{"text": "Hello, how are you?"}] + } + + # Transform to completion format + completion_request = adapter.translate_generate_content_to_completion( + model=model, + contents=contents + ) + + # Verify the transformation + assert completion_request["model"] == model + assert len(completion_request["messages"]) == 1 + assert completion_request["messages"][0]["role"] == "user" + + # Verify content is a simple string (not an array) for backward compatibility + content = completion_request["messages"][0]["content"] + assert isinstance(content, str), "Content should be a string for text-only messages (backward compatibility)" + assert content == "Hello, how are you?" diff --git a/tests/test_litellm/google_genai/test_google_genai_transformation.py b/tests/test_litellm/google_genai/test_google_genai_transformation.py index c953a504a38..0d5468cf275 100644 --- a/tests/test_litellm/google_genai/test_google_genai_transformation.py +++ b/tests/test_litellm/google_genai/test_google_genai_transformation.py @@ -247,3 +247,122 @@ def test_responses_api_no_reasoning(): # reasoning_effort should not be in result if not provided (filtered out as None) assert "reasoning_effort" not in result or result.get("reasoning_effort") is None + + +def test_transform_generate_content_request_with_system_instruction(): + """Test that systemInstruction parameter is properly included in the request""" + config = GoogleGenAIConfig() + + system_instruction = { + "parts": [{"text": "You are a helpful assistant"}] + } + + contents = [ + { + "role": "user", + "parts": [{"text": "Hello"}] + } + ] + + generate_content_config_dict = { + "temperature": 1.0, + "maxOutputTokens": 100 + } + + # Call transform_generate_content_request + result = config.transform_generate_content_request( + model="gemini-3-flash-preview", + contents=contents, + tools=None, + generate_content_config_dict=generate_content_config_dict, + system_instruction=system_instruction, + ) + + # Verify that systemInstruction is in the request + assert "systemInstruction" in result, "systemInstruction should be in request body" + assert result["systemInstruction"] == system_instruction, "systemInstruction should match input" + assert result["model"] == "gemini-3-flash-preview" + assert result["contents"] == contents + + +def test_transform_generate_content_request_without_system_instruction(): + """Test that request works correctly without systemInstruction""" + config = GoogleGenAIConfig() + + contents = [ + { + "role": "user", + "parts": [{"text": "Hello"}] + } + ] + + generate_content_config_dict = { + "temperature": 1.0 + } + + # Call transform_generate_content_request without system_instruction + result = config.transform_generate_content_request( + model="gemini-3-flash-preview", + contents=contents, + tools=None, + generate_content_config_dict=generate_content_config_dict, + system_instruction=None, + ) + + # Verify that systemInstruction is NOT in the request when not provided + assert "systemInstruction" not in result, "systemInstruction should not be in request when None" + assert result["model"] == "gemini-3-flash-preview" + assert result["contents"] == contents + + +def test_transform_generate_content_request_system_instruction_with_tools(): + """Test that systemInstruction works correctly alongside tools""" + config = GoogleGenAIConfig() + + system_instruction = { + "parts": [{"text": "You are a helpful assistant that uses tools"}] + } + + contents = [ + { + "role": "user", + "parts": [{"text": "What's the weather?"}] + } + ] + + tools = [ + { + "functionDeclarations": [ + { + "name": "get_weather", + "description": "Get weather information", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"} + } + } + } + ] + } + ] + + generate_content_config_dict = { + "temperature": 0.7 + } + + # Call transform_generate_content_request with both system_instruction and tools + result = config.transform_generate_content_request( + model="gemini-3-flash-preview", + contents=contents, + tools=tools, + generate_content_config_dict=generate_content_config_dict, + system_instruction=system_instruction, + ) + + # Verify that both systemInstruction and tools are in the request + assert "systemInstruction" in result, "systemInstruction should be in request body" + assert result["systemInstruction"] == system_instruction + assert "tools" in result, "tools should be in request body" + assert result["tools"] == tools + assert result["model"] == "gemini-3-flash-preview" diff --git a/tests/test_litellm/integrations/levo/__init__.py b/tests/test_litellm/integrations/levo/__init__.py new file mode 100644 index 00000000000..1560e78b7b9 --- /dev/null +++ b/tests/test_litellm/integrations/levo/__init__.py @@ -0,0 +1 @@ +# Levo integration tests diff --git a/tests/test_litellm/integrations/levo/test_levo.py b/tests/test_litellm/integrations/levo/test_levo.py new file mode 100644 index 00000000000..5d042cbc060 --- /dev/null +++ b/tests/test_litellm/integrations/levo/test_levo.py @@ -0,0 +1,407 @@ +import unittest +from unittest.mock import patch + +import pytest + +from litellm.integrations.levo.levo import LevoConfig, LevoLogger +from litellm.integrations.opentelemetry import OpenTelemetryConfig + +# Try to import OpenTelemetry packages, skip tests if not available +try: + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, + ) + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + + OPENTELEMETRY_AVAILABLE = True +except ImportError: + OPENTELEMETRY_AVAILABLE = False + + +class TestLevoConfig(unittest.TestCase): + """Unit tests for LevoLogger configuration.""" + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + "LEVOAI_WORKSPACE_ID": "test-workspace-id", + "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", + }, + ) + def test_get_levo_config_with_all_required_vars(self): + """Test get_levo_config() with all required environment variables.""" + config = LevoLogger.get_levo_config() + + # Verify headers include all three values + self.assertIn("Authorization=Bearer test-api-key", config.otlp_auth_headers) + self.assertIn("x-levo-organization-id=test-org-id", config.otlp_auth_headers) + self.assertIn("x-levo-workspace-id=test-workspace-id", config.otlp_auth_headers) + + # Verify endpoint uses provided collector URL exactly as-is + self.assertEqual(config.endpoint, "https://collector.levo.ai") + + # Verify protocol is otlp_http + self.assertEqual(config.protocol, "otlp_http") + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + "LEVOAI_WORKSPACE_ID": "test-workspace-id", + "LEVOAI_COLLECTOR_URL": "https://custom.collector.com", + }, + ) + def test_get_levo_config_with_custom_collector_url(self): + """Test get_levo_config() with custom collector URL.""" + config = LevoLogger.get_levo_config() + + # Verify endpoint uses custom URL exactly as provided + self.assertEqual(config.endpoint, "https://custom.collector.com") + self.assertEqual(config.protocol, "otlp_http") + + @patch.dict("os.environ", {}, clear=True) + def test_get_levo_config_missing_api_key(self): + """Test get_levo_config() raises ValueError when LEVOAI_API_KEY is missing.""" + with pytest.raises(ValueError, match="LEVOAI_API_KEY"): + LevoLogger.get_levo_config() + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + }, + clear=True, + ) + def test_get_levo_config_missing_org_id(self): + """Test get_levo_config() raises ValueError when LEVOAI_ORG_ID is missing.""" + with pytest.raises(ValueError, match="LEVOAI_ORG_ID"): + LevoLogger.get_levo_config() + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + }, + clear=True, + ) + def test_get_levo_config_missing_workspace_id(self): + """Test get_levo_config() raises ValueError when LEVOAI_WORKSPACE_ID is missing.""" + with pytest.raises(ValueError, match="LEVOAI_WORKSPACE_ID"): + LevoLogger.get_levo_config() + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + "LEVOAI_WORKSPACE_ID": "test-workspace-id", + }, + clear=True, + ) + def test_get_levo_config_missing_collector_url(self): + """Test get_levo_config() raises ValueError when LEVOAI_COLLECTOR_URL is missing.""" + with pytest.raises(ValueError, match="LEVOAI_COLLECTOR_URL"): + LevoLogger.get_levo_config() + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + "LEVOAI_WORKSPACE_ID": "test-workspace-id", + "LEVOAI_COLLECTOR_URL": "http://localhost:4318", + }, + ) + def test_get_levo_config_with_http_endpoint(self): + """Test get_levo_config() with HTTP endpoint.""" + config = LevoLogger.get_levo_config() + + # Should use HTTP endpoint exactly as provided + self.assertEqual(config.endpoint, "http://localhost:4318") + self.assertEqual(config.protocol, "otlp_http") + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + "LEVOAI_WORKSPACE_ID": "test-workspace-id", + "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", + }, + ) + def test_levo_config_headers_format(self): + """Test that OTLP headers are formatted correctly.""" + config = LevoLogger.get_levo_config() + + # Verify headers contain all required parts + self.assertIn("Authorization=Bearer test-api-key", config.otlp_auth_headers) + self.assertIn("x-levo-organization-id=test-org-id", config.otlp_auth_headers) + self.assertIn("x-levo-workspace-id=test-workspace-id", config.otlp_auth_headers) + + # Verify headers are comma-separated + header_parts = config.otlp_auth_headers.split(",") + self.assertEqual(len(header_parts), 3) + + +class TestLevoIntegration(unittest.TestCase): + """Integration tests for LevoLogger.""" + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + "LEVOAI_WORKSPACE_ID": "test-workspace-id", + "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", + }, + ) + @pytest.mark.skipif( + not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed" + ) + @patch( + "litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy" + ) + def test_levo_logger_instantiation(self, mock_init_proxy): + """Test that LevoLogger can be instantiated with proper config.""" + # Mock the proxy initialization to avoid importing proxy code + mock_init_proxy.return_value = None + + config = LevoLogger.get_levo_config() + otel_config = OpenTelemetryConfig( + exporter=config.protocol, + endpoint=config.endpoint, + headers=config.otlp_auth_headers, + ) + + # Create a tracer provider with in-memory exporter to avoid requiring OTLP packages + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + + # Create LevoLogger instance with mocked tracer provider + levo_logger = LevoLogger( + config=otel_config, callback_name="levo", tracer_provider=tracer_provider + ) + + # Verify it's an instance of OpenTelemetry + self.assertIsInstance(levo_logger, LevoLogger) + # Check it extends OpenTelemetry by checking base classes + from litellm.integrations.opentelemetry import OpenTelemetry + + self.assertIsInstance(levo_logger, OpenTelemetry) + + # Verify callback_name is set + self.assertEqual(levo_logger.callback_name, "levo") + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + "LEVOAI_WORKSPACE_ID": "test-workspace-id", + "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", + }, + ) + @pytest.mark.skipif( + not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed" + ) + @patch( + "litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy" + ) + @pytest.mark.asyncio + async def test_levo_logger_health_check_healthy(self, mock_init_proxy): + """Test health check returns healthy status when config is valid.""" + # Mock the proxy initialization to avoid importing proxy code + mock_init_proxy.return_value = None + + config = LevoLogger.get_levo_config() + otel_config = OpenTelemetryConfig( + exporter=config.protocol, + endpoint=config.endpoint, + headers=config.otlp_auth_headers, + ) + + # Create tracer provider with in-memory exporter + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + + levo_logger = LevoLogger( + config=otel_config, callback_name="levo", tracer_provider=tracer_provider + ) + + # Run health check + result = await levo_logger.async_health_check() + + self.assertEqual(result["status"], "healthy") + self.assertIn("message", result) + + @patch.dict("os.environ", {}, clear=True) + def test_levo_logger_health_check_unhealthy(self): + """Test health check returns unhealthy status when required vars are missing.""" + # Try to create logger without required env vars + # This should fail during config, but we can test health check logic + with pytest.raises(ValueError): + LevoLogger.get_levo_config() + + @patch.dict( + "os.environ", + { + "LEVOAI_API_KEY": "test-api-key", + "LEVOAI_ORG_ID": "test-org-id", + "LEVOAI_WORKSPACE_ID": "test-workspace-id", + "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", + }, + ) + @pytest.mark.skipif( + not OPENTELEMETRY_AVAILABLE, reason="OpenTelemetry packages not installed" + ) + @patch( + "litellm.integrations.opentelemetry.OpenTelemetry._init_otel_logger_on_litellm_proxy" + ) + def test_levo_logger_callback_name(self, mock_init_proxy): + """Test that callback_name is properly set and used.""" + # Mock the proxy initialization to avoid importing proxy code + mock_init_proxy.return_value = None + + config = LevoLogger.get_levo_config() + otel_config = OpenTelemetryConfig( + exporter=config.protocol, + endpoint=config.endpoint, + headers=config.otlp_auth_headers, + ) + + # Create tracer provider with in-memory exporter + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(InMemorySpanExporter())) + + levo_logger = LevoLogger( + config=otel_config, callback_name="levo", tracer_provider=tracer_provider + ) + + # Verify callback_name attribute + self.assertEqual(levo_logger.callback_name, "levo") + + +@pytest.mark.parametrize( + "env_vars, expected_headers_contains, expected_endpoint, expected_protocol", + [ + pytest.param( + { + "LEVOAI_API_KEY": "test-key", + "LEVOAI_ORG_ID": "test-org", + "LEVOAI_WORKSPACE_ID": "test-workspace", + "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", + }, + [ + "Authorization=Bearer test-key", + "x-levo-organization-id=test-org", + "x-levo-workspace-id=test-workspace", + ], + "https://collector.levo.ai", + "otlp_http", + id="collector URL with all required vars", + ), + pytest.param( + { + "LEVOAI_API_KEY": "key-123", + "LEVOAI_ORG_ID": "org-456", + "LEVOAI_WORKSPACE_ID": "workspace-789", + "LEVOAI_COLLECTOR_URL": "https://custom.example.com", + }, + [ + "Authorization=Bearer key-123", + "x-levo-organization-id=org-456", + "x-levo-workspace-id=workspace-789", + ], + "https://custom.example.com", + "otlp_http", + id="custom collector URL", + ), + pytest.param( + { + "LEVOAI_API_KEY": "key-123", + "LEVOAI_ORG_ID": "org-456", + "LEVOAI_WORKSPACE_ID": "workspace-789", + "LEVOAI_COLLECTOR_URL": "http://localhost:9999", + }, + ["Authorization=Bearer key-123"], + "http://localhost:9999", + "otlp_http", + id="custom HTTP endpoint", + ), + ], +) +def test_get_levo_config_parametrized( + monkeypatch, + env_vars, + expected_headers_contains, + expected_endpoint, + expected_protocol, +): + """Parametrized tests for get_levo_config() with various configurations.""" + # Clear all Levo-related env vars first to ensure clean state + for key in [ + "LEVOAI_API_KEY", + "LEVOAI_ORG_ID", + "LEVOAI_WORKSPACE_ID", + "LEVOAI_COLLECTOR_URL", + "LEVOAI_ENV_NAME", + ]: + monkeypatch.delenv(key, raising=False) + + for key, value in env_vars.items(): + monkeypatch.setenv(key, value) + + config = LevoLogger.get_levo_config() + + assert isinstance(config, LevoConfig) + assert config.endpoint == expected_endpoint + assert config.protocol == expected_protocol + + # Verify all expected header parts are present + for header_part in expected_headers_contains: + assert header_part in config.otlp_auth_headers + + +@pytest.mark.parametrize( + "missing_var", + [ + pytest.param("LEVOAI_API_KEY", id="missing API key"), + pytest.param("LEVOAI_ORG_ID", id="missing org ID"), + pytest.param("LEVOAI_WORKSPACE_ID", id="missing workspace ID"), + pytest.param("LEVOAI_COLLECTOR_URL", id="missing collector URL"), + ], +) +def test_get_levo_config_missing_required_vars(monkeypatch, missing_var): + """Test that missing required environment variables raise ValueError.""" + # Clear all Levo-related env vars + for key in [ + "LEVOAI_API_KEY", + "LEVOAI_ORG_ID", + "LEVOAI_WORKSPACE_ID", + "LEVOAI_COLLECTOR_URL", + ]: + monkeypatch.delenv(key, raising=False) + + # Set all required vars except the missing one + required_vars = { + "LEVOAI_API_KEY": "test-key", + "LEVOAI_ORG_ID": "test-org", + "LEVOAI_WORKSPACE_ID": "test-workspace", + "LEVOAI_COLLECTOR_URL": "https://collector.levo.ai", + } + required_vars.pop(missing_var) + + for key, value in required_vars.items(): + monkeypatch.setenv(key, value) + + with pytest.raises(ValueError, match=missing_var): + LevoLogger.get_levo_config() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index a719d102a7c..a322dfe9a2b 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -498,7 +498,7 @@ class TestPassthroughCallTypeHandling: def test_get_pre_call_type_with_allm_passthrough_route(self): """ Test that _get_pre_call_type correctly maps allm_passthrough_route. - + This tests Fix #1: allm_passthrough_route was not being handled, causing call_type to be None. """ from litellm.proxy.common_request_processing import ( @@ -509,14 +509,14 @@ class TestPassthroughCallTypeHandling: result = ProxyBaseLLMRequestProcessing._get_pre_call_type( route_type="allm_passthrough_route" ) - + # Should return allm_passthrough_route, not None assert result == "allm_passthrough_route" def test_get_pre_call_type_preserves_standard_mappings(self): """ Test that _get_pre_call_type still correctly maps standard route types. - + Ensures Fix #1 didn't break existing functionality. """ from litellm.proxy.common_request_processing import ( @@ -536,3 +536,235 @@ class TestPassthroughCallTypeHandling: ProxyBaseLLMRequestProcessing._get_pre_call_type(route_type="aresponses") == "responses" ) + + +class TestEventTypeLogging: + """Tests for event_type logging in guardrail information.""" + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_pre_call_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.pre_call + from async_pre_call_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_pre_call_hook(self, data: dict, **kwargs): + return {"result": "pre_call_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_pre_call_hook(data=request_data) + + # Check that the guardrail_mode was set to pre_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_post_call_success_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call + from async_post_call_success_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_post_call_success_hook(self, data: dict, **kwargs): + return {"result": "post_call_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_post_call_success_hook(data=request_data) + + # Check that the guardrail_mode was set to post_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_moderation_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.during_call + from async_moderation_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.during_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_moderation_hook(self, data: dict, **kwargs): + return {"result": "moderation_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_moderation_hook(data=request_data) + + # Check that the guardrail_mode was set to during_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.during_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_infers_event_type_from_async_post_call_streaming_hook( + self, + ): + """ + Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call + from async_post_call_streaming_hook function name. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=[ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + ], + ) + + @log_guardrail_information + async def async_post_call_streaming_hook(self, data: dict, **kwargs): + return {"result": "streaming_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.async_post_call_streaming_hook(data=request_data) + + # Check that the guardrail_mode was set to post_call (not the full list) + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + @pytest.mark.asyncio + async def test_log_guardrail_information_returns_none_for_unknown_function_name( + self, + ): + """ + Test that log_guardrail_information decorator returns None for event_type + when function name doesn't match known patterns, and falls back to self.event_hook. + """ + from litellm.integrations.custom_guardrail import log_guardrail_information + from litellm.types.guardrails import GuardrailEventHooks + + class TestGuardrail(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="test_event_type_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + @log_guardrail_information + async def some_other_hook(self, data: dict, **kwargs): + return {"result": "other_hook_executed"} + + guardrail = TestGuardrail() + request_data = {"metadata": {}} + + await guardrail.some_other_hook(data=request_data) + + # Check that the guardrail_mode falls back to self.event_hook + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call + + def test_add_standard_logging_uses_event_type_over_event_hook(self): + """ + Test that add_standard_logging_guardrail_information_to_request_data + prioritizes event_type parameter over self.event_hook. + """ + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + ) + + request_data = {"metadata": {}} + + # Call with explicit event_type + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"result": "ok"}, + request_data=request_data, + guardrail_status="success", + event_type=GuardrailEventHooks.post_call, + ) + + # Should use the provided event_type (post_call), not the full event_hook list + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call + + def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none( + self, + ): + """ + Test that add_standard_logging_guardrail_information_to_request_data + falls back to self.event_hook when event_type is None. + """ + from litellm.types.guardrails import GuardrailEventHooks + + guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + event_hook=GuardrailEventHooks.pre_call, + ) + + request_data = {"metadata": {}} + + # Call with event_type=None + guardrail.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"result": "ok"}, + request_data=request_data, + guardrail_status="success", + event_type=None, + ) + + # Should fall back to self.event_hook + logged_info = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(logged_info) == 1 + assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index 58fccf5e42d..5d648f601f6 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch # Adds the grandparent directory to sys.path to allow importing project modules sys.path.insert(0, os.path.abspath("../..")) +from opentelemetry import trace from opentelemetry.sdk._logs import LoggerProvider as OTLoggerProvider from opentelemetry.sdk._logs.export import InMemoryLogExporter, SimpleLogRecordProcessor from opentelemetry.sdk.metrics import MeterProvider @@ -1318,3 +1319,456 @@ class TestOpenTelemetryProtocolSelection(unittest.TestCase): "http://collector:4317/v1/traces", "logs" ) self.assertEqual(normalized, "http://collector:4317/v1/logs") + + +class TestOpenTelemetryExternalSpan(unittest.TestCase): + """ + Test suite for external span handling in OpenTelemetry integration. + + These tests verify that LiteLLM correctly handles spans created outside + of LiteLLM (e.g., by Langfuse SDK, user application code, or global context) + without closing them prematurely. + + Background: + - External spans can come from: Langfuse SDK, user code, HTTP traceparent headers, global context + - LiteLLM should NEVER close spans it did not create + - Bug: LiteLLM was reusing and closing external spans in _start_primary_span + """ + + HERE = os.path.dirname(__file__) + + def setUp(self): + """Set up common test fixtures""" + self.span_exporter = InMemorySpanExporter() + self.tracer_provider = TracerProvider() + self.tracer_provider.add_span_processor( + SimpleSpanProcessor(self.span_exporter) + ) + + # Don't set global tracer provider - instead, get tracers directly from our provider + # This avoids "Overriding of current TracerProvider is not allowed" warnings + + # Clear any existing spans + self.span_exporter.clear() + + def _create_test_kwargs_and_response(self): + """Load test data from JSON files""" + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_kwargs.json") + ) as f: + kwargs = json.load(f) + + with open( + os.path.join(self.HERE, "open_telemetry", "data", "captured_response.json") + ) as f: + response_obj = json.load(f) + + return kwargs, response_obj + + def _get_spans_by_name(self, name): + """Get all spans with the given name""" + spans = self.span_exporter.get_finished_spans() + return [s for s in spans if s.name == name] + + @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False) + def test_external_span_not_closed_with_use_otel_litellm_request_span_false(self): + """ + Test that external spans are not closed when USE_OTEL_LITELLM_REQUEST_SPAN=false (default). + + Expected behavior: + - External span remains open (is_recording = True) + - raw_gen_ai_request spans are direct children of external span (shallow hierarchy) + - No litellm_request span is created + - Multiple completions work correctly + """ + # Initialize OpenTelemetry + otel = OpenTelemetry(tracer_provider=self.tracer_provider) + + # Load test data + kwargs, response_obj = self._create_test_kwargs_and_response() + + # Create external parent span using our test TracerProvider + tracer = self.tracer_provider.get_tracer(__name__) + + with tracer.start_as_current_span("external_parent_span") as parent_span: + parent_ctx = parent_span.get_span_context() + parent_trace_id = parent_ctx.trace_id + parent_span_id = parent_ctx.span_id + + self.assertTrue( + parent_span.is_recording(), + "External span should be recording before completion calls" + ) + + # First completion call + start_time = datetime.utcnow() + end_time = start_time + timedelta(seconds=1) + otel._handle_success(kwargs, response_obj, start_time, end_time) + + # Verify parent span is still recording + self.assertTrue( + parent_span.is_recording(), + "External span should still be recording after first completion" + ) + + # Second completion call + start_time2 = end_time + end_time2 = start_time2 + timedelta(seconds=1) + otel._handle_success(kwargs, response_obj, start_time2, end_time2) + + # Verify parent span is still recording + self.assertTrue( + parent_span.is_recording(), + "External span should still be recording after second completion" + ) + + # After exiting context, verify spans + spans = self.span_exporter.get_finished_spans() + + # All spans should have the same trace_id + for span in spans: + self.assertEqual( + span.context.trace_id, + parent_trace_id, + f"Span {span.name} should have same trace_id as parent" + ) + + # Should have external_parent_span + parent_spans = self._get_spans_by_name("external_parent_span") + self.assertEqual(len(parent_spans), 1, "Should have exactly one external_parent_span") + + # Verify LiteLLM set attributes on external parent span + parent_span_finished = parent_spans[0] + self.assertIsNotNone( + parent_span_finished.attributes, + "Parent span should have attributes set by LiteLLM" + ) + self.assertIn( + "gen_ai.request.model", + parent_span_finished.attributes, + "Parent span should have model attribute from LiteLLM" + ) + + # Should have raw_gen_ai_request spans (if message_logging is on) + raw_spans = self._get_spans_by_name("raw_gen_ai_request") + # Note: May be 0 if message_logging is off, or 2 if on + + # Should NOT have litellm_request spans (USE_OTEL_LITELLM_REQUEST_SPAN=false) + litellm_spans = self._get_spans_by_name("litellm_request") + self.assertEqual( + len(litellm_spans), + 0, + "Should NOT have litellm_request spans when USE_OTEL_LITELLM_REQUEST_SPAN=false" + ) + + # Verify raw_gen_ai_request spans are direct children of external span + for raw_span in raw_spans: + self.assertEqual( + raw_span.parent.span_id if raw_span.parent else None, + parent_span_id, + f"raw_gen_ai_request should be direct child of external_parent_span" + ) + + @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "true"}, clear=False) + def test_external_span_not_closed_with_use_otel_litellm_request_span_true(self): + """ + Test that external spans are not closed when USE_OTEL_LITELLM_REQUEST_SPAN=true. + + Expected behavior: + - External span remains open (is_recording = True) + - litellm_request spans are created as children of external span + - raw_gen_ai_request spans are children of litellm_request spans + - Correct hierarchy: external_parent → litellm_request → raw_gen_ai_request + """ + # Initialize OpenTelemetry + otel = OpenTelemetry(tracer_provider=self.tracer_provider) + + # Load test data + kwargs, response_obj = self._create_test_kwargs_and_response() + + # Create external parent span using our test TracerProvider + tracer = self.tracer_provider.get_tracer(__name__) + + with tracer.start_as_current_span("external_parent_span") as parent_span: + parent_ctx = parent_span.get_span_context() + parent_trace_id = parent_ctx.trace_id + parent_span_id = parent_ctx.span_id + + # First completion call + start_time = datetime.utcnow() + end_time = start_time + timedelta(seconds=1) + otel._handle_success(kwargs, response_obj, start_time, end_time) + + # Verify parent span is still recording + self.assertTrue( + parent_span.is_recording(), + "External span should still be recording after first completion" + ) + + # Second completion call + start_time2 = end_time + end_time2 = start_time2 + timedelta(seconds=1) + otel._handle_success(kwargs, response_obj, start_time2, end_time2) + + # Verify parent span is still recording + self.assertTrue( + parent_span.is_recording(), + "External span should still be recording after second completion" + ) + + # After exiting context, verify spans + spans = self.span_exporter.get_finished_spans() + + # All spans should have the same trace_id + for span in spans: + self.assertEqual( + span.context.trace_id, + parent_trace_id, + f"Span {span.name} should have same trace_id as parent" + ) + + # Should have litellm_request spans (USE_OTEL_LITELLM_REQUEST_SPAN=true) + litellm_spans = self._get_spans_by_name("litellm_request") + self.assertEqual( + len(litellm_spans), + 2, + "Should have 2 litellm_request spans when USE_OTEL_LITELLM_REQUEST_SPAN=true" + ) + + # Verify litellm_request spans are children of external span + for litellm_span in litellm_spans: + self.assertEqual( + litellm_span.parent.span_id if litellm_span.parent else None, + parent_span_id, + "litellm_request should be child of external_parent_span" + ) + + # Verify raw_gen_ai_request spans (if present) are children of litellm_request + raw_spans = self._get_spans_by_name("raw_gen_ai_request") + if raw_spans: + litellm_span_ids = {s.context.span_id for s in litellm_spans} + for raw_span in raw_spans: + self.assertIn( + raw_span.parent.span_id if raw_span.parent else None, + litellm_span_ids, + "raw_gen_ai_request should be child of litellm_request" + ) + + @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False) + def test_external_span_with_multiple_completions(self): + """ + Test that multiple completion calls work correctly within external span context. + + Expected behavior: + - Both completion calls succeed + - All spans belong to the same trace + - External span remains open throughout + - No errors or warnings about "ended span" + """ + # Initialize OpenTelemetry + otel = OpenTelemetry(tracer_provider=self.tracer_provider) + + # Load test data + kwargs, response_obj = self._create_test_kwargs_and_response() + + # Create external parent span using our test TracerProvider + tracer = self.tracer_provider.get_tracer(__name__) + + with tracer.start_as_current_span("external_parent_span") as parent_span: + parent_ctx = parent_span.get_span_context() + parent_trace_id = parent_ctx.trace_id + + # Make multiple completion calls + for i in range(3): + start_time = datetime.utcnow() + end_time = start_time + timedelta(seconds=1) + + # This should not raise any exceptions + otel._handle_success(kwargs, response_obj, start_time, end_time) + + # Verify parent span is still recording after each call + self.assertTrue( + parent_span.is_recording(), + f"External span should still be recording after completion #{i+1}" + ) + + # Verify all spans have the same trace_id + spans = self.span_exporter.get_finished_spans() + for span in spans: + self.assertEqual( + span.context.trace_id, + parent_trace_id, + f"All spans should belong to the same trace" + ) + + # Should have the external parent span + parent_spans = self._get_spans_by_name("external_parent_span") + self.assertEqual(len(parent_spans), 1, "Should have exactly one external_parent_span") + + # Verify LiteLLM set attributes on external parent span + parent_span_finished = parent_spans[0] + self.assertIn( + "gen_ai.request.model", + parent_span_finished.attributes, + "Parent span should have model attribute from LiteLLM" + ) + + @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False) + def test_external_span_from_global_context(self): + """ + Test external span detection from global context (Priority 3 in _get_span_context). + + This simulates the case where a span is set in the global context + (e.g., by user code or Langfuse SDK) and LiteLLM detects it via + trace.get_current_span(). + + Expected behavior: + - LiteLLM detects the span from global context + - External span is not closed + - Correct parent-child relationship + """ + # Initialize OpenTelemetry + otel = OpenTelemetry(tracer_provider=self.tracer_provider) + + # Load test data + kwargs, response_obj = self._create_test_kwargs_and_response() + + # Create external parent span and set it as current using our test TracerProvider + tracer = self.tracer_provider.get_tracer(__name__) + + with tracer.start_as_current_span("external_global_span") as parent_span: + parent_ctx = parent_span.get_span_context() + parent_trace_id = parent_ctx.trace_id + + # Verify the span is in global context + current_span = trace.get_current_span() + self.assertEqual(current_span, parent_span, "Span should be in global context") + + # Make completion call + start_time = datetime.utcnow() + end_time = start_time + timedelta(seconds=1) + otel._handle_success(kwargs, response_obj, start_time, end_time) + + # Verify parent span is still recording + self.assertTrue( + parent_span.is_recording(), + "External span from global context should not be closed" + ) + + # Verify trace structure + spans = self.span_exporter.get_finished_spans() + for span in spans: + self.assertEqual( + span.context.trace_id, + parent_trace_id, + "All spans should have the same trace_id" + ) + + @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False) + def test_external_span_hierarchy_preserved(self): + """ + Test that span hierarchy is correctly preserved with external parent. + + Expected behavior: + - Parent span IDs are correct + - Trace structure matches expected hierarchy + - Span names are correct + """ + # Initialize OpenTelemetry + otel = OpenTelemetry(tracer_provider=self.tracer_provider) + otel.message_logging = True # Enable message logging to get raw_gen_ai_request spans + + # Load test data + kwargs, response_obj = self._create_test_kwargs_and_response() + + # Create external parent span using our test TracerProvider + tracer = self.tracer_provider.get_tracer(__name__) + + with tracer.start_as_current_span("external_parent_span") as parent_span: + parent_span_id = parent_span.get_span_context().span_id + + # Make completion call + start_time = datetime.utcnow() + end_time = start_time + timedelta(seconds=1) + otel._handle_success(kwargs, response_obj, start_time, end_time) + + # Verify hierarchy + spans = self.span_exporter.get_finished_spans() + + # Get spans by name + parent_spans = self._get_spans_by_name("external_parent_span") + raw_spans = self._get_spans_by_name("raw_gen_ai_request") + + self.assertEqual(len(parent_spans), 1, "Should have one parent span") + + # Verify parent-child relationship + if raw_spans: # If message_logging is on + for raw_span in raw_spans: + self.assertEqual( + raw_span.parent.span_id if raw_span.parent else None, + parent_span_id, + "raw_gen_ai_request should be child of external_parent_span" + ) + + @patch.dict(os.environ, {"USE_OTEL_LITELLM_REQUEST_SPAN": "false"}, clear=False) + def test_external_span_not_ended_on_failure(self): + """ + Test that external spans are not closed even on failure. + + Expected behavior: + - When _handle_failure is called with external span context + - External span remains open (is_recording = True) + - Error span is created correctly + - External span status is NOT changed by LiteLLM + """ + # Initialize OpenTelemetry + otel = OpenTelemetry(tracer_provider=self.tracer_provider) + + # Load test data + kwargs, response_obj = self._create_test_kwargs_and_response() + + # Create external parent span using our test TracerProvider + tracer = self.tracer_provider.get_tracer(__name__) + + with tracer.start_as_current_span("external_parent_span") as parent_span: + parent_ctx = parent_span.get_span_context() + parent_trace_id = parent_ctx.trace_id + + # Simulate failure + start_time = datetime.utcnow() + end_time = start_time + timedelta(seconds=1) + + # Create error response object + error_response = {"error": "Test error"} + + # Call _handle_failure + otel._handle_failure(kwargs, error_response, start_time, end_time) + + # Verify parent span is still recording + self.assertTrue( + parent_span.is_recording(), + "External span should still be recording even after failure" + ) + + # Verify trace structure + spans = self.span_exporter.get_finished_spans() + + # All spans should have the same trace_id + for span in spans: + self.assertEqual( + span.context.trace_id, + parent_trace_id, + "All spans should have the same trace_id even on failure" + ) + + # Should have external_parent_span + parent_spans = self._get_spans_by_name("external_parent_span") + self.assertEqual(len(parent_spans), 1, "Should have exactly one external_parent_span") + + # Verify LiteLLM set attributes on external parent span even on failure + parent_span_finished = parent_spans[0] + self.assertIn( + "gen_ai.request.model", + parent_span_finished.attributes, + "Parent span should have model attribute from LiteLLM even on failure" + ) diff --git a/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py b/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py new file mode 100644 index 00000000000..9658eff3cc5 --- /dev/null +++ b/tests/test_litellm/integrations/test_prometheus_metric_name_consistency.py @@ -0,0 +1,106 @@ +""" +Unit tests for prometheus metric name consistency + +This test ensures that the metric names used when creating Prometheus metrics +match the names defined in DEFINED_PROMETHEUS_METRICS, so that metric filtering +configuration works correctly. + +Related issue: https://github.com/BerriAI/litellm/issues/18221 +""" +from typing import get_args + +import pytest + + +def test_remaining_requests_metric_name_in_defined_metrics(): + """ + Test that litellm_remaining_requests_metric is defined in DEFINED_PROMETHEUS_METRICS. + + The metric name should include the _metric suffix to be consistent with the + configuration format users specify in prometheus_metrics_config. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + assert ( + "litellm_remaining_requests_metric" in defined_metrics + ), "litellm_remaining_requests_metric should be in DEFINED_PROMETHEUS_METRICS" + + +def test_remaining_tokens_metric_name_in_defined_metrics(): + """ + Test that litellm_remaining_tokens_metric is defined in DEFINED_PROMETHEUS_METRICS. + + The metric name should include the _metric suffix to be consistent with the + configuration format users specify in prometheus_metrics_config. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + assert ( + "litellm_remaining_tokens_metric" in defined_metrics + ), "litellm_remaining_tokens_metric should be in DEFINED_PROMETHEUS_METRICS" + + +def test_prometheus_metric_labels_have_remaining_metrics(): + """ + Test that PrometheusMetricLabels has label definitions for remaining metrics. + + This ensures that the labels can be retrieved when creating the metrics. + """ + from litellm.types.integrations.prometheus import PrometheusMetricLabels + + # Test that labels can be retrieved for remaining metrics + remaining_requests_labels = PrometheusMetricLabels.get_labels( + "litellm_remaining_requests_metric" + ) + remaining_tokens_labels = PrometheusMetricLabels.get_labels( + "litellm_remaining_tokens_metric" + ) + + assert isinstance( + remaining_requests_labels, list + ), "Labels for litellm_remaining_requests_metric should be a list" + assert isinstance( + remaining_tokens_labels, list + ), "Labels for litellm_remaining_tokens_metric should be a list" + + # These metrics should have api_provider and api_base labels + assert ( + "api_provider" in remaining_requests_labels + ), "litellm_remaining_requests_metric should have api_provider label" + assert ( + "api_base" in remaining_requests_labels + ), "litellm_remaining_requests_metric should have api_base label" + assert ( + "api_provider" in remaining_tokens_labels + ), "litellm_remaining_tokens_metric should have api_provider label" + assert ( + "api_base" in remaining_tokens_labels + ), "litellm_remaining_tokens_metric should have api_base label" + + +def test_all_defined_metrics_have_consistent_naming(): + """ + Test that all metrics defined in DEFINED_PROMETHEUS_METRICS follow + a consistent naming convention. + + This helps prevent similar inconsistencies in the future. + """ + from litellm.types.integrations.prometheus import DEFINED_PROMETHEUS_METRICS + + defined_metrics = get_args(DEFINED_PROMETHEUS_METRICS) + + for metric_name in defined_metrics: + # All metrics should start with 'litellm_' + assert metric_name.startswith( + "litellm_" + ), f"Metric {metric_name} should start with 'litellm_'" + + +if __name__ == "__main__": + test_remaining_requests_metric_name_in_defined_metrics() + test_remaining_tokens_metric_name_in_defined_metrics() + test_prometheus_metric_labels_have_remaining_metrics() + test_all_defined_metrics_have_consistent_naming() + print("All prometheus metric name consistency tests passed!") diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 65e3dbec8bd..5ba78d9eed1 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -809,3 +809,54 @@ def test_bedrock_anthropic_prompt_caching(): assert completion_cost >= 0 assert round(prompt_cost, 3) == 0.111 assert round(completion_cost, 5) == 0.00820 + + +def test_reasoning_tokens_without_text_tokens_gpt5_nano(): + """ + Test fix for GitHub issue #18599: + https://github.com/BerriAI/litellm/issues/18599 + + When OpenAI models (gpt-5-nano, o1, o3) return reasoning_tokens but don't provide + text_tokens, LiteLLM should calculate text_tokens as: + text_tokens = completion_tokens - reasoning_tokens - audio_tokens - image_tokens + + This ensures ALL completion tokens are billed, not just reasoning tokens. + """ + model = "gpt-5-nano" + custom_llm_provider = "openai" + + # Simulate OpenAI gpt-5-nano response where text_tokens is NOT provided + # completion_tokens: 977 total + # reasoning_tokens: 768 + # text_tokens: should be calculated as 977 - 768 = 209 + usage = Usage( + prompt_tokens=17, + completion_tokens=977, + total_tokens=994, + completion_tokens_details=CompletionTokensDetailsWrapper( + reasoning_tokens=768, + audio_tokens=0, + # text_tokens NOT provided - this is the key part of the bug + ), + ) + + prompt_cost, completion_cost = generic_cost_per_token( + model=model, + usage=usage, + custom_llm_provider=custom_llm_provider, + ) + + # gpt-5-nano pricing: $0.05/1M input, $0.40/1M output + expected_prompt_cost = 17 * 0.05 / 1_000_000 + expected_completion_cost = 977 * 0.40 / 1_000_000 # ALL tokens, not just reasoning + + assert abs(prompt_cost - expected_prompt_cost) < 1e-10, \ + f"Prompt cost incorrect: {prompt_cost} vs {expected_prompt_cost}" + + assert abs(completion_cost - expected_completion_cost) < 1e-10, \ + f"Completion cost incorrect: {completion_cost} vs {expected_completion_cost}" + + # Verify it's NOT using only reasoning_tokens (the bug) + wrong_cost = 768 * 0.40 / 1_000_000 # Only reasoning tokens + assert abs(completion_cost - wrong_cost) > 1e-6, \ + "Bug detected: Cost calculation is using only reasoning_tokens instead of all completion_tokens!" diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 41ac893b4d7..c8fe6efeaa1 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -497,6 +497,57 @@ def test_convert_gemini_messages(): ) +def test_convert_gemini_tool_call_result_with_image_url(): + """ + Test that image_url content type in tool results is handled correctly for Gemini. + Fixes: https://github.com/BerriAI/litellm/issues/18187 + """ + from litellm.litellm_core_utils.prompt_templates.factory import ( + convert_to_gemini_tool_call_result, + ) + from litellm.types.llms.openai import ChatCompletionToolMessage + + # Test with string image_url format + message_str_format = ChatCompletionToolMessage( + role="tool", + tool_call_id="call_123", + content=[{"type": "image_url", "image_url": "data:image/jpeg;base64,/9j/4AAQ"}], + ) + last_message_with_tool_calls = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "index": 0, + "function": {"name": "get_image", "arguments": "{}"}, + } + ], + } + + result = convert_to_gemini_tool_call_result( + message=message_str_format, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + # Should have inline_data for the image + assert isinstance(result, list) and any("inline_data" in p for p in result) + + # Test with dict image_url format (OpenAI standard) + message_dict_format = ChatCompletionToolMessage( + role="tool", + tool_call_id="call_456", + content=[{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,/9j/4AAQ"}}], + ) + last_message_with_tool_calls["tool_calls"][0]["id"] = "call_456" + + result2 = convert_to_gemini_tool_call_result( + message=message_dict_format, + last_message_with_tool_calls=last_message_with_tool_calls, + ) + assert isinstance(result2, list) and any("inline_data" in p for p in result2) + + def test_bedrock_tools_unpack_defs(): """ Test that the unpack_defs method handles nested $ref inside anyOf items correctly diff --git a/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py new file mode 100644 index 00000000000..1a6ed51afd0 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_codestral_provider_routing.py @@ -0,0 +1,69 @@ +""" +Unit tests for codestral provider routing. + +These tests verify that the chat and FIM endpoints for codestral +are correctly routed to different providers: +- Chat endpoint -> codestral provider +- FIM endpoint -> text-completion-codestral provider + +Related issue: https://github.com/BerriAI/litellm/issues/18464 +""" +import pytest + +import litellm + + +class TestCodestralProviderRouting: + """Tests for codestral endpoint routing in get_llm_provider""" + + def test_codestral_chat_endpoint_routes_to_codestral_provider(self): + """ + Test that the codestral chat endpoint routes to the 'codestral' provider. + + The chat/completions endpoint should be handled by the codestral provider. + """ + model, custom_llm_provider, _, api_base = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/chat/completions", + ) + + assert custom_llm_provider == "codestral" + + def test_codestral_fim_endpoint_routes_to_text_completion_provider(self): + """ + Test that the codestral FIM endpoint routes to 'text-completion-codestral'. + + The fim/completions endpoint should be handled by the + text-completion-codestral provider for fill-in-the-middle completions. + """ + model, custom_llm_provider, _, api_base = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/fim/completions", + ) + + assert custom_llm_provider == "text-completion-codestral" + + def test_codestral_endpoints_are_different_providers(self): + """ + Test that chat and FIM endpoints route to different providers. + + This is the core fix for issue #18464 - previously both endpoints + would route to 'codestral' due to duplicate conditions. + """ + _, chat_provider, _, _ = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/chat/completions", + ) + + _, fim_provider, _, _ = litellm.get_llm_provider( + model="codestral-latest", + api_base="https://codestral.mistral.ai/v1/fim/completions", + ) + + assert chat_provider != fim_provider + assert chat_provider == "codestral" + assert fim_provider == "text-completion-codestral" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/litellm_core_utils/test_core_helpers.py b/tests/test_litellm/litellm_core_utils/test_core_helpers.py index 89cc11c40d2..cd9c401143e 100644 --- a/tests/test_litellm/litellm_core_utils/test_core_helpers.py +++ b/tests/test_litellm/litellm_core_utils/test_core_helpers.py @@ -1,173 +1,45 @@ -import json -import os -import sys -from unittest.mock import MagicMock, patch +"""Tests for litellm_core_utils.core_helpers module.""" -import pytest - -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path - -from litellm.litellm_core_utils.core_helpers import ( - get_litellm_metadata_from_kwargs, - safe_divide, - safe_deep_copy -) +from litellm.litellm_core_utils.core_helpers import reconstruct_model_name -def test_get_litellm_metadata_from_kwargs(): - kwargs = { - "litellm_params": { - "litellm_metadata": {}, - "metadata": {"user_api_key": "1234567890"}, - }, - } - assert get_litellm_metadata_from_kwargs(kwargs) == {"user_api_key": "1234567890"} +def test_reconstruct_model_name_prefers_deployment_value(): + """Ensure deployment metadata wins when reconstructing the model name.""" + metadata = {"deployment": "vertex_ai/gemini-1.5-flash"} -def test_add_missing_spend_metadata_to_litellm_metadata(): - litellm_metadata = {"test_key": "test_value"} - metadata = {"user_api_key_hash_value": "1234567890"} - kwargs = { - "litellm_params": { - "litellm_metadata": litellm_metadata, - "metadata": metadata, - }, - } - assert get_litellm_metadata_from_kwargs(kwargs) == { - "test_key": "test_value", - "user_api_key_hash_value": "1234567890", - } - - -def test_preserve_upstream_non_openai_attributes(): - from litellm.litellm_core_utils.core_helpers import ( - preserve_upstream_non_openai_attributes, - ) - from litellm.types.utils import ModelResponseStream - - model_response = ModelResponseStream( - id="123", - object="text_completion", - created=1715811200, - model="gpt-3.5-turbo", + result = reconstruct_model_name( + model_name="gemini-1.5-flash", + custom_llm_provider="vertex_ai", + metadata=metadata, ) - setattr(model_response, "test_key", "test_value") - preserve_upstream_non_openai_attributes( - model_response=ModelResponseStream(), - original_chunk=model_response, + assert result == "vertex_ai/gemini-1.5-flash" + + +def test_reconstruct_model_name_adds_bedrock_prefix_when_missing(): + """Bedrock model names without prefixes should gain the provider prefix.""" + + metadata = {} + + result = reconstruct_model_name( + model_name="us.anthropic.claude-3-sonnet", + custom_llm_provider="bedrock", + metadata=metadata, ) - assert model_response.test_key == "test_value" + assert result == "bedrock/us.anthropic.claude-3-sonnet" -def test_safe_divide_basic(): - """Test basic safe division functionality""" - # Normal division - result = safe_divide(10, 2) - assert result == 5.0, f"Expected 5.0, got {result}" - - # Division with float - result = safe_divide(7.5, 2.5) - assert result == 3.0, f"Expected 3.0, got {result}" - - # Division by zero with default - result = safe_divide(10, 0) - assert result == 0, f"Expected 0, got {result}" - - # Division by zero with custom default - result = safe_divide(10, 0, default=1) - assert result == 1, f"Expected 1, got {result}" - - # Division by zero with custom default as float - result = safe_divide(10, 0, default=0.5) - assert result == 0.5, f"Expected 0.5, got {result}" +def test_reconstruct_model_name_returns_original_for_other_providers(): + """Non-Bedrock providers should not prepend anything.""" + metadata = {} -def test_safe_divide_edge_cases(): - """Test edge cases for safe division""" - # Zero numerator - result = safe_divide(0, 5) - assert result == 0.0, f"Expected 0.0, got {result}" - - # Negative numbers - result = safe_divide(-10, 2) - assert result == -5.0, f"Expected -5.0, got {result}" - - # Negative denominator - result = safe_divide(10, -2) - assert result == -5.0, f"Expected -5.0, got {result}" - - # Both negative - result = safe_divide(-10, -2) - assert result == 5.0, f"Expected 5.0, got {result}" - - # Float division - result = safe_divide(1, 3) - assert abs(result - 0.3333333333333333) < 1e-10, f"Expected ~0.333..., got {result}" + result = reconstruct_model_name( + model_name="claude-3-sonnet", + custom_llm_provider="anthropic", + metadata=metadata, + ) - -def test_safe_divide_weight_scenario(): - """Test safe division in the context of weight calculations""" - # Simulate weight calculation scenario - weights = [3, 7, 0, 2] - total_weight = sum(weights) # 12 - - # Normal case - normalized_weights = [safe_divide(w, total_weight) for w in weights] - expected = [0.25, 7/12, 0.0, 1/6] - - for i, (actual, exp) in enumerate(zip(normalized_weights, expected)): - assert abs(actual - exp) < 1e-10, f"Weight {i}: Expected {exp}, got {actual}" - - # Zero total weight scenario (division by zero) - zero_weights = [0, 0, 0] - zero_total = sum(zero_weights) # 0 - - # Should return default values (0) for all weights - normalized_zero_weights = [safe_divide(w, zero_total) for w in zero_weights] - expected_zero = [0, 0, 0] - - assert normalized_zero_weights == expected_zero, f"Expected {expected_zero}, got {normalized_zero_weights}" - - -def test_safe_deep_copy_with_non_pickleables_and_span(): - """ - Verify safe_deep_copy: - - does not crash when non-pickleables are present, - - preserves structure/keys, - - deep-copies JSON-y payloads (e.g., messages), - - keeps non-pickleables by reference, - - redacts OTEL span in the copy and restores it in the original. - """ - import threading - rlock = threading.RLock() - data = { - "metadata": {"litellm_parent_otel_span": rlock, "x": 1}, - "messages": [{"role": "user", "content": "hi"}], - "optional_params": {"handle": rlock}, - "ok": True, - } - - copied = safe_deep_copy(data) - - # Structure preserved - assert set(copied.keys()) == set(data.keys()) - - # Messages are deep-copied (new object, same content) - assert copied["messages"] is not data["messages"] - assert copied["messages"][0] == data["messages"][0] - - # Non-pickleable subtree kept by reference (no crash) - assert copied["optional_params"] is data["optional_params"] - assert copied["optional_params"]["handle"] is rlock - - # OTEL span: redacted in the copy, restored in original - assert copied["metadata"]["litellm_parent_otel_span"] == "placeholder" - assert data["metadata"]["litellm_parent_otel_span"] is rlock - - # Other simple fields unchanged - assert copied["ok"] is True - assert copied["metadata"]["x"] == 1 + assert result == "claude-3-sonnet" diff --git a/tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py b/tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py new file mode 100644 index 00000000000..6940a5ea7a5 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_dot_notation_indexing.py @@ -0,0 +1,138 @@ +""" +Tests for litellm.litellm_core_utils.dot_notation_indexing module. +""" + +import pytest + +from litellm.litellm_core_utils.dot_notation_indexing import ( + get_nested_value, + delete_nested_value, +) + + +class TestGetNestedValue: + """Tests for get_nested_value function.""" + + def test_simple_key(self): + """Test accessing a simple top-level key.""" + data = {"name": "test"} + assert get_nested_value(data, "name") == "test" + + def test_nested_key(self): + """Test accessing nested keys with dot notation.""" + data = {"a": {"b": {"c": "value"}}} + assert get_nested_value(data, "a.b.c") == "value" + + def test_missing_key_returns_default(self): + """Test that missing keys return the default value.""" + data = {"a": {"b": "value"}} + assert get_nested_value(data, "a.b", "default") == "value" + assert get_nested_value(data, "a.c", "default") == "default" + assert get_nested_value(data, "x.y.z") is None + + def test_empty_key_path(self): + """Test that empty key path returns default.""" + data = {"a": "value"} + assert get_nested_value(data, "") is None + assert get_nested_value(data, "", "default") == "default" + + def test_metadata_prefix_removal(self): + """Test that metadata. prefix is properly removed.""" + data = {"user": {"email": "test@example.com"}} + assert get_nested_value(data, "metadata.user.email") == "test@example.com" + + def test_escaped_dot_in_key(self): + """Test accessing keys that contain dots using escape sequence.""" + data = {"kubernetes.io": {"namespace": "default"}} + assert get_nested_value(data, "kubernetes\\.io.namespace") == "default" + + def test_escaped_dot_nested(self): + """Test multiple levels with escaped dots.""" + data = { + "kubernetes.io": { + "pod.info": { + "name": "my-pod" + } + } + } + assert get_nested_value(data, "kubernetes\\.io.pod\\.info.name") == "my-pod" + + def test_kubernetes_jwt_example(self): + """Test with a realistic Kubernetes JWT structure.""" + jwt_token = { + "aud": ["https://kubernetes.default.svc"], + "exp": "1234567890", + "iat": "123456789", + "iss": "https://oidc.eks.region.amazonaws.com/id/randomstring", + "jti": "randomstring", + "kubernetes.io": { + "namespace": "namespace", + "node": { + "name": "node-name", + "uid": "node-uid" + }, + "pod": { + "name": "pod-name", + "uid": "pod-uid" + }, + "serviceaccount": { + "name": "serviceaccount-name", + "uid": "serviceaccount-uid" + }, + "warnafter": 1234567880 + }, + "nbf": 123456789, + "sub": "system:serviceaccount:namespace:serviceaccount-name" + } + + # Test accessing kubernetes.io.namespace + assert get_nested_value(jwt_token, "kubernetes\\.io.namespace") == "namespace" + + # Test accessing nested values within kubernetes.io + assert get_nested_value(jwt_token, "kubernetes\\.io.pod.name") == "pod-name" + assert get_nested_value(jwt_token, "kubernetes\\.io.serviceaccount.name") == "serviceaccount-name" + + # Test accessing regular keys still works + assert get_nested_value(jwt_token, "sub") == "system:serviceaccount:namespace:serviceaccount-name" + + def test_mixed_escaped_and_regular_dots(self): + """Test path with both escaped dots (in keys) and regular dots (separators).""" + data = { + "config.v1": { + "settings": { + "feature.enabled": True + } + } + } + assert get_nested_value(data, "config\\.v1.settings.feature\\.enabled") is True + + +class TestDeleteNestedValue: + """Tests for delete_nested_value function.""" + + def test_delete_simple_key(self): + """Test deleting a simple top-level key.""" + data = {"a": 1, "b": 2} + result = delete_nested_value(data, "a") + assert result == {"b": 2} + # Original should be unchanged + assert data == {"a": 1, "b": 2} + + def test_delete_nested_key(self): + """Test deleting a nested key.""" + data = {"a": {"b": {"c": 1, "d": 2}}} + result = delete_nested_value(data, "a.b.c") + assert result == {"a": {"b": {"d": 2}}} + + def test_delete_array_wildcard(self): + """Test deleting a field from all array elements.""" + data = {"tools": [{"name": "t1", "secret": "s1"}, {"name": "t2", "secret": "s2"}]} + result = delete_nested_value(data, "tools[*].secret") + assert result == {"tools": [{"name": "t1"}, {"name": "t2"}]} + + def test_delete_array_index(self): + """Test deleting a field from a specific array element.""" + data = {"items": [{"a": 1, "b": 2}, {"a": 3, "b": 4}]} + result = delete_nested_value(data, "items[0].b") + assert result == {"items": [{"a": 1}, {"a": 3, "b": 4}]} + diff --git a/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py new file mode 100644 index 00000000000..b17c02d7006 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py @@ -0,0 +1,156 @@ +""" +Unit tests for _extract_base64_data and extract_images_from_message functions. + +These tests verify that base64 image data is correctly extracted from data URLs, +which fixes the Ollama error "illegal base64 data at input byte 4". + +Related issue: https://github.com/BerriAI/litellm/issues/18338 +""" +import pytest + +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _extract_base64_data, + extract_images_from_message, +) + + +class TestExtractBase64Data: + """Tests for _extract_base64_data function""" + + def test_extract_base64_from_png_data_url(self): + """Test extracting base64 data from a PNG data URL""" + data_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + expected = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + assert _extract_base64_data(data_url) == expected + + def test_extract_base64_from_jpeg_data_url(self): + """Test extracting base64 data from a JPEG data URL""" + data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD" + expected = "/9j/4AAQSkZJRgABAQAAAQABAAD" + assert _extract_base64_data(data_url) == expected + + def test_extract_base64_from_gif_data_url(self): + """Test extracting base64 data from a GIF data URL""" + data_url = "data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP" + expected = "R0lGODlhAQABAIAAAAAAAP" + assert _extract_base64_data(data_url) == expected + + def test_regular_url_unchanged(self): + """Test that regular HTTP URLs are returned unchanged""" + url = "https://example.com/image.png" + assert _extract_base64_data(url) == url + + def test_file_path_unchanged(self): + """Test that file paths are returned unchanged""" + path = "/path/to/image.png" + assert _extract_base64_data(path) == path + + def test_data_url_without_base64_unchanged(self): + """Test that data URLs without base64 encoding are returned unchanged""" + # This is a data URL with URL encoding, not base64 + url = "data:text/plain,Hello%20World" + assert _extract_base64_data(url) == url + + def test_base64_data_with_special_chars(self): + """Test extracting base64 data that contains valid special characters""" + # Base64 can contain +, /, and = characters + data_url = "data:image/png;base64,abc+def/ghi===" + expected = "abc+def/ghi===" + assert _extract_base64_data(data_url) == expected + + +class TestExtractImagesFromMessage: + """Tests for extract_images_from_message function""" + + def test_extract_from_message_with_data_url_string(self): + """Test extracting images when image_url is a string data URL""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgo", + } + ], + } + result = extract_images_from_message(message) + assert result == ["iVBORw0KGgo"] + + def test_extract_from_message_with_data_url_dict(self): + """Test extracting images when image_url is a dict with url key""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo"}, + } + ], + } + result = extract_images_from_message(message) + assert result == ["iVBORw0KGgo"] + + def test_extract_from_message_with_regular_url(self): + """Test that regular URLs are preserved""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + } + ], + } + result = extract_images_from_message(message) + assert result == ["https://example.com/image.png"] + + def test_extract_multiple_images(self): + """Test extracting multiple images from a single message""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,image1base64", + }, + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,image2base64"}, + }, + { + "type": "image_url", + "image_url": "https://example.com/image3.png", + }, + ], + } + result = extract_images_from_message(message) + assert result == [ + "image1base64", + "image2base64", + "https://example.com/image3.png", + ] + + def test_empty_content(self): + """Test message with empty content""" + message = {"role": "user", "content": []} + result = extract_images_from_message(message) + assert result == [] + + def test_no_images_in_content(self): + """Test message with content but no images""" + message = { + "role": "user", + "content": [{"type": "text", "text": "Hello world"}], + } + result = extract_images_from_message(message) + assert result == [] + + def test_string_content(self): + """Test message with string content (no images possible)""" + message = {"role": "user", "content": "Hello world"} + result = extract_images_from_message(message) + assert result == [] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index 0fe15168467..a44b821db87 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -323,3 +323,40 @@ class TestLoggingWorker: await worker.clear_queue() assert len(processed) >= 4, f"Expected 4+ tasks processed, got {len(processed)}" + + @pytest.mark.asyncio + async def test_event_loop_change_handling(self): + """Test that LoggingWorker handles event loop changes correctly. + + This tests the fix for GitHub issue #17813 where asyncio.Queue + was bound to a different event loop when using multiprocessing. + """ + worker = LoggingWorker(timeout=1.0, max_queue_size=10) + + # Start the worker in the current event loop + worker.start() + + # Verify queue was created and bound to current loop + assert worker._queue is not None + assert worker._bound_loop is not None + original_queue = worker._queue + + await worker.stop() + + # Simulate a new event loop by creating a mock scenario + # In a real multiprocessing case, asyncio.run() creates a new loop + # We test the internal state detection + + # Create a new worker to test the _ensure_queue logic + worker2 = LoggingWorker(timeout=1.0, max_queue_size=10) + worker2._queue = original_queue # Pretend we have an old queue + worker2._bound_loop = None # No bound loop (simulates first call) + + # Calling start should create a new queue since _bound_loop != current + worker2.start() + + # The queue should be reinitialized since bound_loop was None + assert worker2._queue is not None + assert worker2._bound_loop is not None + + await worker2.stop() diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index d78a638fd89..bdced849c7e 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -55,11 +55,12 @@ class TestAzureAnthropicMessagesConfig: assert isinstance(call_args[1]["litellm_params"], GenericLiteLLMParams) assert call_args[1]["litellm_params"].api_key == "test-api-key" assert "anthropic-version" in result - # api-key header is preserved as-is (no conversion to x-api-key) - assert "api-key" in result + assert "x-api-key" in result + assert result["x-api-key"] == "test-api-key" + assert "api-key" not in result - def test_validate_anthropic_messages_environment_preserves_api_key_header(self): - """Test that api-key header is preserved as-is (Azure handles the header internally)""" + def test_validate_anthropic_messages_environment_converts_api_key_to_x_api_key(self): + """Test that api-key header is converted to x-api-key""" config = AzureAnthropicMessagesConfig() headers = {} model = "claude-sonnet-4-5" @@ -79,9 +80,10 @@ class TestAzureAnthropicMessagesConfig: litellm_params=litellm_params, ) - # Verify api-key header is preserved as-is - assert "api-key" in result - assert result["api-key"] == "test-api-key" + # Verify api-key was converted to x-api-key + assert "x-api-key" in result + assert result["x-api-key"] == "test-api-key" + assert "api-key" not in result def test_validate_anthropic_messages_environment_sets_headers(self): """Test that required headers are set""" @@ -108,8 +110,7 @@ class TestAzureAnthropicMessagesConfig: assert result["anthropic-version"] == "2023-06-01" assert "content-type" in result assert result["content-type"] == "application/json" - # api-key header is preserved as-is - assert "api-key" in result + assert "x-api-key" in result def test_get_complete_url_with_base_url(self): """Test get_complete_url with base URL""" diff --git a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py index 2dcec689895..a8ac680908e 100644 --- a/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py +++ b/tests/test_litellm/llms/bedrock/rerank/test_bedrock_rerank_header_forwarding.py @@ -8,12 +8,13 @@ forward_client_headers_to_llm_api were not being passed to Bedrock rerank provid import json import os import sys -from unittest.mock import Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest sys.path.insert(0, os.path.abspath("../../../../..")) # Adds the parent directory to the system path import litellm +from litellm.llms.bedrock.base_aws_llm import Boto3CredentialsInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler # Mock response for Bedrock rerank @@ -47,6 +48,19 @@ test_documents = [ ] +def create_mock_credentials(): + """Create mock AWS credentials for testing""" + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + return Boto3CredentialsInfo( + credentials=mock_credentials, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + ) + + @pytest.mark.parametrize( "model", [ @@ -73,7 +87,17 @@ def test_bedrock_rerank_header_forwarding_sync(model): "X-Test-Header": "test-value", } - with patch.object(client, "post") as mock_post: + # Mock AWS credentials and SigV4 auth + mock_credentials_info = create_mock_credentials() + + with patch.object(client, "post") as mock_post, \ + patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \ + patch("botocore.auth.SigV4Auth") as mock_sigv4: + + # Mock SigV4Auth to not actually sign the request + mock_sigv4_instance = MagicMock() + mock_sigv4.return_value = mock_sigv4_instance + mock_response = Mock() mock_response.status_code = 200 mock_response.text = json.dumps(bedrock_rerank_response) @@ -152,9 +176,17 @@ async def test_bedrock_rerank_header_forwarding_async(model): "X-Test-Header": "test-value", } - from unittest.mock import AsyncMock + # Mock AWS credentials and SigV4 auth + mock_credentials_info = create_mock_credentials() - with patch.object(client, "post", new_callable=AsyncMock) as mock_post: + with patch.object(client, "post", new_callable=AsyncMock) as mock_post, \ + patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \ + patch("botocore.auth.SigV4Auth") as mock_sigv4: + + # Mock SigV4Auth to not actually sign the request + mock_sigv4_instance = MagicMock() + mock_sigv4.return_value = mock_sigv4_instance + mock_response = AsyncMock() mock_response.status_code = 200 mock_response.text = json.dumps(bedrock_rerank_response) @@ -223,7 +255,17 @@ def test_bedrock_rerank_extra_headers_and_headers_merge(): # Explicit extra_headers explicit_headers = {"X-Explicit-Header": "ExplicitValue"} - with patch.object(client, "post") as mock_post: + # Mock AWS credentials and SigV4 auth + mock_credentials_info = create_mock_credentials() + + with patch.object(client, "post") as mock_post, \ + patch("litellm.llms.bedrock.rerank.handler.BedrockRerankHandler._get_boto_credentials_from_optional_params", return_value=mock_credentials_info), \ + patch("botocore.auth.SigV4Auth") as mock_sigv4: + + # Mock SigV4Auth to not actually sign the request + mock_sigv4_instance = MagicMock() + mock_sigv4.return_value = mock_sigv4_instance + mock_response = Mock() mock_response.status_code = 200 mock_response.text = json.dumps(bedrock_rerank_response) diff --git a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py index b4dc6c68bb0..800066ac5bf 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py +++ b/tests/test_litellm/llms/databricks/test_databricks_partner_integration.py @@ -122,10 +122,11 @@ class TestRedactSensitiveData: def test_redact_pat_token(self): """Databricks PAT tokens are redacted.""" + test_token = "dapiTESTTOKENFAKEVALUEFORTESTINGPURPOSESONLY123" result = DatabricksBase.redact_sensitive_data( - "Using token dapi_fake_test_token_value" + f"Using token {test_token}" ) - assert "dapi_fake_test_token_value" not in result + assert test_token not in result assert "[REDACTED_PAT]" in result def test_redact_client_secret(self): @@ -347,17 +348,27 @@ class TestSDKPartnerTelemetry: "Authorization": "Bearer token" } - with patch( - "databricks.sdk.WorkspaceClient", return_value=mock_workspace_client - ): - with patch("databricks.sdk.useragent.with_partner") as mock_with_partner: - databricks_base._get_databricks_credentials( - api_key=None, - api_base=None, - headers=None, - ) + mock_useragent = MagicMock() + # Create a mock databricks.sdk module to simulate the SDK being available + # This allows us to test the partner telemetry registration without requiring + # the actual databricks-sdk package to be installed + mock_sdk_module = MagicMock() + mock_sdk_module.WorkspaceClient = MagicMock(return_value=mock_workspace_client) + mock_sdk_module.useragent = mock_useragent + + # Mock both databricks and databricks.sdk modules to ensure the import works + with patch.dict(sys.modules, { + "databricks": MagicMock(), + "databricks.sdk": mock_sdk_module + }): + databricks_base._get_databricks_credentials( + api_key=None, + api_base=None, + headers=None, + ) - mock_with_partner.assert_called_once_with("litellm") + # Verify that partner telemetry registration was called correctly + mock_useragent.with_partner.assert_called_once_with("litellm") class TestUserAgentFromEnvironment: @@ -592,19 +603,28 @@ class TestAuthenticationPriority: "Authorization": "Bearer sdk-token" } - with patch( - "databricks.sdk.WorkspaceClient", return_value=mock_workspace_client - ): - with patch("databricks.sdk.useragent.with_partner"): - api_base, headers = databricks_base.databricks_validate_environment( - api_key=None, - api_base=None, - endpoint_type="chat_completions", - custom_endpoint=False, - headers=None, - ) + # Create a mock databricks.sdk module to simulate the SDK being available + # This allows us to test the SDK fallback authentication without requiring + # the actual databricks-sdk package to be installed + mock_sdk_module = MagicMock() + mock_sdk_module.WorkspaceClient = MagicMock(return_value=mock_workspace_client) + mock_sdk_module.useragent = MagicMock() + + # Mock both databricks and databricks.sdk modules to ensure the import works + with patch.dict(sys.modules, { + "databricks": MagicMock(), + "databricks.sdk": mock_sdk_module + }): + api_base, headers = databricks_base.databricks_validate_environment( + api_key=None, + api_base=None, + endpoint_type="chat_completions", + custom_endpoint=False, + headers=None, + ) - assert "Authorization" in headers + # Verify that SDK authentication was used (headers contain Authorization) + assert "Authorization" in headers class TestEndpointURLConstruction: diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 783d85f471b..91a28ee6ec9 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -2279,3 +2279,194 @@ def test_partial_json_chunk_on_first_chunk(): assert result is None, "Partial first chunk should return None" assert iterator.chunk_type == "accumulated_json", "Should switch to accumulated_json mode" + +# ==================== Tool Type Separation Tests ==================== +# These tests verify that each Tool object contains exactly one type per Vertex AI API spec +# Ref: https://cloud.google.com/vertex-ai/generative-ai/docs/reference/rest/v1beta1/Tool + + +def test_vertex_ai_multiple_tool_types_separate_objects(): + """ + Test that multiple tool types are placed in separate Tool objects. + + This is required by Vertex AI API spec: + "A Tool object should contain exactly one type of Tool" + + Related error without this fix: + "tools[0].tool_type: one_of 'tool_type' has more than one initialized field: + enterprise_web_search, url_context" + + Input: + value=[ + {"enterpriseWebSearch": {}}, + {"url_context": {}}, + ] + + Expected Output: + tools=[ + {"enterpriseWebSearch": {}}, # First Tool object + {"url_context": {}}, # Second Tool object (separate!) + ] + + NOT (incorrect - causes API error): + tools=[ + {"enterpriseWebSearch": {}, "url_context": {}} # Multiple types in one object + ] + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[ + {"enterpriseWebSearch": {}}, + {"url_context": {}}, + ], + optional_params=optional_params + ) + + # Should have 2 separate Tool objects + assert len(tools) == 2, f"Expected 2 separate Tool objects, got {len(tools)}" + + # Each Tool object should contain exactly ONE type + tool_types_in_first = [k for k in tools[0].keys()] + tool_types_in_second = [k for k in tools[1].keys()] + + assert len(tool_types_in_first) == 1, f"First Tool should have exactly 1 type, got {tool_types_in_first}" + assert len(tool_types_in_second) == 1, f"Second Tool should have exactly 1 type, got {tool_types_in_second}" + + # Verify the correct tool types are present + assert "enterpriseWebSearch" in tools[0], "First Tool should contain enterpriseWebSearch" + assert "url_context" in tools[1], "Second Tool should contain url_context" + + +def test_vertex_ai_function_declarations_with_other_tools_separate(): + """ + Test that function declarations and other tool types are in separate Tool objects. + + This ensures that when using both function calling AND special tools like + google_search or code_execution, they are properly separated per API spec. + + Input: + value=[ + {"type": "function", "function": {"name": "get_weather", "description": "Get weather"}}, + {"googleSearch": {}}, + {"code_execution": {}}, + ] + + Expected Output: + tools=[ + {"function_declarations": [{"name": "get_weather", "description": "Get weather"}]}, + {"googleSearch": {}}, + {"code_execution": {}}, + ] + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[ + {"type": "function", "function": {"name": "get_weather", "description": "Get weather"}}, + {"googleSearch": {}}, + {"code_execution": {}}, + ], + optional_params=optional_params + ) + + # Should have 3 separate Tool objects + assert len(tools) == 3, f"Expected 3 separate Tool objects, got {len(tools)}" + + # Find each tool type + func_tool = None + search_tool = None + code_tool = None + + for tool in tools: + if "function_declarations" in tool: + func_tool = tool + elif "googleSearch" in tool: + search_tool = tool + elif "code_execution" in tool: + code_tool = tool + + # Verify all tools are present and separate + assert func_tool is not None, "function_declarations Tool should be present" + assert search_tool is not None, "googleSearch Tool should be present" + assert code_tool is not None, "code_execution Tool should be present" + + # Verify each Tool has exactly one type + assert len(func_tool.keys()) == 1, "function_declarations Tool should have only one key" + assert len(search_tool.keys()) == 1, "googleSearch Tool should have only one key" + assert len(code_tool.keys()) == 1, "code_execution Tool should have only one key" + + # Verify function declaration content + assert func_tool["function_declarations"][0]["name"] == "get_weather" + + +def test_vertex_ai_single_tool_type_still_works(): + """ + Test that single tool type usage still works correctly (backward compatibility). + + Input: + value=[{"code_execution": {}}] + + Expected Output: + tools=[{"code_execution": {}}] + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[{"code_execution": {}}], + optional_params=optional_params + ) + + assert len(tools) == 1 + assert "code_execution" in tools[0] + assert tools[0]["code_execution"] == {} + + +def test_vertex_ai_multiple_function_declarations_grouped(): + """ + Test that multiple function declarations are grouped in ONE Tool object. + + Function declarations are the exception - they CAN be grouped together + in a single Tool object (up to 512 declarations). + + Input: + value=[ + {"type": "function", "function": {"name": "func1", "description": "First function"}}, + {"type": "function", "function": {"name": "func2", "description": "Second function"}}, + ] + + Expected Output: + tools=[ + { + "function_declarations": [ + {"name": "func1", "description": "First function"}, + {"name": "func2", "description": "Second function"}, + ] + } + ] + """ + v = VertexGeminiConfig() + optional_params = {} + + tools = v._map_function( + value=[ + {"type": "function", "function": {"name": "func1", "description": "First function"}}, + {"type": "function", "function": {"name": "func2", "description": "Second function"}}, + ], + optional_params=optional_params + ) + + # Should have only 1 Tool object (function declarations grouped) + assert len(tools) == 1, f"Expected 1 Tool object for grouped functions, got {len(tools)}" + + # Should contain function_declarations with 2 functions + assert "function_declarations" in tools[0] + assert len(tools[0]["function_declarations"]) == 2 + + # Verify function names + func_names = [f["name"] for f in tools[0]["function_declarations"]] + assert "func1" in func_names + assert "func2" in func_names diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py new file mode 100644 index 00000000000..fca784342d7 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_anthropic_image_url_handling.py @@ -0,0 +1,179 @@ +""" +Tests for Vertex AI Anthropic image URL handling. + +Issue: https://github.com/BerriAI/litellm/issues/18430 +Vertex AI Anthropic models don't support URL sources for images. +LiteLLM should convert image URLs to base64 when using Vertex AI Anthropic. +""" +import os +import sys +from unittest.mock import patch, MagicMock + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../../..") +) # Adds the parent directory to the system path + +from litellm.litellm_core_utils.prompt_templates.factory import ( + anthropic_messages_pt, + create_anthropic_image_param, +) + + +class TestVertexAIAnthropicImageURLHandling: + """Test that Vertex AI Anthropic converts image URLs to base64.""" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_vertex_ai_anthropic_converts_https_url_to_base64( + self, mock_convert_url: MagicMock + ): + """ + Test that HTTPS image URLs are converted to base64 for Vertex AI Anthropic. + + For regular Anthropic, HTTPS URLs are passed through as URL type. + For Vertex AI Anthropic, HTTPS URLs should be converted to base64. + """ + mock_convert_url.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQ==" + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.jpg"}, + }, + ], + } + ] + + # For Vertex AI, image URLs should be converted to base64 + result = anthropic_messages_pt( + messages=messages, + model="claude-sonnet-4", + llm_provider="vertex_ai", + ) + + # Verify convert_url_to_base64 was called + mock_convert_url.assert_called_once_with(url="https://example.com/image.jpg") + + # Check the result has base64 source type + user_message = result[0] + assert user_message["role"] == "user" + image_content = user_message["content"][1] + assert image_content["type"] == "image" + assert image_content["source"]["type"] == "base64" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_regular_anthropic_uses_url_type_for_https( + self, mock_convert_url: MagicMock + ): + """ + Test that regular Anthropic API uses URL type for HTTPS images. + + This confirms the original behavior is preserved for non-Vertex AI. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Describe this image"}, + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.jpg"}, + }, + ], + } + ] + + # For regular Anthropic, HTTPS URLs should NOT be converted + result = anthropic_messages_pt( + messages=messages, + model="claude-sonnet-4", + llm_provider="anthropic", + ) + + # convert_url_to_base64 should NOT be called for regular Anthropic with HTTPS + mock_convert_url.assert_not_called() + + # Check the result has URL source type + user_message = result[0] + assert user_message["role"] == "user" + image_content = user_message["content"][1] + assert image_content["type"] == "image" + assert image_content["source"]["type"] == "url" + assert image_content["source"]["url"] == "https://example.com/image.jpg" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_vertex_ai_beta_also_converts_to_base64( + self, mock_convert_url: MagicMock + ): + """ + Test that vertex_ai_beta provider also converts image URLs to base64. + """ + mock_convert_url.return_value = "data:image/png;base64,iVBORw0KGgo=" + + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + { + "type": "image_url", + "image_url": "https://example.com/photo.png", + }, + ], + } + ] + + result = anthropic_messages_pt( + messages=messages, + model="claude-3-opus", + llm_provider="vertex_ai_beta", + ) + + # Verify convert_url_to_base64 was called + mock_convert_url.assert_called_once() + + # Check the result has base64 source type + user_message = result[0] + image_content = user_message["content"][1] + assert image_content["source"]["type"] == "base64" + + +class TestCreateAnthropicImageParam: + """Test the create_anthropic_image_param function directly.""" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_force_base64_converts_https_url(self, mock_convert_url: MagicMock): + """ + Test that is_bedrock_invoke=True (used for both Bedrock and Vertex AI) + forces conversion of HTTPS URLs to base64. + """ + mock_convert_url.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRg==" + + result = create_anthropic_image_param( + image_url_input="https://example.com/image.jpg", + format=None, + is_bedrock_invoke=True, # This flag is set for both Bedrock and Vertex AI + ) + + mock_convert_url.assert_called_once_with(url="https://example.com/image.jpg") + assert result["source"]["type"] == "base64" + + @patch("litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64") + def test_no_force_uses_url_type(self, mock_convert_url: MagicMock): + """ + Test that without force, HTTPS URLs use URL type. + """ + result = create_anthropic_image_param( + image_url_input="https://example.com/image.jpg", + format=None, + is_bedrock_invoke=False, + ) + + mock_convert_url.assert_not_called() + assert result["source"]["type"] == "url" + assert result["source"]["url"] == "https://example.com/image.jpg" diff --git a/tests/test_litellm/llms/zai/test_zai_provider.py b/tests/test_litellm/llms/zai/test_zai_provider.py index a3d47d666bc..d1e4359d048 100644 --- a/tests/test_litellm/llms/zai/test_zai_provider.py +++ b/tests/test_litellm/llms/zai/test_zai_provider.py @@ -1,6 +1,7 @@ """ Tests for Z.AI (Zhipu AI) provider - GLM models """ + import json import math @@ -50,10 +51,12 @@ def test_zai_in_provider_lists(): def test_zai_models_in_model_cost(): """Test that ZAI models are in the model cost map""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") zai_models = [ + "zai/glm-4.7", "zai/glm-4.6", "zai/glm-4.5", "zai/glm-4.5v", @@ -72,6 +75,7 @@ def test_zai_models_in_model_cost(): def test_zai_glm46_cost_calculation(): """Test the cost calculation for glm-4.6""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -92,6 +96,7 @@ def test_zai_glm46_cost_calculation(): def test_zai_flash_model_is_free(): """Test that glm-4.5-flash has zero cost""" import os + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" litellm.model_cost = litellm.get_model_cost_map(url="") @@ -102,6 +107,38 @@ def test_zai_flash_model_is_free(): assert info["output_cost_per_token"] == 0 +def test_glm47_supports_reasoning(): + """Test that GLM-4.7 supports reasoning""" + import os + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + key = "zai/glm-4.7" + assert key in litellm.model_cost, f"Model {key} not found in model_cost" + + info = litellm.model_cost[key] + assert info["supports_reasoning"] is True + + +def test_glm47_cost_calculation(): + """Test cost calculation for GLM-4.7""" + import os + + os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" + litellm.model_cost = litellm.get_model_cost_map(url="") + + prompt_cost, completion_cost = cost_per_token( + model="zai/glm-4.7", + prompt_tokens=1000000, # 1M tokens + completion_tokens=1000000, + ) + + # GLM-4.7: $0.6/M input, $2.2/M output (same as GLM-4.6) + assert math.isclose(prompt_cost, 0.6, rel_tol=1e-6) + assert math.isclose(completion_cost, 2.2, rel_tol=1e-6) + + @pytest.mark.asyncio async def test_zai_completion_call(respx_mock, zai_response, monkeypatch): """Test completion call with zai provider using mocked response""" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py new file mode 100644 index 00000000000..0e150e064c7 --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py @@ -0,0 +1,78 @@ +"""Tests for the MCP guardrail translation handler.""" + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) + + +class MockGuardrail(CustomGuardrail): + """Simple guardrail mock that records invocations.""" + + def __init__(self, return_texts=None): + super().__init__(guardrail_name="mock-mcp-guardrail") + self.return_texts = return_texts + self.call_count = 0 + self.last_inputs = None + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.call_count += 1 + self.last_inputs = inputs + + if self.return_texts is not None: + return {"texts": self.return_texts} + + texts = inputs.get("texts", []) + return {"texts": [f"{text} [SAFE]" for text in texts]} + + +@pytest.mark.asyncio +async def test_process_input_messages_updates_content(): + """Handler should update the synthetic message content when guardrail modifies text.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MockGuardrail() + + original_content = "Tool: weather\nArguments: {'city': 'tokyo'}" + data = { + "messages": [{"role": "user", "content": original_content}], + "mcp_tool_name": "weather", + } + + result = await handler.process_input_messages(data, guardrail) + + assert result["messages"][0]["content"].endswith("[SAFE]") + assert guardrail.last_inputs == {"texts": [original_content]} + assert guardrail.call_count == 1 + + +@pytest.mark.asyncio +async def test_process_input_messages_skips_when_no_messages(): + """Handler should skip guardrail invocation if messages array is missing or empty.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MockGuardrail() + + data = {"mcp_tool_name": "noop"} + result = await handler.process_input_messages(data, guardrail) + + assert result == data + assert guardrail.call_count == 0 + + +@pytest.mark.asyncio +async def test_process_input_messages_handles_empty_guardrail_result(): + """Handler should leave content untouched when guardrail returns no text updates.""" + handler = MCPGuardrailTranslationHandler() + guardrail = MockGuardrail(return_texts=[]) + + original_content = "Tool: calendar\nArguments: {'date': '2024-12-25'}" + data = { + "messages": [{"role": "user", "content": original_content}], + "mcp_tool_name": "calendar", + } + + result = await handler.process_input_messages(data, guardrail) + + assert result["messages"][0]["content"] == original_content + assert guardrail.call_count == 1 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py index 5581070be71..a2425cc659a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_custom_fields.py @@ -71,28 +71,29 @@ class TestMCPCustomFields: manager = MCPServerManager() # Mock database record with custom fields - mock_server = Mock(spec=LiteLLM_MCPServerTable) - mock_server.server_id = "test-server-id" - mock_server.server_name = "Test Server" - mock_server.description = "A test server" - mock_server.url = "http://localhost:3000" - mock_server.transport = "http" - mock_server.auth_type = MCPAuth.bearer_token - mock_server.alias = None - mock_server.mcp_info = { - "server_name": "Test Server", - "description": "A test server", - "custom_db_field": "database_value", - "metadata": {"source": "database"}, - "version": "1.0.0" - } - mock_server.command = None - mock_server.args = None - mock_server.env = None - mock_server.mcp_access_groups = None + mock_server = LiteLLM_MCPServerTable( + server_id="test-server-id", + server_name="Test Server", + alias=None, + description="A test server", + url="http://localhost:3000", + transport="http", + auth_type=MCPAuth.bearer_token, + mcp_info={ + "server_name": "Test Server", + "description": "A test server", + "custom_db_field": "database_value", + "metadata": {"source": "database"}, + "version": "1.0.0", + }, + command=None, + args=[], + env={}, + mcp_access_groups=[], + ) # Add server to manager - await manager.add_update_server(mock_server) + await manager.add_server(mock_server) # Get the added server server = manager.get_mcp_server_by_id("test-server-id") @@ -209,4 +210,4 @@ class TestMCPCustomFields: # Should use mcp_info description, not config level assert mcp_info["description"] == "MCP info description" - assert mcp_info["custom_field"] == "custom_value" \ No newline at end of file + assert mcp_info["custom_field"] == "custom_value" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index ff016a1a130..6491e11024a 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -65,7 +65,7 @@ class TestMCPServerManager: updated_at=datetime.now(), ) - await manager.add_update_server(stdio_server) + await manager.add_server(stdio_server) # Verify server was added assert "stdio-server-1" in manager.registry @@ -641,37 +641,31 @@ class TestMCPServerManager: manager = MCPServerManager() # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock successful _get_tools_from_server - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - tool1 = MagicMock() - tool1.name = "tool1" - tool2 = MagicMock() - tool2.name = "tool2" - return [tool1, tool2] - - manager._get_tools_from_server = mock_get_tools_from_server + # Mock successful client.run_with_session + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + manager._create_mcp_client = MagicMock(return_value=mock_client) # Perform health check result = await manager.health_check_server("test-server") - # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "healthy" - assert result["tools_count"] == 2 - assert result["error"] is None - assert "last_health_check" in result - assert "response_time_ms" in result - assert result["response_time_ms"] >= 0 # Allow 0 for very fast mocks + # Verify results - result is now LiteLLM_MCPServerTable + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "healthy" + assert result.health_check_error is None + assert result.last_health_check is not None @pytest.mark.asyncio async def test_health_check_server_unhealthy(self): @@ -679,32 +673,33 @@ class TestMCPServerManager: manager = MCPServerManager() # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock failed _get_tools_from_server - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - raise Exception("Connection timeout") - - manager._get_tools_from_server = mock_get_tools_from_server + # Mock failed client.run_with_session + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock( + side_effect=Exception("Connection timeout") + ) + manager._create_mcp_client = MagicMock(return_value=mock_client) # Perform health check result = await manager.health_check_server("test-server") # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "unhealthy" - assert result["error"] == "Connection timeout" - assert "last_health_check" in result - assert "response_time_ms" in result - assert result["response_time_ms"] >= 0 # Allow 0 for very fast mocks + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "unhealthy" + assert result.health_check_error == "Connection timeout" + assert result.last_health_check is not None @pytest.mark.asyncio async def test_health_check_server_not_found(self): @@ -718,104 +713,121 @@ class TestMCPServerManager: result = await manager.health_check_server("non-existent-server") # Verify results - assert result["server_id"] == "non-existent-server" - assert result["status"] == "unknown" - assert result["error"] == "Server not found" - assert result["response_time_ms"] is None - assert "last_health_check" in result + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "non-existent-server" + assert result.server_name is None + assert result.status == "unknown" + assert result.health_check_error == "Server not found" + assert result.last_health_check is not None @pytest.mark.asyncio - async def test_health_check_all_servers(self): - """Test health check for all servers""" + async def test_health_check_server_oauth2_skips_check(self): + """Test that health check is skipped for OAuth2 servers and returns unknown status""" manager = MCPServerManager() - # Mock servers - server1 = MagicMock() - server1.server_id = "server1" - server1.name = "server1" - - server2 = MagicMock() - server2.server_id = "server2" - server2.name = "server2" - - # Mock registry - manager.registry = {"server1": server1, "server2": server2} - - # Mock get_mcp_server_by_id - def mock_get_server_by_id(server_id): - if server_id == "server1": - return server1 - elif server_id == "server2": - return server2 - return None - - manager.get_mcp_server_by_id = mock_get_server_by_id - - # Mock _get_tools_from_server with different results - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - if server.server_id == "server1": - tool = MagicMock() - tool.name = "tool1" - return [tool] - elif server.server_id == "server2": - raise Exception("Connection failed") - return [] - - manager._get_tools_from_server = mock_get_tools_from_server - - # Perform health check for all servers - result = await manager.health_check_all_servers() - - # Verify results - assert len(result) == 2 - assert "server1" in result - assert "server2" in result - - # Check server1 (healthy) - assert result["server1"]["status"] == "healthy" - assert result["server1"]["tools_count"] == 1 - assert result["server1"]["error"] is None - - # Check server2 (unhealthy) - assert result["server2"]["status"] == "unhealthy" - assert result["server2"]["error"] == "Connection failed" - - @pytest.mark.asyncio - async def test_health_check_server_with_auth_header(self): - """Test health check with authentication header""" - manager = MCPServerManager() - - # Mock server - server = MagicMock() - server.server_id = "test-server" - server.name = "test-server" + # Mock OAuth2 server + server = MCPServer( + server_id="oauth2-server", + name="oauth2-server", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + url="http://oauth2-server.com", + ) manager.get_mcp_server_by_id = MagicMock(return_value=server) - # Mock _get_tools_from_server to verify auth header is passed - async def mock_get_tools_from_server( - server, - mcp_auth_header=None, - raw_headers=None, - ): - assert mcp_auth_header == "test-token" - tool = MagicMock() - tool.name = "tool1" - return [tool] + # _create_mcp_client should not be called for OAuth2 servers + manager._create_mcp_client = MagicMock() - manager._get_tools_from_server = mock_get_tools_from_server + # Perform health check + result = await manager.health_check_server("oauth2-server") - # Perform health check with auth header - result = await manager.health_check_server("test-server", "test-token") + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() # Verify results - assert result["server_id"] == "test-server" - assert result["status"] == "healthy" - assert result["tools_count"] == 1 + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "oauth2-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_server_no_token_skips_check(self): + """Test that health check is skipped when auth_type is set but authentication_token is missing""" + manager = MCPServerManager() + + # Mock server with auth_type but no authentication_token + server = MCPServer( + server_id="no-token-server", + name="no-token-server", + transport=MCPTransport.http, + auth_type=MCPAuth.bearer_token, + authentication_token=None, # No token + url="http://no-token-server.com", + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # _create_mcp_client should not be called + manager._create_mcp_client = MagicMock() + + # Perform health check + result = await manager.health_check_server("no-token-server") + + # Verify that client was not created (health check was skipped) + manager._create_mcp_client.assert_not_called() + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "no-token-server" + assert result.status == "unknown" + assert result.health_check_error is None + assert result.last_health_check is not None + + @pytest.mark.asyncio + async def test_health_check_server_with_static_headers(self): + """Test health check with static headers configured""" + manager = MCPServerManager() + + # Mock server with static_headers + server = MCPServer( + server_id="test-server", + name="test-server", + transport=MCPTransport.http, + auth_type=None, + authentication_token="test-token", + url="http://test-server.com", + static_headers={"X-Custom-Header": "custom-value"}, + ) + + manager.get_mcp_server_by_id = MagicMock(return_value=server) + + # Mock successful client + mock_client = AsyncMock() + mock_client.run_with_session = AsyncMock(return_value="ok") + + # Capture the extra_headers passed to _create_mcp_client + captured_extra_headers = None + + def capture_create_mcp_client(server, mcp_auth_header, extra_headers, stdio_env): + nonlocal captured_extra_headers + captured_extra_headers = extra_headers + return mock_client + + manager._create_mcp_client = MagicMock(side_effect=capture_create_mcp_client) + + # Perform health check + result = await manager.health_check_server("test-server") + + # Verify static headers were passed + assert captured_extra_headers == {"X-Custom-Header": "custom-value"} + + # Verify results + assert isinstance(result, LiteLLM_MCPServerTable) + assert result.server_id == "test-server" + assert result.status == "healthy" + assert result.health_check_error is None @pytest.mark.asyncio async def test_pre_call_tool_check_allowed_tools_list_allows_tool(self): @@ -1352,7 +1364,7 @@ class TestMCPServerManager: "env": {}, }, ) - await manager.add_update_server(server) + await manager.add_server(server) assert server.server_id in manager.get_registry() @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py index cb48a940b57..1044ad3ed27 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_openapi_to_mcp_generator.py @@ -5,6 +5,8 @@ This test suite ensures that: 1. Parameter names with invalid Python identifiers are handled safely 2. No exec() is used (security) 3. All edge cases (hyphens, dots, keywords, special chars) work correctly +4. Path traversal attacks are prevented +5. Path parameters are properly URL encoded """ import pytest @@ -452,3 +454,69 @@ class TestExtractParameters: if __name__ == "__main__": pytest.main([__file__, "-v"]) + + +class TestPathSecurity: + """Test path traversal security and URL encoding.""" + + @pytest.mark.asyncio + async def test_should_reject_path_traversal_inputs(self): + """Test that path traversal attacks (../admin) are rejected.""" + operation = { + "parameters": [ + { + "name": "filename", + "in": "path", + "required": True, + "schema": {"type": "string"}, + } + ] + } + + tool_function = create_tool_function( + path="/files/{filename}", + method="GET", + operation=operation, + base_url="https://example.com", + ) + + response = await tool_function(**{"filename": "../admin"}) + + assert "Invalid path parameter" in response + + @pytest.mark.asyncio + async def test_should_encode_and_request_safe_path_parameters(self): + """Test that path parameters are properly URL encoded.""" + operation = { + "parameters": [ + { + "name": "filename", + "in": "path", + "required": True, + "schema": {"type": "string"}, + } + ] + } + + tool_function = create_tool_function( + path="/files/{filename}", + method="GET", + operation=operation, + base_url="https://example.com", + ) + + with patch("httpx.AsyncClient") as mock_client: + mock_response = AsyncMock() + mock_response.text = "dummy-response" + mock_client.return_value.__aenter__.return_value.get = AsyncMock( + return_value=mock_response + ) + + response = await tool_function(**{"filename": "report 2024.json"}) + + assert response == "dummy-response" + + # Verify URL was properly encoded + call_args = mock_client.return_value.__aenter__.return_value.get.call_args + url = call_args[0][0] + assert url == "https://example.com/files/report%202024.json" diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index ce38ee59e4d..0c6d0921952 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -7,6 +7,7 @@ from litellm.proxy._experimental.mcp_server import rest_endpoints from litellm.proxy._experimental.mcp_server.auth import ( user_api_key_auth_mcp as auth_mcp, ) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy._types import NewMCPServerRequest, UserAPIKeyAuth from litellm.types.mcp import MCPAuth @@ -31,6 +32,60 @@ def _build_request(headers: Optional[Dict[str, str]] = None) -> Request: return Request(scope, receive=receive) +def _get_route(path: str, method: str): + for route in rest_endpoints.router.routes: + if getattr(route, "path", None) == path and method in getattr( + route, "methods", set() + ): + return route + raise AssertionError(f"Route {method} {path} not found") + + +def _route_has_dependency(route, dependency) -> bool: + if any( + getattr(dep, "dependency", None) == dependency + for dep in getattr(route, "dependencies", []) + ): + return True + dependant = getattr(route, "dependant", None) + if dependant is None: + return False + return any(getattr(dep, "call", None) == dependency for dep in dependant.dependencies) + + +@pytest.mark.asyncio +async def test_execute_with_mcp_client_redacts_stack_trace(monkeypatch): + def fake_create_client(*args, **kwargs): + return object() + + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "_create_mcp_client", + fake_create_client, + ) + + async def failing_operation(client): + raise RuntimeError("boom") + + payload = NewMCPServerRequest( + server_name="example", + url="https://example.com", + auth_type=MCPAuth.none, + ) + + result = await rest_endpoints._execute_with_mcp_client( + payload, failing_operation + ) + + assert result["status"] == "error" + assert "stack_trace" not in result + + +def test_test_connection_requires_auth_dependency(): + route = _get_route("/mcp-rest/test/connection", "POST") + assert _route_has_dependency(route, user_api_key_auth) + + @pytest.mark.asyncio async def test_test_tools_list_forwards_mcp_auth_header(monkeypatch): """Ensure credential-based auth forwards the auth_value to the MCP client.""" diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 8ecbaced21e..b56d13bb932 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1146,4 +1146,343 @@ async def test_auth_builder_uses_team_from_header_e2e(): ) assert result["team_id"] == "team-2" - assert result["team_object"] == team_object \ No newline at end of file + assert result["team_object"] == team_object + + +@pytest.mark.asyncio +async def test_get_team_alias_with_nested_fields(): + """ + Test get_team_alias() method with nested JWT fields + """ + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTHandler + + jwt_handler = JWTHandler() + + # Test token with nested team name + nested_token = { + "organization": { + "team": { + "name": "engineering-team" + } + }, + "team_name": "flat-team" + } + + # Test nested access + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="organization.team.name") + assert jwt_handler.get_team_alias(nested_token, None) == "engineering-team" + + # Test flat access (backward compatibility) + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="team_name") + assert jwt_handler.get_team_alias(nested_token, None) == "flat-team" + + # Test missing field returns default + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="nonexistent.field") + assert jwt_handler.get_team_alias(nested_token, "default-team") == "default-team" + + # Test with team_alias_jwt_field not configured + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() # team_alias_jwt_field is None + assert jwt_handler.get_team_alias(nested_token, "default") is None + + +@pytest.mark.asyncio +async def test_is_required_team_id_with_team_alias_field(): + """ + Test that is_required_team_id() returns True when team_alias_jwt_field is set + """ + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTHandler + + jwt_handler = JWTHandler() + + # Neither field set - should return False + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + assert jwt_handler.is_required_team_id() is False + + # Only team_id_jwt_field set - should return True + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_id_jwt_field="team_id") + assert jwt_handler.is_required_team_id() is True + + # Only team_alias_jwt_field set - should return True + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_alias_jwt_field="team_name") + assert jwt_handler.is_required_team_id() is True + + # Both fields set - should return True + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth( + team_id_jwt_field="team_id", + team_alias_jwt_field="team_name" + ) + assert jwt_handler.is_required_team_id() is True + + +@pytest.mark.asyncio +async def test_find_and_validate_specific_team_id_with_team_alias(): + """ + Test that find_and_validate_specific_team_id resolves team by name when team_id is not found + """ + from unittest.mock import MagicMock + + from litellm.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_TeamTable + from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler + from litellm.proxy.utils import ProxyLogging + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + team_alias_jwt_field="team_alias" + ), + ) + + # Token with team name (no team_id) + jwt_token = { + "sub": "user-1", + "team_alias": "my-team" + } + + # Mock team object returned by get_team_object_by_alias + team_object = LiteLLM_TeamTable(team_id="resolved-team-id", team_alias="my-team") + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", + new_callable=AsyncMock + ) as mock_get_by_alias: + mock_get_by_alias.return_value = team_object + + team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=jwt_token, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + # Should have resolved team_id from team name + assert team_id == "resolved-team-id" + assert result_team == team_object + mock_get_by_alias.assert_called_once_with( + team_alias="my-team", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + +@pytest.mark.asyncio +async def test_find_and_validate_team_id_takes_precedence_over_name(): + """ + Test that team_id_jwt_field takes precedence over team_alias_jwt_field + """ + from unittest.mock import MagicMock + + from litellm.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_TeamTable + from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler + from litellm.proxy.utils import ProxyLogging + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + team_id_jwt_field="team_id", + team_alias_jwt_field="team_alias" + ), + ) + + # Token with both team_id and team name + jwt_token = { + "sub": "user-1", + "team_id": "direct-team-id", + "team_alias": "my-team" + } + + # Mock team object returned by get_team_object (by ID) + team_object = LiteLLM_TeamTable(team_id="direct-team-id") + + with patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock + ) as mock_get_by_id, patch( + "litellm.proxy.auth.handle_jwt.get_team_object_by_alias", + new_callable=AsyncMock + ) as mock_get_by_alias: + mock_get_by_id.return_value = team_object + + team_id, result_team = await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=jwt_token, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + # Should use team_id directly, not resolve by name + assert team_id == "direct-team-id" + assert result_team == team_object + mock_get_by_id.assert_called_once() + mock_get_by_alias.assert_not_called() + + +@pytest.mark.asyncio +async def test_find_and_validate_raises_when_required_team_not_found(): + """ + Test that an exception is raised when team is required but neither team_id nor team_name is found + """ + from litellm.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler + from litellm.proxy.utils import ProxyLogging + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + team_alias_jwt_field="team_alias" # Required, but not in token + ), + ) + + # Token without team info + jwt_token = { + "sub": "user-1" + } + + with pytest.raises(Exception) as exc_info: + await JWTAuthManager.find_and_validate_specific_team_id( + jwt_handler=jwt_handler, + jwt_valid_token=jwt_token, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + assert "No team found in token" in str(exc_info.value) + assert "team_alias field 'team_alias'" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_get_org_alias_with_nested_fields(): + """ + Test get_org_alias() method with nested JWT fields + """ + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTHandler + + jwt_handler = JWTHandler() + + # Test token with nested org name + nested_token = { + "company": { + "organization": { + "name": "acme-corp" + } + }, + "org_name": "flat-org" + } + + # Test nested access + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="company.organization.name") + assert jwt_handler.get_org_alias(nested_token, None) == "acme-corp" + + # Test flat access + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="org_name") + assert jwt_handler.get_org_alias(nested_token, None) == "flat-org" + + # Test missing field returns default + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(org_alias_jwt_field="nonexistent.field") + assert jwt_handler.get_org_alias(nested_token, "default-org") == "default-org" + + # Test with org_alias_jwt_field not configured + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth() + assert jwt_handler.get_org_alias(nested_token, "default") is None + + +@pytest.mark.asyncio +async def test_get_objects_resolves_org_by_name(): + """ + Test that get_objects resolves organization by name when org_id is not provided + """ + from litellm.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth, LiteLLM_OrganizationTable + from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler + from litellm.proxy.utils import ProxyLogging + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + proxy_logging_obj = ProxyLogging(user_api_key_cache=user_api_key_cache) + + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + org_alias_jwt_field="org_alias" + ), + ) + + # Mock org object returned by get_org_object_by_alias + org_object = LiteLLM_OrganizationTable( + organization_id="resolved-org-id", + organization_alias="my-org", + budget_id="budget-1", + created_by="admin", + updated_by="admin", + models=[] + ) + + with patch( + "litellm.proxy.auth.handle_jwt.get_org_object_by_alias", + new_callable=AsyncMock + ) as mock_get_by_alias: + mock_get_by_alias.return_value = org_object + + ( + result_user_obj, + result_org_obj, + result_end_user_obj, + result_team_membership, + ) = await JWTAuthManager.get_objects( + user_id=None, + user_email=None, + org_id=None, # No org_id provided + end_user_id=None, + team_id=None, + valid_user_email=None, + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + route="/chat/completions", + org_alias="my-org", + ) + + # Should resolve org by alias - org_id can be derived from org_object.organization_id + assert result_org_obj == org_object + assert result_org_obj.organization_id == "resolved-org-id" + mock_get_by_alias.assert_called_once_with( + org_alias="my-org", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=proxy_logging_obj, + ) + + + diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py index 201461dc8b5..c04b8114939 100644 --- a/tests/test_litellm/proxy/auth/test_login_utils.py +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -6,6 +6,7 @@ to login_utils.py for better reusability. """ import os +from datetime import datetime, timezone, timedelta from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -21,6 +22,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.login_utils import ( LoginResult, authenticate_user, + expire_previous_ui_session_tokens, get_ui_credentials, ) @@ -282,3 +284,268 @@ async def test_authenticate_user_database_required_for_admin(): finally: if original_db_url: os.environ["DATABASE_URL"] = original_db_url + + +@pytest.mark.asyncio +async def test_expire_previous_ui_session_tokens_none_prisma_client(): + """Test that function returns early when prisma_client is None""" + await expire_previous_ui_session_tokens("test-user", None) + # Should not raise any exception + + +@pytest.mark.asyncio +async def test_expire_previous_ui_session_tokens_only_litellm_dashboard_team(): + """Test that only tokens with team_id='litellm-dashboard' are expired""" + user_id = "test-user" + current_time = datetime.now(timezone.utc) + + # Create mock tokens with proper attributes + token1 = MagicMock() + token1.token = "token1" + token1.user_id = user_id + token1.team_id = "litellm-dashboard" + token1.blocked = None + token1.expires = current_time + timedelta(hours=1) + + token2 = MagicMock() + token2.token = "token2" + token2.user_id = user_id + token2.team_id = "other-team" + token2.blocked = None + token2.expires = current_time + timedelta(hours=1) + + def mock_find_many(**kwargs): + """Mock find_many that filters tokens based on query criteria""" + where_clause = kwargs.get("where", {}) + filtered_tokens = [] + + for token in [token1, token2]: + # Check user_id match + if token.user_id != where_clause.get("user_id"): + continue + # Check team_id match + if token.team_id != where_clause.get("team_id"): + continue + # Check blocked condition (None or False) + if token.blocked is not None and token.blocked is not False: + continue + # Check expires > current_time + if token.expires <= where_clause.get("expires", {}).get("gt"): + continue + filtered_tokens.append(token) + + return filtered_tokens + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=mock_find_many) + mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock() + + await expire_previous_ui_session_tokens(user_id, mock_prisma_client) + + # Should only call update_many with the litellm-dashboard token + mock_prisma_client.db.litellm_verificationtoken.update_many.assert_called_once_with( + where={"token": {"in": ["token1"]}}, + data={"blocked": True} + ) + + +@pytest.mark.asyncio +async def test_expire_previous_ui_session_tokens_blocks_null_and_false(): + """Test that tokens with blocked=None and blocked=False are both processed""" + user_id = "test-user" + current_time = datetime.now(timezone.utc) + + # Create mock tokens with proper attributes + token1 = MagicMock() + token1.token = "token1" + token1.user_id = user_id + token1.team_id = "litellm-dashboard" + token1.blocked = None + token1.expires = current_time + timedelta(hours=1) + + token2 = MagicMock() + token2.token = "token2" + token2.user_id = user_id + token2.team_id = "litellm-dashboard" + token2.blocked = False + token2.expires = current_time + timedelta(hours=1) + + token3 = MagicMock() + token3.token = "token3" + token3.user_id = user_id + token3.team_id = "litellm-dashboard" + token3.blocked = True # This should be ignored + token3.expires = current_time + timedelta(hours=1) + + def mock_find_many(**kwargs): + """Mock find_many that filters tokens based on query criteria""" + where_clause = kwargs.get("where", {}) + filtered_tokens = [] + + for token in [token1, token2, token3]: + # Check user_id match + if token.user_id != where_clause.get("user_id"): + continue + # Check team_id match + if token.team_id != where_clause.get("team_id"): + continue + # Check blocked condition (None or False) + if token.blocked is not None and token.blocked is not False: + continue + # Check expires > current_time + if token.expires <= where_clause.get("expires", {}).get("gt"): + continue + filtered_tokens.append(token) + + return filtered_tokens + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=mock_find_many) + mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock() + + await expire_previous_ui_session_tokens(user_id, mock_prisma_client) + + # Should only block token1 and token2 (not token3 which is already blocked) + mock_prisma_client.db.litellm_verificationtoken.update_many.assert_called_once_with( + where={"token": {"in": ["token1", "token2"]}}, + data={"blocked": True} + ) + + +@pytest.mark.asyncio +async def test_expire_previous_ui_session_tokens_only_non_expired(): + """Test that only non-expired tokens are processed""" + user_id = "test-user" + current_time = datetime.now(timezone.utc) + + # Create mock tokens with proper attributes + token1 = MagicMock() + token1.token = "token1" + token1.user_id = user_id + token1.team_id = "litellm-dashboard" + token1.blocked = None + token1.expires = current_time + timedelta(hours=1) # Not expired + + token2 = MagicMock() + token2.token = "token2" + token2.user_id = user_id + token2.team_id = "litellm-dashboard" + token2.blocked = None + token2.expires = current_time - timedelta(hours=1) # Already expired + + def mock_find_many(**kwargs): + """Mock find_many that filters tokens based on query criteria""" + where_clause = kwargs.get("where", {}) + filtered_tokens = [] + + for token in [token1, token2]: + # Check user_id match + if token.user_id != where_clause.get("user_id"): + continue + # Check team_id match + if token.team_id != where_clause.get("team_id"): + continue + # Check blocked condition (None or False) + if token.blocked is not None and token.blocked is not False: + continue + # Check expires > current_time + if token.expires <= where_clause.get("expires", {}).get("gt"): + continue + filtered_tokens.append(token) + + return filtered_tokens + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=mock_find_many) + mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock() + + await expire_previous_ui_session_tokens(user_id, mock_prisma_client) + + # Should only block the non-expired token + mock_prisma_client.db.litellm_verificationtoken.update_many.assert_called_once_with( + where={"token": {"in": ["token1"]}}, + data={"blocked": True} + ) + + +@pytest.mark.asyncio +async def test_expire_previous_ui_session_tokens_no_tokens_found(): + """Test behavior when no valid tokens are found""" + user_id = "test-user" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock() + + await expire_previous_ui_session_tokens(user_id, mock_prisma_client) + + # Should not call update_many when no tokens found + mock_prisma_client.db.litellm_verificationtoken.update_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_expire_previous_ui_session_tokens_filters_none_token(): + """Test that tokens with None token value are filtered out""" + user_id = "test-user" + current_time = datetime.now(timezone.utc) + + # Create mock tokens with proper attributes + token1 = MagicMock() + token1.token = "token1" + token1.user_id = user_id + token1.team_id = "litellm-dashboard" + token1.blocked = None + token1.expires = current_time + timedelta(hours=1) + + token2 = MagicMock() + token2.token = None # This should be filtered out in the token collection step + token2.user_id = user_id + token2.team_id = "litellm-dashboard" + token2.blocked = None + token2.expires = current_time + timedelta(hours=1) + + def mock_find_many(**kwargs): + """Mock find_many that filters tokens based on query criteria""" + where_clause = kwargs.get("where", {}) + filtered_tokens = [] + + for token in [token1, token2]: + # Check user_id match + if token.user_id != where_clause.get("user_id"): + continue + # Check team_id match + if token.team_id != where_clause.get("team_id"): + continue + # Check blocked condition (None or False) + if token.blocked is not None and token.blocked is not False: + continue + # Check expires > current_time + if token.expires <= where_clause.get("expires", {}).get("gt"): + continue + filtered_tokens.append(token) + + return filtered_tokens + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=mock_find_many) + mock_prisma_client.db.litellm_verificationtoken.update_many = AsyncMock() + + await expire_previous_ui_session_tokens(user_id, mock_prisma_client) + + # Should only block token1 (token with None value should be filtered out) + mock_prisma_client.db.litellm_verificationtoken.update_many.assert_called_once_with( + where={"token": {"in": ["token1"]}}, + data={"blocked": True} + ) + + +@pytest.mark.asyncio +async def test_expire_previous_ui_session_tokens_exception_handling(): + """Test that exceptions during token expiry are silently handled""" + user_id = "test-user" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(side_effect=Exception("Database error")) + + # Should not raise exception despite database error + await expire_previous_ui_session_tokens(user_id, mock_prisma_client) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py new file mode 100644 index 00000000000..f4aa28d98cc --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -0,0 +1,84 @@ +"""Tests for unified guardrail.""" + +import pytest + +from litellm.caching import DualCache +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( + MCPGuardrailTranslationHandler, +) +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import unified_guardrail as unified_module +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypes + + +class RecordingGuardrail(CustomGuardrail): + """Records the event types it is asked to run for.""" + + def __init__(self): + super().__init__(guardrail_name="recording-guardrail") + self.event_history = [] + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + self.event_history.append(event_type) + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + return {"texts": inputs.get("texts", [])} + + +@pytest.fixture(autouse=True) +def _inject_mcp_handler_mapping(): + """Inject MCP handler mapping so the unified guardrail can run inside tests.""" + unified_module.endpoint_guardrail_translation_mappings = { + CallTypes.call_mcp_tool: MCPGuardrailTranslationHandler, + } + yield + unified_module.endpoint_guardrail_translation_mappings = None + + +@pytest.mark.asyncio +async def test_pre_call_hook_uses_mcp_event_type(): + """pre_call hook should swap to GuardrailEventHooks.pre_mcp_call for MCP calls.""" + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + cache = DualCache() + + data = { + "guardrail_to_apply": guardrail, + "messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}], + "model": "mcp-tool-call", + } + + await handler.async_pre_call_hook( + user_api_key_dict=None, + cache=cache, + data=data, + call_type=CallTypes.call_mcp_tool.value, + ) + + assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call] + + +@pytest.mark.asyncio +async def test_moderation_hook_uses_mcp_event_type(): + """moderation hook should request GuardrailEventHooks.during_mcp_call for MCP calls.""" + handler = UnifiedLLMGuardrails() + guardrail = RecordingGuardrail() + + data = { + "guardrail_to_apply": guardrail, + "messages": [{"role": "user", "content": "Tool: test\nArguments: {}"}], + "model": "mcp-tool-call", + } + + await handler.async_moderation_hook( + data=data, + user_api_key_dict=None, + call_type=CallTypes.call_mcp_tool.value, + ) + + assert guardrail.event_history == [GuardrailEventHooks.during_mcp_call] diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index b4dcc33c747..d5c3ecae7d6 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -163,3 +163,78 @@ async def test_update_budget_allows_null_max_budget(client_and_mocks): assert captured_data["max_budget"] is None, "max_budget should be None" mock_table.update.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_new_budget_negative_max_budget(client_and_mocks): + """ + Test that /budget/new rejects negative max_budget values. + + This prevents the issue where negative budgets would always trigger + budget exceeded errors. + """ + client, _, _ = client_and_mocks + + payload = { + "budget_id": "budget_negative", + "max_budget": -7.0, + } + resp = client.post("/budget/new", json=payload) + assert resp.status_code == 400, resp.text + + detail = resp.json()["detail"] + assert "max_budget cannot be negative" in str(detail) + + +@pytest.mark.asyncio +async def test_new_budget_negative_soft_budget(client_and_mocks): + """ + Test that /budget/new rejects negative soft_budget values. + """ + client, _, _ = client_and_mocks + + payload = { + "budget_id": "budget_negative_soft", + "soft_budget": -10.0, + } + resp = client.post("/budget/new", json=payload) + assert resp.status_code == 400, resp.text + + detail = resp.json()["detail"] + assert "soft_budget cannot be negative" in str(detail) + + +@pytest.mark.asyncio +async def test_update_budget_negative_max_budget(client_and_mocks): + """ + Test that /budget/update rejects negative max_budget values. + """ + client, _, _ = client_and_mocks + + payload = { + "budget_id": "budget_update_negative", + "max_budget": -5.0, + } + resp = client.post("/budget/update", json=payload) + assert resp.status_code == 400, resp.text + + detail = resp.json()["detail"] + assert "max_budget cannot be negative" in str(detail) + + +@pytest.mark.asyncio +async def test_update_budget_negative_soft_budget(client_and_mocks): + """ + Test that /budget/update rejects negative soft_budget values. + """ + client, _, _ = client_and_mocks + + payload = { + "budget_id": "budget_update_negative_soft", + "soft_budget": -15.0, + } + resp = client.post("/budget/update", json=payload) + assert resp.status_code == 400, resp.text + + detail = resp.json()["detail"] + assert "soft_budget cannot be negative" in str(detail) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 648045a7ea6..3e69b4e0caf 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -3405,3 +3405,161 @@ async def test_can_modify_verification_token_personal_key_no_user_id(monkeypatch ) assert result is False + + +@pytest.mark.asyncio +async def test_list_keys_with_expand_user(): + """ + Test that expand=user parameter correctly includes user information in the response. + """ + mock_prisma_client = AsyncMock() + + # Create mock keys with user_ids + mock_key1 = MagicMock() + mock_key1.token = "token1" + mock_key1.user_id = "user123" + mock_key1.dict.return_value = { + "token": "token1", + "user_id": "user123", + "key_alias": "key1", + "models": ["gpt-4"], + } + + mock_key2 = MagicMock() + mock_key2.token = "token2" + mock_key2.user_id = "user456" + mock_key2.dict.return_value = { + "token": "token2", + "user_id": "user456", + "key_alias": "key2", + "models": ["gpt-3.5-turbo"], + } + + mock_find_many_keys = AsyncMock(return_value=[mock_key1, mock_key2]) + mock_count_keys = AsyncMock(return_value=2) + + # Create mock users + mock_user1 = MagicMock() + mock_user1.user_id = "user123" + mock_user1.user_email = "user1@example.com" + mock_user1.dict.return_value = { + "user_id": "user123", + "user_email": "user1@example.com", + "user_alias": "User One", + } + + mock_user2 = MagicMock() + mock_user2.user_id = "user456" + mock_user2.user_email = "user2@example.com" + mock_user2.dict.return_value = { + "user_id": "user456", + "user_email": "user2@example.com", + "user_alias": "User Two", + } + + mock_find_many_users = AsyncMock(return_value=[mock_user1, mock_user2]) + + mock_prisma_client.db.litellm_verificationtoken.find_many = mock_find_many_keys + mock_prisma_client.db.litellm_verificationtoken.count = mock_count_keys + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many_users + + args = { + "prisma_client": mock_prisma_client, + "page": 1, + "size": 50, + "user_id": None, + "team_id": None, + "organization_id": None, + "key_alias": None, + "key_hash": None, + "exclude_team_id": None, + "return_full_object": False, # This should be overridden by expand=user + "admin_team_ids": None, + "include_created_by_keys": False, + "expand": ["user"], # Test the expand parameter + } + + result = await _list_key_helper(**args) + + # Verify that keys were fetched + mock_find_many_keys.assert_called_once() + mock_count_keys.assert_called_once() + + # Verify that users were fetched + # Note: Order doesn't matter for the 'in' query, so we just check that both user_ids are present + call_args = mock_find_many_users.call_args + assert call_args is not None + where_clause = call_args.kwargs["where"] + assert "user_id" in where_clause + assert "in" in where_clause["user_id"] + user_ids_in_query = set(where_clause["user_id"]["in"]) + assert user_ids_in_query == {"user123", "user456"} + + # Verify response structure + assert len(result["keys"]) == 2 + assert result["total_count"] == 2 + assert result["current_page"] == 1 + assert result["total_pages"] == 1 + + # Verify that user data is included in the response + # Since expand=user is specified, keys should be full objects + assert isinstance(result["keys"][0], UserAPIKeyAuth) + assert isinstance(result["keys"][1], UserAPIKeyAuth) + + # Verify user data is attached to keys + assert result["keys"][0].user == { + "user_id": "user123", + "user_email": "user1@example.com", + "user_alias": "User One", + } + assert result["keys"][1].user == { + "user_id": "user456", + "user_email": "user2@example.com", + "user_alias": "User Two", + } + + +@pytest.mark.asyncio +async def test_generate_key_negative_max_budget(): + """ + Test that GenerateKeyRequest model allows negative max_budget values. + Validation is done at API level, not model level. + + This prevents GET requests from breaking when they receive data with negative budgets. + """ + # Should not raise any errors at model level + request = GenerateKeyRequest(max_budget=-7.0) + assert request.max_budget == -7.0 + + +@pytest.mark.asyncio +async def test_generate_key_negative_soft_budget(): + """ + Test that GenerateKeyRequest model allows negative soft_budget values. + Validation is done at API level, not model level. + """ + # Should not raise any errors at model level + request = GenerateKeyRequest(soft_budget=-10.0) + assert request.soft_budget == -10.0 + + +@pytest.mark.asyncio +async def test_generate_key_positive_budgets_accepted(): + """ + Test that GenerateKeyRequest accepts positive budget values. + """ + # Should not raise any errors + request = GenerateKeyRequest(max_budget=100.0, soft_budget=50.0) + assert request.max_budget == 100.0 + assert request.soft_budget == 50.0 + + +@pytest.mark.asyncio +async def test_update_key_negative_max_budget(): + """ + Test that UpdateKeyRequest model allows negative max_budget values. + Validation is done at API level, not model level. + """ + # Should not raise any errors at model level + request = UpdateKeyRequest(key="test-key", max_budget=-5.0) + assert request.max_budget == -5.0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 61342e8025b..cc268ab9925 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -169,8 +169,8 @@ class TestListMCPServers: return_value=["config_server_1", "config_server_2"] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ generate_mock_mcp_server_db_record( server_id="config_server_1", alias="Zapier MCP", @@ -184,11 +184,11 @@ class TestListMCPServers: transport="http", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -200,6 +200,9 @@ class TestListMCPServers: ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", return_value=mock_prisma_client, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -300,8 +303,8 @@ class TestListMCPServers: ] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ db_server_1, db_server_2, generate_mock_mcp_server_db_record( @@ -317,11 +320,11 @@ class TestListMCPServers: transport="http", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -333,6 +336,9 @@ class TestListMCPServers: ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", return_value=mock_prisma_client, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -425,8 +431,8 @@ class TestListMCPServers: return_value=["db_server_allowed", "config_server_allowed"] ) - # Mock the new method that returns servers with health and team data - mock_servers_with_health = [ + # Mock the new method that returns servers without health check + mock_servers = [ db_server_allowed, generate_mock_mcp_server_db_record( server_id="config_server_allowed", @@ -434,11 +440,11 @@ class TestListMCPServers: url="https://actions.zapier.com/mcp/sse", ), ] - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=mock_servers_with_health + mock_manager.get_all_allowed_mcp_servers = AsyncMock( + return_value=mock_servers ) - for idx, server in enumerate(mock_servers_with_health): + for idx, server in enumerate(mock_servers): server.credentials = {"auth_value": f"secret_{idx}"} with patch( @@ -450,6 +456,9 @@ class TestListMCPServers: ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", return_value=mock_prisma_client, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), ): # Import and call the function from litellm.proxy.management_endpoints.mcp_management_endpoints import ( @@ -486,11 +495,14 @@ class TestListMCPServers: mock_server.credentials = {"auth_value": "top-secret"} mock_prisma_client = MagicMock() - mock_health_result = { - "status": "healthy", - "last_health_check": datetime.now().isoformat(), - "error": None, - } + + # Mock health check result as LiteLLM_MCPServerTable + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", alias="Server 1" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None mock_user_auth = generate_mock_user_api_key_auth( user_role=LitellmUserRoles.PROXY_ADMIN @@ -531,11 +543,14 @@ class TestListMCPServers: delattr(mock_server, "credentials") mock_prisma_client = MagicMock() - mock_health_result = { - "status": "healthy", - "last_health_check": datetime.now().isoformat(), - "error": None, - } + + # Mock health check result as LiteLLM_MCPServerTable + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-2", alias="Server 2" + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None mock_user_auth = generate_mock_user_api_key_auth( user_role=LitellmUserRoles.PROXY_ADMIN @@ -568,296 +583,6 @@ class TestListMCPServers: assert result.status == "healthy" -class TestMCPHealthCheckEndpoints: - """Test MCP health check endpoints""" - - @pytest.mark.asyncio - async def test_health_check_mcp_server_success(self): - """Test successful health check for a specific MCP server""" - # Mock server - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - - # Mock dependencies - mock_prisma_client = MagicMock() - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.health_check_server = AsyncMock( - return_value={ - "server_id": "test-server", - "server_name": "Test Server", - "status": "healthy", - "tools_count": 3, - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 150.5, - "error": None, - } - ) - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=mock_server), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - result = await health_check_mcp_server( - server_id="test-server", user_api_key_dict=mock_user_auth - ) - - # Verify results - assert result["server_id"] == "test-server" - assert result["server_name"] == "Test Server" - assert result["status"] == "healthy" - assert result["tools_count"] == 3 - assert result["response_time_ms"] == 150.5 - assert result["error"] is None - - @pytest.mark.asyncio - async def test_health_check_mcp_server_not_found(self): - """Test health check for a server that doesn't exist""" - # Mock dependencies - mock_prisma_client = MagicMock() - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=None), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - # Should raise HTTPException - with pytest.raises(Exception) as exc_info: - await health_check_mcp_server( - server_id="non-existent-server", user_api_key_dict=mock_user_auth - ) - - assert "not found" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_health_check_mcp_server_unauthorized(self): - """Test health check for a server user doesn't have access to""" - # Mock server - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - - # Mock dependencies - mock_prisma_client = MagicMock() - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.INTERNAL_USER # Non-admin user - ) - - # Mock user doesn't have access to this server - mock_user_servers = [] - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_all_mcp_servers_for_user", - return_value=mock_user_servers, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_mcp_server", - AsyncMock(return_value=mock_server), - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_mcp_server, - ) - - # Should raise HTTPException - with pytest.raises(Exception) as exc_info: - await health_check_mcp_server( - server_id="test-server", user_api_key_dict=mock_user_auth - ) - - assert "permission" in str(exc_info.value) - - @pytest.mark.asyncio - async def test_health_check_all_mcp_servers(self): - """Test health check for all accessible MCP servers""" - # Mock team records - team_records = [ - generate_mock_team_record( - team_id="team1", - team_alias="Team 1", - organization_id="org1", - mcp_servers=["server1", "server2"], - ) - ] - - # Mock DB servers - db_servers = [ - generate_mock_mcp_server_db_record(server_id="server1"), - generate_mock_mcp_server_db_record(server_id="server2"), - ] - - # Mock dependencies - mock_prisma_client = MagicMock() - mock_prisma_client = setup_mock_prisma_client( - mock_prisma_client=mock_prisma_client, - team_records=team_records, - mcp_servers=db_servers, - ) - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.health_check_allowed_servers = AsyncMock( - return_value={ - "server1": { - "server_id": "server1", - "server_name": "Test DB Server", - "status": "healthy", - "tools_count": 2, - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 100.0, - "error": None, - }, - "server2": { - "server_id": "server2", - "server_name": "Test DB Server", - "status": "unhealthy", - "last_health_check": "2024-01-01T12:00:00", - "response_time_ms": 5000.0, - "error": "Connection timeout", - }, - } - ) - mock_manager.get_allowed_mcp_servers = AsyncMock( - return_value=["server1", "server2"] - ) - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.INTERNAL_USER - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=False, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - health_check_all_mcp_servers, - ) - - result = await health_check_all_mcp_servers( - user_api_key_dict=mock_user_auth - ) - - # Verify results - assert result["total_servers"] == 2 - assert result["healthy_count"] == 1 - assert result["unhealthy_count"] == 1 - assert result["unknown_count"] == 0 - assert "server1" in result["servers"] - assert "server2" in result["servers"] - - # Check individual server results - assert result["servers"]["server1"]["status"] == "healthy" - assert result["servers"]["server1"]["tools_count"] == 2 - assert result["servers"]["server1"]["server_name"] == "Test DB Server" - assert result["servers"]["server2"]["status"] == "unhealthy" - assert result["servers"]["server2"]["error"] == "Connection timeout" - assert result["servers"]["server2"]["server_name"] == "Test DB Server" - - @pytest.mark.asyncio - async def test_fetch_all_mcp_servers_with_health_status(self): - """Test that fetch_all_mcp_servers includes health check status""" - # Mock server with health status - mock_server = generate_mock_mcp_server_db_record( - server_id="test-server", alias="Test Server" - ) - # Add health status to the mock server - mock_server.status = "healthy" - mock_server.last_health_check = datetime.now() - mock_server.health_check_error = None - - # Mock dependencies - mock_prisma_client = MagicMock() - mock_prisma_client = setup_mock_prisma_client( - mock_prisma_client=mock_prisma_client, - team_records=[], - mcp_servers=[], # Don't add servers here since we're mocking get_all_mcp_servers - ) - - # Mock global MCP server manager - mock_manager = MagicMock() - mock_manager.config_mcp_servers = {} - mock_manager.get_allowed_mcp_servers = AsyncMock(return_value=[]) - mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( - return_value=[mock_server] - ) - - mock_server.credentials = {"auth_value": "secret"} - - mock_user_auth = generate_mock_user_api_key_auth( - user_role=LitellmUserRoles.PROXY_ADMIN - ) - - with patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.get_prisma_client_or_throw", - return_value=mock_prisma_client, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", - return_value=True, - ), patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", - mock_manager, - ): - # Import and call the function - from litellm.proxy.management_endpoints.mcp_management_endpoints import ( - fetch_all_mcp_servers, - ) - - result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth) - - # Verify health check status is included - assert len(result) == 1 - server = result[0] - assert server.server_id == "test-server" - assert server.status == "healthy" - assert server.last_health_check is not None - assert server.health_check_error is None - assert server.credentials is None - - class TestTemporaryMCPSessionEndpoints: def test_inherit_credentials_from_existing_server(self): payload = NewMCPServerRequest( @@ -1170,7 +895,6 @@ class TestTemporaryMCPSessionEndpoints: fallback_client_id="server-1", ) - class TestUpdateMCPServer: """Test suite for update MCP server functionality""" @@ -1233,7 +957,7 @@ class TestUpdateMCPServer: "litellm.proxy.management_endpoints.mcp_management_endpoints.update_mcp_server", AsyncMock(return_value=updated_server), ) as update_mock, patch( - "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_update_server", + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.add_server", AsyncMock(), ), patch( "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager.reload_servers_from_database", @@ -1260,3 +984,163 @@ class TestUpdateMCPServer: # Verify the result includes extra_headers assert result.extra_headers == ["X-Custom-Header", "X-Another-Header"] assert result.alias == "Updated Test Server" + + +class TestHealthCheckServers: + """Test suite for health check servers endpoint""" + + @pytest.mark.asyncio + async def test_health_check_all_servers(self): + """ + Test health check for all accessible servers + + Scenario: User has access to 2 servers, checks all + Expected: Returns health status for both servers + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check results + mock_health_result_1 = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result_1.status = "healthy" + mock_health_result_1.last_health_check = datetime.now() + mock_health_result_1.health_check_error = None + + mock_health_result_2 = generate_mock_mcp_server_db_record( + server_id="server-2", + alias="Server 2", + url="https://server2.example.com", + ) + mock_health_result_2.status = "unhealthy" + mock_health_result_2.last_health_check = datetime.now() + mock_health_result_2.health_check_error = "Connection timeout" + + # Mock manager + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result_1, mock_health_result_2] + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=None, + user_api_key_dict=mock_user_auth, + ) + + # Verify results + assert len(result) == 2 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" + assert result[1]["server_id"] == "server-2" + assert result[1]["status"] == "unhealthy" + + @pytest.mark.asyncio + async def test_health_check_specific_servers(self): + """ + Test health check for specific servers + + Scenario: User requests health check for specific server IDs + Expected: Returns health status only for requested servers + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check result + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + # Mock manager + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result] + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=["server-1"], + user_api_key_dict=mock_user_auth, + ) + + # Verify results + assert len(result) == 1 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" + + @pytest.mark.asyncio + async def test_health_check_unauthorized_servers(self): + """ + Test health check with unauthorized servers + + Scenario: User requests health check for servers they don't have access to + Expected: Only checks accessible servers, unauthorized servers are filtered out + """ + from litellm.proxy.management_endpoints.mcp_management_endpoints import ( + health_check_servers, + ) + + # Mock user auth + mock_user_auth = generate_mock_user_api_key_auth() + + # Mock health check result for authorized server + mock_health_result = generate_mock_mcp_server_db_record( + server_id="server-1", + alias="Server 1", + url="https://server1.example.com", + ) + mock_health_result.status = "healthy" + mock_health_result.last_health_check = datetime.now() + mock_health_result.health_check_error = None + + # Mock manager - server_ids filter is applied inside get_all_mcp_servers_with_health_and_teams + # So it only returns servers the user has access to + mock_manager = MagicMock() + mock_manager.get_all_mcp_servers_with_health_and_teams = AsyncMock( + return_value=[mock_health_result] # Only server-1 is returned (accessible) + ) + + with patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), patch( + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ): + result = await health_check_servers( + server_ids=["server-1", "server-unauthorized"], + user_api_key_dict=mock_user_auth, + ) + + # Verify results - only accessible server is returned + assert len(result) == 1 + assert result[0]["server_id"] == "server-1" + assert result[0]["status"] == "healthy" diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 83b4fc35a0d..6cf8f745e07 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -3958,3 +3958,74 @@ async def test_update_team_guardrails_with_org_id(): assert "include" in first_call_kwargs assert "teams" in first_call_kwargs["include"] assert first_call_kwargs["include"]["teams"] is True + + +@pytest.mark.asyncio +async def test_new_team_negative_max_budget(): + """ + Test that NewTeamRequest model allows negative max_budget values. + Validation is done at API level, not model level. + + This prevents GET requests from breaking when they receive data with negative budgets. + """ + from litellm.proxy._types import NewTeamRequest + + # Should not raise any errors at model level + request = NewTeamRequest(team_alias="test-team", max_budget=-7.0) + assert request.max_budget == -7.0 + + +@pytest.mark.asyncio +async def test_new_team_negative_team_member_budget(): + """ + Test that NewTeamRequest model allows negative team_member_budget values. + Validation is done at API level, not model level. + """ + from litellm.proxy._types import NewTeamRequest + + # Should not raise any errors at model level + request = NewTeamRequest(team_alias="test-team", team_member_budget=-10.0) + assert request.team_member_budget == -10.0 + + +@pytest.mark.asyncio +async def test_update_team_negative_max_budget(): + """ + Test that UpdateTeamRequest model allows negative max_budget values. + Validation is done at API level, not model level. + """ + from litellm.proxy._types import UpdateTeamRequest + + # Should not raise any errors at model level + request = UpdateTeamRequest(team_id="test-team-id", max_budget=-5.0) + assert request.max_budget == -5.0 + + +@pytest.mark.asyncio +async def test_update_team_negative_team_member_budget(): + """ + Test that UpdateTeamRequest model allows negative team_member_budget values. + Validation is done at API level, not model level. + """ + from litellm.proxy._types import UpdateTeamRequest + + # Should not raise any errors at model level + request = UpdateTeamRequest(team_id="test-team-id", team_member_budget=-15.0) + assert request.team_member_budget == -15.0 + + +@pytest.mark.asyncio +async def test_new_team_positive_budgets_accepted(): + """ + Test that NewTeamRequest accepts positive budget values. + """ + from litellm.proxy._types import NewTeamRequest + + # Should not raise any errors + request = NewTeamRequest( + team_alias="test-team", + max_budget=100.0, + team_member_budget=50.0 + ) + assert request.max_budget == 100.0 + assert request.team_member_budget == 50.0 diff --git a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py index f6a9b5bddb3..213297fc80a 100644 --- a/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/test_litellm/proxy/public_endpoints/test_public_endpoints.py @@ -78,3 +78,26 @@ def test_get_litellm_model_cost_map_returns_cost_map(): # Check for common cost fields that should be present assert "input_cost_per_token" in sample_model_data or "output_cost_per_token" in sample_model_data + +def test_watsonx_provider_fields(): + """Test that Watsonx provider has all required credential fields including multiple auth options.""" + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + response = client.get("/public/providers/fields") + providers = response.json() + + watsonx = next((p for p in providers if p["provider"] == "WATSONX"), None) + assert watsonx is not None + + field_keys = [f["key"] for f in watsonx["credential_fields"]] + # Core fields + assert "api_base" in field_keys + assert "project_id" in field_keys + assert "space_id" in field_keys + # Multiple auth methods supported + assert "api_key" in field_keys + assert "token" in field_keys + assert "zen_api_key" in field_keys + diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 5e3652c6d9d..56bba39e6c3 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -202,6 +202,7 @@ ignored_keys = [ "metadata.cold_storage_object_key", "metadata.additional_usage_values.prompt_tokens_details.cache_creation_tokens", "metadata.litellm_overhead_time_ms", + "metadata.cost_breakdown", ] MODEL_LIST = [ diff --git a/tests/test_litellm/proxy/test_empty_model_list.py b/tests/test_litellm/proxy/test_empty_model_list.py new file mode 100644 index 00000000000..6b3e59d3194 --- /dev/null +++ b/tests/test_litellm/proxy/test_empty_model_list.py @@ -0,0 +1,155 @@ +""" +Tests for graceful handling of empty model list scenarios. + +These tests verify that /v2/model/info and /model_group/info endpoints +return empty data arrays instead of 500 errors when no models are configured. +""" + +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system-path + +from litellm.proxy.proxy_server import app + + +@pytest.fixture +def client(): + """Create a test client for the FastAPI app.""" + return TestClient(app) + + +class TestEmptyModelListHandling: + """Test suite for empty model list scenarios.""" + + def test_v2_model_info_returns_empty_data_when_router_is_none( + self, client, monkeypatch + ): + """ + Test that /v2/model/info returns {"data": []} instead of 500 + when llm_router is None. + """ + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + return_value=MagicMock( + user_id="test-user", + team_id=None, + team_models=[], + models=[], + user_role="proxy_admin", + ), + ): + response = client.get( + "/v2/model/info", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert response.json() == {"data": []} + + def test_v2_model_info_returns_empty_data_when_model_list_empty( + self, client, monkeypatch + ): + """ + Test that /v2/model/info returns {"data": []} instead of 500 + when llm_router exists but model_list is empty. + """ + mock_router = MagicMock() + mock_router.model_list = [] + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", []) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + return_value=MagicMock( + user_id="test-user", + team_id=None, + team_models=[], + models=[], + user_role="proxy_admin", + ), + ): + response = client.get( + "/v2/model/info", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert response.json() == {"data": []} + + def test_model_group_info_returns_empty_data_when_model_list_none( + self, client, monkeypatch + ): + """ + Test that /model_group/info returns {"data": []} instead of 500 + when llm_model_list is None. + """ + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", None) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + return_value=MagicMock( + user_id="test-user", + team_id=None, + team_models=[], + models=[], + user_role="proxy_admin", + ), + ): + response = client.get( + "/model_group/info", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert response.json() == {"data": []} + + def test_model_group_info_returns_empty_data_when_model_list_empty( + self, client, monkeypatch + ): + """ + Test that /model_group/info returns {"data": []} instead of 500 + when llm_model_list is empty. + """ + mock_router = MagicMock() + mock_router.model_list = [] + + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_model_list", []) + monkeypatch.setattr("litellm.proxy.proxy_server.user_model", None) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + + with patch( + "litellm.proxy.auth.user_api_key_auth.user_api_key_auth", + return_value=MagicMock( + user_id="test-user", + team_id=None, + team_models=[], + models=[], + user_role="proxy_admin", + ), + ): + response = client.get( + "/model_group/info", + headers={"Authorization": "Bearer sk-test"}, + ) + + assert response.status_code == 200 + assert response.json() == {"data": []} diff --git a/tests/test_litellm/test_gpt_image_cost_calculator.py b/tests/test_litellm/test_gpt_image_cost_calculator.py new file mode 100644 index 00000000000..0a2a62b6c97 --- /dev/null +++ b/tests/test_litellm/test_gpt_image_cost_calculator.py @@ -0,0 +1,239 @@ +""" +Tests for OpenAI gpt-image-1 cost calculator + +This tests the fix for GitHub issue #13847: +https://github.com/BerriAI/litellm/issues/13847 + +gpt-image-1 uses token-based pricing: +- Text Input: $5.00/1M tokens +- Image Input: $10.00/1M tokens +- Image Output: $40.00/1M tokens +""" + +import os +import sys + +sys.path.insert(0, os.path.abspath("../..")) + +import pytest + +import litellm +from litellm.types.utils import ( + ImageResponse, + ImageObject, + ImageUsage, + ImageUsageInputTokensDetails, +) + + +class TestGPTImageCostCalculator: + """Test the OpenAI gpt-image-1 cost calculator""" + + def test_gpt_image_1_cost_with_text_only(self): + """Test cost calculation with only text input tokens""" + from litellm.llms.openai.image_generation.cost_calculator import cost_calculator + + usage = ImageUsage( + input_tokens=100, + output_tokens=5000, + total_tokens=5100, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=100, + image_tokens=0, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + + cost = cost_calculator( + model="gpt-image-1", + image_response=image_response, + custom_llm_provider="openai", + ) + + # Expected cost: + # Text input: 100 * $5/1M = 0.0005 + # Image output: 5000 * $40/1M = 0.2 + # Total: 0.2005 + expected_cost = 0.0005 + 0.2 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + + def test_gpt_image_1_cost_with_image_input(self): + """Test cost calculation with both text and image input tokens (for edits)""" + from litellm.llms.openai.image_generation.cost_calculator import cost_calculator + + usage = ImageUsage( + input_tokens=600, + output_tokens=5000, + total_tokens=5600, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=100, + image_tokens=500, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + + cost = cost_calculator( + model="gpt-image-1", + image_response=image_response, + custom_llm_provider="openai", + ) + + # Expected cost: + # Text input: 100 * $5/1M = 0.0005 + # Image input: 500 * $10/1M = 0.005 + # Image output: 5000 * $40/1M = 0.2 + # Total: 0.2055 + expected_cost = 0.0005 + 0.005 + 0.2 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + + def test_gpt_image_1_mini_cost(self): + """Test cost calculation for gpt-image-1-mini model""" + from litellm.llms.openai.image_generation.cost_calculator import cost_calculator + + usage = ImageUsage( + input_tokens=100, + output_tokens=5000, + total_tokens=5100, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=100, + image_tokens=0, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + + cost = cost_calculator( + model="gpt-image-1-mini", + image_response=image_response, + custom_llm_provider="openai", + ) + + # Expected cost for gpt-image-1-mini: + # Text input: 100 * $2/1M = 0.0002 + # Image output: 5000 * $8/1M = 0.04 + # Total: 0.0402 + expected_cost = 0.0002 + 0.04 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + + def test_gpt_image_1_cost_no_usage(self): + """Test that cost returns 0 when no usage data is available""" + from litellm.llms.openai.image_generation.cost_calculator import cost_calculator + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + + cost = cost_calculator( + model="gpt-image-1", + image_response=image_response, + custom_llm_provider="openai", + ) + + assert cost == 0.0 + + +class TestGPTImageCostRouting: + """Test that gpt-image models are properly routed to the token-based calculator""" + + def test_openai_gpt_image_routes_to_token_calculator(self): + """Test that OpenAI gpt-image-1 routes to token-based calculator""" + from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils + + usage = ImageUsage( + input_tokens=100, + output_tokens=5000, + total_tokens=5100, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=100, + image_tokens=0, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="gpt-image-1", + completion_response=image_response, + custom_llm_provider="openai", + ) + + expected_cost = 0.0005 + 0.2 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + + def test_openai_dalle_routes_to_pixel_calculator(self): + """Test that OpenAI DALL-E still routes to pixel-based calculator""" + from litellm.litellm_core_utils.llm_cost_calc.utils import CostCalculatorUtils + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.size = "1024x1024" + image_response.quality = "standard" + + cost = CostCalculatorUtils.route_image_generation_cost_calculator( + model="dall-e-3", + completion_response=image_response, + custom_llm_provider="openai", + size="1024x1024", + quality="standard", + n=1, + ) + + assert cost >= 0 + + +class TestCompletionCostIntegration: + """Test the full completion_cost integration for gpt-image-1""" + + def test_completion_cost_gpt_image_1(self): + """Test completion_cost correctly calculates gpt-image-1 costs""" + usage = ImageUsage( + input_tokens=100, + output_tokens=5000, + total_tokens=5100, + input_tokens_details=ImageUsageInputTokensDetails( + text_tokens=100, + image_tokens=0, + ), + ) + + image_response = ImageResponse( + created=1234567890, + data=[ImageObject(url="http://example.com/image.jpg")], + ) + image_response.usage = usage + image_response._hidden_params = {"custom_llm_provider": "openai"} + + cost = litellm.completion_cost( + completion_response=image_response, + model="gpt-image-1", + call_type="image_generation", + custom_llm_provider="openai", + ) + + expected_cost = 0.0005 + 0.2 + assert abs(cost - expected_cost) < 1e-6, f"Expected {expected_cost}, got {cost}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index 0eaedaab601..7de99495b79 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -33,6 +33,8 @@ from litellm._lazy_imports import ( _lazy_import_llm_configs, TYPES_NAMES, _lazy_import_types, + LLM_PROVIDER_LOGIC_NAMES, + _lazy_import_llm_provider_logic, ) @@ -218,6 +220,9 @@ def test_unknown_attribute_raises_error(): with pytest.raises(AttributeError): _lazy_import_types("unknown") + with pytest.raises(AttributeError): + _lazy_import_llm_provider_logic("unknown") + def test_llm_config_lazy_imports(): """Test that LLM config classes can be lazy imported.""" @@ -246,3 +251,16 @@ def test_types_lazy_imports(): _verify_only_requested_name_imported(name, TYPES_NAMES) + +def test_llm_provider_logic_lazy_imports(): + """Test that LLM provider logic functions can be lazy imported.""" + for name in LLM_PROVIDER_LOGIC_NAMES: + _clear_names_from_globals(LLM_PROVIDER_LOGIC_NAMES) + + func = _lazy_import_llm_provider_logic(name) + assert func is not None + assert callable(func) + assert name in litellm.__dict__ + + _verify_only_requested_name_imported(name, LLM_PROVIDER_LOGIC_NAMES) + diff --git a/tests/test_litellm/types/llms/test_types_llms_openai.py b/tests/test_litellm/types/llms/test_types_llms_openai.py index 05dec06d469..87cc9586665 100644 --- a/tests/test_litellm/types/llms/test_types_llms_openai.py +++ b/tests/test_litellm/types/llms/test_types_llms_openai.py @@ -35,3 +35,137 @@ def test_output_item_added_event(): assert event.sequence_number == 4 assert event.output_index == 1 assert event.item is None + + +class TestResponsesAPIResponseOutputText: + """Tests for the output_text property on ResponsesAPIResponse""" + + def test_output_text_with_single_message(self): + """Test output_text with a single message containing text output""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Hello, world!", + } + ], + } + ], + ) + + assert response.output_text == "Hello, world!" + + def test_output_text_with_multiple_messages(self): + """Test output_text with multiple messages aggregates all text""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "First part. ", + } + ], + }, + { + "type": "message", + "id": "msg_2", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "Second part.", + } + ], + }, + ], + ) + + assert response.output_text == "First part. Second part." + + def test_output_text_with_no_text_content(self): + """Test output_text returns empty string when no output_text content exists""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + } + ], + ) + + assert response.output_text == "" + + def test_output_text_with_mixed_content(self): + """Test output_text only aggregates output_text type content""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + { + "type": "output_text", + "text": "The weather is sunny. ", + }, + { + "type": "refusal", + "refusal": "I cannot do that.", + }, + ], + }, + { + "type": "function_call", + "id": "call_123", + "status": "completed", + "name": "get_weather", + "arguments": "{}", + }, + ], + ) + + assert response.output_text == "The weather is sunny. " + + def test_output_text_with_empty_output(self): + """Test output_text returns empty string with empty output list""" + from litellm.types.llms.openai import ResponsesAPIResponse + + response = ResponsesAPIResponse( + id="resp_123", + created_at=1234567890, + output=[], + ) + + assert response.output_text == "" diff --git a/ui/litellm-dashboard/e2e_tests/constants.ts b/ui/litellm-dashboard/e2e_tests/constants.ts new file mode 100644 index 00000000000..b07bd68fcf1 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/constants.ts @@ -0,0 +1 @@ +export const ADMIN_STORAGE_PATH = "admin.storageState.json"; diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts b/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts new file mode 100644 index 00000000000..913230ad44b --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/roles.ts @@ -0,0 +1,6 @@ +export enum Role { + ProxyAdmin = "proxy_admin", + ProxyAdminViewer = "proxy_admin_viewer", + InternalUser = "internal_user", + InternalUserViewer = "internal_user_viewer", +} diff --git a/ui/litellm-dashboard/e2e_tests/fixtures/users.ts b/ui/litellm-dashboard/e2e_tests/fixtures/users.ts new file mode 100644 index 00000000000..d1f1eab00e5 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/fixtures/users.ts @@ -0,0 +1,10 @@ +import { Role } from "./roles"; + +const isCI = !!process.env.CI; + +export const users = { + [Role.ProxyAdmin]: { + email: "admin", + password: isCI ? "gm" : "sk-1234", + }, +}; diff --git a/ui/litellm-dashboard/e2e_tests/globalSetup.ts b/ui/litellm-dashboard/e2e_tests/globalSetup.ts new file mode 100644 index 00000000000..a725c58f35b --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/globalSetup.ts @@ -0,0 +1,18 @@ +import { chromium } from "@playwright/test"; +import { users } from "./fixtures/users"; +import { Role } from "./fixtures/roles"; + +async function globalSetup() { + const browser = await chromium.launch(); + const page = await browser.newPage(); + await page.goto("http://localhost:4000/ui/login"); + await page.getByPlaceholder("Enter your username").fill(users[Role.ProxyAdmin].email); + await page.getByPlaceholder("Enter your password").fill(users[Role.ProxyAdmin].password); + const loginButton = page.getByRole("button", { name: "Login" }); + await loginButton.click(); + await page.waitForSelector("text=AI Gateway"); + await page.context().storageState({ path: "admin.storageState.json" }); + await browser.close(); +} + +export default globalSetup; diff --git a/ui/litellm-dashboard/e2e_tests/playwright.config.ts b/ui/litellm-dashboard/e2e_tests/playwright.config.ts new file mode 100644 index 00000000000..329bb7f7afc --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/playwright.config.ts @@ -0,0 +1,48 @@ +import { defineConfig, devices } from "@playwright/test"; + +/** + * See https://playwright.dev/docs/test-configuration. + */ +export default defineConfig({ + testDir: ".", + testMatch: ["**/*.spec.ts", "**/*.setup.ts"], + testIgnore: ["**/*.test.*"], + /* Run tests in files in parallel */ + fullyParallel: true, + /* Fail the build on CI if you accidentally left test.only in the source code. */ + forbidOnly: !!process.env.CI, + /* Retry on CI only */ + retries: process.env.CI ? 2 : 0, + /* Opt out of parallel tests on CI. */ + workers: process.env.CI ? 1 : undefined, + /* Reporter to use. See https://playwright.dev/docs/test-reporters */ + reporter: "html", + /* Shared settings for all the projects below. See https://playwright.dev/docs/api/class-testoptions. */ + use: { + /* Base URL to use in actions like `await page.goto('/')`. */ + baseURL: "http://localhost:4000", + + /* Collect trace when retrying the failed test. See https://playwright.dev/docs/trace-viewer */ + trace: "on-first-retry", + }, + + /* Configure projects for major browsers */ + projects: [ + { + name: "chromium", + use: { ...devices["Desktop Chrome"] }, + }, + + { + name: "firefox", + use: { ...devices["Desktop Firefox"] }, + }, + ], + + /* Timeout settings */ + timeout: 4 * 60 * 1000, + expect: { + timeout: 10 * 1000, + }, + globalSetup: require.resolve("./globalSetup"), +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts new file mode 100644 index 00000000000..d8cc26f8642 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/auth/unauthenticatedRedirect.spec.ts @@ -0,0 +1,11 @@ +import { test, expect } from "@playwright/test"; + +test.describe("Authentication Checks", () => { + test("should redirect unauthenticated user from a protected page", async ({ page }) => { + const protectedPageUrl = "http://localhost:4000/ui?page=llm-playground"; + const expectedRedirectUrl = "http://localhost:4000/ui/login/"; + await page.goto(protectedPageUrl, { waitUntil: "domcontentloaded" }); + await expect(page).toHaveURL(expectedRedirectUrl); + await expect(page.getByRole("heading", { name: "Login" })).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts new file mode 100644 index 00000000000..5ac977ff0c8 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/login/login.spec.ts @@ -0,0 +1,13 @@ +import { expect, test } from "@playwright/test"; +import { users } from "../../fixtures/users"; +import { Role } from "../../fixtures/roles"; + +test("user can log in", async ({ page }) => { + await page.goto("http://localhost:4000/ui/login"); + await page.getByPlaceholder("Enter your username").fill(users[Role.ProxyAdmin].email); + await page.getByPlaceholder("Enter your password").fill(users[Role.ProxyAdmin].password); + const loginButton = page.getByRole("button", { name: "Login" }); + await expect(loginButton).toBeEnabled(); + await loginButton.click(); + await expect(page.getByText("AI Gateway")).toBeVisible(); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts new file mode 100644 index 00000000000..5fa11a98ef6 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/modelsPage/addModel.spec.ts @@ -0,0 +1,23 @@ +import { test, expect } from "@playwright/test"; +import { ADMIN_STORAGE_PATH } from "../../constants"; + +test.describe("Add Model", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Able to see all models for a specific provider in the model dropdown", async ({ page }) => { + await page.goto("http://localhost:4000/ui"); + + await page.getByText("Models + Endpoints").click(); + await page.getByRole("tab", { name: "Add Model" }).click(); + + const providerInputDropdown = page.getByRole("combobox", { name: /Provider/i }); + await providerInputDropdown.fill("Anthropic"); + await page.waitForTimeout(1000); + await providerInputDropdown.press("Enter"); + await page.waitForTimeout(1000); + + const providerModelsDropdown = page.locator(".ant-select-selection-overflow").first(); + await providerModelsDropdown.click(); + await expect(page.getByTitle("claude-haiku-4-5", { exact: true })).toBeVisible(); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts new file mode 100644 index 00000000000..6801f891e87 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/navigation/sidebar.spec.ts @@ -0,0 +1,35 @@ +import test, { expect } from "@playwright/test"; +import { Role } from "../../fixtures/roles"; +import { ADMIN_STORAGE_PATH } from "../../constants"; + +const sidebarButtons = { + [Role.ProxyAdmin]: [ + "Virtual Keys", + "Playground", + "Models", + "Usage", + "Teams", + "Internal User", + "Settings", + "Experimental", + "API Reference", + "AI Hub", + ], +}; + +const roles = [{ role: Role.ProxyAdmin, storage: ADMIN_STORAGE_PATH }]; + +for (const { role, storage } of roles) { + test.describe(`${role} sidebar`, () => { + test.use({ storageState: storage }); + + test("can see and navigate all sidebar buttons", async ({ page }) => { + await page.goto("http://localhost:4000/ui"); + for (const button of sidebarButtons[role as keyof typeof sidebarButtons]) { + const tab = page.getByRole("menuitem", { name: button }); + await expect(tab).toBeVisible(); + await tab.click(); + } + }); + }); +} diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts new file mode 100644 index 00000000000..5873bb3125c --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/users/searchUsers.spec.ts @@ -0,0 +1,90 @@ +import { test, expect, Page } from "@playwright/test"; +test.describe("Internal Users Search", () => { + test.use({ storageState: "admin.storageState.json" }); + + async function goToInternalUsers(page: Page) { + await page.goto("http://localhost:4000/ui"); + + const tab = page.getByRole("menuitem", { name: "Internal User" }); + await expect(tab).toBeVisible(); + await tab.click(); + + await expect(page.locator("tbody tr").first()).toBeVisible(); + await expect(page.locator(".ant-skeleton")).toHaveCount(0); + } + + test("can search users by email", async ({ page }) => { + await goToInternalUsers(page); + + const rows = page.locator("tbody tr"); + const searchInput = page.getByPlaceholder("Search by email..."); + + await expect(searchInput).toBeVisible(); + + // Ensure initial data is loaded + const initialCount = await rows.count(); + expect(initialCount).toBeGreaterThan(0); + + // 🔹 Apply filter + wait for backend response + await Promise.all([ + page.waitForResponse( + (res) => + res.url().includes("/user/list") && + res.url().includes("user_email=test%40") && // encoded "test@" + res.status() === 200, + ), + searchInput.fill("test@"), + ]); + await page.waitForTimeout(5000); + const filteredCount = await rows.count(); + await expect(filteredCount).toBeLessThan(initialCount); + + // 🔹 Clear filter + wait for unfiltered request + await Promise.all([ + page.waitForResponse( + (res) => res.url().includes("/user/list") && !res.url().includes("user_email=") && res.status() === 200, + ), + searchInput.clear(), + ]); + + const resetCount = await rows.count(); + await expect(resetCount).toBe(initialCount); + }); + + test("can filter users by user ID and SSO ID", async ({ page }) => { + await goToInternalUsers(page); + const rows = page.locator("tbody tr"); + + // Ensure initial data is loaded + const initialCount = await rows.count(); + expect(initialCount).toBeGreaterThan(0); + + const filtersButton = page.getByRole("button", { + name: "Filters", + exact: true, + }); + await filtersButton.click(); + + const userIdInput = page.getByPlaceholder("Filter by User ID"); + const ssoIdInput = page.getByPlaceholder("Filter by SSO ID"); + await Promise.all([ + page.waitForResponse( + (res) => res.url().includes("/user/list") && res.url().includes("user_ids=user") && res.status() === 200, + ), + userIdInput.fill("user"), + ]); + + await Promise.all([ + page.waitForResponse( + (res) => + res.url().includes("/user/list") && + res.url().includes("user_ids=user") && + res.url().includes("sso_user_ids=sso") && + res.status() === 200, + ), + ssoIdInput.fill("sso"), + ]); + const combinedFilteredCount = await rows.count(); + await expect(combinedFilteredCount).toBeLessThan(initialCount); + }); +}); diff --git a/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts b/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts new file mode 100644 index 00000000000..980c7233e42 --- /dev/null +++ b/ui/litellm-dashboard/e2e_tests/tests/users/viewInternalUsers.spec.ts @@ -0,0 +1,52 @@ +import { test, expect, Page } from "@playwright/test"; + +test.describe("Internal Users Page", () => { + test.use({ storageState: "admin.storageState.json" }); + + async function goToInternalUsers(page: Page) { + await page.goto("http://localhost:4000/ui"); + + const internalUserTab = page.getByRole("menuitem", { name: "Internal User" }); + await expect(internalUserTab).toBeVisible(); + await internalUserTab.click(); + + const firstRow = page.locator("tbody tr").first(); + await expect(firstRow).toBeVisible(); + await expect(page.locator(".ant-skeleton")).toHaveCount(0); + } + + test("renders internal users table correctly", async ({ page }) => { + await goToInternalUsers(page); + + const rows = page.locator("tbody tr"); + const rowCount = await rows.count(); + expect(rowCount).toBeGreaterThan(0); + + const userIdHeader = page.getByRole("columnheader", { name: "User ID" }); + await expect(userIdHeader).toBeVisible(); + + const virtualKeysHeader = page.getByRole("columnheader", { name: "Virtual Keys" }); + await expect(virtualKeysHeader).toBeVisible(); + }); + + test("pagination controls work correctly", async ({ page }) => { + await goToInternalUsers(page); + + const paginationInfo = page.locator(".text-sm.text-gray-700"); + const prevButton = page.getByRole("button", { name: "Previous" }); + const nextButton = page.getByRole("button", { name: "Next" }); + + const infoText = (await paginationInfo.textContent()) || ""; + + // On first page, Previous should be disabled + if (infoText.includes("1 -")) { + await expect(prevButton).toBeDisabled(); + } + + // Check if there are more pages + const hasMorePages = infoText.includes("of") && !infoText.endsWith("25 of 25"); + if (hasMorePages) { + await expect(nextButton).toBeEnabled(); + } + }); +}); diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index b65c35e78d8..0e9675f40a1 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -39,6 +39,7 @@ "uuid": "^11.1.0" }, "devDependencies": { + "@playwright/test": "^1.57.0", "@tailwindcss/forms": "^0.5.7", "@testing-library/dom": "^10.4.1", "@testing-library/jest-dom": "^6.8.0", @@ -91,7 +92,6 @@ "version": "5.2.0", "resolved": "https://registry.npmjs.org/@alloc/quick-lru/-/quick-lru-5.2.0.tgz", "integrity": "sha512-UrcABB+4bUrFABwbluTIBErXwvbsU/V7TZWfmbgJfbkwiBuziS9gxdODUyuiecfdGQ85jglMW6juS3+z5TsKLw==", - "dev": true, "license": "MIT", "engines": { "node": ">=10" @@ -325,6 +325,7 @@ "resolved": "https://registry.npmjs.org/@babel/core/-/core-7.28.5.tgz", "integrity": "sha512-e7jT4DxYvIDLk1ZHmU/m/mB19rex9sv0c2ftBtjSBv+kVM/902eh0fINUzD7UwLLNR+jU585GxUJ8/EBfAM5fw==", "license": "MIT", + "peer": true, "dependencies": { "@babel/code-frame": "^7.27.1", "@babel/generator": "^7.28.5", @@ -2186,6 +2187,7 @@ } ], "license": "MIT", + "peer": true, "engines": { "node": ">=18" }, @@ -2228,6 +2230,7 @@ } ], "license": "MIT", + "peer": true, "engines": { "node": ">=18" } @@ -2337,6 +2340,7 @@ "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz", "integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==", "license": "MIT", + "peer": true, "dependencies": { "cssesc": "^3.0.0", "util-deprecate": "^1.0.2" @@ -2758,6 +2762,7 @@ "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz", "integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==", "license": "MIT", + "peer": true, "dependencies": { "cssesc": "^3.0.0", "util-deprecate": "^1.0.2" @@ -3541,6 +3546,40 @@ "react-dom": "*" } }, + "node_modules/@docusaurus/plugin-content-docs": { + "version": "3.9.2", + "resolved": "https://registry.npmjs.org/@docusaurus/plugin-content-docs/-/plugin-content-docs-3.9.2.tgz", + "integrity": "sha512-C5wZsGuKTY8jEYsqdxhhFOe1ZDjH0uIYJ9T/jebHwkyxqnr4wW0jTkB72OMqNjsoQRcb0JN3PcSeTwFlVgzCZg==", + "license": "MIT", + "peer": true, + "dependencies": { + "@docusaurus/core": "3.9.2", + "@docusaurus/logger": "3.9.2", + "@docusaurus/mdx-loader": "3.9.2", + "@docusaurus/module-type-aliases": "3.9.2", + "@docusaurus/theme-common": "3.9.2", + "@docusaurus/types": "3.9.2", + "@docusaurus/utils": "3.9.2", + "@docusaurus/utils-common": "3.9.2", + "@docusaurus/utils-validation": "3.9.2", + "@types/react-router-config": "^5.0.7", + "combine-promises": "^1.1.0", + "fs-extra": "^11.1.1", + "js-yaml": "^4.1.0", + "lodash": "^4.17.21", + "schema-dts": "^1.1.2", + "tslib": "^2.6.0", + "utility-types": "^3.10.0", + "webpack": "^5.88.1" + }, + "engines": { + "node": ">=20.0" + }, + "peerDependencies": { + "react": "^18.0.0 || ^19.0.0", + "react-dom": "^18.0.0 || ^19.0.0" + } + }, "node_modules/@docusaurus/theme-common": { "version": "3.9.2", "resolved": "https://registry.npmjs.org/@docusaurus/theme-common/-/theme-common-3.9.2.tgz", @@ -4700,6 +4739,24 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/@mdx-js/react": { + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/@mdx-js/react/-/react-3.1.1.tgz", + "integrity": "sha512-f++rKLQgUVYDAtECQ6fn/is15GkEH9+nZPM3MS0RcxVqoTfawHvDlSCH7JbMhAM6uJ32v3eXLvLmLvjGu7PTQw==", + "license": "MIT", + "peer": true, + "dependencies": { + "@types/mdx": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + }, + "peerDependencies": { + "@types/react": ">=16", + "react": ">=16" + } + }, "node_modules/@mermaid-js/parser": { "version": "0.6.3", "resolved": "https://registry.npmjs.org/@mermaid-js/parser/-/parser-0.6.3.tgz", @@ -4927,6 +4984,23 @@ "node": ">=12.4.0" } }, + "node_modules/@playwright/test": { + "version": "1.57.0", + "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.57.0.tgz", + "integrity": "sha512-6TyEnHgd6SArQO8UO2OMTxshln3QMWBtPGrOCgs3wVEmQmwyuNtB10IZMfmYDE0riwNR1cu4q+pPcxMVtaG3TA==", + "devOptional": true, + "license": "Apache-2.0", + "peer": true, + "dependencies": { + "playwright": "1.57.0" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, "node_modules/@pnpm/config.env-replace": { "version": "1.1.0", "resolved": "https://registry.npmjs.org/@pnpm/config.env-replace/-/config.env-replace-1.1.0.tgz", @@ -5773,6 +5847,7 @@ "integrity": "sha512-o4PXJQidqJl82ckFaXUeoAW+XysPLauYI43Abki5hABd853iMhitooc6znOnczgbTYmEP6U6/y1ZyKAIsvMKGg==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@babel/code-frame": "^7.10.4", "@babel/runtime": "^7.12.5", @@ -6566,6 +6641,7 @@ "resolved": "https://registry.npmjs.org/@types/react/-/react-18.2.48.tgz", "integrity": "sha512-qboRCl6Ie70DQQG9hhNREz81jqC1cs9EVNcjQ1AU+jH6NFfSAhVVbrrY/+nSF+Bsk4AOwm9Qa61InvMCyV+H3w==", "license": "MIT", + "peer": true, "dependencies": { "@types/prop-types": "*", "@types/scheduler": "*", @@ -6588,6 +6664,7 @@ "integrity": "sha512-MEe3UeoENYVFXzoXEWsvcpg6ZvlrFNlOQ7EOsvhI3CfAXwzPfO8Qwuxd40nepsYKqyyVQnTdEfv68q91yLcKrQ==", "dev": true, "license": "MIT", + "peer": true, "peerDependencies": { "@types/react": "^18.0.0" } @@ -6784,6 +6861,7 @@ "integrity": "sha512-lJi3PfxVmo0AkEY93ecfN+r8SofEqZNGByvHAI3GBLrvt1Cw6H5k1IM02nSzu0RfUafr2EvFSw0wAsZgubNplQ==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@typescript-eslint/scope-manager": "8.47.0", "@typescript-eslint/types": "8.47.0", @@ -7445,6 +7523,7 @@ "integrity": "sha512-hGISOaP18plkzbWEcP/QvtRW1xDXF2+96HbEX6byqQhAUbiS5oH6/9JwW+QsQCIYON2bI6QZBF+2PvOmrRZ9wA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@vitest/utils": "3.2.4", "fflate": "^0.8.2", @@ -7673,6 +7752,7 @@ "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz", "integrity": "sha512-NZyJarBfL7nWwIq+FDL6Zp/yHEhePMNnnJ0y3qfieCrmNvYct8uvtiV41UvlSe6apAfk0fY1FbWx+NwfmpvtTg==", "license": "MIT", + "peer": true, "bin": { "acorn": "bin/acorn" }, @@ -7762,6 +7842,7 @@ "resolved": "https://registry.npmjs.org/ajv/-/ajv-6.12.6.tgz", "integrity": "sha512-j3fVLgvTo527anyYyJOGTYJbG+vnnQYvE0m5mmkc1TK+nxAppkCLMIL0aZ4dblVCNoGShhm+kzE4ZUykBoMg4g==", "license": "MIT", + "peer": true, "dependencies": { "fast-deep-equal": "^3.1.1", "fast-json-stable-stringify": "^2.0.0", @@ -7982,7 +8063,6 @@ "version": "1.3.0", "resolved": "https://registry.npmjs.org/any-promise/-/any-promise-1.3.0.tgz", "integrity": "sha512-7UvmKalWRt1wgjL1RrGxoSJW/0QZFIegpeGvZG9kjp8vrRu55XTHbwnqq2GpXm9uLbcuhxm3IqX9OB4MZR1b2A==", - "dev": true, "license": "MIT" }, "node_modules/anymatch": { @@ -8002,7 +8082,6 @@ "version": "5.0.2", "resolved": "https://registry.npmjs.org/arg/-/arg-5.0.2.tgz", "integrity": "sha512-PYjyFOLKQ9y57JvQ6QLo8dAgNqswh8M1RMJYdQduT6xbWSgK36P/Z/v+p888pM69jMMfS8Xd8F6I1kQ/I9HUGg==", - "dev": true, "license": "MIT" }, "node_modules/argparse": { @@ -8479,23 +8558,23 @@ } }, "node_modules/body-parser": { - "version": "1.20.3", - "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.3.tgz", - "integrity": "sha512-7rAxByjUMqQ3/bHJy7D6OGXvx/MMc4IqBn/X0fcM1QUcAItpZrBEYhWGem+tzXH90c+G01ypMcYJBO9Y30203g==", + "version": "1.20.4", + "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.4.tgz", + "integrity": "sha512-ZTgYYLMOXY9qKU/57FAo8F+HA2dGX7bqGc71txDRC1rS4frdFI5R7NhluHxH6M0YItAP0sHB4uqAOcYKxO6uGA==", "license": "MIT", "dependencies": { - "bytes": "3.1.2", + "bytes": "~3.1.2", "content-type": "~1.0.5", "debug": "2.6.9", "depd": "2.0.0", - "destroy": "1.2.0", - "http-errors": "2.0.0", - "iconv-lite": "0.4.24", - "on-finished": "2.4.1", - "qs": "6.13.0", - "raw-body": "2.5.2", + "destroy": "~1.2.0", + "http-errors": "~2.0.1", + "iconv-lite": "~0.4.24", + "on-finished": "~2.4.1", + "qs": "~6.14.0", + "raw-body": "~2.5.3", "type-is": "~1.6.18", - "unpipe": "1.0.0" + "unpipe": "~1.0.0" }, "engines": { "node": ">= 0.8", @@ -8520,6 +8599,26 @@ "ms": "2.0.0" } }, + "node_modules/body-parser/node_modules/http-errors": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", + "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", + "license": "MIT", + "dependencies": { + "depd": "~2.0.0", + "inherits": "~2.0.4", + "setprototypeof": "~1.2.0", + "statuses": "~2.0.2", + "toidentifier": "~1.0.1" + }, + "engines": { + "node": ">= 0.8" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, "node_modules/body-parser/node_modules/iconv-lite": { "version": "0.4.24", "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz", @@ -8538,6 +8637,15 @@ "integrity": "sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A==", "license": "MIT" }, + "node_modules/body-parser/node_modules/statuses": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", + "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, "node_modules/bonjour-service": { "version": "1.3.0", "resolved": "https://registry.npmjs.org/bonjour-service/-/bonjour-service-1.3.0.tgz", @@ -8617,6 +8725,7 @@ } ], "license": "MIT", + "peer": true, "dependencies": { "baseline-browser-mapping": "^2.8.25", "caniuse-lite": "^1.0.30001754", @@ -8797,7 +8906,6 @@ "version": "2.0.1", "resolved": "https://registry.npmjs.org/camelcase-css/-/camelcase-css-2.0.1.tgz", "integrity": "sha512-QOSvevhslijgYwRx6Rv7zKdMF8lbRmx+uQGx2+vDc+KI/eBnsy9kit5aj23AgGu3pa4t9AgwbnXWqS+iOY+2aA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -8942,6 +9050,7 @@ "resolved": "https://registry.npmjs.org/chevrotain/-/chevrotain-11.0.3.tgz", "integrity": "sha512-ci2iJH6LeIkvP9eJW6gpueU8cnZhv85ELY8w8WiFtNjMHA5ad6pQLaJo9mEly/9qUyCpvqX8/POVUTf18/HFdw==", "license": "Apache-2.0", + "peer": true, "dependencies": { "@chevrotain/cst-dts-gen": "11.0.3", "@chevrotain/gast": "11.0.3", @@ -9670,6 +9779,7 @@ "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz", "integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==", "license": "MIT", + "peer": true, "dependencies": { "cssesc": "^3.0.0", "util-deprecate": "^1.0.2" @@ -10032,6 +10142,7 @@ "resolved": "https://registry.npmjs.org/cytoscape/-/cytoscape-3.33.1.tgz", "integrity": "sha512-iJc4TwyANnOGR1OmWhsS9ayRS3s+XQ185FmuHObThD+5AeJCakAAbWv8KimMTt08xCCLNgneQwFp+JRJOr9qGQ==", "license": "MIT", + "peer": true, "engines": { "node": ">=0.10" } @@ -10441,6 +10552,7 @@ "resolved": "https://registry.npmjs.org/d3-selection/-/d3-selection-3.0.0.tgz", "integrity": "sha512-fmTRWbNMmsmWq6xJV8D19U/gw/bwrHfNXxrIN+HfZgnzqTHp9jOmKMhsTUjXOJnZOdZY9Q28y4yebKzqDKlxlQ==", "license": "ISC", + "peer": true, "engines": { "node": ">=12" } @@ -10615,6 +10727,7 @@ "resolved": "https://registry.npmjs.org/date-fns/-/date-fns-3.6.0.tgz", "integrity": "sha512-fRHTG8g/Gif+kSh50gaGEdToemgfj74aRX3swtiouboip5JDLAyDE9F11nHMIcvOaXeOC6D7SpNhi7uFyB7Uww==", "license": "MIT", + "peer": true, "funding": { "type": "github", "url": "https://github.com/sponsors/kossnocorp" @@ -10894,7 +11007,6 @@ "version": "1.2.2", "resolved": "https://registry.npmjs.org/didyoumean/-/didyoumean-1.2.2.tgz", "integrity": "sha512-gxtyfqMg7GKyhQmb056K7M3xszy/myH8w+B4RT+QXBQsvAOdc3XymqDDPHx1BgPgsdAA5SIifona89YtRATDzw==", - "dev": true, "license": "Apache-2.0" }, "node_modules/dir-glob": { @@ -10913,7 +11025,6 @@ "version": "1.1.3", "resolved": "https://registry.npmjs.org/dlv/-/dlv-1.1.3.tgz", "integrity": "sha512-+HlytyjlPKnIG8XuRG8WvmBP8xs8P71y+SKKS6ZXWoEgLuePxtDoUEiH7WkdePWrQ5JBpE6aoVqfZfJUQkjXwA==", - "dev": true, "license": "MIT" }, "node_modules/dns-packet": { @@ -11494,6 +11605,7 @@ "deprecated": "This version is no longer supported. Please see https://eslint.org/version-support for other options.", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@eslint-community/eslint-utils": "^4.2.0", "@eslint-community/regexpp": "^4.6.1", @@ -11679,6 +11791,7 @@ "integrity": "sha512-whOE1HFo/qJDyX4SnXzP4N6zOWn79WhnCUY/iDR0mPfQZO8wcYE4JClzI2oZrhBnnMUCBCHZhO6VQyoBU95mZA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@rtsao/scc": "^1.1.0", "array-includes": "^3.1.9", @@ -12181,39 +12294,39 @@ } }, "node_modules/express": { - "version": "4.21.2", - "resolved": "https://registry.npmjs.org/express/-/express-4.21.2.tgz", - "integrity": "sha512-28HqgMZAmih1Czt9ny7qr6ek2qddF4FclbMzwhCREB6OFfH+rXAnuNCwo1/wFvrtbgsQDb4kSbX9de9lFbrXnA==", + "version": "4.22.1", + "resolved": "https://registry.npmjs.org/express/-/express-4.22.1.tgz", + "integrity": "sha512-F2X8g9P1X7uCPZMA3MVf9wcTqlyNp7IhH5qPCI0izhaOIYXaW9L535tGA3qmjRzpH+bZczqq7hVKxTR4NWnu+g==", "license": "MIT", "dependencies": { "accepts": "~1.3.8", "array-flatten": "1.1.1", - "body-parser": "1.20.3", - "content-disposition": "0.5.4", + "body-parser": "~1.20.3", + "content-disposition": "~0.5.4", "content-type": "~1.0.4", - "cookie": "0.7.1", - "cookie-signature": "1.0.6", + "cookie": "~0.7.1", + "cookie-signature": "~1.0.6", "debug": "2.6.9", "depd": "2.0.0", "encodeurl": "~2.0.0", "escape-html": "~1.0.3", "etag": "~1.8.1", - "finalhandler": "1.3.1", - "fresh": "0.5.2", - "http-errors": "2.0.0", + "finalhandler": "~1.3.1", + "fresh": "~0.5.2", + "http-errors": "~2.0.0", "merge-descriptors": "1.0.3", "methods": "~1.1.2", - "on-finished": "2.4.1", + "on-finished": "~2.4.1", "parseurl": "~1.3.3", - "path-to-regexp": "0.1.12", + "path-to-regexp": "~0.1.12", "proxy-addr": "~2.0.7", - "qs": "6.13.0", + "qs": "~6.14.0", "range-parser": "~1.2.1", "safe-buffer": "5.2.1", - "send": "0.19.0", - "serve-static": "1.16.2", + "send": "~0.19.0", + "serve-static": "~1.16.2", "setprototypeof": "1.2.0", - "statuses": "2.0.1", + "statuses": "~2.0.1", "type-is": "~1.6.18", "utils-merge": "1.0.1", "vary": "~1.1.2" @@ -14942,6 +15055,7 @@ "integrity": "sha512-454TI39PeRDW1LgpyLPyURtB4Zx1tklSr6+OFOipsxGUH1WMTvk6C65JQdrj455+DP2uJ1+veBEHTGFKWVLFoA==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@acemir/cssom": "^0.9.23", "@asamuzakjp/dom-selector": "^6.7.4", @@ -18069,6 +18183,7 @@ "resolved": "https://registry.npmjs.org/moment/-/moment-2.30.1.tgz", "integrity": "sha512-uEmtNhbDOrWPFS+hdjFCBfy9f2YoyzRpwcl+DqpC6taX21FzsTLQVbMV/W7PzNSX6x/bhC1zA3c2UQ5NzH6how==", "license": "MIT", + "peer": true, "engines": { "node": "*" } @@ -18105,7 +18220,6 @@ "version": "2.7.0", "resolved": "https://registry.npmjs.org/mz/-/mz-2.7.0.tgz", "integrity": "sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q==", - "dev": true, "license": "MIT", "dependencies": { "any-promise": "^1.0.0", @@ -18454,7 +18568,6 @@ "version": "3.0.0", "resolved": "https://registry.npmjs.org/object-hash/-/object-hash-3.0.0.tgz", "integrity": "sha512-RSn9F68PjH9HqtltsSnqYC1XXoWe9Bju5+213R98cNGttag9q9yAOTzdbsqvIa7aNm5WffBZFpWYr2aWrklWAw==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -19095,7 +19208,6 @@ "version": "2.3.0", "resolved": "https://registry.npmjs.org/pify/-/pify-2.3.0.tgz", "integrity": "sha512-udgsAY+fTnvv7kI7aaxbqwWNb0AHiB0qBO89PZKPkoTmGOgdbrHDKD+0B2X4uTfJ/FT1R09r9gTsjUjNJotuog==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.10.0" @@ -19105,7 +19217,6 @@ "version": "4.0.7", "resolved": "https://registry.npmjs.org/pirates/-/pirates-4.0.7.tgz", "integrity": "sha512-TfySrs/5nm8fQJDcBDuUng3VOUKsd7S+zqvbOTiGXHfxX4wK31ard+hoNuvkicM/2YFzlpDgABOevKSsB4G/FA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -19219,6 +19330,53 @@ "pathe": "^2.0.3" } }, + "node_modules/playwright": { + "version": "1.57.0", + "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.57.0.tgz", + "integrity": "sha512-ilYQj1s8sr2ppEJ2YVadYBN0Mb3mdo9J0wQ+UuDhzYqURwSoW4n1Xs5vs7ORwgDGmyEh33tRMeS8KhdkMoLXQw==", + "devOptional": true, + "license": "Apache-2.0", + "dependencies": { + "playwright-core": "1.57.0" + }, + "bin": { + "playwright": "cli.js" + }, + "engines": { + "node": ">=18" + }, + "optionalDependencies": { + "fsevents": "2.3.2" + } + }, + "node_modules/playwright-core": { + "version": "1.57.0", + "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.57.0.tgz", + "integrity": "sha512-agTcKlMw/mjBWOnD6kFZttAAGHgi/Nw0CZ2o6JqWSbMlI219lAFLZZCyqByTsvVAJq5XA5H8cA6PrvBRpBWEuQ==", + "devOptional": true, + "license": "Apache-2.0", + "bin": { + "playwright-core": "cli.js" + }, + "engines": { + "node": ">=18" + } + }, + "node_modules/playwright/node_modules/fsevents": { + "version": "2.3.2", + "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", + "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", + "dev": true, + "hasInstallScript": true, + "license": "MIT", + "optional": true, + "os": [ + "darwin" + ], + "engines": { + "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + } + }, "node_modules/points-on-curve": { "version": "0.2.0", "resolved": "https://registry.npmjs.org/points-on-curve/-/points-on-curve-0.2.0.tgz", @@ -19264,6 +19422,7 @@ } ], "license": "MIT", + "peer": true, "dependencies": { "nanoid": "^3.3.11", "picocolors": "^1.1.1", @@ -19820,7 +19979,6 @@ "version": "15.1.0", "resolved": "https://registry.npmjs.org/postcss-import/-/postcss-import-15.1.0.tgz", "integrity": "sha512-hpr+J05B2FVYUAXHeK1YyI267J/dDDhMU6B6civm8hSY1jYJnBXxzKDKDswzJmtLHryrjhnDjqqp/49t8FALew==", - "dev": true, "license": "MIT", "dependencies": { "postcss-value-parser": "^4.0.0", @@ -19838,7 +19996,6 @@ "version": "4.1.0", "resolved": "https://registry.npmjs.org/postcss-js/-/postcss-js-4.1.0.tgz", "integrity": "sha512-oIAOTqgIo7q2EOwbhb8UalYePMvYoIeRY2YKntdpFQXNosSu3vLrniGgmH9OKs/qAkfoj5oB3le/7mINW1LCfw==", - "dev": true, "funding": [ { "type": "opencollective", @@ -19893,7 +20050,6 @@ "version": "6.0.1", "resolved": "https://registry.npmjs.org/postcss-load-config/-/postcss-load-config-6.0.1.tgz", "integrity": "sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g==", - "dev": true, "funding": [ { "type": "opencollective", @@ -20182,7 +20338,6 @@ "version": "6.2.0", "resolved": "https://registry.npmjs.org/postcss-nested/-/postcss-nested-6.2.0.tgz", "integrity": "sha512-HQbt28KulC5AJzG+cZtj9kvKB93CFCdLvog1WFLf1D+xmMvPGlBstkpTEZfK5+AN9hfJocyBFCNiqyS48bpgzQ==", - "dev": true, "funding": [ { "type": "opencollective", @@ -20280,6 +20435,7 @@ "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-7.1.0.tgz", "integrity": "sha512-8sLjZwK0R+JlxlYcTuVnyT2v+htpdrjDOKuMcOVdYjt52Lh8hWRYpxBPoKx/Zg+bcjc3wx6fmQevMmUztS/ccA==", "license": "MIT", + "peer": true, "dependencies": { "cssesc": "^3.0.0", "util-deprecate": "^1.0.2" @@ -21011,12 +21167,12 @@ } }, "node_modules/qs": { - "version": "6.13.0", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.13.0.tgz", - "integrity": "sha512-+38qI9SOr8tfZ4QmJNplMUxqjbe7LKvvZgWdExBOmd+egZTtjLB67Gu0HRX3u/XOq7UU2Nx6nsjvS16Z9uwfpg==", + "version": "6.14.1", + "resolved": "https://registry.npmjs.org/qs/-/qs-6.14.1.tgz", + "integrity": "sha512-4EK3+xJl8Ts67nLYNwqw/dsFVnCf+qR7RgXSK9jEEm9unao3njwMDdmsdvoKBKHzxd7tCYz5e5M+SnMjdtXGQQ==", "license": "BSD-3-Clause", "dependencies": { - "side-channel": "^1.0.6" + "side-channel": "^1.1.0" }, "engines": { "node": ">=0.6" @@ -21092,15 +21248,15 @@ } }, "node_modules/raw-body": { - "version": "2.5.2", - "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.2.tgz", - "integrity": "sha512-8zGqypfENjCIqGhgXToC8aB2r7YrBX+AQAfIPs/Mlk+BtPTztOvTS01NRW/3Eh60J+a48lt8qsCzirQ6loCVfA==", + "version": "2.5.3", + "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.3.tgz", + "integrity": "sha512-s4VSOf6yN0rvbRZGxs8Om5CWj6seneMwK3oDb4lWDH0UPhWcxwOWw5+qk24bxq87szX1ydrwylIOp2uG1ojUpA==", "license": "MIT", "dependencies": { - "bytes": "3.1.2", - "http-errors": "2.0.0", - "iconv-lite": "0.4.24", - "unpipe": "1.0.0" + "bytes": "~3.1.2", + "http-errors": "~2.0.1", + "iconv-lite": "~0.4.24", + "unpipe": "~1.0.0" }, "engines": { "node": ">= 0.8" @@ -21115,6 +21271,26 @@ "node": ">= 0.8" } }, + "node_modules/raw-body/node_modules/http-errors": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", + "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", + "license": "MIT", + "dependencies": { + "depd": "~2.0.0", + "inherits": "~2.0.4", + "setprototypeof": "~1.2.0", + "statuses": "~2.0.2", + "toidentifier": "~1.0.1" + }, + "engines": { + "node": ">= 0.8" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, "node_modules/raw-body/node_modules/iconv-lite": { "version": "0.4.24", "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz", @@ -21127,6 +21303,15 @@ "node": ">=0.10.0" } }, + "node_modules/raw-body/node_modules/statuses": { + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", + "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", + "license": "MIT", + "engines": { + "node": ">= 0.8" + } + }, "node_modules/rc": { "version": "1.2.8", "resolved": "https://registry.npmjs.org/rc/-/rc-1.2.8.tgz", @@ -21774,6 +21959,7 @@ "resolved": "https://registry.npmjs.org/react/-/react-18.3.1.tgz", "integrity": "sha512-wS+hAgJShR0KhEvPJArfuPVN1+Hz1t0Y6n5jLrGQbkb4urgPE/0Rve+1kMB1v/oWgHgm4WIcV+i7F2pTVj+2iQ==", "license": "MIT", + "peer": true, "dependencies": { "loose-envify": "^1.1.0" }, @@ -21813,6 +21999,7 @@ "resolved": "https://registry.npmjs.org/react-dom/-/react-dom-18.3.1.tgz", "integrity": "sha512-5m4nQKp+rZRb09LNH59GM4BxTh9251/ylbKIbpe7TpGxfJ+9kv6BLkLBXIjjspbgbnIBNqlI23tRnTWT0snUIw==", "license": "MIT", + "peer": true, "dependencies": { "loose-envify": "^1.1.0", "scheduler": "^0.23.2" @@ -21870,6 +22057,7 @@ "resolved": "https://registry.npmjs.org/@docusaurus/react-loadable/-/react-loadable-6.0.0.tgz", "integrity": "sha512-YMMxTUQV/QFSnbgrP3tjDzLHRg7vsbMn8e9HAa8o/1iXoiomo48b7sk/kkmWEuWNDPJVlKSJRB6Y2fHqdJk+SQ==", "license": "MIT", + "peer": true, "dependencies": { "@types/react": "*" }, @@ -21935,6 +22123,7 @@ "resolved": "https://registry.npmjs.org/react-router/-/react-router-5.3.4.tgz", "integrity": "sha512-Ys9K+ppnJah3QuaRiLxk+jDWOR1MekYQrlytiXxC1RyfbdsZkS5pvKAzCCr031xHixZwpnsYNT5xysdFHQaYsA==", "license": "MIT", + "peer": true, "dependencies": { "@babel/runtime": "^7.12.13", "history": "^4.9.0", @@ -22049,7 +22238,6 @@ "version": "1.0.0", "resolved": "https://registry.npmjs.org/read-cache/-/read-cache-1.0.0.tgz", "integrity": "sha512-Owdv/Ft7IjOgm/i0xvNDZ1LrRANRfew4b2prF3OWMQLxLfu3bS8FVhCsrSCMK4lR56Y9ya+AThoTpDCTxCmpRA==", - "dev": true, "license": "MIT", "dependencies": { "pify": "^2.3.0" @@ -22969,6 +23157,12 @@ "loose-envify": "^1.1.0" } }, + "node_modules/schema-dts": { + "version": "1.1.5", + "resolved": "https://registry.npmjs.org/schema-dts/-/schema-dts-1.1.5.tgz", + "integrity": "sha512-RJr9EaCmsLzBX2NDiO5Z3ux2BVosNZN5jo0gWgsyKvxKIUL5R3swNvoorulAeL9kLB0iTSX7V6aokhla2m7xbg==", + "license": "Apache-2.0" + }, "node_modules/schema-utils": { "version": "4.3.3", "resolved": "https://registry.npmjs.org/schema-utils/-/schema-utils-4.3.3.tgz", @@ -22993,6 +23187,7 @@ "resolved": "https://registry.npmjs.org/ajv/-/ajv-8.17.1.tgz", "integrity": "sha512-B/gBuNg5SiMTrPkC+A2+cW0RszwxYmn6VYxB/inlBStS5nx6xHIt/ehKRhIMhqusl7a8LjQoZnjCs5vhwxOQ1g==", "license": "MIT", + "peer": true, "dependencies": { "fast-deep-equal": "^3.1.3", "fast-uri": "^3.0.1", @@ -24033,7 +24228,6 @@ "version": "3.35.1", "resolved": "https://registry.npmjs.org/sucrase/-/sucrase-3.35.1.tgz", "integrity": "sha512-DhuTmvZWux4H1UOnWMB3sk0sbaCVOoQZjv8u1rDoTV0HTdGem9hkAZtl4JZy8P2z4Bg0nT+YMeOFyVr4zcG5Tw==", - "dev": true, "license": "MIT", "dependencies": { "@jridgewell/gen-mapping": "^0.3.2", @@ -24056,7 +24250,6 @@ "version": "4.1.1", "resolved": "https://registry.npmjs.org/commander/-/commander-4.1.1.tgz", "integrity": "sha512-NOKm8xhkzAjzFx8B2v5OAHT+u5pRQc2UCa2Vq9jYL/31o2wi9mxBA7LIFs3sV5VSC49z6pEhfbMULvShKj26WA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -24225,8 +24418,8 @@ "version": "3.4.18", "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-3.4.18.tgz", "integrity": "sha512-6A2rnmW5xZMdw11LYjhcI5846rt9pbLSabY5XPxo+XWdxwZaFEn47Go4NzFiHu9sNNmr/kXivP1vStfvMaK1GQ==", - "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@alloc/quick-lru": "^5.2.0", "arg": "^5.0.2", @@ -24263,7 +24456,6 @@ "version": "6.0.2", "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-6.0.2.tgz", "integrity": "sha512-XxwI8EOhVQgWp6iDL+3b0r86f4d6AX6zSU55HfB4ydCEuXLXc5FcYeOu+nnGftS4TEju/11rt4KJPTMgbfmv4A==", - "dev": true, "license": "ISC", "dependencies": { "is-glob": "^4.0.3" @@ -24424,7 +24616,6 @@ "version": "3.3.1", "resolved": "https://registry.npmjs.org/thenify/-/thenify-3.3.1.tgz", "integrity": "sha512-RVZSIV5IG10Hk3enotrhvz0T9em6cyHBLkH/YAZuKqd8hRkKhSfCGIcP2KUY0EPxndzANBmNllzWPwak+bheSw==", - "dev": true, "license": "MIT", "dependencies": { "any-promise": "^1.0.0" @@ -24434,7 +24625,6 @@ "version": "1.6.0", "resolved": "https://registry.npmjs.org/thenify-all/-/thenify-all-1.6.0.tgz", "integrity": "sha512-RNxQH/qI8/t3thXJDwcstUO4zeqo64+Uy/+sNVRBx4Xn2OX+OZ9oP+iJnNFqplFra2ZUVeKCSa2oVWi3T4uVmA==", - "dev": true, "license": "MIT", "dependencies": { "thenify": ">= 3.1.0 < 4" @@ -24506,7 +24696,6 @@ "version": "0.2.15", "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.15.tgz", "integrity": "sha512-j2Zq4NyQYG5XMST4cbs02Ak8iJUdxRM0XI5QyxXuZOzKOINmWurp3smXu3y5wDcJrptwpSjgXHzIQxR0omXljQ==", - "dev": true, "license": "MIT", "dependencies": { "fdir": "^6.5.0", @@ -24523,7 +24712,6 @@ "version": "6.5.0", "resolved": "https://registry.npmjs.org/fdir/-/fdir-6.5.0.tgz", "integrity": "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==", - "dev": true, "license": "MIT", "engines": { "node": ">=12.0.0" @@ -24541,8 +24729,8 @@ "version": "4.0.3", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", - "dev": true, "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -24723,7 +24911,6 @@ "version": "0.1.13", "resolved": "https://registry.npmjs.org/ts-interface-checker/-/ts-interface-checker-0.1.13.tgz", "integrity": "sha512-Y/arvbn+rrz3JCKl9C4kVNfTfSm2/mEp5FSz5EsZSANGPSlQrpRI5M4PKF+mJnE52jOO90PnPSc3Ur3bTQw0gA==", - "dev": true, "license": "Apache-2.0" }, "node_modules/tsconfig-paths": { @@ -24756,7 +24943,8 @@ "version": "2.8.1", "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", - "license": "0BSD" + "license": "0BSD", + "peer": true }, "node_modules/type-check": { "version": "0.4.0", @@ -24887,8 +25075,9 @@ "version": "5.3.3", "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.3.3.tgz", "integrity": "sha512-pXWcraxM0uxAS+tN0AG/BF2TyqmHO014Z070UsJ+pFvYuRSq8KH8DmWpnbXe0pEPDHXZV3FcAbJkijJ5oNEnWw==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", + "peer": true, "bin": { "tsc": "bin/tsc", "tsserver": "bin/tsserver" @@ -25431,6 +25620,7 @@ "integrity": "sha512-NL8jTlbo0Tn4dUEXEsUg8KeyG/Lkmc4Fnzb8JXN/Ykm9G4HNImjtABMJgkQoVjOBN/j2WAwDTRytdqJbZsah7w==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "esbuild": "^0.25.0", "fdir": "^6.5.0", @@ -25547,6 +25737,7 @@ "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", "dev": true, "license": "MIT", + "peer": true, "engines": { "node": ">=12" }, @@ -25560,6 +25751,7 @@ "integrity": "sha512-LUCP5ev3GURDysTWiP47wRRUpLKMOfPh+yKTx3kVIEiu5KOMeqzpnYNsKyOoVrULivR8tLcks4+lga33Whn90A==", "dev": true, "license": "MIT", + "peer": true, "dependencies": { "@types/chai": "^5.2.2", "@vitest/expect": "3.2.4", @@ -25765,6 +25957,7 @@ "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.103.0.tgz", "integrity": "sha512-HU1JOuV1OavsZ+mfigY0j8d1TgQgbZ6M+J75zDkpEAwYeXjWSqrGJtgnPblJjd/mAyTNQ7ygw0MiKOn6etz8yw==", "license": "MIT", + "peer": true, "dependencies": { "@types/eslint-scope": "^3.7.7", "@types/estree": "^1.0.8", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ce42d0ba41a..0ecba62140f 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -9,8 +9,11 @@ "lint": "next lint", "test": "vitest", "test:watch": "vitest -w", + "test:coverage": "vitest run --coverage", "format": "prettier --write .", - "format:check": "prettier --check ." + "format:check": "prettier --check .", + "e2e": "playwright test --config e2e_tests/playwright.config.ts", + "e2e:ui": "playwright test --ui --config e2e_tests/playwright.config.ts" }, "dependencies": { "@anthropic-ai/sdk": "^0.54.0", @@ -44,6 +47,7 @@ "uuid": "^11.1.0" }, "devDependencies": { + "@playwright/test": "^1.57.0", "@tailwindcss/forms": "^0.5.7", "@testing-library/dom": "^10.4.1", "@testing-library/jest-dom": "^6.8.0", diff --git a/ui/litellm-dashboard/public/assets/logos/minimax.svg b/ui/litellm-dashboard/public/assets/logos/minimax.svg new file mode 100644 index 00000000000..59b741bbcb7 --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/minimax.svg @@ -0,0 +1 @@ +资源 2 \ No newline at end of file diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts new file mode 100644 index 00000000000..44fc6a96836 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/agents/useAgents.test.ts @@ -0,0 +1,332 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useAgents } from "./useAgents"; +import { getAgentsList } from "@/components/networking"; +import type { AgentsResponse, Agent } from "@/components/agents/types"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + getAgentsList: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Import actual roles instead of mocking them + +// Mock data +const mockAgents: Agent[] = [ + { + agent_id: "agent-1", + agent_name: "Test Agent 1", + litellm_params: { + model: "gpt-3.5-turbo", + api_key: "test-key-1", + }, + agent_card_params: { + description: "A test agent for unit testing", + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_by: "user-1", + }, + { + agent_id: "agent-2", + agent_name: "Test Agent 2", + litellm_params: { + model: "claude-3", + api_key: "test-key-2", + }, + agent_card_params: { + description: "Another test agent", + }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + created_by: "user-2", + updated_by: "user-2", + }, +]; + +const mockAgentsResponse: AgentsResponse = { + agents: mockAgents, +}; + +describe("useAgents", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return agents data when query is successful", async () => { + // Mock successful API call + (getAgentsList as any).mockResolvedValue(mockAgentsResponse); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockAgentsResponse); + expect(result.current.error).toBeNull(); + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + expect(getAgentsList).toHaveBeenCalledTimes(1); + }); + + it("should handle error when getAgentsList fails", async () => { + const errorMessage = "Failed to fetch agents"; + const testError = new Error(errorMessage); + + // Mock failed API call + (getAgentsList as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + expect(getAgentsList).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is not an admin role", async () => { + // Mock non-admin userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "member", // Not in all_admin_roles + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is null", async () => { + // Mock null userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: null, + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is empty string", async () => { + // Mock empty string userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when both accessToken and userRole are missing", async () => { + // Mock both auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: null, + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getAgentsList).not.toHaveBeenCalled(); + }); + + it("should execute query when accessToken is present and userRole is Admin", async () => { + // Mock successful API call + (getAgentsList as any).mockResolvedValue(mockAgentsResponse); + + // Ensure auth values are set (already done in beforeEach) + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + expect(getAgentsList).toHaveBeenCalledTimes(1); + }); + + it("should execute query when accessToken is present and userRole is proxy_admin", async () => { + // Mock successful API call + (getAgentsList as any).mockResolvedValue(mockAgentsResponse); + + // Mock proxy_admin role + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "proxy_admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + expect(getAgentsList).toHaveBeenCalledTimes(1); + }); + + it("should return empty agents array when API returns empty data", async () => { + // Mock API returning empty agents array + (getAgentsList as any).mockResolvedValue({ agents: [] }); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({ agents: [] }); + expect(getAgentsList).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (getAgentsList as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useAgents(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts new file mode 100644 index 00000000000..ee903628f08 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.test.ts @@ -0,0 +1,194 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useCredentials } from "./useCredentials"; +import { credentialListCall, CredentialsResponse, CredentialItem } from "@/components/networking"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + credentialListCall: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockCredentialItems: CredentialItem[] = [ + { + credential_name: "openai-api-key", + credential_values: { api_key: "sk-test123" }, + credential_info: { + custom_llm_provider: "openai", + description: "OpenAI API Key for GPT models", + required: true, + }, + }, + { + credential_name: "anthropic-api-key", + credential_values: { api_key: "sk-ant-test456" }, + credential_info: { + custom_llm_provider: "anthropic", + description: "Anthropic API Key for Claude models", + required: true, + }, + }, +]; + +const mockCredentialsResponse: CredentialsResponse = { + credentials: mockCredentialItems, +}; + +describe("useCredentials", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return credentials data when query is successful", async () => { + // Mock successful API call + (credentialListCall as any).mockResolvedValue(mockCredentialsResponse); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockCredentialsResponse); + expect(result.current.error).toBeNull(); + expect(credentialListCall).toHaveBeenCalledWith("test-access-token"); + expect(credentialListCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when credentialListCall fails", async () => { + const errorMessage = "Failed to fetch credentials"; + const testError = new Error(errorMessage); + + // Mock failed API call + (credentialListCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(credentialListCall).toHaveBeenCalledWith("test-access-token"); + expect(credentialListCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(credentialListCall).not.toHaveBeenCalled(); + }); + + it("should return empty credentials array when API returns empty data", async () => { + // Mock API returning empty credentials array + (credentialListCall as any).mockResolvedValue({ credentials: [] }); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({ credentials: [] }); + expect(credentialListCall).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (credentialListCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should execute query when accessToken is present", async () => { + // Mock successful API call + (credentialListCall as any).mockResolvedValue(mockCredentialsResponse); + + // Ensure auth values are set (already done in beforeEach) + const { result } = renderHook(() => useCredentials(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(credentialListCall).toHaveBeenCalledWith("test-access-token"); + expect(credentialListCall).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts new file mode 100644 index 00000000000..716d6f75399 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/customers/useCustomers.test.ts @@ -0,0 +1,334 @@ +import { allEndUsersCall } from "@/components/networking"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { Customer, CustomersResponse } from "./useCustomers"; +import { useCustomers } from "./useCustomers"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + allEndUsersCall: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Import actual roles instead of mocking them + +// Mock data +const mockCustomers: Customer[] = [ + { + user_id: "customer-1", + alias: "Test Customer 1", + spend: 150.5, + blocked: false, + allowed_model_region: "us-east-1", + default_model: "gpt-3.5-turbo", + budget_id: "budget-1", + litellm_budget_table: { + budget_id: "budget-1", + max_budget: 1000, + soft_budget: 800, + max_parallel_requests: 10, + tpm_limit: 1000, + rpm_limit: 100, + model_max_budget: { "gpt-4": 500 }, + budget_duration: "monthly", + budget_reset_at: "2024-02-01T00:00:00Z", + created_at: "2024-01-01T00:00:00Z", + created_by: "admin-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "admin-1", + }, + }, + { + user_id: "customer-2", + alias: null, + spend: 0, + blocked: true, + allowed_model_region: null, + default_model: null, + budget_id: null, + litellm_budget_table: null, + }, +]; + +const mockCustomersResponse: CustomersResponse = mockCustomers; + +describe("useCustomers", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return customers data when query is successful", async () => { + // Mock successful API call + (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockCustomersResponse); + expect(result.current.error).toBeNull(); + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + expect(allEndUsersCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when allEndUsersCall fails", async () => { + const errorMessage = "Failed to fetch customers"; + const testError = new Error(errorMessage); + + // Mock failed API call + (allEndUsersCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + expect(allEndUsersCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is not an admin role", async () => { + // Mock non-admin userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "member", // Not in all_admin_roles + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is null", async () => { + // Mock null userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: null, + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is empty string", async () => { + // Mock empty string userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when both accessToken and userRole are missing", async () => { + // Mock both auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: null, + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(allEndUsersCall).not.toHaveBeenCalled(); + }); + + it("should execute query when accessToken is present and userRole is Admin", async () => { + // Mock successful API call + (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse); + + // Ensure auth values are set (already done in beforeEach) + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + expect(allEndUsersCall).toHaveBeenCalledTimes(1); + }); + + it("should execute query when accessToken is present and userRole is proxy_admin", async () => { + // Mock successful API call + (allEndUsersCall as any).mockResolvedValue(mockCustomersResponse); + + // Mock proxy_admin role + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "proxy_admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + expect(allEndUsersCall).toHaveBeenCalledTimes(1); + }); + + it("should return empty customers array when API returns empty data", async () => { + // Mock API returning empty customers array + (allEndUsersCall as any).mockResolvedValue([]); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual([]); + expect(allEndUsersCall).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (allEndUsersCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useCustomers(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts new file mode 100644 index 00000000000..d9e96a5308c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.test.ts @@ -0,0 +1,273 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useGuardrails } from "./useGuardrails"; +import { getGuardrailsList } from "@/components/networking"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + getGuardrailsList: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockGuardrailsResponse = { + guardrails: [ + { guardrail_name: "content-safety" }, + { guardrail_name: "toxicity-filter" }, + { guardrail_name: "pii-detection" }, + ], +}; + +const expectedGuardrailNames = ["content-safety", "toxicity-filter", "pii-detection"]; + +describe("useGuardrails", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return guardrail names when query is successful", async () => { + // Mock successful API call + (getGuardrailsList as any).mockResolvedValue(mockGuardrailsResponse); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(expectedGuardrailNames); + expect(result.current.error).toBeNull(); + expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token"); + expect(getGuardrailsList).toHaveBeenCalledTimes(1); + }); + + it("should handle error when getGuardrailsList fails", async () => { + const errorMessage = "Failed to fetch guardrails"; + const testError = new Error(errorMessage); + + // Mock failed API call + (getGuardrailsList as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token"); + expect(getGuardrailsList).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getGuardrailsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", async () => { + // Mock missing userId + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getGuardrailsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", async () => { + // Mock missing userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getGuardrailsList).not.toHaveBeenCalled(); + }); + + it("should not execute query when all auth values are missing", async () => { + // Mock all auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: null, + userRole: null, + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getGuardrailsList).not.toHaveBeenCalled(); + }); + + it("should execute query when all auth values are present", async () => { + // Mock successful API call + (getGuardrailsList as any).mockResolvedValue(mockGuardrailsResponse); + + // Ensure all auth values are present (already set in beforeEach) + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token"); + expect(getGuardrailsList).toHaveBeenCalledTimes(1); + }); + + it("should return empty array when API returns empty guardrails", async () => { + // Mock API returning empty guardrails array + (getGuardrailsList as any).mockResolvedValue({ guardrails: [] }); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual([]); + expect(getGuardrailsList).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (getGuardrailsList as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should correctly transform guardrail objects to names array", async () => { + const customGuardrailsResponse = { + guardrails: [{ guardrail_name: "custom-guardrail-1" }, { guardrail_name: "custom-guardrail-2" }], + }; + const expectedNames = ["custom-guardrail-1", "custom-guardrail-2"]; + + // Mock API call with custom data + (getGuardrailsList as any).mockResolvedValue(customGuardrailsResponse); + + const { result } = renderHook(() => useGuardrails(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(expectedNames); + expect(result.current.data).toHaveLength(2); + expect(result.current.data).toContain("custom-guardrail-1"); + expect(result.current.data).toContain("custom-guardrail-2"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts new file mode 100644 index 00000000000..9786b7fa359 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/guardrails/useGuardrails.ts @@ -0,0 +1,18 @@ +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { getGuardrailsList } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +const guardrailKeys = createQueryKeys("guardrails"); + +export const useGuardrails = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: guardrailKeys.list({}), + queryFn: async () => { + const response = await getGuardrailsList(accessToken!); + return response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name); + }, + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts new file mode 100644 index 00000000000..db394b9f7f8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/healthReadiness/useHealthReadiness.ts @@ -0,0 +1,27 @@ +import { getProxyBaseUrl } from "@/components/networking"; +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const healthReadinessKeys = createQueryKeys("healthReadiness"); + +interface HealthReadinessResponse { + litellm_version?: string; + [key: string]: any; +} + +const fetchHealthReadiness = async (): Promise => { + const baseUrl = getProxyBaseUrl(); + const response = await fetch(`${baseUrl}/health/readiness`); + if (!response.ok) { + throw new Error(`Failed to fetch health readiness: ${response.statusText}`); + } + return response.json(); +}; + +export const useHealthReadiness = (): UseQueryResult => { + return useQuery({ + queryKey: healthReadinessKeys.detail("readiness"), + queryFn: fetchHealthReadiness, + staleTime: 5 * 60 * 1000, // 5 minutes + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts new file mode 100644 index 00000000000..c4ffb7041aa --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.test.ts @@ -0,0 +1,362 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useKeys } from "./useKeys"; +import { keyListCall } from "@/components/networking"; +import type { KeyResponse } from "@/components/key_team_helpers/key_list"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + keyListCall: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockKeys: KeyResponse[] = [ + { + token: "sk-test-key-1", + token_id: "key-1", + key_name: "Test Key 1", + key_alias: "test-key-1", + spend: 10.5, + max_budget: 100, + expires: "2024-12-31T23:59:59Z", + models: ["gpt-3.5-turbo"], + aliases: {}, + config: {}, + user_id: "user-1", + team_id: null, + max_parallel_requests: 10, + metadata: {}, + tpm_limit: 1000, + rpm_limit: 100, + duration: "30d", + budget_duration: "1mo", + budget_reset_at: "2024-02-01T00:00:00Z", + allowed_cache_controls: [], + allowed_routes: [], + permissions: {}, + model_spend: { "gpt-3.5-turbo": 10.5 }, + model_max_budget: { "gpt-3.5-turbo": 100 }, + soft_budget_cooldown: false, + blocked: false, + litellm_budget_table: {}, + organization_id: null, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + team_spend: 0, + team_alias: "", + team_tpm_limit: 0, + team_rpm_limit: 0, + team_max_budget: 0, + team_models: [], + team_blocked: false, + soft_budget: 0, + team_model_aliases: {}, + team_member_spend: 0, + team_metadata: {}, + end_user_id: "", + end_user_tpm_limit: 0, + end_user_rpm_limit: 0, + end_user_max_budget: 0, + last_refreshed_at: 0, + api_key: "", + user_role: "user", + rpm_limit_per_model: {}, + tpm_limit_per_model: {}, + user_tpm_limit: 0, + user_rpm_limit: 0, + user_email: "", + }, + { + token: "sk-test-key-2", + token_id: "key-2", + key_name: "Test Key 2", + key_alias: "test-key-2", + spend: 25.0, + max_budget: 200, + expires: "2024-12-31T23:59:59Z", + models: ["claude-3"], + aliases: {}, + config: {}, + user_id: "user-2", + team_id: "team-1", + max_parallel_requests: 5, + metadata: {}, + tpm_limit: 500, + rpm_limit: 50, + duration: "30d", + budget_duration: "1mo", + budget_reset_at: "2024-02-01T00:00:00Z", + allowed_cache_controls: [], + allowed_routes: [], + permissions: {}, + model_spend: { "claude-3": 25.0 }, + model_max_budget: { "claude-3": 200 }, + soft_budget_cooldown: false, + blocked: false, + litellm_budget_table: {}, + organization_id: null, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + team_spend: 0, + team_alias: "test-team", + team_tpm_limit: 1000, + team_rpm_limit: 100, + team_max_budget: 500, + team_models: ["claude-3"], + team_blocked: false, + soft_budget: 0, + team_model_aliases: {}, + team_member_spend: 0, + team_metadata: {}, + end_user_id: "", + end_user_tpm_limit: 0, + end_user_rpm_limit: 0, + end_user_max_budget: 0, + last_refreshed_at: 0, + api_key: "", + user_role: "user", + rpm_limit_per_model: {}, + tpm_limit_per_model: {}, + user_tpm_limit: 0, + user_rpm_limit: 0, + user_email: "", + }, +]; + +const mockKeysResponse = { + keys: mockKeys, + total_count: 2, + current_page: 1, + total_pages: 1, +}; + +describe("useKeys", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return keys data when query is successful", async () => { + // Mock successful API call + (keyListCall as any).mockResolvedValue(mockKeysResponse); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockKeysResponse); + expect(result.current.error).toBeNull(); + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + 1, // page + 10, // pageSize + ); + expect(keyListCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when keyListCall fails", async () => { + const errorMessage = "Failed to fetch keys"; + const testError = new Error(errorMessage); + + // Mock failed API call + (keyListCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + 1, // page + 10, // pageSize + ); + expect(keyListCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(keyListCall).not.toHaveBeenCalled(); + }); + + it("should pass correct page and pageSize parameters to the API", async () => { + // Mock successful API call + (keyListCall as any).mockResolvedValue(mockKeysResponse); + + const page = 2; + const pageSize = 20; + + const { result } = renderHook(() => useKeys(page, pageSize), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + page, // page + pageSize, // pageSize + ); + }); + + it("should return empty keys array when API returns empty data", async () => { + // Mock API returning empty keys array + const emptyResponse = { + keys: [], + total_count: 0, + current_page: 1, + total_pages: 0, + }; + (keyListCall as any).mockResolvedValue(emptyResponse); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(emptyResponse); + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + 1, // page + 10, // pageSize + ); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (keyListCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useKeys(1, 10), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should handle pagination correctly", async () => { + const paginatedResponse = { + keys: [mockKeys[0]], // Only first key + total_count: 15, + current_page: 2, + total_pages: 2, + }; + (keyListCall as any).mockResolvedValue(paginatedResponse); + + const { result } = renderHook(() => useKeys(2, 10), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(result.current.data).toEqual(paginatedResponse); + expect(keyListCall).toHaveBeenCalledWith( + "test-access-token", + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + 2, // page + 10, // pageSize + ); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts new file mode 100644 index 00000000000..8ae4d76ff5d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/keys/useKeys.ts @@ -0,0 +1,36 @@ +import { keepPreviousData, useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { keyListCall } from "@/components/networking"; +import { KeyResponse } from "@/components/key_team_helpers/key_list"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +const keyKeys = createQueryKeys("keys"); + +export interface KeysResponse { + keys: KeyResponse[]; + total_count: number; + current_page: number; + total_pages: number; +} + +export const useKeys = (page: number, pageSize: number): UseQueryResult => { + const { accessToken } = useAuthorized(); + + return useQuery({ + queryKey: keyKeys.list({ page, limit: pageSize }), + queryFn: async () => + await keyListCall( + accessToken!, + null, // organizationID + null, // teamID + null, // selectedKeyAlias + null, // userID + null, // keyHash + page, + pageSize, + ), + enabled: Boolean(accessToken), + staleTime: 30000, // 30 seconds + placeholderData: keepPreviousData, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts new file mode 100644 index 00000000000..be910acf7e4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.test.ts @@ -0,0 +1,127 @@ +/* @vitest-environment jsdom */ +import React from "react"; +import { renderHook, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { useMCPServerHealth } from "./useMCPServerHealth"; +import * as networking from "@/components/networking"; + +// Mock the networking module +vi.mock("@/components/networking", () => ({ + fetchMCPServerHealth: vi.fn(), +})); + +// Mock useAuthorized hook +vi.mock("../useAuthorized", () => ({ + default: vi.fn(() => ({ + accessToken: "test-token-123", + })), +})); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +const wrapper = ({ children }: { children: React.ReactNode }) => { + const queryClient = createQueryClient(); + return React.createElement(QueryClientProvider, { client: queryClient }, children); +}; + +describe("useMCPServerHealth", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should fetch health status for given server IDs", async () => { + const mockHealthStatuses = [ + { server_id: "server-1", status: "healthy" }, + { server_id: "server-2", status: "unhealthy" }, + ]; + + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); + + const { result } = renderHook(() => useMCPServerHealth(["server-1", "server-2"]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", ["server-1", "server-2"]); + expect(result.current.data).toEqual(mockHealthStatuses); + }); + + it("should fetch health status for all servers when no server IDs provided", async () => { + const mockHealthStatuses = [ + { server_id: "server-1", status: "healthy" }, + { server_id: "server-2", status: "healthy" }, + { server_id: "server-3", status: "unhealthy" }, + ]; + + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue(mockHealthStatuses); + + const { result } = renderHook(() => useMCPServerHealth(), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", undefined); + expect(result.current.data).toEqual(mockHealthStatuses); + }); + + it("should handle empty server list", async () => { + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + + const { result } = renderHook(() => useMCPServerHealth([]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(networking.fetchMCPServerHealth).toHaveBeenCalledWith("test-token-123", []); + expect(result.current.data).toEqual([]); + }); + + it("should handle errors when fetching health status", async () => { + const mockError = new Error("Failed to fetch health status"); + vi.mocked(networking.fetchMCPServerHealth).mockRejectedValue(mockError); + + const { result } = renderHook(() => useMCPServerHealth(["server-1"]), { + wrapper, + }); + + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(mockError); + }); + + it("should not fetch when accessToken is not available", async () => { + // Mock useAuthorized to return no token + const useAuthorizedModule = await import("../useAuthorized"); + vi.mocked(useAuthorizedModule.default).mockReturnValue({ + accessToken: null, + } as any); + + const { result } = renderHook(() => useMCPServerHealth(["server-1"]), { + wrapper, + }); + + // Should remain in idle state since query is not enabled + expect(result.current.status).toBe("pending"); + expect(networking.fetchMCPServerHealth).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts new file mode 100644 index 00000000000..95d7f3bcee0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/mcpServers/useMCPServerHealth.ts @@ -0,0 +1,22 @@ +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { fetchMCPServerHealth } from "@/components/networking"; +import useAuthorized from "../useAuthorized"; + +const mcpServerHealthKeys = createQueryKeys("mcpServerHealth"); + +interface MCPServerHealth { + server_id: string; + status: string; +} + +export const useMCPServerHealth = (serverIds?: string[]) => { + const { accessToken } = useAuthorized(); + return useQuery({ + queryKey: [...mcpServerHealthKeys.lists(), { serverIds }], + queryFn: async () => await fetchMCPServerHealth(accessToken!, serverIds), + enabled: !!accessToken, + // Refetch health status every 30 seconds to keep it up to date + refetchInterval: 30000, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts new file mode 100644 index 00000000000..f79ca33bc5d --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.test.ts @@ -0,0 +1,144 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useModelCostMap } from "./useModelCostMap"; +import { modelCostMap } from "@/components/networking"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + modelCostMap: vi.fn(), +})); + +// Mock data +const mockModelCostData: Record = { + "gpt-3.5-turbo": { + litellm_provider: "openai", + input_cost_per_token: 0.0015, + output_cost_per_token: 0.002, + }, + "claude-3-sonnet-20240229": { + litellm_provider: "anthropic", + input_cost_per_token: 0.003, + output_cost_per_token: 0.015, + }, +}; + +describe("useModelCostMap", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return model cost map data when query is successful", async () => { + // Mock successful API call + (modelCostMap as any).mockResolvedValue(mockModelCostData); + + const { result } = renderHook(() => useModelCostMap(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockModelCostData); + expect(result.current.error).toBeNull(); + expect(modelCostMap).toHaveBeenCalledTimes(1); + }); + + it("should handle error when modelCostMap fails", async () => { + const errorMessage = "Failed to fetch model cost map"; + const testError = new Error(errorMessage); + + // Mock failed API call + (modelCostMap as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useModelCostMap(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(modelCostMap).toHaveBeenCalledTimes(1); + }); + + it("should return empty object when API returns empty data", async () => { + // Mock API returning empty object + (modelCostMap as any).mockResolvedValue({}); + + const { result } = renderHook(() => useModelCostMap(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({}); + expect(modelCostMap).toHaveBeenCalledTimes(1); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (modelCostMap as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useModelCostMap(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should have correct query configuration", async () => { + // Mock successful API call + (modelCostMap as any).mockResolvedValue(mockModelCostData); + + const { result } = renderHook(() => useModelCostMap(), { wrapper }); + + // Wait for query to complete + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + // Verify the query was called + expect(modelCostMap).toHaveBeenCalledTimes(1); + + // The hook should have the expected properties from useQuery + expect(result.current).toHaveProperty("data"); + expect(result.current).toHaveProperty("isLoading"); + expect(result.current).toHaveProperty("isError"); + expect(result.current).toHaveProperty("isSuccess"); + expect(result.current).toHaveProperty("error"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.ts new file mode 100644 index 00000000000..2d82eedf25c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/models/useModelCostMap.ts @@ -0,0 +1,14 @@ +import { modelCostMap } from "@/components/networking"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const modelCostMapKeys = createQueryKeys("modelCostMap"); + +export const useModelCostMap = () => { + return useQuery>({ + queryKey: modelCostMapKeys.list({}), + queryFn: async () => await modelCostMap(), + staleTime: 60 * 1000, // 1 minute + gcTime: 60 * 1000, // 1 minute + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts new file mode 100644 index 00000000000..66c005f37c4 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.test.ts @@ -0,0 +1,282 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useOrganizations } from "./useOrganizations"; +import { organizationListCall } from "@/components/networking"; +import type { Organization } from "@/components/networking"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + organizationListCall: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockOrganizations: Organization[] = [ + { + organization_id: "org-1", + organization_alias: "Test Organization 1", + budget_id: "budget-1", + metadata: {}, + models: ["gpt-3.5-turbo", "gpt-4"], + spend: 100.5, + model_spend: { "gpt-3.5-turbo": 50.25, "gpt-4": 50.25 }, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-1", + litellm_budget_table: null, + teams: null, + users: null, + members: [ + { user_id: "user-1", user_role: "admin" }, + { user_id: "user-2", user_role: "member" }, + ], + }, + { + organization_id: "org-2", + organization_alias: "Test Organization 2", + budget_id: "budget-2", + metadata: {}, + models: ["claude-3"], + spend: 250.75, + model_spend: { "claude-3": 250.75 }, + created_at: "2024-01-01T00:00:00Z", + created_by: "user-3", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user-3", + litellm_budget_table: null, + teams: null, + users: null, + members: [{ user_id: "user-3", user_role: "admin" }], + }, +]; + +describe("useOrganizations", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return organizations data when query is successful", async () => { + // Mock successful API call + (organizationListCall as any).mockResolvedValue(mockOrganizations); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockOrganizations); + expect(result.current.error).toBeNull(); + expect(organizationListCall).toHaveBeenCalledWith("test-access-token"); + expect(organizationListCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when organizationListCall fails", async () => { + const errorMessage = "Failed to fetch organizations"; + const testError = new Error(errorMessage); + + // Mock failed API call + (organizationListCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(organizationListCall).toHaveBeenCalledWith("test-access-token"); + expect(organizationListCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(organizationListCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", async () => { + // Mock missing userId + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(organizationListCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", async () => { + // Mock missing userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(organizationListCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when all auth values are missing", async () => { + // Mock all auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: null, + userRole: null, + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(organizationListCall).not.toHaveBeenCalled(); + }); + + it("should execute query when all auth values are present", async () => { + // Mock successful API call + (organizationListCall as any).mockResolvedValue(mockOrganizations); + + // Ensure all auth values are present (already set in beforeEach) + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(organizationListCall).toHaveBeenCalledWith("test-access-token"); + expect(organizationListCall).toHaveBeenCalledTimes(1); + }); + + it("should return empty array when API returns empty data", async () => { + // Mock API returning empty array + (organizationListCall as any).mockResolvedValue([]); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual([]); + expect(organizationListCall).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (organizationListCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useOrganizations(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts index 57c9c057652..27a946d112a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/organizations/useOrganizations.ts @@ -6,8 +6,7 @@ import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; const organizationKeys = createQueryKeys("organizations"); export const useOrganizations = (): UseQueryResult => { - const { accessToken } = useAuthorized(); - const { userId, userRole } = useAuthorized(); + const { accessToken, userId, userRole } = useAuthorized(); return useQuery({ queryKey: organizationKeys.list({}), queryFn: async () => await organizationListCall(accessToken!), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts new file mode 100644 index 00000000000..33242e0452f --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/providers/useProviderFields.test.ts @@ -0,0 +1,182 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useProviderFields } from "./useProviderFields"; +import { getProviderCreateMetadata } from "@/components/networking"; +import type { ProviderCreateInfo } from "@/components/networking"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + getProviderCreateMetadata: vi.fn(), +})); + +// Mock data +const mockProviderFields: ProviderCreateInfo[] = [ + { + provider: "OpenAI", + provider_display_name: "OpenAI", + litellm_provider: "openai", + default_model_placeholder: "gpt-3.5-turbo", + credential_fields: [], + }, + { + provider: "Anthropic", + provider_display_name: "Anthropic", + litellm_provider: "anthropic", + default_model_placeholder: "claude-3-sonnet-20240229", + credential_fields: [], + }, + { + provider: "Azure", + provider_display_name: "Azure OpenAI", + litellm_provider: "azure", + default_model_placeholder: "gpt-35-turbo", + credential_fields: [], + }, +]; + +describe("useProviderFields", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return provider fields data when query is successful", async () => { + // Mock successful API call + (getProviderCreateMetadata as any).mockResolvedValue(mockProviderFields); + + const { result } = renderHook(() => useProviderFields(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockProviderFields); + expect(result.current.error).toBeNull(); + expect(getProviderCreateMetadata).toHaveBeenCalledTimes(1); + }); + + it("should handle error when getProviderCreateMetadata fails", async () => { + const errorMessage = "Failed to fetch provider fields"; + const testError = new Error(errorMessage); + + // Mock failed API call + (getProviderCreateMetadata as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useProviderFields(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(getProviderCreateMetadata).toHaveBeenCalledTimes(1); + }); + + it("should return empty array when API returns empty data", async () => { + // Mock API returning empty array + (getProviderCreateMetadata as any).mockResolvedValue([]); + + const { result } = renderHook(() => useProviderFields(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual([]); + expect(getProviderCreateMetadata).toHaveBeenCalledTimes(1); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (getProviderCreateMetadata as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useProviderFields(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should have correct query configuration", async () => { + // Mock successful API call + (getProviderCreateMetadata as any).mockResolvedValue(mockProviderFields); + + const { result } = renderHook(() => useProviderFields(), { wrapper }); + + // Wait for query to complete + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + // Verify the query was called + expect(getProviderCreateMetadata).toHaveBeenCalledTimes(1); + + // The hook should have the expected properties from useQuery + expect(result.current).toHaveProperty("data"); + expect(result.current).toHaveProperty("isLoading"); + expect(result.current).toHaveProperty("isError"); + expect(result.current).toHaveProperty("isSuccess"); + expect(result.current).toHaveProperty("error"); + }); + + it("should return provider fields with populated credential fields", async () => { + const mockFieldsWithCredentials: ProviderCreateInfo[] = [ + { + provider: "TestProvider", + provider_display_name: "Test Provider", + litellm_provider: "test", + default_model_placeholder: "test-model", + credential_fields: [], // Keeping empty as per existing test patterns + }, + ]; + + // Mock successful API call with provider that has credential fields + (getProviderCreateMetadata as any).mockResolvedValue(mockFieldsWithCredentials); + + const { result } = renderHook(() => useProviderFields(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockFieldsWithCredentials); + expect(result.current.data?.[0].provider).toBe("TestProvider"); + expect(result.current.data?.[0].litellm_provider).toBe("test"); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts new file mode 100644 index 00000000000..69e52d0ff25 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useEditSSOSettings.ts @@ -0,0 +1,38 @@ +import { useMutation, UseMutationResult } from "@tanstack/react-query"; +import { updateSSOSettings } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export interface EditSSOSettingsParams { + google_client_id?: string | null; + google_client_secret?: string | null; + microsoft_client_id?: string | null; + microsoft_client_secret?: string | null; + microsoft_tenant?: string | null; + generic_client_id?: string | null; + generic_client_secret?: string | null; + generic_authorization_endpoint?: string | null; + generic_token_endpoint?: string | null; + generic_userinfo_endpoint?: string | null; + proxy_base_url?: string | null; + user_email?: string | null; + sso_provider?: string | null; + role_mappings?: any; + [key: string]: any; +} + +export interface EditSSOSettingsResponse { + [key: string]: any; +} + +export const useEditSSOSettings = (): UseMutationResult => { + const { accessToken } = useAuthorized(); + + return useMutation({ + mutationFn: async (params: EditSSOSettingsParams) => { + if (!accessToken) { + throw new Error("Access token is required"); + } + return await updateSSOSettings(accessToken, params); + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts new file mode 100644 index 00000000000..3e09c3c2ca8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/sso/useSSOSettings.ts @@ -0,0 +1,54 @@ +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { getSSOSettings } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +export interface SSOFieldSchema { + description: string; + properties: { + [key: string]: { + description: string; + type: string; + }; + }; +} + +export interface SSOSettingsValues { + google_client_id: string | null; + google_client_secret: string | null; + microsoft_client_id: string | null; + microsoft_client_secret: string | null; + microsoft_tenant: string | null; + generic_client_id: string | null; + generic_client_secret: string | null; + generic_authorization_endpoint: string | null; + generic_token_endpoint: string | null; + generic_userinfo_endpoint: string | null; + proxy_base_url: string | null; + user_email: string | null; + ui_access_mode: string | null; + role_mappings: { + provider: string; + group_claim: string; + default_role: "internal_user" | "internal_user_viewer" | "proxy_admin" | "proxy_admin_viewer"; + roles: { + [key: string]: string[]; + }; + }; +} + +export interface SSOSettingsResponse { + values: SSOSettingsValues; + field_schema: SSOFieldSchema; +} + +const ssoKeys = createQueryKeys("sso"); + +export const useSSOSettings = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: ssoKeys.detail("settings"), + queryFn: async () => await getSSOSettings(accessToken!), + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts new file mode 100644 index 00000000000..a1751339568 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.test.ts @@ -0,0 +1,283 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useTags } from "./useTags"; +import { tagListCall } from "@/components/networking"; +import type { TagListResponse } from "@/components/tag_management/types"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + tagListCall: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockTags: TagListResponse = { + "tag-1": { + name: "tag-1", + description: "Test tag 1 description", + models: ["gpt-3.5-turbo", "gpt-4"], + model_info: { "gpt-3.5-turbo": "GPT-3.5 Turbo", "gpt-4": "GPT-4" }, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + created_by: "user-1", + updated_by: "user-1", + litellm_budget_table: { + max_budget: 1000, + soft_budget: 800, + tpm_limit: 100000, + rpm_limit: 1000, + max_parallel_requests: 10, + budget_duration: "monthly", + model_max_budget: { "gpt-3.5-turbo": 500, "gpt-4": 500 }, + }, + }, + "tag-2": { + name: "tag-2", + description: "Test tag 2 description", + models: ["claude-3"], + model_info: { "claude-3": "Claude 3" }, + created_at: "2024-01-02T00:00:00Z", + updated_at: "2024-01-02T00:00:00Z", + created_by: "user-2", + updated_by: "user-2", + litellm_budget_table: { + max_budget: 2000, + soft_budget: 1500, + tpm_limit: 200000, + rpm_limit: 2000, + max_parallel_requests: 20, + budget_duration: "monthly", + model_max_budget: { "claude-3": 2000 }, + }, + }, +}; + +describe("useTags", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return tags data when query is successful", async () => { + // Mock successful API call + (tagListCall as any).mockResolvedValue(mockTags); + + const { result } = renderHook(() => useTags(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockTags); + expect(result.current.error).toBeNull(); + expect(tagListCall).toHaveBeenCalledWith("test-access-token"); + expect(tagListCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when tagListCall fails", async () => { + const errorMessage = "Failed to fetch tags"; + const testError = new Error(errorMessage); + + // Mock failed API call + (tagListCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useTags(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(tagListCall).toHaveBeenCalledWith("test-access-token"); + expect(tagListCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTags(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(tagListCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", async () => { + // Mock missing userId + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTags(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(tagListCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", async () => { + // Mock missing userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTags(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(tagListCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when all auth values are missing", async () => { + // Mock all auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: null, + userRole: null, + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTags(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(tagListCall).not.toHaveBeenCalled(); + }); + + it("should execute query when all auth values are present", async () => { + // Mock successful API call + (tagListCall as any).mockResolvedValue(mockTags); + + // Ensure all auth values are present (already set in beforeEach) + const { result } = renderHook(() => useTags(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(tagListCall).toHaveBeenCalledWith("test-access-token"); + expect(tagListCall).toHaveBeenCalledTimes(1); + }); + + it("should return empty object when API returns empty data", async () => { + // Mock API returning empty object + (tagListCall as any).mockResolvedValue({}); + + const { result } = renderHook(() => useTags(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({}); + expect(tagListCall).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (tagListCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useTags(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts new file mode 100644 index 00000000000..8f82502a74c --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/tags/useTags.ts @@ -0,0 +1,16 @@ +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; +import { tagListCall } from "@/components/networking"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { TagListResponse } from "@/components/tag_management/types"; + +const tagKeys = createQueryKeys("tags"); + +export const useTags = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: tagKeys.list({}), + queryFn: async () => await tagListCall(accessToken!), + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts new file mode 100644 index 00000000000..91ffbcfafa2 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/teams/useTeams.test.ts @@ -0,0 +1,275 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useTeams } from "./useTeams"; +import { fetchTeams } from "@/app/(dashboard)/networking"; +import type { Team } from "@/components/key_team_helpers/key_list"; + +// Mock the networking function +vi.mock("@/app/(dashboard)/networking", () => ({ + fetchTeams: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockTeams: Team[] = [ + { + team_id: "team-1", + team_alias: "Test Team 1", + models: ["gpt-3.5-turbo", "claude-3"], + max_budget: 100.0, + budget_duration: "monthly", + tpm_limit: 1000, + rpm_limit: 100, + organization_id: "org-1", + created_at: "2024-01-01T00:00:00Z", + keys: [], + members_with_roles: [], + }, + { + team_id: "team-2", + team_alias: "Test Team 2", + models: ["gpt-4"], + max_budget: 200.0, + budget_duration: "monthly", + tpm_limit: 2000, + rpm_limit: 200, + organization_id: "org-1", + created_at: "2024-01-02T00:00:00Z", + keys: [], + members_with_roles: [], + }, +]; + +describe("useTeams", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return teams data when query is successful", async () => { + // Mock successful API call + (fetchTeams as any).mockResolvedValue(mockTeams); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockTeams); + expect(result.current.error).toBeNull(); + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null); + expect(fetchTeams).toHaveBeenCalledTimes(1); + }); + + it("should handle error when fetchTeams fails", async () => { + const errorMessage = "Failed to fetch teams"; + const testError = new Error(errorMessage); + + // Mock failed API call + (fetchTeams as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null); + expect(fetchTeams).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(fetchTeams).not.toHaveBeenCalled(); + }); + + it("should not execute query when accessToken is empty string", async () => { + // Mock empty string accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: "", + userId: "test-user-id", + userRole: "Admin", + token: "", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(fetchTeams).not.toHaveBeenCalled(); + }); + + it("should execute query when accessToken is present", async () => { + // Mock successful API call + (fetchTeams as any).mockResolvedValue(mockTeams); + + // Ensure auth values are set (already done in beforeEach) + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null); + expect(fetchTeams).toHaveBeenCalledTimes(1); + }); + + it("should return empty teams array when API returns empty data", async () => { + // Mock API returning empty teams array + (fetchTeams as any).mockResolvedValue([]); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual([]); + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", null); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (fetchTeams as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should pass userId and userRole to fetchTeams", async () => { + // Mock successful API call + (fetchTeams as any).mockResolvedValue(mockTeams); + + // Mock specific userId and userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "custom-user-id", + userRole: "member", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", "custom-user-id", "member", null); + }); + + it("should handle null userId", async () => { + // Mock successful API call + (fetchTeams as any).mockResolvedValue(mockTeams); + + // Mock null userId + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useTeams(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(fetchTeams).toHaveBeenCalledWith("test-access-token", null, "Admin", null); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts new file mode 100644 index 00000000000..6429aeafb5a --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.test.ts @@ -0,0 +1,169 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useUIConfig } from "./useUIConfig"; +import { getUiConfig, LiteLLMWellKnownUiConfig } from "@/components/networking"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + getUiConfig: vi.fn(), +})); + +// Mock the queryKeysFactory - we'll mock the specific return value +vi.mock("../common/queryKeysFactory", () => ({ + createQueryKeys: vi.fn((resource: string) => ({ + all: [resource], + lists: () => [resource, "list"], + list: (params?: any) => [resource, "list", { params }], + details: () => [resource, "detail"], + detail: (uid: string) => [resource, "detail", uid], + })), +})); + +// Mock data +const mockUIConfig: LiteLLMWellKnownUiConfig = { + server_root_path: "/api", + proxy_base_url: "https://proxy.example.com", + auto_redirect_to_sso: true, + admin_ui_disabled: false, +}; + +describe("useUIConfig", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return UI config data when query is successful", async () => { + // Mock successful API call + (getUiConfig as any).mockResolvedValue(mockUIConfig); + + const { result } = renderHook(() => useUIConfig(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockUIConfig); + expect(result.current.error).toBeNull(); + expect(getUiConfig).toHaveBeenCalledWith(); + expect(getUiConfig).toHaveBeenCalledTimes(1); + }); + + it("should handle error when getUiConfig fails", async () => { + const errorMessage = "Failed to fetch UI config"; + const testError = new Error(errorMessage); + + // Mock failed API call + (getUiConfig as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useUIConfig(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(getUiConfig).toHaveBeenCalledWith(); + expect(getUiConfig).toHaveBeenCalledTimes(1); + }); + + it("should return different UI config data correctly", async () => { + const alternativeUIConfig: LiteLLMWellKnownUiConfig = { + server_root_path: "/v1", + proxy_base_url: null, + auto_redirect_to_sso: false, + admin_ui_disabled: true, + }; + + // Mock successful API call with different data + (getUiConfig as any).mockResolvedValue(alternativeUIConfig); + + const { result } = renderHook(() => useUIConfig(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(alternativeUIConfig); + expect(result.current.error).toBeNull(); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (getUiConfig as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useUIConfig(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); + + it("should handle malformed response error", async () => { + const malformedError = new Error("Invalid JSON response"); + + // Mock malformed response + (getUiConfig as any).mockRejectedValue(malformedError); + + const { result } = renderHook(() => useUIConfig(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(malformedError); + expect(result.current.data).toBeUndefined(); + }); + + it("should use correct query key structure", async () => { + // Mock successful API call + (getUiConfig as any).mockResolvedValue(mockUIConfig); + + const { result } = renderHook(() => useUIConfig(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + // The query key should be generated by createQueryKeys("uiConfig").list({}) + // Based on our mock, this should be ["uiConfig", "list", {}] + expect(getUiConfig).toHaveBeenCalledTimes(1); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts new file mode 100644 index 00000000000..785f003d2f8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.test.ts @@ -0,0 +1,185 @@ +import { getUiSettings } from "@/components/networking"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { renderHook, waitFor } from "@testing-library/react"; +import React, { ReactNode } from "react"; +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { useUISettings } from "./useUISettings"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + getUiSettings: vi.fn(), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("../useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data +const mockUISettings: Record = { + theme: "dark", + language: "en", + notifications: true, + dashboard_layout: "compact", + api_keys_visible: false, +}; + +describe("useUISettings", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + userId: "test-user-id", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return UI settings data when query is successful", async () => { + // Mock successful API call + (getUiSettings as any).mockResolvedValue(mockUISettings); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockUISettings); + expect(result.current.error).toBeNull(); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + expect(getUiSettings).toHaveBeenCalledTimes(1); + }); + + it("should handle error when getUiSettings fails", async () => { + const errorMessage = "Failed to fetch UI settings"; + const testError = new Error(errorMessage); + + // Mock failed API call + (getUiSettings as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + expect(getUiSettings).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userRole: "Admin", + userId: "test-user-id", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getUiSettings).not.toHaveBeenCalled(); + }); + + it("should not execute query when accessToken is empty string", async () => { + // Mock empty accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: "", + userRole: "Admin", + userId: "test-user-id", + token: "", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(getUiSettings).not.toHaveBeenCalled(); + }); + + it("should return empty object when API returns empty settings", async () => { + // Mock API returning empty object + (getUiSettings as any).mockResolvedValue({}); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual({}); + expect(getUiSettings).toHaveBeenCalledWith("test-access-token"); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (getUiSettings as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useUISettings(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts index 823c0067b5c..46a0254d0db 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiSettings/useUISettings.ts @@ -1,10 +1,12 @@ import { getUiSettings } from "@/components/networking"; import { useQuery } from "@tanstack/react-query"; import { createQueryKeys } from "../common/queryKeysFactory"; +import useAuthorized from "../useAuthorized"; const uiSettingsKeys = createQueryKeys("uiSettings"); -export const useUISettings = (accessToken: string) => { +export const useUISettings = () => { + const { accessToken } = useAuthorized(); return useQuery>({ queryKey: uiSettingsKeys.list({}), queryFn: async () => await getUiSettings(accessToken), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts new file mode 100644 index 00000000000..d0a618e27ba --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useDisableShowNewBadge.ts @@ -0,0 +1,35 @@ +// hooks/useDisableShowNewBadge.ts +import { useSyncExternalStore } from "react"; +import { getLocalStorageItem } from "@/utils/localStorageUtils"; +import { LOCAL_STORAGE_EVENT } from "@/utils/localStorageUtils"; + +function subscribe(callback: () => void) { + const onStorage = (e: StorageEvent) => { + if (e.key === "disableShowNewBadge") { + callback(); + } + }; + + const onCustom = (e: Event) => { + const { key } = (e as CustomEvent).detail; + if (key === "disableShowNewBadge") { + callback(); + } + }; + + window.addEventListener("storage", onStorage); + window.addEventListener(LOCAL_STORAGE_EVENT, onCustom); + + return () => { + window.removeEventListener("storage", onStorage); + window.removeEventListener(LOCAL_STORAGE_EVENT, onCustom); + }; +} + +function getSnapshot() { + return getLocalStorageItem("disableShowNewBadge") === "true"; +} + +export function useDisableShowNewBadge() { + return useSyncExternalStore(subscribe, getSnapshot); +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts new file mode 100644 index 00000000000..a392a940f98 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.test.ts @@ -0,0 +1,253 @@ +import { describe, it, expect, vi, beforeEach } from "vitest"; +import { renderHook, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import React, { ReactNode } from "react"; +import { useCurrentUser } from "./useCurrentUser"; +import { userInfoCall } from "@/components/networking"; +import type { UserInfo } from "@/components/view_users/types"; + +// Mock the networking function +vi.mock("@/components/networking", () => ({ + userInfoCall: vi.fn(), +})); + +// Mock the queryKeysFactory - we'll mock the specific return value +vi.mock("../common/queryKeysFactory", () => ({ + createQueryKeys: vi.fn((resource: string) => ({ + all: [resource], + lists: () => [resource, "list"], + list: (params?: any) => [resource, "list", { params }], + details: () => [resource, "detail"], + detail: (uid: string) => [resource, "detail", uid], + })), +})); + +// Mock useAuthorized hook - we can override this in individual tests +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + +// Mock data - response from userInfoCall should have user_info property +const mockUserInfoResponse = { + user_info: { + user_id: "test-user-id", + user_email: "test@example.com", + user_alias: "Test User", + user_role: "Admin", + spend: 150.75, + max_budget: 1000.0, + key_count: 5, + created_at: "2024-01-01T00:00:00Z", + updated_at: "2024-01-01T00:00:00Z", + sso_user_id: null, + budget_duration: "monthly", + } as UserInfo, +}; + +describe("useCurrentUser", () => { + let queryClient: QueryClient; + + beforeEach(() => { + queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + + // Reset all mocks + vi.clearAllMocks(); + + // Set default mock for useAuthorized (enabled state) + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + }); + + const wrapper = ({ children }: { children: ReactNode }) => + React.createElement(QueryClientProvider, { client: queryClient }, children); + + it("should return user info data when query is successful", async () => { + // Mock successful API call + (userInfoCall as any).mockResolvedValue(mockUserInfoResponse); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + expect(result.current.data).toBeUndefined(); + + // Wait for success + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isSuccess).toBe(true); + }); + + expect(result.current.data).toEqual(mockUserInfoResponse.user_info); + expect(result.current.error).toBeNull(); + expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null); + expect(userInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should handle error when userInfoCall fails", async () => { + const errorMessage = "Failed to fetch user info"; + const testError = new Error(errorMessage); + + // Mock failed API call + (userInfoCall as any).mockRejectedValue(testError); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Initially loading + expect(result.current.isLoading).toBe(true); + + // Wait for error + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(testError); + expect(result.current.data).toBeUndefined(); + expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null); + expect(userInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should not execute query when accessToken is missing", async () => { + // Mock missing accessToken + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: "test-user-id", + userRole: "Admin", + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(userInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userId is missing", async () => { + // Mock missing userId + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: null, + userRole: "Admin", + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(userInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when userRole is missing", async () => { + // Mock missing userRole + mockUseAuthorized.mockReturnValue({ + accessToken: "test-access-token", + userId: "test-user-id", + userRole: null, + token: "test-token", + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(userInfoCall).not.toHaveBeenCalled(); + }); + + it("should not execute query when all auth values are missing", async () => { + // Mock all auth values missing + mockUseAuthorized.mockReturnValue({ + accessToken: null, + userId: null, + userRole: null, + token: null, + userEmail: "test@example.com", + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Query should not execute + expect(result.current.isLoading).toBe(false); + expect(result.current.data).toBeUndefined(); + expect(result.current.isFetched).toBe(false); + + // API should not be called + expect(userInfoCall).not.toHaveBeenCalled(); + }); + + it("should execute query when all auth values are present", async () => { + // Mock successful API call + (userInfoCall as any).mockResolvedValue(mockUserInfoResponse); + + // Ensure all auth values are present (already set in beforeEach) + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Wait for query to execute + await waitFor(() => { + expect(result.current.isLoading).toBe(false); + }); + + expect(userInfoCall).toHaveBeenCalledWith("test-access-token", "test-user-id", "Admin", false, null, null); + expect(userInfoCall).toHaveBeenCalledTimes(1); + }); + + it("should handle network timeout error", async () => { + const timeoutError = new Error("Network timeout"); + + // Mock network timeout + (userInfoCall as any).mockRejectedValue(timeoutError); + + const { result } = renderHook(() => useCurrentUser(), { wrapper }); + + // Wait for error + await waitFor(() => { + expect(result.current.isError).toBe(true); + }); + + expect(result.current.error).toEqual(timeoutError); + expect(result.current.data).toBeUndefined(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts new file mode 100644 index 00000000000..f4028ada0dc --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/users/useCurrentUser.ts @@ -0,0 +1,19 @@ +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { UserInfo, userInfoCall } from "@/components/networking"; +import { useQuery, UseQueryResult } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const userKeys = createQueryKeys("users"); + +export const useCurrentUser = (): UseQueryResult => { + const { accessToken, userId, userRole } = useAuthorized(); + return useQuery({ + queryKey: userKeys.detail(userId!), + queryFn: async () => { + const data = await userInfoCall(accessToken!, userId!, userRole!, false, null, null); + console.log(`userInfo: ${JSON.stringify(data)}`); + return data.user_info; + }, + enabled: Boolean(accessToken && userId && userRole), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx index 86967b660fd..c37a935976b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/model-hub/page.tsx @@ -1,6 +1,6 @@ "use client"; -import ModelHubTable from "@/components/model_hub_table"; +import ModelHubTable from "@/components/AIHub/ModelHubTable"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; const ModelHubPage = () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx index 8dc7d1ff3d6..1e8eabaea2e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.test.tsx @@ -9,34 +9,24 @@ vi.mock("@/components/networking", () => ({ credentialListCall: vi.fn().mockResolvedValue({ credentials: [] }), modelInfoCall: vi.fn().mockResolvedValue({ data: [] }), modelCostMap: vi.fn().mockResolvedValue({}), - modelMetricsCall: vi.fn().mockResolvedValue({ data: [], all_api_bases: [] }), - streamingModelMetricsCall: vi.fn().mockResolvedValue({ data: [], all_api_bases: [] }), - modelExceptionsCall: vi.fn().mockResolvedValue({ data: [], exception_types: [] }), - modelMetricsSlowResponsesCall: vi.fn().mockResolvedValue([]), + getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: {} }), getCallbacksCall: vi.fn().mockResolvedValue({ router_settings: {} }), setCallbacksCall: vi.fn().mockResolvedValue(undefined), - modelSettingsCall: vi.fn().mockResolvedValue([]), - adminGlobalActivityExceptions: vi.fn().mockResolvedValue({ sum_num_rate_limit_exceptions: 0, daily_data: [] }), - adminGlobalActivityExceptionsPerDeployment: vi.fn().mockResolvedValue([]), - allEndUsersCall: vi.fn().mockResolvedValue([]), - latestHealthChecksCall: vi.fn().mockResolvedValue({ latest_health_checks: {} }), - getPassThroughEndpointsCall: vi.fn().mockResolvedValue({ endpoints: {} }), - getGuardrailsList: vi.fn().mockResolvedValue([]), - tagListCall: vi.fn().mockResolvedValue([]), - modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), - modelHubCall: vi.fn().mockResolvedValue({ data: [] }), - getModelCostMapReloadStatus: vi.fn().mockResolvedValue({ - scheduled: false, - interval_hours: null, - last_run: null, - next_run: null, - }), + getUiSettings: vi.fn().mockResolvedValue({ values: {} }), })); vi.mock("@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab", () => ({ default: () => null, })); +vi.mock("@/components/add_model/add_auto_router_tab", () => ({ + default: () => null, +})); + +vi.mock("@/components/add_model/AddModelForm", () => ({ + default: () => null, +})); + vi.mock("@/app/(dashboard)/hooks/useTeams", () => ({ default: () => ({ teams: [], @@ -54,6 +44,16 @@ vi.mock("@/app/(dashboard)/hooks/uiSettings/useUISettings", () => ({ useUISettings: () => mockUseUISettings(), })); +const mockUseModelCostMap = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/models/useModelCostMap", () => ({ + useModelCostMap: () => mockUseModelCostMap(), +})); + +const mockUseAuthorized = vi.fn(); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); + const createQueryClient = () => new QueryClient({ defaultOptions: { queries: { retry: false, gcTime: 0 } }, @@ -69,6 +69,17 @@ describe("ModelsAndEndpointsView", () => { mockUseUISettings.mockReturnValue({ data: { values: {} }, }); + mockUseModelCostMap.mockReturnValue({ + data: {}, + isLoading: false, + error: null, + }); + mockUseAuthorized.mockReturnValue({ + accessToken: "123", + token: "123", + userRole: "Admin", + userId: "123", + }); // eslint-disable-next-line @typescript-eslint/no-explicit-any (global as any).ResizeObserver = class { observe() {} @@ -82,10 +93,7 @@ describe("ModelsAndEndpointsView", () => { const { findByText } = render( {}} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx index bf001c62126..65418179595 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView.tsx @@ -1,52 +1,35 @@ -import { useQueryClient } from "@tanstack/react-query"; -import { Col, Grid, Text } from "@tremor/react"; -import React, { useEffect, useRef, useState } from "react"; - -import { handleAddModelSubmit } from "@/components/add_model/handle_add_model_submit"; - import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials"; +import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels"; -import { Team } from "@/components/key_team_helpers/key_list"; -import CredentialsPanel from "@/components/model_add/credentials"; -import { - adminGlobalActivityExceptions, - adminGlobalActivityExceptionsPerDeployment, - allEndUsersCall, - getCallbacksCall, - modelCostMap, - modelExceptionsCall, - modelMetricsCall, - modelMetricsSlowResponsesCall, - modelSettingsCall, - setCallbacksCall, - streamingModelMetricsCall, -} from "@/components/networking"; -import { Providers, getPlaceholder, getProviderModels } from "@/components/provider_info_helpers"; -import { getDisplayModelName } from "@/components/view_model/model_name_display"; -import { RefreshIcon } from "@heroicons/react/outline"; -import { DateRangePickerValue, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; -import type { UploadProps } from "antd"; -import { Form, Typography } from "antd"; -import AddModelTab from "../../../components/add_model/add_model_tab"; -import ModelInfoView from "../../../components/model_info_view"; -import TeamInfoView from "../../../components/team/team_info"; - +import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; import AllModelsTab from "@/app/(dashboard)/models-and-endpoints/components/AllModelsTab"; -import ModelAnalyticsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelAnalyticsTab/ModelAnalyticsTab"; import ModelRetrySettingsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab"; import PriceDataManagementTab from "@/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab"; -import { useUISettings } from "@/app/(dashboard)/hooks/uiSettings/useUISettings"; +import { handleAddModelSubmit } from "@/components/add_model/handle_add_model_submit"; +import { Team } from "@/components/key_team_helpers/key_list"; +import CredentialsPanel from "@/components/model_add/credentials"; +import { getCallbacksCall, setCallbacksCall } from "@/components/networking"; +import { Providers, getPlaceholder, getProviderModels } from "@/components/provider_info_helpers"; +import { getDisplayModelName } from "@/components/view_model/model_name_display"; +import { transformModelData } from "./utils/modelDataTransformer"; import { all_admin_roles, internalUserRoles, isProxyAdminRole, isUserTeamAdminForAnyTeam } from "@/utils/roles"; +import { RefreshIcon } from "@heroicons/react/outline"; +import { useQueryClient } from "@tanstack/react-query"; +import { Col, Grid, Icon, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react"; +import type { UploadProps } from "antd"; +import { Form, Typography } from "antd"; +import React, { useEffect, useMemo, useState } from "react"; +import AddModelTab from "../../../components/add_model/add_model_tab"; import HealthCheckComponent from "../../../components/model_dashboard/HealthCheckComponent"; import ModelGroupAliasSettings from "../../../components/model_group_alias_settings"; +import ModelInfoView from "../../../components/model_info_view"; import NotificationsManager from "../../../components/molecules/notifications_manager"; import PassThroughSettings from "../../../components/pass_through_settings"; +import TeamInfoView from "../../../components/team/team_info"; +import useAuthorized from "../hooks/useAuthorized"; interface ModelDashboardProps { - accessToken: string | null; token: string | null; - userRole: string | null; - userID: string | null; modelData: any; keys: any[] | null; setModelData: any; @@ -62,100 +45,70 @@ interface GlobalRetryPolicyObject { [retryPolicyKey: string]: number; } -interface GlobalExceptionActivityData { - sum_num_rate_limit_exceptions: number; - daily_data: { date: string; num_rate_limit_exceptions: number }[]; -} - -//["OpenAI", "Azure OpenAI", "Anthropic", "Gemini (Google AI Studio)", "Amazon Bedrock", "OpenAI-Compatible Endpoints (Groq, Together AI, Mistral AI, etc.)"] - -interface ProviderFields { - field_name: string; - field_type: string; - field_description: string; - field_value: string; -} - -interface ProviderSettings { - name: string; - fields: ProviderFields[]; -} - -const ModelsAndEndpointsView: React.FC = ({ - accessToken, - token, - userRole, - userID, - modelData = { data: [] }, - keys, - setModelData, - premiumUser, - teams, -}) => { +const ModelsAndEndpointsView: React.FC = ({ premiumUser, teams }) => { + const { accessToken, token, userRole, userId: userID } = useAuthorized(); const [addModelForm] = Form.useForm(); - const [modelMap, setModelMap] = useState(null); const [lastRefreshed, setLastRefreshed] = useState(""); - - const [providerModels, setProviderModels] = useState>([]); // Explicitly typing providerModels as a string array - - const [providerSettings, setProviderSettings] = useState([]); + const [providerModels, setProviderModels] = useState>([]); const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); - const [editModalVisible, setEditModalVisible] = useState(false); - - const [selectedModel, setSelectedModel] = useState(null); - const [availableModelGroups, setAvailableModelGroups] = useState>([]); - const [availableModelAccessGroups, setAvailableModelAccessGroups] = useState>([]); const [selectedModelGroup, setSelectedModelGroup] = useState(null); - const [modelMetrics, setModelMetrics] = useState([]); - const [modelMetricsCategories, setModelMetricsCategories] = useState([]); - const [streamingModelMetrics, setStreamingModelMetrics] = useState([]); - const [streamingModelMetricsCategories, setStreamingModelMetricsCategories] = useState([]); - const [modelExceptions, setModelExceptions] = useState([]); - const [allExceptions, setAllExceptions] = useState([]); - const [slowResponsesData, setSlowResponsesData] = useState([]); - const [dateValue, setDateValue] = useState({ - from: new Date(Date.now() - 7 * 24 * 60 * 60 * 1000), - to: new Date(), - }); const [modelGroupRetryPolicy, setModelGroupRetryPolicy] = useState(null); const [globalRetryPolicy, setGlobalRetryPolicy] = useState(null); const [defaultRetry, setDefaultRetry] = useState(0); - - const [globalExceptionData, setGlobalExceptionData] = useState( - {} as GlobalExceptionActivityData, - ); - const [globalExceptionPerDeployment, setGlobalExceptionPerDeployment] = useState([]); - - const [showAdvancedFilters, setShowAdvancedFilters] = useState(false); - const [selectedAPIKey, setSelectedAPIKey] = useState(null); - const [selectedCustomer, setSelectedCustomer] = useState(null); - - const [allEndUsers, setAllEndUsers] = useState([]); - - // Model Group Alias state const [modelGroupAlias, setModelGroupAlias] = useState<{ [key: string]: string }>({}); - - // Add state for advanced settings visibility const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); - - // Add these state variables const [selectedModelId, setSelectedModelId] = useState(null); - const [editModel, setEditModel] = useState(false); - const [selectedTeamId, setSelectedTeamId] = useState(null); - const [selectedTeam, setSelectedTeam] = useState(null); - - const [isDropdownOpen, setIsDropdownOpen] = useState(false); - const dropdownRef = useRef(null); - const [selectedTabIndex, setSelectedTabIndex] = useState(0); const queryClient = useQueryClient(); const { data: modelDataResponse, isLoading: isLoadingModels, refetch: refetchModels } = useModelsInfo(); - const { data: credentialsResponse } = useCredentials(); + const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap(); + const { data: credentialsResponse, isLoading: isLoadingCredentials } = useCredentials(); const credentialsList = credentialsResponse?.credentials || []; - const { data: uiSettings } = useUISettings(accessToken || ""); + const { data: uiSettings, isLoading: isLoadingUISettings } = useUISettings(); + + const availableModelGroups = useMemo(() => { + if (!modelDataResponse?.data) return []; + const allModelGroups = new Set(); + for (const model of modelDataResponse.data) { + allModelGroups.add(model.model_name); + } + return Array.from(allModelGroups).sort(); + }, [modelDataResponse?.data]); + + const availableModelAccessGroups = useMemo(() => { + if (!modelDataResponse?.data) return []; + const allModelAccessGroups = new Set(); + for (const model of modelDataResponse.data) { + const modelInfo = model.model_info; + if (modelInfo?.access_groups) { + for (const group of modelInfo.access_groups) { + allModelAccessGroups.add(group); + } + } + } + return Array.from(allModelAccessGroups); + }, [modelDataResponse?.data]); + + const allModelsOnProxy = useMemo(() => { + return modelDataResponse?.data?.map((model: any) => model.model_name); + }, [modelDataResponse?.data]); + + const getProviderFromModel = (model: string) => { + if (modelCostMapData !== null && modelCostMapData !== undefined) { + if (typeof modelCostMapData == "object" && model in modelCostMapData) { + return modelCostMapData[model]["litellm_provider"]; + } + } + return "openai"; + }; + + const processedModelData = useMemo(() => { + if (!modelDataResponse?.data) return { data: [] }; + return transformModelData(modelDataResponse, getProviderFromModel); + }, [modelDataResponse?.data, getProviderFromModel]); const isProxyAdmin = userRole && isProxyAdminRole(userRole); const isInternalUser = userRole && internalUserRoles.includes(userRole); @@ -166,21 +119,10 @@ const ModelsAndEndpointsView: React.FC = ({ const shouldHideAddModelTab = !isProxyAdmin && (addModelDisabledForInternalUsers || !isUserTeamAdmin); const setProviderModelsFn = (provider: Providers) => { - const _providerModels = getProviderModels(provider, modelMap); + const _providerModels = getProviderModels(provider, modelCostMapData); setProviderModels(_providerModels); }; - useEffect(() => { - const handleClickOutside = (event: MouseEvent) => { - if (dropdownRef.current && !dropdownRef.current.contains(event.target as Node)) { - setIsDropdownOpen(false); - } - }; - - document.addEventListener("mousedown", handleClickOutside); - return () => document.removeEventListener("mousedown", handleClickOutside); - }, []); - const uploadProps: UploadProps = { name: "file", accept: ".json", @@ -196,7 +138,6 @@ const ModelsAndEndpointsView: React.FC = ({ }; reader.readAsText(file); } - // Prevent upload return false; }, onChange(info) { @@ -209,10 +150,8 @@ const ModelsAndEndpointsView: React.FC = ({ }; const handleRefreshClick = () => { - // Update the 'lastRefreshed' state to the current date and time const currentDate = new Date(); setLastRefreshed(currentDate.toLocaleString()); - // Invalidate and refetch models data using React Query queryClient.invalidateQueries({ queryKey: ["models", "list"] }); refetchModels(); }; @@ -228,7 +167,6 @@ const ModelsAndEndpointsView: React.FC = ({ }; if (selectedModelGroup === "global") { - // Only update global retry policy if (globalRetryPolicy) { payload.router_settings.retry_policy = globalRetryPolicy; } @@ -252,114 +190,6 @@ const ModelsAndEndpointsView: React.FC = ({ } const fetchData = async () => { try { - setModelData(modelDataResponse); - const _providerSettings = await modelSettingsCall(accessToken); - if (_providerSettings) { - setProviderSettings(_providerSettings); - } - - // loop through modelDataResponse and get all`model_name` values - let all_model_groups: Set = new Set(); - for (let i = 0; i < modelDataResponse.data.length; i++) { - const model = modelDataResponse.data[i]; - all_model_groups.add(model.model_name); - } - let _array_model_groups = Array.from(all_model_groups); - // sort _array_model_groups alphabetically - _array_model_groups = _array_model_groups.sort(); - - setAvailableModelGroups(_array_model_groups); - - let all_model_access_groups: Set = new Set(); - for (let i = 0; i < modelDataResponse.data.length; i++) { - const model = modelDataResponse.data[i]; - let model_info: any | null = model.model_info; - if (model_info) { - let access_groups = model_info.access_groups; - if (access_groups) { - for (let j = 0; j < access_groups.length; j++) { - all_model_access_groups.add(access_groups[j]); - } - } - } - } - - setAvailableModelAccessGroups(Array.from(all_model_access_groups)); - - let _initial_model_group = "all"; - if (_array_model_groups.length > 0) { - _initial_model_group = _array_model_groups[_array_model_groups.length - 1]; - } - - const modelMetricsResponse = await modelMetricsCall( - accessToken, - userID, - userRole, - _initial_model_group, - dateValue.from?.toISOString(), - dateValue.to?.toISOString(), - selectedAPIKey?.token, - selectedCustomer, - ); - - setModelMetrics(modelMetricsResponse.data); - setModelMetricsCategories(modelMetricsResponse.all_api_bases); - - const streamingModelMetricsResponse = await streamingModelMetricsCall( - accessToken, - _initial_model_group, - dateValue.from?.toISOString(), - dateValue.to?.toISOString(), - ); - - // Assuming modelMetricsResponse now contains the metric data for the specified model group - setStreamingModelMetrics(streamingModelMetricsResponse.data); - setStreamingModelMetricsCategories(streamingModelMetricsResponse.all_api_bases); - - const modelExceptionsResponse = await modelExceptionsCall( - accessToken, - userID, - userRole, - _initial_model_group, - dateValue.from?.toISOString(), - dateValue.to?.toISOString(), - selectedAPIKey?.token, - selectedCustomer, - ); - setModelExceptions(modelExceptionsResponse.data); - setAllExceptions(modelExceptionsResponse.exception_types); - - const slowResponses = await modelMetricsSlowResponsesCall( - accessToken, - userID, - userRole, - _initial_model_group, - dateValue.from?.toISOString(), - dateValue.to?.toISOString(), - selectedAPIKey?.token, - selectedCustomer, - ); - - const dailyExceptions = await adminGlobalActivityExceptions( - accessToken, - dateValue.from?.toISOString().split("T")[0], - dateValue.to?.toISOString().split("T")[0], - _initial_model_group, - ); - - setGlobalExceptionData(dailyExceptions); - - const dailyExceptionsPerDeplyment = await adminGlobalActivityExceptionsPerDeployment( - accessToken, - dateValue.from?.toISOString().split("T")[0], - dateValue.to?.toISOString().split("T")[0], - _initial_model_group, - ); - - setGlobalExceptionPerDeployment(dailyExceptionsPerDeplyment); - setSlowResponsesData(slowResponses); - let all_end_users_data = await allEndUsersCall(accessToken); - setAllEndUsers(all_end_users_data?.map((u: any) => u.user_id)); const routerSettingsInfo = await getCallbacksCall(accessToken, userID, userRole); let router_settings = routerSettingsInfo.router_settings; @@ -370,7 +200,6 @@ const ModelsAndEndpointsView: React.FC = ({ setGlobalRetryPolicy(router_settings.retry_policy); setDefaultRetry(default_retries); - // Set model group alias const model_group_alias = router_settings.model_group_alias || {}; setModelGroupAlias(model_group_alias); } catch (error) { @@ -381,110 +210,9 @@ const ModelsAndEndpointsView: React.FC = ({ if (accessToken && token && userRole && userID && modelDataResponse) { fetchData(); } - - const fetchModelMap = async () => { - const data = await modelCostMap(); - console.log(`received model cost map data: ${Object.keys(data)}`); - setModelMap(data); - }; - if (modelMap == null) { - fetchModelMap(); - } }, [accessToken, token, userRole, userID, modelDataResponse]); - if (!modelData || isLoadingModels) { - return
Loading...
; - } - - if (!accessToken || !token || !userRole || !userID) { - return
Loading...
; - } - let all_models_on_proxy: any[] = []; - let all_providers: string[] = []; - - // loop through model data and edit each row - for (let i = 0; i < modelData.data.length; i++) { - let curr_model = modelData.data[i]; - let litellm_model_name = curr_model?.litellm_params?.model; - let custom_llm_provider = curr_model?.litellm_params?.custom_llm_provider; - let model_info = curr_model?.model_info; - - let defaultProvider = "openai"; - let provider = ""; - let input_cost = "Undefined"; - let output_cost = "Undefined"; - let max_tokens = "Undefined"; - let max_input_tokens = "Undefined"; - let cleanedLitellmParams = {}; - - const getProviderFromModel = (model: string) => { - /** - * Use model map - * - check if model in model map - * - return it's litellm_provider, if so - */ - if (modelMap !== null && modelMap !== undefined) { - if (typeof modelMap == "object" && model in modelMap) { - return modelMap[model]["litellm_provider"]; - } - } - return "openai"; - }; - - // Check if litellm_model_name is null or undefined - if (litellm_model_name) { - // Split litellm_model_name based on "/" - let splitModel = litellm_model_name.split("/"); - - // Get the first element in the split - let firstElement = splitModel[0]; - - // If there is only one element, default provider to openai - provider = custom_llm_provider; - if (!provider) { - provider = splitModel.length === 1 ? getProviderFromModel(litellm_model_name) : firstElement; - } - } else { - // litellm_model_name is null or undefined, default provider to openai - provider = "-"; - } - - if (model_info) { - input_cost = model_info?.input_cost_per_token; - output_cost = model_info?.output_cost_per_token; - max_tokens = model_info?.max_tokens; - max_input_tokens = model_info?.max_input_tokens; - } - - if (curr_model?.litellm_params) { - cleanedLitellmParams = Object.fromEntries( - Object.entries(curr_model?.litellm_params).filter(([key]) => key !== "model" && key !== "api_base"), - ); - } - - modelData.data[i].provider = provider; - modelData.data[i].input_cost = input_cost; - modelData.data[i].output_cost = output_cost; - modelData.data[i].litellm_model_name = litellm_model_name; - all_providers.push(provider); - - // Convert Cost in terms of Cost per 1M tokens - if (modelData.data[i].input_cost) { - modelData.data[i].input_cost = (Number(modelData.data[i].input_cost) * 1000000).toFixed(2); - } - - if (modelData.data[i].output_cost) { - modelData.data[i].output_cost = (Number(modelData.data[i].output_cost) * 1000000).toFixed(2); - } - - modelData.data[i].max_tokens = max_tokens; - modelData.data[i].max_input_tokens = max_input_tokens; - modelData.data[i].api_base = curr_model?.litellm_params?.api_base; - modelData.data[i].cleanedLitellmParams = cleanedLitellmParams; - - all_models_on_proxy.push(curr_model.model_name); - } - // when users click request access show pop up to allow them to request access + const isLoading = isLoadingModels || isLoadingModelCostMap || isLoadingCredentials || isLoadingUISettings; if (userRole && userRole == "Admin Viewer") { const { Title, Paragraph } = Typography; @@ -495,62 +223,20 @@ const ModelsAndEndpointsView: React.FC = ({ ); } - const customTooltip = (props: any) => { - const { payload, active } = props; - if (!active || !payload) return null; - // Extract the date from the first item in the payload array - const date = payload[0]?.payload?.date; - - // Sort the payload array by category.value in descending order - let sortedPayload = payload.sort((a: any, b: any) => b.value - a.value); - - // Only show the top 5, the 6th one should be called "X other categories" depending on how many categories were not shown - if (sortedPayload.length > 5) { - let remainingItems = sortedPayload.length - 5; - sortedPayload = sortedPayload.slice(0, 5); - sortedPayload.push({ - dataKey: `${remainingItems} other deployments`, - value: payload.slice(5).reduce((acc: number, curr: any) => acc + curr.value, 0), - color: "gray", - }); + const handleOk = async () => { + try { + const values = await addModelForm.validateFields(); + await handleAddModelSubmit(values, accessToken, addModelForm, handleRefreshClick); + } catch (error: any) { + const errorMessages = + error.errorFields + ?.map((field: any) => { + return `${field.name.join(".")}: ${field.errors.join(", ")}`; + }) + .join(" | ") || "Unknown validation error"; + NotificationsManager.fromBackend(`Please fill in the following required fields: ${errorMessages}`); } - - return ( -
- {date &&

Date: {date}

} - {sortedPayload.map((category: any, idx: number) => { - const roundedValue = parseFloat(category.value.toFixed(5)); - const displayValue = roundedValue === 0 && category.value > 0 ? "<0.00001" : roundedValue.toFixed(5); - return ( -
-
-
-

{category.dataKey}

-
-

{displayValue}

-
- ); - })} -
- ); - }; - - const handleOk = () => { - addModelForm - .validateFields() - .then((values: any) => { - handleAddModelSubmit(values, accessToken, addModelForm, handleRefreshClick); - }) - .catch((error: any) => { - const errorMessages = - error.errorFields - ?.map((field: any) => { - return `${field.name.join(".")}: ${field.errors.join(", ")}`; - }) - .join(" | ") || "Unknown validation error"; - NotificationsManager.fromBackend(`Please fill in the following required fields: ${errorMessages}`); - }); }; Object.keys(Providers).find((key) => (Providers as { [index: string]: any })[key] === selectedProvider); @@ -564,7 +250,7 @@ const ModelsAndEndpointsView: React.FC = ({ accessToken={accessToken} is_team_admin={userRole === "Admin"} is_proxy_admin={userRole === "Proxy Admin"} - userModels={all_models_on_proxy} + userModels={allModelsOnProxy} editTeam={false} onUpdate={handleRefreshClick} premiumUser={premiumUser} @@ -588,39 +274,17 @@ const ModelsAndEndpointsView: React.FC = ({ )}
- {selectedModelId ? ( + {selectedModelId && !isLoading ? ( { setSelectedModelId(null); - setEditModel(false); }} - modelData={modelData.data.find((model: any) => model.model_info.id === selectedModelId)} + modelData={processedModelData.data.find((model: any) => model.model_info.id === selectedModelId)} accessToken={accessToken} userID={userID} userRole={userRole} - setEditModalVisible={setEditModalVisible} - setSelectedModel={setSelectedModel} onModelUpdate={(updatedModel) => { - // Handle model deletion - if (updatedModel.deleted) { - const updatedModelData = { - ...modelData, - data: modelData.data.filter((model: any) => model.model_info.id !== updatedModel.model_info.id), - }; - setModelData(updatedModelData); - } else { - // Update the model in the modelData.data array - const updatedModelData = { - ...modelData, - data: modelData.data.map((model: any) => - model.model_info.id === updatedModel.model_info.id ? updatedModel : model, - ), - }; - setModelData(updatedModelData); - } - // Invalidate cache and trigger a refresh to update UI queryClient.invalidateQueries({ queryKey: ["models", "list"] }); handleRefreshClick(); }} @@ -635,7 +299,6 @@ const ModelsAndEndpointsView: React.FC = ({ {all_admin_roles.includes(userRole) && LLM Credentials} {all_admin_roles.includes(userRole) && Pass-Through Endpoints} {all_admin_roles.includes(userRole) && Health Status} - {all_admin_roles.includes(userRole) && Model Analytics} {all_admin_roles.includes(userRole) && Model Retry Settings} {all_admin_roles.includes(userRole) && Model Group Alias} {all_admin_roles.includes(userRole) && Price Data Reload} @@ -660,7 +323,6 @@ const ModelsAndEndpointsView: React.FC = ({ availableModelAccessGroups={availableModelAccessGroups} setSelectedModelId={setSelectedModelId} setSelectedTeamId={setSelectedTeamId} - setEditModel={setEditModel} /> {!shouldHideAddModelTab && ( @@ -679,7 +341,6 @@ const ModelsAndEndpointsView: React.FC = ({ credentials={credentialsList} accessToken={accessToken} userRole={userRole} - premiumUser={premiumUser} /> )} @@ -691,54 +352,19 @@ const ModelsAndEndpointsView: React.FC = ({ accessToken={accessToken} userRole={userRole} userID={userID} - modelData={modelData} + modelData={processedModelData} premiumUser={premiumUser} /> - = ({ onAliasUpdate={setModelGroupAlias} /> - + )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx index dfa400e6ea9..ae376701a9a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.test.tsx @@ -10,6 +10,28 @@ vi.mock("../../hooks/models/useModels", () => ({ useModelsInfo: () => mockUseModelsInfo(), })); +// Mock the useModelCostMap hook +const mockUseModelCostMap = vi.fn(() => ({ + data: { + "gpt-4": { litellm_provider: "openai" }, + "gpt-3.5-turbo": { litellm_provider: "openai" }, + "gpt-4-accessible": { litellm_provider: "openai" }, + "gpt-3.5-turbo-blocked": { litellm_provider: "openai" }, + "gpt-4-sales": { litellm_provider: "openai" }, + "gpt-4-engineering": { litellm_provider: "openai" }, + "gpt-4-personal": { litellm_provider: "openai" }, + "gpt-4-team-only": { litellm_provider: "openai" }, + "gpt-4-config": { litellm_provider: "openai" }, + "gpt-4-db": { litellm_provider: "openai" }, + }, + isLoading: false, + error: null, +})) as any; + +vi.mock("../../hooks/models/useModelCostMap", () => ({ + useModelCostMap: () => mockUseModelCostMap(), +})); + // Mock the useTeams hook (react-query implementation) const mockUseTeams = vi.fn(() => ({ data: [], @@ -22,6 +44,13 @@ vi.mock("../../hooks/teams/useTeams", () => ({ useTeams: () => mockUseTeams(), })); +// Helper function to create model cost map mock return value +const createModelCostMapMock = (data: Record) => ({ + data, + isLoading: false, + error: null, +}); + describe("AllModelsTab", () => { const mockSetSelectedModelGroup = vi.fn(); const mockSetSelectedModelId = vi.fn(); @@ -64,6 +93,8 @@ describe("AllModelsTab", () => { refetch: vi.fn(), }); + mockUseModelCostMap.mockReturnValueOnce(createModelCostMapMock({})); + render(); expect(screen.getByText("Current Team:")).toBeInTheDocument(); }); @@ -92,6 +123,13 @@ describe("AllModelsTab", () => { refetch: vi.fn(), }); + mockUseModelCostMap.mockReturnValueOnce( + createModelCostMapMock({ + "gpt-4-accessible": { litellm_provider: "openai" }, + "gpt-3.5-turbo-blocked": { litellm_provider: "openai" }, + }), + ); + const modelData = { data: [ { @@ -146,6 +184,13 @@ describe("AllModelsTab", () => { refetch: vi.fn(), }); + mockUseModelCostMap.mockReturnValueOnce( + createModelCostMapMock({ + "gpt-4-sales": { litellm_provider: "openai" }, + "gpt-4-engineering": { litellm_provider: "openai" }, + }), + ); + const modelData = { data: [ { @@ -184,6 +229,13 @@ describe("AllModelsTab", () => { refetch: vi.fn(), }); + mockUseModelCostMap.mockReturnValueOnce( + createModelCostMapMock({ + "gpt-4-personal": { litellm_provider: "openai" }, + "gpt-4-team-only": { litellm_provider: "openai" }, + }), + ); + const modelData = { data: [ { @@ -224,6 +276,13 @@ describe("AllModelsTab", () => { refetch: vi.fn(), }); + mockUseModelCostMap.mockReturnValueOnce( + createModelCostMapMock({ + "gpt-4-config": { litellm_provider: "openai" }, + "gpt-4-db": { litellm_provider: "openai" }, + }), + ); + const modelData = { data: [ { @@ -277,6 +336,12 @@ describe("AllModelsTab", () => { refetch: vi.fn(), }); + mockUseModelCostMap.mockReturnValueOnce( + createModelCostMapMock({ + "gpt-4-config": { litellm_provider: "openai" }, + }), + ); + const modelData = { data: [ { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx index 04c05ede5c2..a85e6516585 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AllModelsTab.tsx @@ -1,3 +1,4 @@ +import { useModelCostMap } from "@/app/(dashboard)/hooks/models/useModelCostMap"; import { useTeams } from "@/app/(dashboard)/hooks/teams/useTeams"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { Team } from "@/components/key_team_helpers/key_list"; @@ -9,7 +10,8 @@ import { PaginationState } from "@tanstack/react-table"; import { Grid, Select, SelectItem, TabPanel, Text } from "@tremor/react"; import { useEffect, useMemo, useState } from "react"; import { useModelsInfo } from "../../hooks/models/useModels"; - +import { transformModelData } from "../utils/modelDataTransformer"; +import { Skeleton } from "antd"; type ModelViewMode = "all" | "current_team"; interface AllModelsTabProps { @@ -19,7 +21,6 @@ interface AllModelsTabProps { availableModelAccessGroups: string[]; setSelectedModelId: (id: string) => void; setSelectedTeamId: (id: string) => void; - setEditModel: (edit: boolean) => void; } const AllModelsTab = ({ @@ -29,12 +30,26 @@ const AllModelsTab = ({ availableModelAccessGroups, setSelectedModelId, setSelectedTeamId, - setEditModel, }: AllModelsTabProps) => { - const { data: modelData } = useModelsInfo(); + const { data: rawModelData, isLoading: isLoadingModelsInfo } = useModelsInfo(); + const { data: modelCostMapData, isLoading: isLoadingModelCostMap } = useModelCostMap(); const { userId, userRole, premiumUser } = useAuthorized(); const { data: teams } = useTeams(); + const getProviderFromModel = (model: string) => { + if (modelCostMapData !== null && modelCostMapData !== undefined) { + if (typeof modelCostMapData == "object" && model in modelCostMapData) { + return modelCostMapData[model]["litellm_provider"]; + } + } + return "openai"; + }; + + const modelData = useMemo(() => { + if (!rawModelData) return { data: [] }; + return transformModelData(rawModelData, getProviderFromModel); + }, [rawModelData, modelCostMapData]); + const [modelNameSearch, setModelNameSearch] = useState(""); const [modelViewMode, setModelViewMode] = useState("current_team"); const [currentTeam, setCurrentTeam] = useState("personal"); @@ -46,6 +61,8 @@ const AllModelsTab = ({ pageSize: 50, }); + const isLoading = isLoadingModelsInfo || isLoadingModelCostMap; + const filteredData = useMemo(() => { if (!modelData || !modelData.data || modelData.data.length === 0) { return []; @@ -110,63 +127,71 @@ const AllModelsTab = ({
Current Team: - + {isLoading ? ( + + ) : ( + + )}
View: - + {isLoading ? ( + + ) : ( + + )}
@@ -304,18 +329,23 @@ const AllModelsTab = ({ {/* Results Count and Pagination Controls */}
- - {filteredData.length > 0 - ? `Showing ${pagination.pageIndex * pagination.pageSize + 1} - ${Math.min( - (pagination.pageIndex + 1) * pagination.pageSize, - filteredData.length, - )} of ${filteredData.length} results` - : "Showing 0 results"} - + {isLoading ? ( + + ) : ( + + {filteredData.length > 0 + ? `Showing ${pagination.pageIndex * pagination.pageSize + 1} - ${Math.min( + (pagination.pageIndex + 1) * pagination.pageSize, + filteredData.length, + )} of ${filteredData.length} results` + : "Showing 0 results"} + + )} - {/* Pagination Controls */} - {filteredData.length > pagination.pageSize && ( -
+
+ {isLoading ? ( + + ) : ( + )} + {isLoading ? ( + + ) : ( -
- )} + )} +
@@ -359,7 +393,6 @@ const AllModelsTab = ({ getDisplayModelName, () => {}, () => {}, - setEditModel, expandedRows, setExpandedRows, )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx index 4076c19c665..d44d19879d5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab.tsx @@ -1,15 +1,12 @@ import { TabPanel, Text, Title } from "@tremor/react"; import PriceDataReload from "@/components/price_data_reload"; -import { modelCostMap } from "@/components/networking"; import React from "react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { useModelCostMap } from "../../hooks/models/useModelCostMap"; -interface PriceDataManagementPanelProps { - setModelMap: (data: any) => void; -} - -const PriceDataManagementTab = ({ setModelMap }: PriceDataManagementPanelProps) => { +const PriceDataManagementTab = () => { const { accessToken } = useAuthorized(); + const { refetch: refetchModelCostMap } = useModelCostMap(); return ( @@ -23,12 +20,7 @@ const PriceDataManagementTab = ({ setModelMap }: PriceDataManagementPanelProps) { - // Refresh the model map after successful reload - const fetchModelMap = async () => { - const data = await modelCostMap(); - setModelMap(data); - }; - fetchModelMap(); + refetchModelCostMap(); }} buttonText="Reload Price Data" size="middle" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx index 01dd97505c8..77496aef3e6 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/page.tsx @@ -6,17 +6,14 @@ import { useState } from "react"; import ModelsAndEndpointsView from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; const ModelsAndEndpointsPage = () => { - const { token, accessToken, userRole, userId, premiumUser } = useAuthorized(); + const { token, premiumUser } = useAuthorized(); const [keys, setKeys] = useState([]); const { teams } = useTeams(); return ( {}} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.test.ts new file mode 100644 index 00000000000..eb7aecaa679 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.test.ts @@ -0,0 +1,53 @@ +import { transformModelData } from "./modelDataTransformer"; +import { describe, it, expect } from "vitest"; +describe("transformModelData", () => { + const mockGetProviderFromModel = (model: string) => { + if (model.includes("gpt")) return "openai"; + if (model.includes("claude")) return "anthropic"; + return "openai"; + }; + + it("should transform raw model data correctly", () => { + const rawData = { + data: [ + { + model_name: "gpt-4", + litellm_params: { + model: "gpt-4", + api_base: "https://api.openai.com", + api_key: "sk-123", + }, + model_info: { + input_cost_per_token: 0.0000015, + output_cost_per_token: 0.000002, + max_tokens: 8192, + max_input_tokens: 128000, + }, + }, + ], + }; + + const result = transformModelData(rawData, mockGetProviderFromModel); + + expect(result.data[0]).toHaveProperty("provider", "openai"); + expect(result.data[0]).toHaveProperty("input_cost", "1.50"); + expect(result.data[0]).toHaveProperty("output_cost", "2.00"); + expect(result.data[0]).toHaveProperty("max_tokens", 8192); + expect(result.data[0]).toHaveProperty("max_input_tokens", 128000); + expect(result.data[0]).toHaveProperty("api_base", "https://api.openai.com"); + expect(result.data[0]).toHaveProperty("litellm_model_name", "gpt-4"); + expect(result.data[0]).toHaveProperty("cleanedLitellmParams"); + expect(result.data[0].cleanedLitellmParams).not.toHaveProperty("model"); + expect(result.data[0].cleanedLitellmParams).not.toHaveProperty("api_base"); + }); + + it("should handle empty data", () => { + const result = transformModelData({ data: [] }, mockGetProviderFromModel); + expect(result).toEqual({ data: [] }); + }); + + it("should handle null/undefined data", () => { + const result = transformModelData(null, mockGetProviderFromModel); + expect(result).toEqual({ data: [] }); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.ts b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.ts new file mode 100644 index 00000000000..3ebf9ddd72b --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/utils/modelDataTransformer.ts @@ -0,0 +1,76 @@ +/** + * Utility function to transform raw model data into the format expected by UI components + * This creates a new transformed data object without mutating the original + */ +export const transformModelData = (rawModelData: any, getProviderFromModel: (model: string) => string) => { + if (!rawModelData?.data) return { data: [] }; + + // Deep copy the data to avoid mutating the original + const transformedData = JSON.parse(JSON.stringify(rawModelData.data)); + + for (let i = 0; i < transformedData.length; i++) { + let curr_model = transformedData[i]; + let litellm_model_name = curr_model?.litellm_params?.model; + let custom_llm_provider = curr_model?.litellm_params?.custom_llm_provider; + let model_info = curr_model?.model_info; + + let provider = ""; + let input_cost = "Undefined"; + let output_cost = "Undefined"; + let max_tokens = "Undefined"; + let max_input_tokens = "Undefined"; + let cleanedLitellmParams = {}; + + // Check if litellm_model_name is null or undefined + if (litellm_model_name) { + // Split litellm_model_name based on "/" + let splitModel = litellm_model_name.split("/"); + + // Get the first element in the split + let firstElement = splitModel[0]; + + // If there is only one element, default provider to openai + provider = custom_llm_provider; + if (!provider) { + provider = splitModel.length === 1 ? getProviderFromModel(litellm_model_name) : firstElement; + } + } else { + // litellm_model_name is null or undefined, default provider to openai + provider = "-"; + } + + if (model_info) { + input_cost = model_info?.input_cost_per_token; + output_cost = model_info?.output_cost_per_token; + max_tokens = model_info?.max_tokens; + max_input_tokens = model_info?.max_input_tokens; + } + + if (curr_model?.litellm_params) { + cleanedLitellmParams = Object.fromEntries( + Object.entries(curr_model?.litellm_params).filter(([key]) => key !== "model" && key !== "api_base"), + ); + } + + transformedData[i].provider = provider; + transformedData[i].input_cost = input_cost; + transformedData[i].output_cost = output_cost; + transformedData[i].litellm_model_name = litellm_model_name; + + // Convert Cost in terms of Cost per 1M tokens + if (transformedData[i].input_cost) { + transformedData[i].input_cost = (Number(transformedData[i].input_cost) * 1000000).toFixed(2); + } + + if (transformedData[i].output_cost) { + transformedData[i].output_cost = (Number(transformedData[i].output_cost) * 1000000).toFixed(2); + } + + transformedData[i].max_tokens = max_tokens; + transformedData[i].max_input_tokens = max_input_tokens; + transformedData[i].api_base = curr_model?.litellm_params?.api_base; + transformedData[i].cleanedLitellmParams = cleanedLitellmParams; + } + + return { data: transformedData }; +}; diff --git a/ui/litellm-dashboard/src/app/model_hub_table/page.tsx b/ui/litellm-dashboard/src/app/model_hub_table/page.tsx index fb83f28fc1b..1df7019ad25 100644 --- a/ui/litellm-dashboard/src/app/model_hub_table/page.tsx +++ b/ui/litellm-dashboard/src/app/model_hub_table/page.tsx @@ -1,12 +1,13 @@ "use client"; import React, { useEffect, useState } from "react"; import { useSearchParams } from "next/navigation"; -import ModelHubTable from "@/components/model_hub_table"; +import ModelHubTable from "@/components/AIHub/ModelHubTable"; export default function PublicModelHubTable() { const searchParams = useSearchParams()!; const key = searchParams.get("key"); const [accessToken, setAccessToken] = useState(null); + console.log("PublicModelHubTable accessToken:", accessToken); useEffect(() => { if (!key) { diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 6b94f514d91..598266df9a2 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -15,7 +15,7 @@ import GeneralSettings from "@/components/general_settings"; import GuardrailsPanel from "@/components/guardrails"; import { Team } from "@/components/key_team_helpers/key_list"; import { MCPServers } from "@/components/mcp_tools"; -import ModelHubTable from "@/components/model_hub_table"; +import ModelHubTable from "@/components/AIHub/ModelHubTable"; import Navbar from "@/components/navbar"; import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName } from "@/components/networking"; import NewUsagePage from "@/components/UsagePage/components/UsagePageView"; @@ -323,11 +323,8 @@ export default function CreateKeyPage() { /> ) : page == "models" ? ( void, copyToClipboard: (text: string) => void, publicPage: boolean = false, @@ -69,11 +69,7 @@ export const agentHubColumns = ( cell: ({ row }) => { const agent = row.original; - return ( - - {agent.description || "-"} - - ); + return {agent.description || "-"}; }, meta: { className: "hidden md:table-cell", @@ -105,11 +101,7 @@ export const agentHubColumns = ( cell: ({ row }) => { const agent = row.original; - return ( - - {agent.protocolVersion || "-"} - - ); + return {agent.protocolVersion || "-"}; }, meta: { className: "hidden lg:table-cell", @@ -135,9 +127,7 @@ export const agentHubColumns = ( {skill.name} ))} - {skills.length > 2 && ( - +{skills.length - 2} - )} + {skills.length > 2 && +{skills.length - 2}} )} @@ -240,4 +230,3 @@ export const agentHubColumns = ( return allColumns; }; - diff --git a/ui/litellm-dashboard/src/components/model_hub_table.test.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx similarity index 88% rename from ui/litellm-dashboard/src/components/model_hub_table.test.tsx rename to ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx index be0b9a113f4..a88ce0d7938 100644 --- a/ui/litellm-dashboard/src/components/model_hub_table.test.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.test.tsx @@ -1,7 +1,7 @@ import * as networking from "@/components/networking"; import { render, screen, waitFor } from "@testing-library/react"; import { afterEach, describe, expect, it, vi } from "vitest"; -import ModelHubTable from "./model_hub_table"; +import ModelHubTable from "./ModelHubTable"; vi.mock("@/components/networking", () => ({ getUiConfig: vi.fn(), @@ -19,7 +19,7 @@ vi.mock("next/navigation", () => ({ }), })); -vi.mock("./public_model_hub", () => ({ +vi.mock("@/components/public_model_hub", () => ({ default: () =>
Public Model Hub
, })); @@ -51,7 +51,12 @@ describe("ModelHubTable", () => { const getUiConfigMock = vi.mocked(networking.getUiConfig); const modelHubPublicModelsCallMock = vi.mocked(networking.modelHubPublicModelsCall); - getUiConfigMock.mockResolvedValue({ server_root_path: "/", proxy_base_url: "http://localhost:4000" }); + getUiConfigMock.mockResolvedValue({ + server_root_path: "/", + proxy_base_url: "http://localhost:4000", + auto_redirect_to_sso: false, + admin_ui_disabled: false, + }); modelHubPublicModelsCallMock.mockResolvedValue([]); render(); diff --git a/ui/litellm-dashboard/src/components/model_hub_table.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/model_hub_table.tsx rename to ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index f45e44ce905..7c538b618f1 100644 --- a/ui/litellm-dashboard/src/components/model_hub_table.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -1,20 +1,13 @@ -import { CopyOutlined } from "@ant-design/icons"; -import { Badge, Button, Card, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; -import { Modal } from "antd"; -import { Copy } from "lucide-react"; -import { useRouter } from "next/navigation"; -import React, { useCallback, useEffect, useState } from "react"; -import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; -import { isAdminRole } from "../utils/roles"; -import { agentHubColumns, AgentHubData } from "./agent_hub_table_columns"; -import MakeAgentPublicForm from "./make_agent_public_form"; -import MakeMCPPublicForm from "./make_mcp_public_form"; -import MakeModelPublicForm from "./make_model_public_form"; -import { mcpHubColumns, MCPServerData } from "./mcp_hub_table_columns"; -import { ModelDataTable } from "./model_dashboard/table"; -import ModelFilters from "./model_filters"; -import { modelHubColumns } from "./model_hub_table_columns"; -import NotificationsManager from "./molecules/notifications_manager"; +import { AgentHubData, getAgentHubTableColumns } from "@/components/AIHub/AgentHubTableColumns"; +import MakeAgentPublicForm from "@/components/AIHub/forms/MakeAgentPublicForm"; +import MakeMCPPublicForm from "@/components/AIHub/forms/MakeMCPPublicForm"; +import MakeModelPublicForm from "@/components/AIHub/forms/MakeModelPublicForm"; +import { mcpHubColumns, MCPServerData } from "@/components/mcp_hub_table_columns"; +import { modelHubColumns } from "@/components/model_hub_table_columns"; +import UsefulLinksManagement from "@/components/AIHub/UsefulLinksManagement"; +import { ModelDataTable } from "@/components/model_dashboard/table"; +import ModelFilters from "@/components/model_filters"; +import NotificationsManager from "@/components/molecules/notifications_manager"; import { fetchMCPServers, getAgentsList, @@ -23,9 +16,16 @@ import { getUiConfig, modelHubCall, modelHubPublicModelsCall, -} from "./networking"; -import PublicModelHub from "./public_model_hub"; -import UsefulLinksManagement from "./useful_links_management"; +} from "@/components/networking"; +import PublicModelHub from "@/components/public_model_hub"; +import { isAdminRole } from "@/utils/roles"; +import { CopyOutlined } from "@ant-design/icons"; +import { Badge, Button, Card, Tab, TabGroup, TabList, TabPanel, TabPanels, Text, Title } from "@tremor/react"; +import { Modal } from "antd"; +import { Copy } from "lucide-react"; +import { useRouter } from "next/navigation"; +import React, { useCallback, useEffect, useState } from "react"; +import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; interface ModelHubTableProps { accessToken: string | null; @@ -423,7 +423,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, {/* Agent Table */} ({ + getPublicModelHubInfo: vi.fn(), + updateUsefulLinksCall: vi.fn(), + getProxyBaseUrl: vi.fn(), +})); + +vi.mock("@/components/molecules/notifications_manager", () => ({ + __esModule: true, + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +const mockedGetPublicModelHubInfo = vi.mocked(getPublicModelHubInfo); +const mockedUpdateUsefulLinksCall = vi.mocked(updateUsefulLinksCall); +const mockedGetProxyBaseUrl = vi.mocked(getProxyBaseUrl); +const mockedNotifications = vi.mocked(NotificationsManager); + +describe("UsefulLinksManagement", () => { + beforeEach(() => { + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: {}, + }); + mockedUpdateUsefulLinksCall.mockResolvedValue({}); + mockedGetProxyBaseUrl.mockReturnValue("https://proxy.example.com"); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + it("should render link management for admin users", async () => { + render(); + + expect(await screen.findByText("Link Management")).toBeInTheDocument(); + await waitFor(() => expect(mockedGetPublicModelHubInfo).toHaveBeenCalled()); + }); + + it("should add a new link when fields are valid", async () => { + const user = userEvent.setup(); + render(); + + const displayNameInput = await screen.findByPlaceholderText("Friendly name"); + const urlInput = screen.getByPlaceholderText("https://example.com"); + + await user.type(displayNameInput, "Docs"); + await user.type(urlInput, "https://docs.example.com"); + await user.click(screen.getByRole("button", { name: /add link/i })); + + await waitFor(() => + expect(mockedUpdateUsefulLinksCall).toHaveBeenCalledWith("token", { + Docs: { url: "https://docs.example.com", index: 0 }, + }), + ); + + expect(await screen.findByText("Docs")).toBeInTheDocument(); + expect(screen.getByText("https://docs.example.com")).toBeInTheDocument(); + expect(mockedNotifications.success).toHaveBeenCalledWith("Link added successfully"); + }); + + it("should rearrange links and save the new order", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "First Link": "https://first.example.com", + "Second Link": "https://second.example.com", + "Third Link": "https://third.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("First Link")).toBeInTheDocument()); + + await user.click(screen.getByRole("button", { name: /rearrange order/i })); + + const secondLinkMoveUpButton = screen.getByTestId("move-up-1-Second Link"); + await user.click(secondLinkMoveUpButton); + + await user.click(screen.getByRole("button", { name: /save order/i })); + + await waitFor(() => + expect(mockedUpdateUsefulLinksCall).toHaveBeenCalledWith("token", { + "Second Link": { url: "https://second.example.com", index: 0 }, + "First Link": { url: "https://first.example.com", index: 1 }, + "Third Link": { url: "https://third.example.com", index: 2 }, + }), + ); + + expect(mockedNotifications.success).toHaveBeenCalledWith("Link order saved successfully"); + }); + + it("should display the Model Hub link", async () => { + render(); + + expect(await screen.findByRole("link", { name: /public model hub/i })).toBeInTheDocument(); + }); + + it("should edit a link when edit button is clicked", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Should show input fields in edit mode + expect(screen.getByDisplayValue("Test Link")).toBeInTheDocument(); + expect(screen.getByDisplayValue("https://test.example.com")).toBeInTheDocument(); + }); + + it("should update a link when save is clicked in edit mode", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Update the display name + const displayNameInput = screen.getByDisplayValue("Test Link"); + await user.clear(displayNameInput); + await user.type(displayNameInput, "Updated Link"); + + // Click save + await user.click(screen.getByRole("button", { name: /save/i })); + + await waitFor(() => + expect(mockedUpdateUsefulLinksCall).toHaveBeenCalledWith("token", { + "Updated Link": { url: "https://test.example.com", index: 0 }, + }), + ); + + expect(mockedNotifications.success).toHaveBeenCalledWith("Link updated successfully"); + }); + + it("should cancel editing when cancel button is clicked", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "Test Link": "https://test.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("Test Link")).toBeInTheDocument()); + + // Click edit button + const editButton = screen.getByTestId("edit-link-0-Test Link"); + await user.click(editButton); + + // Update the display name + const displayNameInput = screen.getByDisplayValue("Test Link"); + await user.clear(displayNameInput); + await user.type(displayNameInput, "Updated Link"); + + // Click cancel + await user.click(screen.getByRole("button", { name: /cancel/i })); + + // Should go back to normal view + expect(screen.getByText("Test Link")).toBeInTheDocument(); + expect(screen.queryByDisplayValue("Updated Link")).not.toBeInTheDocument(); + }); + + it("should not move down the last item in rearrange mode", async () => { + const user = userEvent.setup(); + mockedGetPublicModelHubInfo.mockResolvedValue({ + docs_title: "Docs", + custom_docs_description: null, + litellm_version: "1.0.0", + useful_links: { + "First Link": "https://first.example.com", + "Second Link": "https://second.example.com", + }, + }); + + render(); + + await waitFor(() => expect(screen.getByText("First Link")).toBeInTheDocument()); + + // Enter rearrange mode + await user.click(screen.getByRole("button", { name: /rearrange order/i })); + + // Try to move down the last item (should not do anything) + const secondLinkMoveDownButton = screen.getByTestId("move-down-1-Second Link"); + await user.click(secondLinkMoveDownButton); + + // Links should remain in same order + const linksAfter = screen.getAllByText(/First Link|Second Link/); + expect(linksAfter[0]).toHaveTextContent("First Link"); + expect(linksAfter[1]).toHaveTextContent("Second Link"); + }); + + it("should expand and collapse the component", async () => { + const user = userEvent.setup(); + render(); + + await waitFor(() => expect(screen.getByText("Link Management")).toBeInTheDocument()); + + // Initially expanded + expect(screen.getByText("Manage Existing Links")).toBeInTheDocument(); + + // Click to collapse + await user.click(screen.getByText("Link Management")); + + // Should be collapsed + expect(screen.queryByText("Manage Existing Links")).not.toBeInTheDocument(); + + // Click to expand again + await user.click(screen.getByText("Link Management")); + + // Should be expanded + expect(screen.getByText("Manage Existing Links")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/useful_links_management.tsx b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx similarity index 87% rename from ui/litellm-dashboard/src/components/useful_links_management.tsx rename to ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx index 19ef4605d87..c73eaf52384 100644 --- a/ui/litellm-dashboard/src/components/useful_links_management.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/UsefulLinksManagement.tsx @@ -1,11 +1,11 @@ -import React, { useState, useEffect } from "react"; -import { Modal } from "antd"; -import { PlusCircleIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; -import { isAdminRole } from "../utils/roles"; -import { getPublicModelHubInfo, updateUsefulLinksCall, getProxyBaseUrl } from "./networking"; -import { Card, Title, Text, Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell } from "@tremor/react"; -import NotificationsManager from "./molecules/notifications_manager"; -import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; +import TableIconActionButton from "@/components/common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { isAdminRole } from "@/utils/roles"; +import { ChevronDownIcon, ChevronRightIcon, ExternalLinkIcon, PlusCircleIcon } from "@heroicons/react/outline"; +import { Card, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, Text, Title } from "@tremor/react"; +import Link from "next/link"; +import React, { useEffect, useState } from "react"; +import { getProxyBaseUrl, getPublicModelHubInfo, updateUsefulLinksCall } from "../networking"; interface UsefulLinksManagementProps { accessToken: string | null; @@ -102,32 +102,6 @@ const UsefulLinksManagement: React.FC = ({ accessTok }); await updateUsefulLinksCall(accessToken, linksObject); - // show success modal with public model hub link - Modal.success({ - title: "Links Saved Successfully", - content: ( -
-

- Your useful links have been saved and are now visible on the public model hub. -

-
-

View your updated model hub:

- - Open Public Model Hub → - -
-
- ), - width: 500, - okText: "Close", - maskClosable: true, - keyboard: true, - }); return true; } catch (error) { @@ -319,29 +293,41 @@ const UsefulLinksManagement: React.FC = ({ accessTok
Manage Existing Links - {!isRearranging ? ( - - ) : ( -
+ Public Model Hub + + + {!isRearranging ? ( - -
- )} + ) : ( +
+ + +
+ )} +
diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.test.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.test.tsx new file mode 100644 index 00000000000..67c6d7d6cc9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.test.tsx @@ -0,0 +1,505 @@ +import { render, screen, fireEvent, act, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import MakeAgentPublicForm from "./MakeAgentPublicForm"; +import { AgentHubData } from "@/components/AIHub/AgentHubTableColumns"; + +// Mock the networking function +vi.mock("../../networking", () => ({ + makeAgentsPublicCall: vi.fn(), +})); + +// Import the mocked function +import { makeAgentsPublicCall } from "../../networking"; +const mockMakeAgentsPublicCall = vi.mocked(makeAgentsPublicCall); + +// Mock antd components +vi.mock("antd", () => ({ + Modal: ({ open, title, children, onCancel, footer }: any) => + open ? ( +
+
{title}
+ {children} + {footer} +
+ ) : null, + Form: Object.assign(({ children, form }: any) =>
{children}
, { + useForm: () => [ + { + resetFields: vi.fn(), + validateFields: vi.fn(), + getFieldsValue: vi.fn(), + setFieldsValue: vi.fn(), + }, + vi.fn(), + ], + Item: ({ children }: any) =>
{children}
, + }), + Steps: Object.assign( + ({ children, current, className }: any) => ( +
+ {children} +
+ ), + { + Step: ({ title }: any) =>
{title}
, + }, + ), + Button: ({ children, onClick, disabled, loading, ...props }: any) => ( + + ), + Checkbox: ({ checked, indeterminate, onChange, children, disabled }: any) => ( + + ), +})); + +// Mock @tremor/react components +vi.mock("@tremor/react", () => ({ + Text: ({ children, className }: any) => {children}, + Title: ({ children }: any) =>

{children}

, + Badge: ({ children, color, size }: any) => ( + + {children} + + ), +})); + +describe("MakeAgentPublicForm", () => { + const mockProps = { + visible: true, + onClose: vi.fn(), + accessToken: "test-token", + agentHubData: [ + { + agent_id: "agent-1", + name: "Test Agent 1", + description: "Description 1", + version: "1.0", + is_public: false, + skills: [ + { id: "skill-1", name: "Skill 1", description: "Skill desc" }, + { id: "skill-2", name: "Skill 2", description: "Skill desc" }, + ], + protocolVersion: "1.0", + }, + { + agent_id: "agent-2", + name: "Test Agent 2", + description: "Description 2", + version: "2.0", + is_public: true, + skills: [], + protocolVersion: "1.0", + }, + ] as AgentHubData[], + onSuccess: vi.fn(), + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.resetAllMocks(); + }); + + it("should render the component", () => { + render(); + + expect(screen.getByText("Make Agents Public")).toBeInTheDocument(); + expect(screen.getByText("Select Agents to Make Public")).toBeInTheDocument(); + }); + + it("should initialize with correct state", () => { + render(); + + // Check that the component renders with the correct title and content + expect(screen.getByText("Make Agents Public")).toBeInTheDocument(); + expect(screen.getByText("Select Agents to Make Public")).toBeInTheDocument(); + + // Check that all agent checkboxes are present + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(3); // Select all + 2 agents + + // Check that the Next button is enabled (agents are preselected) + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).not.toBeDisabled(); + }); + + it("should handle agent selection and navigation", async () => { + render(); + + // Initially on step 1 + expect(screen.getByText("Select Agents to Make Public")).toBeInTheDocument(); + + // Select all agents using the select all checkbox + const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // Verify Next button is enabled + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).not.toBeDisabled(); + + // Click Next + await act(async () => { + fireEvent.click(nextButton); + }); + + // Should move to step 2 + await waitFor(() => { + expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); + }); + }); + + it("should submit selected agents successfully", async () => { + mockMakeAgentsPublicCall.mockResolvedValueOnce({}); + + render(); + + // Select all agents + const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Wait for navigation to complete + await waitFor(() => { + expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + await waitFor(() => { + expect(mockMakeAgentsPublicCall).toHaveBeenCalledWith("test-token", ["agent-1", "agent-2"]); + expect(mockProps.onSuccess).toHaveBeenCalled(); + expect(mockProps.onClose).toHaveBeenCalled(); + }); + }); + + it("should handle select all functionality", async () => { + render(); + + const checkboxes = screen.getAllByRole("checkbox"); + const selectAllCheckbox = checkboxes[0]; + + // Select all + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // All checkboxes should be checked + checkboxes.forEach((checkbox) => { + expect(checkbox).toBeChecked(); + }); + + // Deselect all + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // All checkboxes should be unchecked except the indeterminate state + expect(checkboxes[0]).not.toBeChecked(); + expect(checkboxes[1]).not.toBeChecked(); + expect(checkboxes[2]).not.toBeChecked(); + }); + + it("should show error when no agents selected", async () => { + render(); + + // Deselect all agents first + const checkboxes = screen.getAllByRole("checkbox"); + await act(async () => { + fireEvent.click(checkboxes[0]); // Click select all to select all + fireEvent.click(checkboxes[0]); // Click select all again to deselect all + }); + + // Try to go to next step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Should stay on same step + expect(screen.getByText("Select Agents to Make Public")).toBeInTheDocument(); + }); + + it("should display empty state when no agents are available", () => { + const emptyProps = { + ...mockProps, + agentHubData: [] as AgentHubData[], + }; + + render(); + + expect(screen.getByText("No agents available.")).toBeInTheDocument(); + + // Select All checkbox should be disabled + const selectAllCheckbox = screen.getByLabelText("Select All"); + expect(selectAllCheckbox).toBeDisabled(); + + // Next button should be disabled + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).toBeDisabled(); + }); + + it("should handle Cancel button functionality", async () => { + render(); + + // Click Cancel button + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + await act(async () => { + fireEvent.click(cancelButton); + }); + + // Should call onClose + expect(mockProps.onClose).toHaveBeenCalled(); + }); + + it("should handle Previous button functionality", async () => { + render(); + + // Navigate to step 1 + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Verify we're on step 1 + await waitFor(() => { + expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); + }); + + // Click Previous button + const previousButton = screen.getByRole("button", { name: "Previous" }); + await act(async () => { + fireEvent.click(previousButton); + }); + + // Should go back to step 0 + expect(screen.getByText("Select Agents to Make Public")).toBeInTheDocument(); + }); + + it("should handle individual agent selection", async () => { + render(); + + // Get all checkboxes (select all + individual agents) + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(3); // Select all + 2 agents + + // Initially, agent-2 should be selected (it's already public) + const agent1Checkbox = checkboxes[1]; // First agent checkbox + const agent2Checkbox = checkboxes[2]; // Second agent checkbox + + expect(agent2Checkbox).toBeChecked(); // agent-2 is already public + + // Select agent-1 + await act(async () => { + fireEvent.click(agent1Checkbox); + }); + + expect(agent1Checkbox).toBeChecked(); + expect(agent2Checkbox).toBeChecked(); + + // Deselect agent-2 + await act(async () => { + fireEvent.click(agent2Checkbox); + }); + + expect(agent1Checkbox).toBeChecked(); + expect(agent2Checkbox).not.toBeChecked(); + + // Select all should be indeterminate now + const selectAllCheckbox = checkboxes[0]; + expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + }); + + it("should display skills overflow text when agent has more than 3 skills", () => { + const agentWithManySkills = { + ...mockProps.agentHubData[0], + skills: [ + { id: "skill-1", name: "Skill 1", description: "Skill desc" }, + { id: "skill-2", name: "Skill 2", description: "Skill desc" }, + { id: "skill-3", name: "Skill 3", description: "Skill desc" }, + { id: "skill-4", name: "Skill 4", description: "Skill desc" }, + { id: "skill-5", name: "Skill 5", description: "Skill desc" }, + ], + }; + + const propsWithManySkills = { + ...mockProps, + agentHubData: [agentWithManySkills], + }; + + render(); + + // Should show first 3 skills as badges + expect(screen.getByText("Skill 1")).toBeInTheDocument(); + expect(screen.getByText("Skill 2")).toBeInTheDocument(); + expect(screen.getByText("Skill 3")).toBeInTheDocument(); + + // Should show "+2 more" text for the remaining skills + expect(screen.getByText("+2 more")).toBeInTheDocument(); + }); + + it("should handle submit error properly", async () => { + const errorMessage = "Network error"; + mockMakeAgentsPublicCall.mockRejectedValueOnce(new Error(errorMessage)); + + render(); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + await waitFor(() => { + expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + // Should handle error and show error notification + await waitFor(() => { + expect(mockMakeAgentsPublicCall).toHaveBeenCalledWith("test-token", ["agent-2"]); + }); + + // Should not call onSuccess or onClose on error + expect(mockProps.onSuccess).not.toHaveBeenCalled(); + expect(mockProps.onClose).not.toHaveBeenCalled(); + }); + + it("should show loading state during submit", async () => { + let resolvePromise: (value: any) => void = () => {}; + const pendingPromise = new Promise((resolve) => { + resolvePromise = resolve; + }); + mockMakeAgentsPublicCall.mockReturnValueOnce(pendingPromise); + + render(); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + await waitFor(() => { + expect(screen.getByText("Confirm Making Agents Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + // Check loading state + expect(submitButton).toHaveAttribute("data-loading", "true"); + expect(submitButton).toBeDisabled(); + + // Resolve the promise + resolvePromise({}); + await waitFor(() => { + expect(mockProps.onSuccess).toHaveBeenCalled(); + expect(mockProps.onClose).toHaveBeenCalled(); + }); + }); + + it("should not render modal when visible is false", () => { + const invisibleProps = { + ...mockProps, + visible: false, + }; + + render(); + + // Modal should not be rendered + expect(screen.queryByTestId("modal")).not.toBeInTheDocument(); + expect(screen.queryByText("Make Agents Public")).not.toBeInTheDocument(); + }); + + it("should preselect already public agents when modal opens", () => { + // Test data where one agent is public and one is not + const mixedPublicProps = { + ...mockProps, + agentHubData: [ + { + agent_id: "agent-1", + name: "Test Agent 1", + description: "Description 1", + url: "http://example.com/agent1", + version: "1.0", + is_public: false, // Not public + skills: [], + protocolVersion: "1.0", + }, + { + agent_id: "agent-2", + name: "Test Agent 2", + description: "Description 2", + url: "http://example.com/agent2", + version: "2.0", + is_public: true, // Already public + skills: [], + protocolVersion: "1.0", + }, + { + agent_id: "agent-3", + name: "Test Agent 3", + description: "Description 3", + url: "http://example.com/agent3", + version: "3.0", + is_public: true, // Already public + skills: [], + protocolVersion: "1.0", + }, + ] as AgentHubData[], + }; + + render(); + + // Check that the correct checkboxes are selected + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(4); // Select all + 3 agents + + // agent-2 and agent-3 should be checked (they're already public) + const agent1Checkbox = checkboxes[1]; + const agent2Checkbox = checkboxes[2]; + const agent3Checkbox = checkboxes[3]; + + expect(agent1Checkbox).not.toBeChecked(); // agent-1 is not public + expect(agent2Checkbox).toBeChecked(); // agent-2 is public + expect(agent3Checkbox).toBeChecked(); // agent-3 is public + + // Select all should be indeterminate + const selectAllCheckbox = checkboxes[0]; + expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/make_agent_public_form.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.tsx similarity index 97% rename from ui/litellm-dashboard/src/components/make_agent_public_form.tsx rename to ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.tsx index 54548ddba07..a38950b8fb7 100644 --- a/ui/litellm-dashboard/src/components/make_agent_public_form.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeAgentPublicForm.tsx @@ -1,9 +1,9 @@ import React, { useState, useEffect } from "react"; import { Modal, Form, Steps, Button, Checkbox } from "antd"; import { Text, Title, Badge } from "@tremor/react"; -import { makeAgentsPublicCall } from "./networking"; -import NotificationsManager from "./molecules/notifications_manager"; -import { AgentHubData } from "./agent_hub_table_columns"; +import { makeAgentsPublicCall } from "../../networking"; +import NotificationsManager from "../../molecules/notifications_manager"; +import { AgentHubData } from "@/components/AIHub/AgentHubTableColumns"; const { Step } = Steps; diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.test.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.test.tsx new file mode 100644 index 00000000000..b0228e9e868 --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.test.tsx @@ -0,0 +1,562 @@ +import { render, screen, fireEvent, act, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import MakeMCPPublicForm from "./MakeMCPPublicForm"; +import { MCPServerData } from "../../mcp_hub_table_columns"; + +// Mock the networking function +vi.mock("../../networking", () => ({ + makeMCPPublicCall: vi.fn(), +})); + +// Import the mocked function +import { makeMCPPublicCall } from "../../networking"; +const mockMakeMCPPublicCall = vi.mocked(makeMCPPublicCall); + +// Mock antd components +vi.mock("antd", () => ({ + Modal: ({ open, title, children, onCancel, footer }: any) => + open ? ( +
+
{title}
+ {children} + {footer} +
+ ) : null, + Form: Object.assign(({ children, form }: any) =>
{children}
, { + useForm: () => [ + { + resetFields: vi.fn(), + validateFields: vi.fn(), + getFieldsValue: vi.fn(), + setFieldsValue: vi.fn(), + }, + vi.fn(), + ], + Item: ({ children }: any) =>
{children}
, + }), + Steps: Object.assign( + ({ children, current, className }: any) => ( +
+ {children} +
+ ), + { + Step: ({ title }: any) =>
{title}
, + }, + ), + Button: ({ children, onClick, disabled, loading, ...props }: any) => ( + + ), + Checkbox: ({ checked, indeterminate, onChange, children, disabled }: any) => ( + + ), +})); + +// Additional @tremor/react mocks (Button is already mocked globally) +vi.mock("@tremor/react", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + Text: ({ children, className }: any) => {children}, + Title: ({ children }: any) =>

{children}

, + Badge: ({ children, color, size }: any) => ( + + {children} + + ), + }; +}); + +describe("MakeMCPPublicForm", () => { + const mockProps = { + visible: true, + onClose: vi.fn(), + accessToken: "test-token", + mcpHubData: [ + { + server_id: "server-1", + server_name: "Test Server 1", + description: "Description 1", + url: "http://example.com/server1", + transport: "http", + status: "active", + mcp_info: { is_public: false }, + allowed_tools: ["tool-1", "tool-2"], + auth_type: "bearer", + credentials: {}, + created_at: "2024-01-01T00:00:00Z", + created_by: "user1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user1", + teams: [], + mcp_access_groups: [], + extra_headers: [], + static_headers: {}, + args: [], + env: {}, + }, + { + server_id: "server-2", + server_name: "Test Server 2", + description: "Description 2", + url: "http://example.com/server2", + transport: "websocket", + status: "inactive", + mcp_info: { is_public: true }, + allowed_tools: [], + auth_type: "none", + credentials: {}, + created_at: "2024-01-01T00:00:00Z", + created_by: "user2", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user2", + teams: [], + mcp_access_groups: [], + extra_headers: [], + static_headers: {}, + args: [], + env: {}, + }, + ] as MCPServerData[], + onSuccess: vi.fn(), + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.resetAllMocks(); + }); + + it("should render the component", () => { + render(); + + expect(screen.getByText("Make MCP Servers Public")).toBeInTheDocument(); + expect(screen.getByText("Select MCP Servers to Make Public")).toBeInTheDocument(); + }); + + it("should initialize with correct state", () => { + render(); + + // Check that the component renders with the correct title and content + expect(screen.getByText("Make MCP Servers Public")).toBeInTheDocument(); + expect(screen.getByText("Select MCP Servers to Make Public")).toBeInTheDocument(); + + // Check that all server checkboxes are present + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(3); // Select all + 2 servers + + // Check that the Next button is enabled (servers are preselected) + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).not.toBeDisabled(); + }); + + it("should handle server selection and navigation", async () => { + render(); + + // Initially on step 1 + expect(screen.getByText("Select MCP Servers to Make Public")).toBeInTheDocument(); + + // Select all servers using the select all checkbox + const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // Verify Next button is enabled + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).not.toBeDisabled(); + + // Click Next + await act(async () => { + fireEvent.click(nextButton); + }); + + // Should move to step 2 + await waitFor(() => { + expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); + }); + }); + + it("should submit selected servers successfully", async () => { + mockMakeMCPPublicCall.mockResolvedValueOnce({}); + + render(); + + // Select all servers + const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Wait for navigation to complete + await waitFor(() => { + expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + await waitFor(() => { + expect(mockMakeMCPPublicCall).toHaveBeenCalledWith("test-token", ["server-1", "server-2"]); + expect(mockProps.onSuccess).toHaveBeenCalled(); + expect(mockProps.onClose).toHaveBeenCalled(); + }); + }); + + it("should handle select all functionality", async () => { + render(); + + const checkboxes = screen.getAllByRole("checkbox"); + const selectAllCheckbox = checkboxes[0]; + + // Select all + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // All checkboxes should be checked + checkboxes.forEach((checkbox) => { + expect(checkbox).toBeChecked(); + }); + + // Deselect all + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // All checkboxes should be unchecked except the indeterminate state + expect(checkboxes[0]).not.toBeChecked(); + expect(checkboxes[1]).not.toBeChecked(); + expect(checkboxes[2]).not.toBeChecked(); + }); + + it("should show error when no servers selected", async () => { + render(); + + // Deselect all servers first + const checkboxes = screen.getAllByRole("checkbox"); + await act(async () => { + fireEvent.click(checkboxes[0]); // Click select all to select all + fireEvent.click(checkboxes[0]); // Click select all again to deselect all + }); + + // Try to go to next step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Should stay on same step + expect(screen.getByText("Select MCP Servers to Make Public")).toBeInTheDocument(); + }); + + it("should display empty state when no servers are available", () => { + const emptyProps = { + ...mockProps, + mcpHubData: [] as MCPServerData[], + }; + + render(); + + expect(screen.getByText("No MCP servers available.")).toBeInTheDocument(); + + // Select All checkbox should be disabled + const selectAllCheckbox = screen.getByLabelText("Select All"); + expect(selectAllCheckbox).toBeDisabled(); + + // Next button should be disabled + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).toBeDisabled(); + }); + + it("should handle Cancel button functionality", async () => { + render(); + + // Click Cancel button + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + await act(async () => { + fireEvent.click(cancelButton); + }); + + // Should call onClose + expect(mockProps.onClose).toHaveBeenCalled(); + }); + + it("should handle Previous button functionality", async () => { + render(); + + // Navigate to step 1 + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Verify we're on step 1 + await waitFor(() => { + expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); + }); + + // Click Previous button + const previousButton = screen.getByRole("button", { name: "Previous" }); + await act(async () => { + fireEvent.click(previousButton); + }); + + // Should go back to step 0 + expect(screen.getByText("Select MCP Servers to Make Public")).toBeInTheDocument(); + }); + + it("should handle individual server selection", async () => { + render(); + + // Get all checkboxes (select all + individual servers) + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(3); // Select all + 2 servers + + // Initially, server-2 should be selected (it's already public) + const server1Checkbox = checkboxes[1]; // First server checkbox + const server2Checkbox = checkboxes[2]; // Second server checkbox + + expect(server2Checkbox).toBeChecked(); // server-2 is already public + + // Select server-1 + await act(async () => { + fireEvent.click(server1Checkbox); + }); + + expect(server1Checkbox).toBeChecked(); + expect(server2Checkbox).toBeChecked(); + + // Deselect server-2 + await act(async () => { + fireEvent.click(server2Checkbox); + }); + + expect(server1Checkbox).toBeChecked(); + expect(server2Checkbox).not.toBeChecked(); + + // Select all should be indeterminate now + const selectAllCheckbox = checkboxes[0]; + expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + }); + + it("should display tools overflow text when server has more than 3 tools", () => { + const serverWithManyTools = { + ...mockProps.mcpHubData[0], + allowed_tools: ["tool-1", "tool-2", "tool-3", "tool-4", "tool-5"], + }; + + const propsWithManyTools = { + ...mockProps, + mcpHubData: [serverWithManyTools], + }; + + render(); + + // Should show first 3 tools as badges + expect(screen.getByText("tool-1")).toBeInTheDocument(); + expect(screen.getByText("tool-2")).toBeInTheDocument(); + expect(screen.getByText("tool-3")).toBeInTheDocument(); + + // Should show "+2 more" text for the remaining tools + expect(screen.getByText("+2 more")).toBeInTheDocument(); + }); + + it("should handle submit error properly", async () => { + const errorMessage = "Network error"; + mockMakeMCPPublicCall.mockRejectedValueOnce(new Error(errorMessage)); + + render(); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + await waitFor(() => { + expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + // Should handle error and show error notification + await waitFor(() => { + expect(mockMakeMCPPublicCall).toHaveBeenCalledWith("test-token", ["server-2"]); + }); + + // Should not call onSuccess or onClose on error + expect(mockProps.onSuccess).not.toHaveBeenCalled(); + expect(mockProps.onClose).not.toHaveBeenCalled(); + }); + + it("should show loading state during submit", async () => { + let resolvePromise: (value: any) => void = () => {}; + const pendingPromise = new Promise((resolve) => { + resolvePromise = resolve; + }); + mockMakeMCPPublicCall.mockReturnValueOnce(pendingPromise); + + render(); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + await waitFor(() => { + expect(screen.getByText("Confirm Making MCP Servers Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + // Check loading state + expect(submitButton).toHaveAttribute("data-loading", "true"); + expect(submitButton).toBeDisabled(); + + // Resolve the promise + resolvePromise({}); + await waitFor(() => { + expect(mockProps.onSuccess).toHaveBeenCalled(); + expect(mockProps.onClose).toHaveBeenCalled(); + }); + }); + + it("should not render modal when visible is false", () => { + const invisibleProps = { + ...mockProps, + visible: false, + }; + + render(); + + // Modal should not be rendered + expect(screen.queryByTestId("modal")).not.toBeInTheDocument(); + expect(screen.queryByText("Make MCP Servers Public")).not.toBeInTheDocument(); + }); + + it("should preselect already public servers when modal opens", () => { + // Test data where one server is public and one is not + const mixedPublicProps = { + ...mockProps, + mcpHubData: [ + { + server_id: "server-1", + server_name: "Test Server 1", + description: "Description 1", + url: "http://example.com/server1", + transport: "http", + status: "active", + mcp_info: { is_public: false }, // Not public + allowed_tools: [], + auth_type: "bearer", + credentials: {}, + created_at: "2024-01-01T00:00:00Z", + created_by: "user1", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user1", + teams: [], + mcp_access_groups: [], + extra_headers: [], + static_headers: {}, + args: [], + env: {}, + }, + { + server_id: "server-2", + server_name: "Test Server 2", + description: "Description 2", + url: "http://example.com/server2", + transport: "websocket", + status: "inactive", + mcp_info: { is_public: true }, // Already public + allowed_tools: [], + auth_type: "none", + credentials: {}, + created_at: "2024-01-01T00:00:00Z", + created_by: "user2", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user2", + teams: [], + mcp_access_groups: [], + extra_headers: [], + static_headers: {}, + args: [], + env: {}, + }, + { + server_id: "server-3", + server_name: "Test Server 3", + description: "Description 3", + url: "http://example.com/server3", + transport: "sse", + status: "healthy", + mcp_info: { is_public: true }, // Already public + allowed_tools: [], + auth_type: "oauth", + credentials: {}, + created_at: "2024-01-01T00:00:00Z", + created_by: "user3", + updated_at: "2024-01-01T00:00:00Z", + updated_by: "user3", + teams: [], + mcp_access_groups: [], + extra_headers: [], + static_headers: {}, + args: [], + env: {}, + }, + ] as MCPServerData[], + }; + + render(); + + // Check that the correct checkboxes are selected + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(4); // Select all + 3 servers + + // server-2 and server-3 should be checked (they're already public) + const server1Checkbox = checkboxes[1]; + const server2Checkbox = checkboxes[2]; + const server3Checkbox = checkboxes[3]; + + expect(server1Checkbox).not.toBeChecked(); // server-1 is not public + expect(server2Checkbox).toBeChecked(); // server-2 is public + expect(server3Checkbox).toBeChecked(); // server-3 is public + + // Select all should be indeterminate + const selectAllCheckbox = checkboxes[0]; + expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/make_mcp_public_form.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/make_mcp_public_form.tsx rename to ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.tsx index f7bba175800..d7103da9ed7 100644 --- a/ui/litellm-dashboard/src/components/make_mcp_public_form.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeMCPPublicForm.tsx @@ -1,9 +1,9 @@ import React, { useState, useEffect } from "react"; import { Modal, Form, Steps, Button, Checkbox } from "antd"; import { Text, Title, Badge } from "@tremor/react"; -import { makeMCPPublicCall } from "./networking"; -import NotificationsManager from "./molecules/notifications_manager"; -import { MCPServerData } from "./mcp_hub_table_columns"; +import { makeMCPPublicCall } from "../../networking"; +import NotificationsManager from "../../molecules/notifications_manager"; +import { MCPServerData } from "@/components/mcp_hub_table_columns"; const { Step } = Steps; diff --git a/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.test.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.test.tsx new file mode 100644 index 00000000000..2b57535f3ad --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.test.tsx @@ -0,0 +1,557 @@ +import { render, screen, fireEvent, act, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import MakeModelPublicForm from "./MakeModelPublicForm"; + +interface ModelGroupInfo { + model_group: string; + providers: string[]; + max_input_tokens?: number; + max_output_tokens?: number; + input_cost_per_token?: number; + output_cost_per_token?: number; + mode?: string; + tpm?: number; + rpm?: number; + supports_parallel_function_calling: boolean; + supports_vision: boolean; + supports_function_calling: boolean; + supported_openai_params?: string[]; + is_public_model_group: boolean; + [key: string]: any; +} + +// Mock the networking function +vi.mock("../../networking", () => ({ + makeModelGroupPublic: vi.fn(), +})); + +// Import the mocked function +import { makeModelGroupPublic } from "../../networking"; +const mockMakeModelGroupPublic = vi.mocked(makeModelGroupPublic); + +// Mock antd components +vi.mock("antd", () => ({ + Modal: ({ open, title, children, onCancel, footer }: any) => + open ? ( +
+
{title}
+ {children} + {footer} +
+ ) : null, + Form: Object.assign(({ children, form }: any) =>
{children}
, { + useForm: () => [ + { + resetFields: vi.fn(), + validateFields: vi.fn(), + getFieldsValue: vi.fn(), + setFieldsValue: vi.fn(), + }, + vi.fn(), + ], + Item: ({ children }: any) =>
{children}
, + }), + Steps: Object.assign( + ({ children, current, className }: any) => ( +
+ {children} +
+ ), + { + Step: ({ title }: any) =>
{title}
, + }, + ), + Button: ({ children, onClick, disabled, loading, ...props }: any) => ( + + ), + Checkbox: ({ checked, indeterminate, onChange, children, disabled }: any) => ( + + ), +})); + +// Mock @tremor/react components +vi.mock("@tremor/react", () => ({ + Text: ({ children, className }: any) => {children}, + Title: ({ children }: any) =>

{children}

, + Badge: ({ children, color, size }: any) => ( + + {children} + + ), +})); + +// Mock ModelFilters component +vi.mock("../../model_filters", () => ({ + default: ({ onFilteredDataChange, modelHubData }: any) => ( +
+ +
+ ), +})); + +// Mock NotificationsManager +vi.mock("../../molecules/notifications_manager", () => ({ + default: { + fromBackend: vi.fn(), + success: vi.fn(), + }, +})); + +describe("MakeModelPublicForm", () => { + const mockProps = { + visible: true, + onClose: vi.fn(), + accessToken: "test-token", + modelHubData: [ + { + model_group: "gpt-4", + providers: ["openai"], + max_input_tokens: 8192, + max_output_tokens: 4096, + input_cost_per_token: 0.03, + output_cost_per_token: 0.06, + mode: "chat", + tpm: 10000, + rpm: 200, + supports_parallel_function_calling: true, + supports_vision: false, + supports_function_calling: true, + supported_openai_params: ["temperature", "max_tokens"], + is_public_model_group: false, + }, + { + model_group: "gpt-3.5-turbo", + providers: ["openai"], + max_input_tokens: 4096, + max_output_tokens: 2048, + input_cost_per_token: 0.0015, + output_cost_per_token: 0.002, + mode: "chat", + tpm: 60000, + rpm: 3500, + supports_parallel_function_calling: false, + supports_vision: false, + supports_function_calling: true, + supported_openai_params: ["temperature", "max_tokens"], + is_public_model_group: true, + }, + ] as ModelGroupInfo[], + onSuccess: vi.fn(), + }; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.resetAllMocks(); + }); + + it("should render the component", () => { + render(); + + expect(screen.getByText("Make Models Public")).toBeInTheDocument(); + expect(screen.getByText("Select Models to Make Public")).toBeInTheDocument(); + }); + + it("should initialize with correct state", () => { + render(); + + // Check that the component renders with the correct title and content + expect(screen.getByText("Make Models Public")).toBeInTheDocument(); + expect(screen.getByText("Select Models to Make Public")).toBeInTheDocument(); + + // Check that all model checkboxes are present + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(3); // Select all + 2 models + + // Check that the Next button is enabled (models are preselected) + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).not.toBeDisabled(); + }); + + it("should handle model selection and navigation", async () => { + render(); + + // Initially on step 1 + expect(screen.getByText("Select Models to Make Public")).toBeInTheDocument(); + + // Select all models using the select all checkbox + const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // Verify Next button is enabled + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).not.toBeDisabled(); + + // Click Next + await act(async () => { + fireEvent.click(nextButton); + }); + + // Should move to step 2 + await waitFor(() => { + expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); + }); + }); + + it("should submit selected models successfully", async () => { + mockMakeModelGroupPublic.mockResolvedValueOnce({}); + + render(); + + // Select all models + const selectAllCheckbox = screen.getByLabelText("Select All (2)"); + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Wait for navigation to complete + await waitFor(() => { + expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + await waitFor(() => { + expect(mockMakeModelGroupPublic).toHaveBeenCalledWith("test-token", ["gpt-4", "gpt-3.5-turbo"]); + expect(mockProps.onSuccess).toHaveBeenCalled(); + expect(mockProps.onClose).toHaveBeenCalled(); + }); + }); + + it("should handle select all functionality", async () => { + render(); + + const checkboxes = screen.getAllByRole("checkbox"); + const selectAllCheckbox = checkboxes[0]; + + // Select all + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // All checkboxes should be checked + checkboxes.forEach((checkbox) => { + expect(checkbox).toBeChecked(); + }); + + // Deselect all + await act(async () => { + fireEvent.click(selectAllCheckbox); + }); + + // All checkboxes should be unchecked except the indeterminate state + expect(checkboxes[0]).not.toBeChecked(); + expect(checkboxes[1]).not.toBeChecked(); + expect(checkboxes[2]).not.toBeChecked(); + }); + + it("should show error when no models selected", async () => { + render(); + + // Deselect all models first + const checkboxes = screen.getAllByRole("checkbox"); + await act(async () => { + fireEvent.click(checkboxes[0]); // Click select all to select all + fireEvent.click(checkboxes[0]); // Click select all again to deselect all + }); + + // Try to go to next step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Should stay on same step + expect(screen.getByText("Select Models to Make Public")).toBeInTheDocument(); + }); + + it("should display empty state when no models are available", () => { + const emptyProps = { + ...mockProps, + modelHubData: [] as ModelGroupInfo[], + }; + + render(); + + expect(screen.getByText("No models match the current filters.")).toBeInTheDocument(); + + // Select All checkbox should be disabled + const selectAllCheckbox = screen.getByLabelText("Select All"); + expect(selectAllCheckbox).toBeDisabled(); + + // Next button should be disabled + const nextButton = screen.getByRole("button", { name: "Next" }); + expect(nextButton).toBeDisabled(); + }); + + it("should handle Cancel button functionality", async () => { + render(); + + // Click Cancel button + const cancelButton = screen.getByRole("button", { name: "Cancel" }); + await act(async () => { + fireEvent.click(cancelButton); + }); + + // Should call onClose + expect(mockProps.onClose).toHaveBeenCalled(); + }); + + it("should handle Previous button functionality", async () => { + render(); + + // Navigate to step 1 + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + // Verify we're on step 1 + await waitFor(() => { + expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); + }); + + // Click Previous button + const previousButton = screen.getByRole("button", { name: "Previous" }); + await act(async () => { + fireEvent.click(previousButton); + }); + + // Should go back to step 0 + expect(screen.getByText("Select Models to Make Public")).toBeInTheDocument(); + }); + + it("should handle individual model selection", async () => { + render(); + + // Get all checkboxes (select all + individual models) + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(3); // Select all + 2 models + + // Initially, gpt-3.5-turbo should be selected (it's already public) + const gpt4Checkbox = checkboxes[1]; // First model checkbox + const gpt35Checkbox = checkboxes[2]; // Second model checkbox + + expect(gpt35Checkbox).toBeChecked(); // gpt-3.5-turbo is already public + + // Select gpt-4 + await act(async () => { + fireEvent.click(gpt4Checkbox); + }); + + expect(gpt4Checkbox).toBeChecked(); + expect(gpt35Checkbox).toBeChecked(); + + // Deselect gpt-3.5-turbo + await act(async () => { + fireEvent.click(gpt35Checkbox); + }); + + expect(gpt4Checkbox).toBeChecked(); + expect(gpt35Checkbox).not.toBeChecked(); + + // Select all should be indeterminate now + const selectAllCheckbox = checkboxes[0]; + expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + }); + + it("should display model badges and information", () => { + render(); + + // Should show model names + expect(screen.getByText("gpt-4")).toBeInTheDocument(); + expect(screen.getByText("gpt-3.5-turbo")).toBeInTheDocument(); + + // Should show mode badges + expect(screen.getAllByText("chat")).toHaveLength(2); + + // Should show provider badges + expect(screen.getAllByText("openai")).toHaveLength(2); + }); + + it("should handle submit error properly", async () => { + const errorMessage = "Network error"; + mockMakeModelGroupPublic.mockRejectedValueOnce(new Error(errorMessage)); + + render(); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + await waitFor(() => { + expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + // Should handle error and show error notification + await waitFor(() => { + expect(mockMakeModelGroupPublic).toHaveBeenCalledWith("test-token", ["gpt-3.5-turbo"]); + }); + + // Should not call onSuccess or onClose on error + expect(mockProps.onSuccess).not.toHaveBeenCalled(); + expect(mockProps.onClose).not.toHaveBeenCalled(); + }); + + it("should show loading state during submit", async () => { + let resolvePromise: (value: any) => void = () => {}; + const pendingPromise = new Promise((resolve) => { + resolvePromise = resolve; + }); + mockMakeModelGroupPublic.mockReturnValueOnce(pendingPromise); + + render(); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + await waitFor(() => { + expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); + }); + + // Submit + const submitButton = screen.getByRole("button", { name: "Make Public" }); + await act(async () => { + fireEvent.click(submitButton); + }); + + // Check loading state + expect(submitButton).toHaveAttribute("data-loading", "true"); + expect(submitButton).toBeDisabled(); + + // Resolve the promise + resolvePromise({}); + await waitFor(() => { + expect(mockProps.onSuccess).toHaveBeenCalled(); + expect(mockProps.onClose).toHaveBeenCalled(); + }); + }); + + it("should not render modal when visible is false", () => { + const invisibleProps = { + ...mockProps, + visible: false, + }; + + render(); + + // Modal should not be rendered + expect(screen.queryByTestId("modal")).not.toBeInTheDocument(); + expect(screen.queryByText("Make Models Public")).not.toBeInTheDocument(); + }); + + it("should preselect already public models when modal opens", () => { + // Test data where one model is public and one is not + const mixedPublicProps = { + ...mockProps, + modelHubData: [ + { + model_group: "private-model", + providers: ["openai"], + is_public_model_group: false, + mode: "chat", + }, + { + model_group: "public-model", + providers: ["anthropic"], + is_public_model_group: true, + mode: "completion", + }, + { + model_group: "another-public-model", + providers: ["cohere"], + is_public_model_group: true, + mode: "chat", + }, + ] as ModelGroupInfo[], + }; + + render(); + + // Check that the correct checkboxes are selected + const checkboxes = screen.getAllByRole("checkbox"); + expect(checkboxes).toHaveLength(4); // Select all + 3 models + + // private-model should not be checked, public models should be checked + const privateModelCheckbox = checkboxes[1]; + const publicModelCheckbox = checkboxes[2]; + const anotherPublicModelCheckbox = checkboxes[3]; + + expect(privateModelCheckbox).not.toBeChecked(); // private-model is not public + expect(publicModelCheckbox).toBeChecked(); // public-model is public + expect(anotherPublicModelCheckbox).toBeChecked(); // another-public-model is public + + // Select all should be indeterminate + const selectAllCheckbox = checkboxes[0]; + expect(selectAllCheckbox).toHaveAttribute("data-indeterminate", "true"); + }); + + it("should show selected count", () => { + render(); + + // Should show that 1 model is selected (gpt-3.5-turbo is preselected) + expect(screen.getByText("1")).toBeInTheDocument(); + expect(screen.getByText("model selected")).toBeInTheDocument(); + }); + + it("should show confirmation step with selected models", async () => { + render(); + + // Navigate to confirm step + const nextButton = screen.getByRole("button", { name: "Next" }); + await act(async () => { + fireEvent.click(nextButton); + }); + + await waitFor(() => { + expect(screen.getByText("Confirm Making Models Public")).toBeInTheDocument(); + }); + + // Should show the selected model + expect(screen.getByText("gpt-3.5-turbo")).toBeInTheDocument(); + + // Should show the warning message + expect(screen.getByText(/Warning:/)).toBeInTheDocument(); + expect(screen.getByText(/model_hub_table/)).toBeInTheDocument(); + + // Should show total count (already verified by checking the presence of the confirmation step) + }); +}); diff --git a/ui/litellm-dashboard/src/components/make_model_public_form.tsx b/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.tsx similarity index 98% rename from ui/litellm-dashboard/src/components/make_model_public_form.tsx rename to ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.tsx index 750bdc24eeb..16ed04c1779 100644 --- a/ui/litellm-dashboard/src/components/make_model_public_form.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/forms/MakeModelPublicForm.tsx @@ -1,9 +1,9 @@ import React, { useState, useCallback, useEffect } from "react"; import { Modal, Form, Steps, Button, Checkbox } from "antd"; import { Text, Title, Badge } from "@tremor/react"; -import { makeModelGroupPublic } from "./networking"; -import ModelFilters from "./model_filters"; -import NotificationsManager from "./molecules/notifications_manager"; +import { makeModelGroupPublic } from "../../networking"; +import ModelFilters from "../../model_filters"; +import NotificationsManager from "../../molecules/notifications_manager"; const { Step } = Steps; diff --git a/ui/litellm-dashboard/src/components/SSOModals.test.tsx b/ui/litellm-dashboard/src/components/SSOModals.test.tsx index e9d2389b69e..365d23f4036 100644 --- a/ui/litellm-dashboard/src/components/SSOModals.test.tsx +++ b/ui/litellm-dashboard/src/components/SSOModals.test.tsx @@ -1,25 +1,21 @@ -import { render, fireEvent, waitFor } from "@testing-library/react"; -import { describe, expect, it, beforeAll } from "vitest"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; import { Form } from "antd"; +import { describe, expect, it, vi } from "vitest"; import SSOModals from "./SSOModals"; -import React from "react"; -// Mock window.matchMedia for Ant Design components -beforeAll(() => { - Object.defineProperty(window, "matchMedia", { - writable: true, - value: (query: string) => ({ - matches: false, - media: query, - onchange: null, - addListener: () => {}, // deprecated - removeListener: () => {}, // deprecated - addEventListener: () => {}, - removeEventListener: () => {}, - dispatchEvent: () => true, - }), - }); -}); +// Mock the networking functions +vi.mock("./networking", () => ({ + getSSOSettings: vi.fn(), + updateSSOSettings: vi.fn(), +})); + +// Mock parseErrorMessage +vi.mock("./shared/errorUtils", () => ({ + parseErrorMessage: vi.fn((error) => error?.message || "An error occurred"), +})); + +import NotificationsManager from "./molecules/notifications_manager"; +import { getSSOSettings, updateSSOSettings } from "./networking"; describe("SSOModals", () => { it("should render the SSOModals component", () => { @@ -42,11 +38,11 @@ describe("SSOModals", () => { ); }; - const { getByText } = render(); - expect(getByText("Add SSO")).toBeInTheDocument(); + render(); + expect(screen.getByText("Add SSO")).toBeInTheDocument(); }); - it("should have a validation error if the proxy base url is not a valid URL", async () => { + it("should show validation error if proxy base url is not a valid URL", async () => { const TestWrapper = () => { const [form] = Form.useForm(); return ( @@ -65,42 +61,40 @@ describe("SSOModals", () => { ); }; - const { getByLabelText, getByText, container } = render(); + render(); // Find and interact with the SSO provider select - const ssoProviderSelect = container.querySelector("#sso_provider"); - if (ssoProviderSelect) { - fireEvent.mouseDown(ssoProviderSelect); - // Wait for dropdown and select Google - await waitFor(() => { - const googleOption = getByText("Google SSO"); - fireEvent.click(googleOption); - }); - } + const ssoProviderSelect = screen.getByLabelText("SSO Provider"); + fireEvent.mouseDown(ssoProviderSelect); + // Wait for dropdown and select Google + await waitFor(() => { + const googleOption = screen.getByText("Google SSO"); + fireEvent.click(googleOption); + }); // Fill in the email field - const emailInput = getByLabelText("Proxy Admin Email"); + const emailInput = screen.getByLabelText("Proxy Admin Email"); fireEvent.change(emailInput, { target: { value: "test@example.com" } }); // Fill in an invalid URL - const urlInput = getByLabelText("Proxy Base URL"); + const urlInput = screen.getByLabelText("Proxy Base URL"); fireEvent.change(urlInput, { target: { value: "invalid-url" } }); // Submit the form - const saveButton = getByText("Save"); + const saveButton = screen.getByText("Save"); fireEvent.click(saveButton); // Check for validation error await waitFor( () => { - expect(getByText("URL must start with http:// or https://")).toBeInTheDocument(); + expect(screen.getByText("URL must start with http:// or https://")).toBeInTheDocument(); }, // The validation is based on a Promise, so we need to wait for it to resolve { timeout: 5000 }, ); }); - it("should show validation error if the proxy base url ends with a trailing slash", async () => { + it("should show validation error if proxy base url ends with trailing slash", async () => { const TestWrapper = () => { const [form] = Form.useForm(); return ( @@ -119,33 +113,31 @@ describe("SSOModals", () => { ); }; - const { getByLabelText, getByText, findByText, container } = render(); + render(); // Find and interact with the SSO provider select - const ssoProviderSelect = container.querySelector("#sso_provider"); - if (ssoProviderSelect) { - fireEvent.mouseDown(ssoProviderSelect); - // Wait for dropdown and select Google - await waitFor(() => { - const googleOption = getByText("Google SSO"); - fireEvent.click(googleOption); - }); - } + const ssoProviderSelect = screen.getByLabelText("SSO Provider"); + fireEvent.mouseDown(ssoProviderSelect); + // Wait for dropdown and select Google + await waitFor(() => { + const googleOption = screen.getByText("Google SSO"); + fireEvent.click(googleOption); + }); // Fill in the email field - const emailInput = getByLabelText("Proxy Admin Email"); + const emailInput = screen.getByLabelText("Proxy Admin Email"); fireEvent.change(emailInput, { target: { value: "test@example.com" } }); // Fill in a URL with trailing slash - const urlInput = getByLabelText("Proxy Base URL") as HTMLInputElement; + const urlInput = screen.getByLabelText("Proxy Base URL") as HTMLInputElement; fireEvent.change(urlInput, { target: { value: "https://example.com/" } }); // Submit the form - const saveButton = getByText("Save"); + const saveButton = screen.getByText("Save"); fireEvent.click(saveButton); // Check for validation error using findByText for async rendering - const errorMessage = await findByText("URL must not end with a trailing slash", {}, { timeout: 5000 }); + const errorMessage = await screen.findByText("URL must not end with a trailing slash", {}, { timeout: 5000 }); expect(errorMessage).toBeInTheDocument(); }); @@ -168,9 +160,9 @@ describe("SSOModals", () => { ); }; - const { getByLabelText } = render(); + render(); - const urlInput = getByLabelText("Proxy Base URL") as HTMLInputElement; + const urlInput = screen.getByLabelText("Proxy Base URL") as HTMLInputElement; // Simulate user typing "https://" fireEvent.change(urlInput, { target: { value: "h" } }); @@ -218,36 +210,266 @@ describe("SSOModals", () => { ); }; - const { getByLabelText, getByText, queryByText, container, findByText } = render(); + render(); // Find and interact with the SSO provider select - const ssoProviderSelect = container.querySelector("#sso_provider"); - if (ssoProviderSelect) { - fireEvent.mouseDown(ssoProviderSelect); - // Wait for dropdown and select Google - await waitFor(() => { - const googleOption = getByText("Google SSO"); - fireEvent.click(googleOption); - }); - } + const ssoProviderSelect = screen.getByLabelText("SSO Provider"); + fireEvent.mouseDown(ssoProviderSelect); + // Wait for dropdown and select Google + await waitFor(() => { + const googleOption = screen.getByText("Google SSO"); + fireEvent.click(googleOption); + }); // Fill in the email field - const emailInput = getByLabelText("Proxy Admin Email"); + const emailInput = screen.getByLabelText("Proxy Admin Email"); fireEvent.change(emailInput, { target: { value: "test@example.com" } }); // Fill in an incomplete URL like "http:" - const urlInput = getByLabelText("Proxy Base URL"); + const urlInput = screen.getByLabelText("Proxy Base URL"); fireEvent.change(urlInput, { target: { value: "http:" } }); // Submit the form - const saveButton = getByText("Save"); + const saveButton = screen.getByText("Save"); fireEvent.click(saveButton); // Check that only the URL format error appears (use findByText for async rendering) - const errorMessage = await findByText("URL must start with http:// or https://", {}, { timeout: 3000 }); + const errorMessage = await screen.findByText("URL must start with http:// or https://", {}, { timeout: 3000 }); expect(errorMessage).toBeInTheDocument(); // Verify the trailing slash error does NOT appear - expect(queryByText("URL must not end with a trailing slash")).not.toBeInTheDocument(); + expect(screen.queryByText("URL must not end with a trailing slash")).not.toBeInTheDocument(); + }); + + it("should load existing SSO settings when modal opens", async () => { + const mockSSOData = { + values: { + google_client_id: "test-client-id", + google_client_secret: "test-client-secret", + proxy_base_url: "https://example.com", + user_email: "admin@example.com", + role_mappings: { + group_claim: "groups", + default_role: "internal_user", + roles: { + proxy_admin: ["admin-group"], + proxy_admin_viewer: ["viewer-group"], + internal_user: ["user-group"], + internal_user_viewer: ["readonly-group"], + }, + }, + }, + }; + + (getSSOSettings as any).mockResolvedValue(mockSSOData); + + const TestWrapper = () => { + const [form] = Form.useForm(); + + return ( + {}} + handleAddSSOCancel={() => {}} + handleShowInstructions={() => {}} + handleInstructionsOk={() => {}} + handleInstructionsCancel={() => {}} + form={form} + accessToken="test-token" + ssoConfigured={false} + /> + ); + }; + + render(); + + // Wait for the useEffect to load data and populate form + await waitFor(() => { + expect(getSSOSettings).toHaveBeenCalledWith("test-token"); + }); + + // Check that form fields are populated with loaded data + await waitFor(() => { + const emailInput = screen.getByLabelText("Proxy Admin Email") as HTMLInputElement; + expect(emailInput.value).toBe("admin@example.com"); + }); + + const urlInput = screen.getByLabelText("Proxy Base URL") as HTMLInputElement; + expect(urlInput.value).toBe("https://example.com"); + + // Check that role mappings are populated + const groupClaimInput = screen.getByLabelText("Group Claim") as HTMLInputElement; + expect(groupClaimInput.value).toBe("groups"); + }); + + it("should submit form with role mappings enabled", async () => { + const mockHandleShowInstructions = vi.fn(); + (updateSSOSettings as any).mockResolvedValue({}); + // Mock getSSOSettings to return empty data so form starts clean + (getSSOSettings as any).mockResolvedValue({ values: {} }); + + let formInstance: any = null; + + const TestWrapper = () => { + const [form] = Form.useForm(); + formInstance = form; + + return ( + {}} + handleAddSSOCancel={() => {}} + handleShowInstructions={mockHandleShowInstructions} + handleInstructionsOk={() => {}} + handleInstructionsCancel={() => {}} + form={form} + accessToken="test-token" + ssoConfigured={false} + /> + ); + }; + + render(); + + // Wait for any initial loading to complete + await waitFor(() => { + expect(getSSOSettings).toHaveBeenCalledWith("test-token"); + }); + + // Set the provider directly using the form to trigger conditional rendering + formInstance.setFieldsValue({ sso_provider: "okta" }); + + // Wait for the "Use Role Mappings" checkbox to appear + await waitFor(() => { + expect(screen.getByLabelText("Use Role Mappings")).toBeInTheDocument(); + }); + + // Enable role mappings + const roleMappingsCheckbox = screen.getByLabelText("Use Role Mappings"); + fireEvent.click(roleMappingsCheckbox); + + // Fill required fields + const emailInput = screen.getByLabelText("Proxy Admin Email"); + fireEvent.change(emailInput, { target: { value: "admin@example.com" } }); + + const urlInput = screen.getByLabelText("Proxy Base URL"); + fireEvent.change(urlInput, { target: { value: "https://example.com" } }); + + // Fill Okta specific fields + const clientIdInput = screen.getByLabelText("Generic Client ID"); + fireEvent.change(clientIdInput, { target: { value: "test-client-id" } }); + + const clientSecretInput = screen.getByLabelText("Generic Client Secret"); + fireEvent.change(clientSecretInput, { target: { value: "test-client-secret" } }); + + const authEndpointInput = screen.getByLabelText("Authorization Endpoint"); + fireEvent.change(authEndpointInput, { target: { value: "https://example.okta.com/authorize" } }); + + const tokenEndpointInput = screen.getByLabelText("Token Endpoint"); + fireEvent.change(tokenEndpointInput, { target: { value: "https://example.okta.com/token" } }); + + const userinfoEndpointInput = screen.getByLabelText("Userinfo Endpoint"); + fireEvent.change(userinfoEndpointInput, { target: { value: "https://example.okta.com/userinfo" } }); + + // Fill role mapping fields + const groupClaimInput = screen.getByLabelText("Group Claim"); + fireEvent.change(groupClaimInput, { target: { value: "groups" } }); + + const proxyAdminTeamsInput = screen.getByLabelText("Proxy Admin Teams"); + fireEvent.change(proxyAdminTeamsInput, { target: { value: "admin-group, super-admin" } }); + + // Submit the form + const saveButton = screen.getByText("Save"); + fireEvent.click(saveButton); + + // Verify the API was called with correct payload including role mappings + await waitFor(() => { + expect(updateSSOSettings).toHaveBeenCalledWith("test-token", { + sso_provider: "okta", + user_email: "admin@example.com", + proxy_base_url: "https://example.com", + generic_client_id: "test-client-id", + generic_client_secret: "test-client-secret", + generic_authorization_endpoint: "https://example.okta.com/authorize", + generic_token_endpoint: "https://example.okta.com/token", + generic_userinfo_endpoint: "https://example.okta.com/userinfo", + role_mappings: { + provider: "generic", + group_claim: "groups", + default_role: "internal_user", + roles: { + proxy_admin: ["admin-group", "super-admin"], + proxy_admin_viewer: [], + internal_user: [], + internal_user_viewer: [], + }, + }, + }); + }); + + expect(mockHandleShowInstructions).toHaveBeenCalled(); + }); + + it("should show Clear button and clear SSO settings when configured", async () => { + const mockHandleAddSSOOk = vi.fn(); + (updateSSOSettings as any).mockResolvedValue({}); + (NotificationsManager.success as any).mockImplementation(() => {}); + + const TestWrapper = () => { + const [form] = Form.useForm(); + + return ( + {}} + handleShowInstructions={() => {}} + handleInstructionsOk={() => {}} + handleInstructionsCancel={() => {}} + form={form} + accessToken="test-token" + ssoConfigured={true} + /> + ); + }; + + render(); + + // Check that Clear button is visible when SSO is configured + const clearButton = screen.getByText("Clear"); + expect(clearButton).toBeInTheDocument(); + + // Click Clear button to open confirmation modal + fireEvent.click(clearButton); + + // Confirm the clear action in the modal + const confirmButton = screen.getByText("Yes, Clear"); + fireEvent.click(confirmButton); + + // Verify the clear API was called with null values + await waitFor(() => { + expect(updateSSOSettings).toHaveBeenCalledWith("test-token", { + google_client_id: null, + google_client_secret: null, + microsoft_client_id: null, + microsoft_client_secret: null, + microsoft_tenant: null, + generic_client_id: null, + generic_client_secret: null, + generic_authorization_endpoint: null, + generic_token_endpoint: null, + generic_userinfo_endpoint: null, + proxy_base_url: null, + user_email: null, + sso_provider: null, + role_mappings: null, + }); + }); + + expect(NotificationsManager.success).toHaveBeenCalledWith("SSO settings cleared successfully"); + expect(mockHandleAddSSOOk).toHaveBeenCalled(); }); }); diff --git a/ui/litellm-dashboard/src/components/SSOModals.tsx b/ui/litellm-dashboard/src/components/SSOModals.tsx index 26e33ace2d7..6cb57f41736 100644 --- a/ui/litellm-dashboard/src/components/SSOModals.tsx +++ b/ui/litellm-dashboard/src/components/SSOModals.tsx @@ -1,5 +1,5 @@ import React, { useEffect, useState } from "react"; -import { Modal, Form, Input, Button as Button2, Select } from "antd"; +import { Modal, Form, Input, Button as Button2, Select, Checkbox } from "antd"; import { Text, TextInput } from "@tremor/react"; import { getSSOSettings, updateSSOSettings } from "./networking"; import NotificationsManager from "./molecules/notifications_manager"; @@ -144,12 +144,35 @@ const SSOModals: React.FC = ({ } } + // Extract role mappings if they exist + let roleMappingFields = {}; + if (ssoData.values.role_mappings) { + const roleMappings = ssoData.values.role_mappings; + + // Helper function to join arrays into comma-separated strings + const joinTeams = (teams: string[] | undefined): string => { + if (!teams || teams.length === 0) return ""; + return teams.join(", "); + }; + + roleMappingFields = { + use_role_mappings: true, + group_claim: roleMappings.group_claim, + default_role: roleMappings.default_role || "internal_user", + proxy_admin_teams: joinTeams(roleMappings.roles?.proxy_admin), + admin_viewer_teams: joinTeams(roleMappings.roles?.proxy_admin_viewer), + internal_user_teams: joinTeams(roleMappings.roles?.internal_user), + internal_viewer_teams: joinTeams(roleMappings.roles?.internal_user_viewer), + }; + } + // Set form values with existing data (excluding UI access control fields) const formValues = { sso_provider: selectedProvider, proxy_base_url: ssoData.values.proxy_base_url, user_email: ssoData.values.user_email, ...ssoData.values, + ...roleMappingFields, }; console.log("Setting form values:", formValues); // Debug log @@ -178,8 +201,55 @@ const SSOModals: React.FC = ({ } try { + const { + proxy_admin_teams, + admin_viewer_teams, + internal_user_teams, + internal_viewer_teams, + default_role, + group_claim, + use_role_mappings, + ...rest + } = formValues; + + const payload: any = { + ...rest, + }; + + // Add role mappings if use_role_mappings is checked + if (use_role_mappings) { + // Helper function to split comma-separated string into array + const splitTeams = (teams: string | undefined): string[] => { + if (!teams || teams.trim() === "") return []; + return teams + .split(",") + .map((team) => team.trim()) + .filter((team) => team.length > 0); + }; + + // Map default role display values to backend values + const defaultRoleMapping: Record = { + internal_user_viewer: "internal_user_viewer", + internal_user: "internal_user", + proxy_admin_viewer: "proxy_admin_viewer", + proxy_admin: "proxy_admin", + }; + + payload.role_mappings = { + provider: "generic", + group_claim, + default_role: defaultRoleMapping[default_role] || "internal_user", + roles: { + proxy_admin: splitTeams(proxy_admin_teams), + proxy_admin_viewer: splitTeams(admin_viewer_teams), + internal_user: splitTeams(internal_user_teams), + internal_user_viewer: splitTeams(internal_viewer_teams), + }, + }; + } + // Save SSO settings using the new API - await updateSSOSettings(accessToken, formValues); + await updateSSOSettings(accessToken, payload); // Continue with the original flow (show instructions) handleShowInstructions(formValues); @@ -211,6 +281,7 @@ const SSOModals: React.FC = ({ proxy_base_url: null, user_email: null, sso_provider: null, + role_mappings: null, }; await updateSSOSettings(accessToken, clearSettings); @@ -334,6 +405,79 @@ const SSOModals: React.FC = ({ > + + prevValues.sso_provider !== currentValues.sso_provider} + > + {({ getFieldValue }) => { + const provider = getFieldValue("sso_provider"); + return provider === "okta" || provider === "generic" ? ( + + + + ) : null; + }} + + + + prevValues.use_role_mappings !== currentValues.use_role_mappings + } + > + {({ getFieldValue }) => { + const useRoleMappings = getFieldValue("use_role_mappings"); + return useRoleMappings ? ( + + + + ) : null; + }} + + + + prevValues.use_role_mappings !== currentValues.use_role_mappings + } + > + {({ getFieldValue }) => { + const useRoleMappings = getFieldValue("use_role_mappings"); + return useRoleMappings ? ( + <> + + + + + + + + + + + + + + + + + + + + + ) : null; + }} +
({ + updateSSOSettings: vi.fn(), +})); + +// Mock error utils +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn((error) => error?.message || "Unknown error"), +})); + +// Mock the useAuthorized hook +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + accessToken: "test-access-token", + userId: "test-user-id", + userEmail: "test@example.com", + userRole: "admin", + }), +})); + +// Mock NotificationsManager +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +describe("AddSSOSettingsModal", () => { + it("should render", () => { + const onCancel = vi.fn(); + const onSuccess = vi.fn(); + + renderWithProviders(); + + expect(screen.getByText("SSO Provider")).toBeInTheDocument(); + expect(screen.getByText("Cancel")).toBeInTheDocument(); + expect(screen.getAllByText("Add SSO")).toHaveLength(2); // Title and button + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx new file mode 100644 index 00000000000..7af6240b19e --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/AddSSOSettingsModal.tsx @@ -0,0 +1,63 @@ +"use client"; + +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { Button, Form, Modal, Space } from "antd"; +import React from "react"; +import BaseSSOSettingsForm from "./BaseSSOSettingsForm"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; +import { processSSOSettingsPayload } from "../utils"; + +interface AddSSOSettingsModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess: () => void; +} + +const AddSSOSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { + const [form] = Form.useForm(); + const { mutateAsync, isPending } = useEditSSOSettings(); + + // Enhanced form submission handler + const handleFormSubmit = async (formValues: Record) => { + const payload = processSSOSettingsPayload(formValues); + + await mutateAsync(payload, { + onSuccess: () => { + NotificationsManager.success("SSO settings added successfully"); + onSuccess(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); + }, + }); + }; + + const handleCancel = () => { + form.resetFields(); + onCancel(); + }; + + return ( + + + + + } + onCancel={handleCancel} + > + + + ); +}; + +export default AddSSOSettingsModal; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx new file mode 100644 index 00000000000..a4b36e5190e --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/BaseSSOSettingsForm.tsx @@ -0,0 +1,249 @@ +"use client"; + +import { TextInput } from "@tremor/react"; +import { Checkbox, Form, Input, Select } from "antd"; +import React from "react"; +import { ssoProviderLogoMap, ssoProviderDisplayNames } from "../constants"; + +export interface BaseSSOSettingsFormProps { + form: any; // Replace with proper Form type if available + onFormSubmit: (formValues: Record) => Promise; +} + +// Define the SSO provider configuration type +export interface SSOProviderConfig { + envVarMap: Record; + fields: Array<{ + label: string; + name: string; + placeholder?: string; + }>; +} + +// Define configurations for each SSO provider +export const ssoProviderConfigs: Record = { + google: { + envVarMap: { + google_client_id: "GOOGLE_CLIENT_ID", + google_client_secret: "GOOGLE_CLIENT_SECRET", + }, + fields: [ + { label: "Google Client ID", name: "google_client_id" }, + { label: "Google Client Secret", name: "google_client_secret" }, + ], + }, + microsoft: { + envVarMap: { + microsoft_client_id: "MICROSOFT_CLIENT_ID", + microsoft_client_secret: "MICROSOFT_CLIENT_SECRET", + microsoft_tenant: "MICROSOFT_TENANT", + }, + fields: [ + { label: "Microsoft Client ID", name: "microsoft_client_id" }, + { label: "Microsoft Client Secret", name: "microsoft_client_secret" }, + { label: "Microsoft Tenant", name: "microsoft_tenant" }, + ], + }, + okta: { + envVarMap: { + generic_client_id: "GENERIC_CLIENT_ID", + generic_client_secret: "GENERIC_CLIENT_SECRET", + generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", + generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", + generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", + }, + fields: [ + { label: "Generic Client ID", name: "generic_client_id" }, + { label: "Generic Client Secret", name: "generic_client_secret" }, + { + label: "Authorization Endpoint", + name: "generic_authorization_endpoint", + placeholder: "https://your-domain/authorize", + }, + { label: "Token Endpoint", name: "generic_token_endpoint", placeholder: "https://your-domain/token" }, + { + label: "Userinfo Endpoint", + name: "generic_userinfo_endpoint", + placeholder: "https://your-domain/userinfo", + }, + ], + }, + generic: { + envVarMap: { + generic_client_id: "GENERIC_CLIENT_ID", + generic_client_secret: "GENERIC_CLIENT_SECRET", + generic_authorization_endpoint: "GENERIC_AUTHORIZATION_ENDPOINT", + generic_token_endpoint: "GENERIC_TOKEN_ENDPOINT", + generic_userinfo_endpoint: "GENERIC_USERINFO_ENDPOINT", + }, + fields: [ + { label: "Generic Client ID", name: "generic_client_id" }, + { label: "Generic Client Secret", name: "generic_client_secret" }, + { label: "Authorization Endpoint", name: "generic_authorization_endpoint" }, + { label: "Token Endpoint", name: "generic_token_endpoint" }, + { label: "Userinfo Endpoint", name: "generic_userinfo_endpoint" }, + ], + }, +}; + +// Helper function to render provider fields +export const renderProviderFields = (provider: string) => { + const config = ssoProviderConfigs[provider]; + if (!config) return null; + + return config.fields.map((field) => ( + + {field.name.includes("client") ? : } + + )); +}; + +const BaseSSOSettingsForm: React.FC = ({ form, onFormSubmit }) => { + return ( +
+
+ + + + + prevValues.sso_provider !== currentValues.sso_provider} + > + {({ getFieldValue }) => { + const provider = getFieldValue("sso_provider"); + return provider ? renderProviderFields(provider) : null; + }} + + + + + + value?.trim()} + rules={[ + { required: true, message: "Please enter the proxy base url" }, + { + pattern: /^https?:\/\/.+/, + message: "URL must start with http:// or https://", + }, + { + validator: (_, value) => { + // Only check for trailing slash if the URL starts with http:// or https:// + if (value && /^https?:\/\/.+/.test(value) && value.endsWith("/")) { + return Promise.reject("URL must not end with a trailing slash"); + } + return Promise.resolve(); + }, + }, + ]} + > + + + + prevValues.sso_provider !== currentValues.sso_provider} + > + {({ getFieldValue }) => { + const provider = getFieldValue("sso_provider"); + return provider === "okta" || provider === "generic" ? ( + + + + ) : null; + }} + + + prevValues.use_role_mappings !== currentValues.use_role_mappings} + > + {({ getFieldValue }) => { + const useRoleMappings = getFieldValue("use_role_mappings"); + return useRoleMappings ? ( + + + + ) : null; + }} + + + prevValues.use_role_mappings !== currentValues.use_role_mappings} + > + {({ getFieldValue }) => { + const useRoleMappings = getFieldValue("use_role_mappings"); + return useRoleMappings ? ( + <> + + + + + + + + + + + + + + + + + + + + + ) : null; + }} + +
+
+ ); +}; + +export default BaseSSOSettingsForm; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx new file mode 100644 index 00000000000..ef6ec6c7055 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.test.tsx @@ -0,0 +1,20 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import DeleteSSOSettingsModal from "./DeleteSSOSettingsModal"; + +describe("DeleteSSOSettingsModal", () => { + it("should render", () => { + const onCancel = vi.fn(); + const onSuccess = vi.fn(); + + render( + , + ); + + expect(screen.getByText("Confirm Clear SSO Settings")).toBeInTheDocument(); + expect( + screen.getByText("Are you sure you want to clear all SSO settings? This action cannot be undone."), + ).toBeInTheDocument(); + expect(screen.getByText("Users will no longer be able to login using SSO after this change.")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.tsx new file mode 100644 index 00000000000..6a28b6490e0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/DeleteSSOSettingsModal.tsx @@ -0,0 +1,80 @@ +import { Modal } from "antd"; +import React from "react"; +import NotificationsManager from "../../../../molecules/notifications_manager"; +import { updateSSOSettings } from "../../../../networking"; +import { parseErrorMessage } from "../../../../shared/errorUtils"; + +interface DeleteSSOSettingsModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess: () => void; + accessToken: string | null; +} + +const DeleteSSOSettingsModal: React.FC = ({ + isVisible, + onCancel, + onSuccess, + accessToken, +}) => { + // Handle clearing SSO settings + const handleClearSSO = async () => { + if (!accessToken) { + NotificationsManager.fromBackend("No access token available"); + return; + } + + try { + // Clear all SSO settings + const clearSettings = { + google_client_id: null, + google_client_secret: null, + microsoft_client_id: null, + microsoft_client_secret: null, + microsoft_tenant: null, + generic_client_id: null, + generic_client_secret: null, + generic_authorization_endpoint: null, + generic_token_endpoint: null, + generic_userinfo_endpoint: null, + proxy_base_url: null, + user_email: null, + sso_provider: null, + }; + + await updateSSOSettings(accessToken, clearSettings); + + NotificationsManager.success("SSO settings cleared successfully"); + + // Close modal and trigger success callback + onCancel(); + onSuccess(); + } catch (error) { + console.error("Failed to clear SSO settings:", error); + NotificationsManager.fromBackend("Failed to clear SSO settings: " + parseErrorMessage(error)); + } + }; + + return ( + +

Are you sure you want to clear all SSO settings? This action cannot be undone.

+

Users will no longer be able to login using SSO after this change.

+
+ ); +}; + +export default DeleteSSOSettingsModal; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx new file mode 100644 index 00000000000..559d837b409 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.test.tsx @@ -0,0 +1,620 @@ +import { render, screen, fireEvent, waitFor } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach, Mock } from "vitest"; +import EditSSOSettingsModal from "./EditSSOSettingsModal"; +import { useSSOSettings } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { processSSOSettingsPayload } from "../utils"; + +// Constants +const SSO_PROVIDERS = { + GOOGLE: "google", + MICROSOFT: "microsoft", + OKTA: "okta", + AUTH0: "auth0", + GENERIC: "generic", +} as const; + +const TEST_DATA = { + MODAL_TITLE: "Edit SSO Settings", + MODAL_WIDTH: "800", + SUCCESS_MESSAGE: "SSO settings updated successfully", + ERROR_MESSAGE_PREFIX: "Failed to save SSO settings:", + BUTTON_TEXT: { + CANCEL: "Cancel", + SAVE: "Save", + SAVING: "Saving...", + }, +} as const; + +const TEST_IDS = { + MODAL: "modal", + BUTTON: "button", + BASE_SSO_FORM: "base-sso-form", + TRIGGER_FORM_SUBMIT: "trigger-form-submit", +} as const; + +// Mock form instance +const mockForm = { + resetFields: vi.fn(), + setFieldsValue: vi.fn(), + getFieldsValue: vi.fn(), + submit: vi.fn(), +}; + +// Types +type SSOData = { + values: Record; +} & Record; + +type SSOSettingsHookReturn = { + data: SSOData | null; + isLoading: boolean; + error: any; +}; + +type EditSSOSettingsHookReturn = { + mutateAsync: ReturnType; + isPending: boolean; +}; + +// Test data factories +const createSSOData = (overrides: Record = {}): SSOData => ({ + values: { + user_email: "test@example.com", + ...overrides, + }, +}); + +const createGoogleSSOData = (overrides: Record = {}) => + createSSOData({ + google_client_id: "test-google-id", + google_client_secret: "test-google-secret", + ...overrides, + }); + +const createMicrosoftSSOData = (overrides: Record = {}) => + createSSOData({ + microsoft_client_id: "test-microsoft-id", + microsoft_client_secret: "test-microsoft-secret", + microsoft_tenant: "test-tenant", + ...overrides, + }); + +const createGenericSSOData = (overrides: Record = {}) => + createSSOData({ + generic_client_id: "test-generic-id", + generic_client_secret: "test-generic-secret", + generic_authorization_endpoint: overrides.authorization_endpoint || "https://custom.example.com/oauth", + ...overrides, + }); + +const createRoleMappingsSSOData = (overrides: Record = {}) => + createGoogleSSOData({ + role_mappings: { + group_claim: "groups", + default_role: "internal_user", + roles: { + proxy_admin: overrides.proxy_admin || ["admin-group"], + proxy_admin_viewer: overrides.proxy_admin_viewer || ["viewer-group"], + internal_user: overrides.internal_user || ["user-group"], + internal_user_viewer: overrides.internal_user_viewer || ["readonly-group"], + }, + }, + ...overrides, + }); + +// Mock utilities +const createMockHooks = (): { + useSSOSettings: SSOSettingsHookReturn; + useEditSSOSettings: EditSSOSettingsHookReturn; +} => ({ + useSSOSettings: { + data: null, + isLoading: false, + error: null, + }, + useEditSSOSettings: { + mutateAsync: vi.fn(), + isPending: false, + }, +}); + +vi.mock("antd", () => ({ + Modal: ({ children, open, title, footer, onCancel, width, ...props }: any) => ( +
+
{children}
+
{footer}
+
+ ), + Button: ({ children, onClick, loading, disabled, ...props }: any) => ( + + ), + Form: { + useForm: () => [mockForm], + }, + Space: ({ children, ...props }: any) => ( +
+ {children} +
+ ), +})); + +vi.mock("./BaseSSOSettingsForm", () => ({ + default: ({ form, onFormSubmit }: any) => ( +
+ +
+ ), +})); + +vi.mock("@/app/(dashboard)/hooks/sso/useSSOSettings", () => ({ + useSSOSettings: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/sso/useEditSSOSettings", () => ({ + useEditSSOSettings: vi.fn(), +})); + +vi.mock("@/components/molecules/notifications_manager", () => ({ + default: { + success: vi.fn(), + fromBackend: vi.fn(), + }, +})); + +vi.mock("@/components/shared/errorUtils", () => ({ + parseErrorMessage: vi.fn(), +})); + +vi.mock("../utils", () => ({ + processSSOSettingsPayload: vi.fn(), +})); + +// Test helpers +const setupMocks = ( + overrides: Partial<{ + useSSOSettings: Partial; + useEditSSOSettings: Partial; + }> = {}, +) => { + const defaultMocks = createMockHooks(); + const mocks = { + useSSOSettings: { ...defaultMocks.useSSOSettings, ...overrides.useSSOSettings }, + useEditSSOSettings: { ...defaultMocks.useEditSSOSettings, ...overrides.useEditSSOSettings }, + }; + + (useSSOSettings as Mock).mockReturnValue(mocks.useSSOSettings); + (useEditSSOSettings as Mock).mockReturnValue(mocks.useEditSSOSettings); + + return mocks; +}; + +const renderComponent = (props: Partial> = {}) => { + const defaultProps = { + isVisible: true, + onCancel: vi.fn(), + onSuccess: vi.fn(), + }; + + return { + ...render(), + mockOnCancel: defaultProps.onCancel, + mockOnSuccess: defaultProps.onSuccess, + }; +}; + +const getButtons = () => screen.getAllByTestId(TEST_IDS.BUTTON); +const getCancelButton = () => getButtons()[0]; +const getSaveButton = () => getButtons()[1]; + +describe("EditSSOSettingsModal", () => { + beforeEach(() => { + vi.clearAllMocks(); + setupMocks(); + }); + + describe("Rendering", () => { + it("renders without crashing", () => { + expect(() => renderComponent()).not.toThrow(); + }); + + it("displays modal with correct configuration", () => { + renderComponent(); + + const modal = screen.getByTestId(TEST_IDS.MODAL); + expect(modal).toHaveAttribute("data-open", "true"); + expect(modal).toHaveAttribute("data-title", TEST_DATA.MODAL_TITLE); + expect(modal).toHaveAttribute("data-width", TEST_DATA.MODAL_WIDTH); + }); + + it("displays modal as closed when not visible", () => { + renderComponent({ isVisible: false }); + + const modal = screen.getByTestId(TEST_IDS.MODAL); + expect(modal).toHaveAttribute("data-open", "false"); + }); + }); + + describe("Footer Actions", () => { + it("renders cancel and save buttons", () => { + renderComponent(); + + const buttons = getButtons(); + expect(buttons).toHaveLength(2); + expect(buttons[0]).toHaveTextContent(TEST_DATA.BUTTON_TEXT.CANCEL); + expect(buttons[1]).toHaveTextContent(TEST_DATA.BUTTON_TEXT.SAVE); + }); + + it("calls onCancel and resets form when cancel button is clicked", () => { + const { mockOnCancel } = renderComponent(); + + fireEvent.click(getCancelButton()); + + expect(mockForm.resetFields).toHaveBeenCalled(); + expect(mockOnCancel).toHaveBeenCalled(); + }); + + it("calls form.submit when save button is clicked", () => { + renderComponent(); + + fireEvent.click(getSaveButton()); + + expect(mockForm.submit).toHaveBeenCalled(); + }); + + describe("Loading States", () => { + it("disables cancel button during submission", () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: true }, + }); + + renderComponent(); + + expect(getCancelButton()).toBeDisabled(); + }); + + it("shows loading state on save button during submission", () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: true }, + }); + + renderComponent(); + + expect(getSaveButton()).toHaveAttribute("data-loading", "true"); + expect(getSaveButton()).toHaveTextContent(TEST_DATA.BUTTON_TEXT.SAVING); + }); + }); + }); + + describe("Form Submission", () => { + const formValues = { testField: "testValue" }; + const processedPayload = { processed: "payload" }; + + beforeEach(() => { + (processSSOSettingsPayload as any).mockReturnValue(processedPayload); + }); + + it("processes form values and submits successfully", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onSuccess(); + return Promise.resolve({ success: true }); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + const { mockOnSuccess } = renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(processSSOSettingsPayload).toHaveBeenCalledWith(formValues); + expect(mockMutateAsync).toHaveBeenCalledWith( + processedPayload, + expect.objectContaining({ + onSuccess: expect.any(Function), + onError: expect.any(Function), + }), + ); + }); + + it("shows success notification and calls onSuccess callback", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onSuccess(); + return Promise.resolve({ success: true }); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + const { mockOnSuccess } = renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(NotificationsManager.success).toHaveBeenCalledWith(TEST_DATA.SUCCESS_MESSAGE); + expect(mockOnSuccess).toHaveBeenCalled(); + }); + + it("handles submission errors gracefully", async () => { + const error = new Error("Submission failed"); + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(error); + return Promise.reject(error); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (parseErrorMessage as any).mockReturnValue("Parsed error message"); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(parseErrorMessage).toHaveBeenCalledWith(error); + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith( + `${TEST_DATA.ERROR_MESSAGE_PREFIX} Parsed error message`, + ); + }); + }); + + describe("Form Initialization", () => { + describe("Provider Detection", () => { + const testProviderDetection = (testName: string, ssoData: SSOData, expectedProvider: string) => { + it(`detects ${testName} provider`, async () => { + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: expectedProvider, + ...ssoData.values, + }); + }); + }); + }; + + testProviderDetection("Google", createGoogleSSOData(), SSO_PROVIDERS.GOOGLE); + + testProviderDetection("Microsoft", createMicrosoftSSOData(), SSO_PROVIDERS.MICROSOFT); + + testProviderDetection( + "Okta", + createGenericSSOData({ + authorization_endpoint: "https://okta.example.com/oauth2/authorize", + }), + SSO_PROVIDERS.OKTA, + ); + + testProviderDetection( + "Auth0 (detected as Okta)", + createGenericSSOData({ + authorization_endpoint: "https://auth0.example.com/authorize", + }), + SSO_PROVIDERS.OKTA, // Auth0 URLs are detected as Okta provider + ); + + testProviderDetection("generic", createGenericSSOData(), SSO_PROVIDERS.GENERIC); + }); + + describe("Role Mappings", () => { + it("processes role mappings with all roles assigned", async () => { + const ssoData = createRoleMappingsSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "admin-group", + admin_viewer_teams: "viewer-group", + internal_user_teams: "user-group", + internal_viewer_teams: "readonly-group", + }); + }); + }); + + it("handles empty role mapping arrays", async () => { + const ssoData = createRoleMappingsSSOData({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user_viewer: [], + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "user-group", + internal_viewer_teams: "", + }); + }); + }); + }); + + describe("Initialization Guards", () => { + it("resets form before setting values", async () => { + const ssoData = createGoogleSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.resetFields).toHaveBeenCalled(); + expect(mockForm.setFieldsValue).toHaveBeenCalled(); + }); + }); + + it("skips initialization when modal is not visible", () => { + const ssoData = createGoogleSSOData(); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent({ isVisible: false }); + + expect(mockForm.setFieldsValue).not.toHaveBeenCalled(); + }); + + it("skips initialization when SSO data is unavailable", () => { + setupMocks({ + useSSOSettings: { data: null, isLoading: false, error: null }, + }); + + renderComponent(); + + expect(mockForm.setFieldsValue).not.toHaveBeenCalled(); + }); + }); + }); + + describe("Error Handling", () => { + it("handles form submission errors with undefined error message", async () => { + const error = new Error("Network error"); + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(error); + return Promise.reject(error); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (parseErrorMessage as any).mockReturnValue(undefined); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(NotificationsManager.fromBackend).toHaveBeenCalledWith(`${TEST_DATA.ERROR_MESSAGE_PREFIX} undefined`); + }); + + it("handles form submission with malformed data", async () => { + const mockMutateAsync = vi.fn().mockImplementation((payload, options) => { + options.onError(new Error("Invalid data")); + return Promise.reject(new Error("Invalid data")); + }); + + setupMocks({ + useEditSSOSettings: { mutateAsync: mockMutateAsync, isPending: false }, + }); + + (processSSOSettingsPayload as any).mockImplementation(() => { + throw new Error("Processing failed"); + }); + + renderComponent(); + + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + + expect(processSSOSettingsPayload).toHaveBeenCalled(); + expect(mockMutateAsync).not.toHaveBeenCalled(); + }); + }); + + describe("Edge Cases", () => { + it("handles role mappings with undefined roles object", async () => { + const ssoData = createGoogleSSOData({ + role_mappings: { + group_claim: "groups", + default_role: "internal_user", + // roles is undefined + }, + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GOOGLE, + ...ssoData.values, + use_role_mappings: true, + group_claim: "groups", + default_role: "internal_user", + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "", + internal_viewer_teams: "", + }); + }); + }); + + it("handles provider detection with partial SSO data", async () => { + const ssoData = createSSOData({ + // Only has generic fields, no specific provider identifiers + generic_client_id: "test-id", + generic_authorization_endpoint: "https://unknown.provider.com/auth", + }); + + setupMocks({ + useSSOSettings: { data: ssoData, isLoading: false, error: null }, + }); + + renderComponent(); + + await waitFor(() => { + expect(mockForm.setFieldsValue).toHaveBeenCalledWith({ + sso_provider: SSO_PROVIDERS.GENERIC, + ...ssoData.values, + }); + }); + }); + + it("handles form submission when processing throws error", async () => { + setupMocks({ + useEditSSOSettings: { mutateAsync: vi.fn(), isPending: false }, + }); + + (processSSOSettingsPayload as any).mockImplementation(() => { + throw new Error("Processing error"); + }); + + renderComponent(); + + expect(() => { + fireEvent.click(screen.getByTestId(TEST_IDS.TRIGGER_FORM_SUBMIT)); + }).not.toThrow(); + + expect(processSSOSettingsPayload).toHaveBeenCalled(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx new file mode 100644 index 00000000000..a731af68ff1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/Modals/EditSSOSettingsModal.tsx @@ -0,0 +1,136 @@ +"use client"; + +import { Button, Form, Modal, Space } from "antd"; +import React, { useEffect } from "react"; +import BaseSSOSettingsForm from "./BaseSSOSettingsForm"; +import NotificationsManager from "@/components/molecules/notifications_manager"; +import { parseErrorMessage } from "@/components/shared/errorUtils"; +import { processSSOSettingsPayload } from "../utils"; +import { useSSOSettings } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; +import { useEditSSOSettings } from "@/app/(dashboard)/hooks/sso/useEditSSOSettings"; + +interface EditSSOSettingsModalProps { + isVisible: boolean; + onCancel: () => void; + onSuccess: () => void; +} + +const EditSSOSettingsModal: React.FC = ({ isVisible, onCancel, onSuccess }) => { + const [form] = Form.useForm(); + + // Use react-query hooks for SSO settings + const ssoSettings = useSSOSettings(); + const { mutateAsync, isPending } = useEditSSOSettings(); + useEffect(() => { + if (isVisible && ssoSettings.data && ssoSettings.data.values) { + const ssoData = ssoSettings.data; + console.log("Raw SSO data received:", ssoData); // Debug log + console.log("SSO values:", ssoData.values); // Debug log + console.log("user_email from API:", ssoData.values.user_email); // Debug log + + // Determine which SSO provider is configured + let selectedProvider = null; + if (ssoData.values.google_client_id) { + selectedProvider = "google"; + } else if (ssoData.values.microsoft_client_id) { + selectedProvider = "microsoft"; + } else if (ssoData.values.generic_client_id) { + // Check if it looks like Okta based on endpoints + if ( + ssoData.values.generic_authorization_endpoint?.includes("okta") || + ssoData.values.generic_authorization_endpoint?.includes("auth0") + ) { + selectedProvider = "okta"; + } else { + selectedProvider = "generic"; + } + } + + // Extract role mappings if they exist + let roleMappingFields = {}; + if (ssoData.values.role_mappings) { + const roleMappings = ssoData.values.role_mappings; + + // Helper function to join arrays into comma-separated strings + const joinTeams = (teams: string[] | undefined): string => { + if (!teams || teams.length === 0) return ""; + return teams.join(", "); + }; + + roleMappingFields = { + use_role_mappings: true, + group_claim: roleMappings.group_claim, + default_role: roleMappings.default_role || "internal_user", + proxy_admin_teams: joinTeams(roleMappings.roles?.proxy_admin), + admin_viewer_teams: joinTeams(roleMappings.roles?.proxy_admin_viewer), + internal_user_teams: joinTeams(roleMappings.roles?.internal_user), + internal_viewer_teams: joinTeams(roleMappings.roles?.internal_user_viewer), + }; + } + + // Set form values with existing data (excluding UI access control fields) + const formValues = { + sso_provider: selectedProvider, + ...ssoData.values, + ...roleMappingFields, + }; + + console.log("Setting form values:", formValues); // Debug log + + // Clear form first, then set values with a small delay to ensure proper initialization + form.resetFields(); + setTimeout(() => { + form.setFieldsValue(formValues); + console.log("Form values set, current form values:", form.getFieldsValue()); // Debug log + }, 100); + } + }, [isVisible, ssoSettings.data, form]); + + // Enhanced form submission handler + const handleFormSubmit = async (formValues: Record) => { + try { + const payload = processSSOSettingsPayload(formValues); + + await mutateAsync(payload, { + onSuccess: () => { + NotificationsManager.success("SSO settings updated successfully"); + onSuccess(); + }, + onError: (error) => { + NotificationsManager.fromBackend("Failed to save SSO settings: " + parseErrorMessage(error)); + }, + }); + } catch (error) { + // Handle processing errors gracefully + NotificationsManager.fromBackend("Failed to process SSO settings: " + parseErrorMessage(error)); + } + }; + + const handleCancel = () => { + form.resetFields(); + onCancel(); + }; + + return ( + + + + + } + onCancel={handleCancel} + > + + + ); +}; + +export default EditSSOSettingsModal; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx new file mode 100644 index 00000000000..a047d7aea4f --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.test.tsx @@ -0,0 +1,108 @@ +import { render, screen, fireEvent } from "@testing-library/react"; +import { describe, expect, it } from "vitest"; +import RedactableField from "./RedactableField"; + +describe("RedactableField", () => { + describe("when value is null", () => { + it("should display 'Not configured' text", () => { + render(); + + expect(screen.getByText("Not configured")).toBeInTheDocument(); + }); + + it("should not display toggle button", () => { + render(); + + // There should be no button elements + const buttons = screen.queryAllByRole("button"); + expect(buttons).toHaveLength(0); + }); + }); + + describe("when value is provided", () => { + const testValue = "secret-password"; + + it("should be hidden by default and show redacted dots", () => { + render(); + + // Should show dots equal to the length of the value + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + }); + + it("should show actual value when defaultHidden is false", () => { + render(); + + expect(screen.getByText(testValue)).toBeInTheDocument(); + expect(screen.queryByText("•".repeat(testValue.length))).not.toBeInTheDocument(); + }); + + it("should display toggle button with eye icon when hidden", () => { + render(); + + const button = screen.getByRole("button"); + expect(button).toBeInTheDocument(); + + // Check that the Eye icon is rendered (we can check by title or by the presence of the icon) + // The button should contain the Eye icon when hidden + const eyeIcon = button.querySelector("svg"); + expect(eyeIcon).toBeInTheDocument(); + }); + + it("should display toggle button with eye-off icon when shown", () => { + render(); + + const button = screen.getByRole("button"); + expect(button).toBeInTheDocument(); + + // The button should contain the EyeOff icon when shown + const eyeOffIcon = button.querySelector("svg"); + expect(eyeOffIcon).toBeInTheDocument(); + }); + + it("should toggle visibility when button is clicked", () => { + render(); + + // Initially hidden + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + + // Click to show + const button = screen.getByRole("button"); + fireEvent.click(button); + + // Should now show the actual value + expect(screen.getByText(testValue)).toBeInTheDocument(); + expect(screen.queryByText("•".repeat(testValue.length))).not.toBeInTheDocument(); + + // Click again to hide + fireEvent.click(button); + + // Should be hidden again + expect(screen.getByText("•".repeat(testValue.length))).toBeInTheDocument(); + expect(screen.queryByText(testValue)).not.toBeInTheDocument(); + }); + + it("should handle empty string value", () => { + render(); + + // Empty string should show "Not configured" since value is falsy + expect(screen.getByText("Not configured")).toBeInTheDocument(); + + // No toggle button for empty string + const buttons = screen.queryAllByRole("button"); + expect(buttons).toHaveLength(0); + }); + + it("should handle different value lengths correctly", () => { + const shortValue = "hi"; + const longValue = "this-is-a-very-long-secret-value"; + + const { rerender } = render(); + expect(screen.getByText("••")).toBeInTheDocument(); + + rerender(); + expect(screen.getByText("•".repeat(longValue.length))).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx new file mode 100644 index 00000000000..44fef5cc7f8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/RedactableField.tsx @@ -0,0 +1,38 @@ +import { useState } from "react"; +import { Button } from "antd"; +import { Eye, EyeOff } from "lucide-react"; + +export default function RedactableField({ + defaultHidden = true, + value, +}: { + defaultHidden?: boolean; + value: string | null; +}) { + const [isHidden, setIsHidden] = useState(defaultHidden); + + return ( +
+ + {value ? ( + isHidden ? ( + "•".repeat(value.length) + ) : ( + value + ) + ) : ( + Not configured + )} + + {value && ( +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx new file mode 100644 index 00000000000..5e7908a872b --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.test.tsx @@ -0,0 +1,37 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import SSOSettings from "./SSOSettings"; + +// Mock the useSSOSettings hook +vi.mock("@/app/(dashboard)/hooks/sso/useSSOSettings", () => ({ + useSSOSettings: () => ({ + data: null, + refetch: vi.fn(), + }), +})); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +describe("SSOSettings", () => { + it("should render", () => { + const queryClient = createQueryClient(); + + render( + + + , + ); + + expect(screen.getByText("SSO Configuration")).toBeInTheDocument(); + expect(screen.getByText("Manage Single Sign-On authentication settings")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx new file mode 100644 index 00000000000..27ff96af05f --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettings.tsx @@ -0,0 +1,249 @@ +"use client"; + +import { useSSOSettings, type SSOSettingsValues } from "@/app/(dashboard)/hooks/sso/useSSOSettings"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { Button, Card, Descriptions, Space, Typography } from "antd"; +import { Edit, Shield, Trash2 } from "lucide-react"; +import { useState } from "react"; +import AddSSOSettingsModal from "./Modals/AddSSOSettingsModal"; +import DeleteSSOSettingsModal from "./Modals/DeleteSSOSettingsModal"; +import EditSSOSettingsModal from "./Modals/EditSSOSettingsModal"; +import RedactableField from "./RedactableField"; +import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; +import SSOSettingsLoadingSkeleton from "./SSOSettingsLoadingSkeleton"; +import { ssoProviderDisplayNames, ssoProviderLogoMap } from "./constants"; + +const { Title, Text } = Typography; + +export default function SSOSettings() { + const { data: ssoSettings, refetch, isLoading } = useSSOSettings(); + const { accessToken } = useAuthorized(); + const [isDeleteModalVisible, setIsDeleteModalVisible] = useState(false); + const [isAddModalVisible, setIsAddModalVisible] = useState(false); + const [isEditModalVisible, setIsEditModalVisible] = useState(false); + const isSSOConfigured = + Boolean(ssoSettings?.values.google_client_id) || + Boolean(ssoSettings?.values.microsoft_client_id) || + Boolean(ssoSettings?.values.generic_client_id); + + // Determine the SSO provider based on the configuration + const detectSSOProvider = (values: SSOSettingsValues): string | null => { + if (values.google_client_id) return "google"; + if (values.microsoft_client_id) return "microsoft"; + if (values.generic_client_id) { + // Check if it looks like Okta/Auth0 based on endpoints + if ( + values.generic_authorization_endpoint?.includes("okta") || + values.generic_authorization_endpoint?.includes("auth0") + ) { + return "okta"; + } + return "generic"; + } + return null; + }; + + const selectedProvider = ssoSettings?.values ? detectSSOProvider(ssoSettings.values) : null; + + const renderEndpointValue = (value?: string | null) => ( + + {value || "-"} + + ); + + const renderSimpleValue = (value?: string | null) => + value ? value : Not configured; + + const descriptionsConfig = { + column: { + xxl: 1, + xl: 1, + lg: 1, + md: 1, + sm: 1, + xs: 1, + }, + }; + + const providerConfigs = { + google: { + providerText: ssoProviderDisplayNames.google, + fields: [ + { + label: "Client ID", + render: (values: SSOSettingsValues) => , + }, + { + label: "Client Secret", + render: (values: SSOSettingsValues) => , + }, + { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, + ], + }, + microsoft: { + providerText: ssoProviderDisplayNames.microsoft, + fields: [ + { + label: "Client ID", + render: (values: SSOSettingsValues) => , + }, + { + label: "Client Secret", + render: (values: SSOSettingsValues) => , + }, + { label: "Tenant", render: (values: any) => renderSimpleValue(values.microsoft_tenant) }, + { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, + ], + }, + okta: { + providerText: ssoProviderDisplayNames.okta, + fields: [ + { + label: "Client ID", + render: (values: SSOSettingsValues) => , + }, + { + label: "Client Secret", + render: (values: SSOSettingsValues) => , + }, + { + label: "Authorization Endpoint", + render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_authorization_endpoint), + }, + { + label: "Token Endpoint", + render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_token_endpoint), + }, + { + label: "User Info Endpoint", + render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_userinfo_endpoint), + }, + { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, + ], + }, + generic: { + providerText: ssoProviderDisplayNames.generic, + fields: [ + { + label: "Client ID", + render: (values: SSOSettingsValues) => , + }, + { + label: "Client Secret", + render: (values: SSOSettingsValues) => , + }, + { + label: "Authorization Endpoint", + render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_authorization_endpoint), + }, + { + label: "Token Endpoint", + render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_token_endpoint), + }, + { + label: "User Info Endpoint", + render: (values: SSOSettingsValues) => renderEndpointValue(values.generic_userinfo_endpoint), + }, + { label: "Proxy Base URL", render: (values: SSOSettingsValues) => renderSimpleValue(values.proxy_base_url) }, + ], + }, + }; + + const renderSSOSettings = () => { + if (!ssoSettings?.values || !selectedProvider) return null; + + const { values } = ssoSettings; + const config = providerConfigs[selectedProvider as keyof typeof providerConfigs]; + + if (!config) return null; + + return ( + + +
+ {ssoProviderLogoMap[selectedProvider] && ( + {selectedProvider} + )} + {config.providerText} +
+
+ {config.fields.map((field, index) => ( + + {field.render(values)} + + ))} +
+ ); + }; + + return ( + <> + {isLoading ? ( + + ) : ( + + + {/* Header Section */} +
+
+ +
+ SSO Configuration + Manage Single Sign-On authentication settings +
+
+ +
+ {isSSOConfigured && ( + <> + + + + )} +
+
+ + {isSSOConfigured ? ( + renderSSOSettings() + ) : ( + setIsAddModalVisible(true)} /> + )} +
+
+ )} + + setIsDeleteModalVisible(false)} + onSuccess={() => refetch()} + accessToken={accessToken} + /> + + setIsAddModalVisible(false)} + onSuccess={() => { + setIsAddModalVisible(false); + refetch(); + }} + /> + + setIsEditModalVisible(false)} + onSuccess={() => { + setIsEditModalVisible(false); + refetch(); + }} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx new file mode 100644 index 00000000000..6676ba1c2c9 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.test.tsx @@ -0,0 +1,14 @@ +import { render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import SSOSettingsEmptyPlaceholder from "./SSOSettingsEmptyPlaceholder"; + +describe("SSOSettingsEmptyPlaceholder", () => { + it("should render", () => { + const onAdd = vi.fn(); + + render(); + + expect(screen.getByText("No SSO Configuration Found")).toBeInTheDocument(); + expect(screen.getByText("Configure SSO")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx new file mode 100644 index 00000000000..fc315493a54 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsEmptyPlaceholder.tsx @@ -0,0 +1,30 @@ +import { Empty, Typography, Button } from "antd"; + +const { Title, Paragraph } = Typography; + +interface SSOSettingsEmptyPlaceholderProps { + onAdd: () => void; +} + +export default function SSOSettingsEmptyPlaceholder({ onAdd }: SSOSettingsEmptyPlaceholderProps) { + return ( +
+ + No SSO Configuration Found + + Configure Single Sign-On (SSO) to enable seamless authentication for your team members using your identity + provider. + +
+ } + > + + +
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx new file mode 100644 index 00000000000..fd4fde69588 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.test.tsx @@ -0,0 +1,222 @@ +import { render, screen } from "@testing-library/react"; +import { describe, it, expect, vi } from "vitest"; +import SSOSettingsLoadingSkeleton from "./SSOSettingsLoadingSkeleton"; + +// Mock lucide-react icons +vi.mock("lucide-react", () => ({ + Shield: ({ className }: any) =>
, +})); + +// Mock Ant Design components +vi.mock("antd", () => ({ + Card: ({ children, ...props }: any) => ( +
+ {children} +
+ ), + Descriptions: Object.assign( + ({ children, bordered, column, ...props }: any) => ( +
+ {children} +
+ ), + { + Item: ({ children, label, ...props }: any) => ( +
+
{label}
+
{children}
+
+ ), + }, + ), + Typography: { + Title: ({ children, level, ...props }: any) => ( +
+ {children} +
+ ), + Text: ({ children, type, ...props }: any) => ( +
+ {children} +
+ ), + }, + Space: ({ children, direction, size, className, ...props }: any) => ( +
+ {children} +
+ ), + Skeleton: { + Button: ({ active, size, style, ...props }: any) => ( +
+ Button Skeleton +
+ ), + Node: ({ active, style, ...props }: any) => ( +
+ Node Skeleton +
+ ), + }, +})); + +describe("SSOSettingsLoadingSkeleton", () => { + it("should render without crashing", () => { + expect(() => render()).not.toThrow(); + }); + + it("should render Card component", () => { + render(); + expect(screen.getByTestId("card")).toBeInTheDocument(); + }); + + it("should render Space component with correct props", () => { + render(); + const space = screen.getByTestId("space"); + expect(space).toBeInTheDocument(); + expect(space).toHaveAttribute("data-direction", "vertical"); + expect(space).toHaveAttribute("data-size", "large"); + expect(space).toHaveClass("w-full"); + }); + + describe("Header Section", () => { + it("should render Shield icon", () => { + render(); + const shieldIcon = screen.getByTestId("shield-icon"); + expect(shieldIcon).toBeInTheDocument(); + expect(shieldIcon).toHaveClass("w-6 h-6 text-gray-400"); + }); + + it("should render title with correct text and level", () => { + render(); + const title = screen.getByTestId("typography-title"); + expect(title).toBeInTheDocument(); + expect(title).toHaveAttribute("data-level", "3"); + expect(title).toHaveTextContent("SSO Configuration"); + }); + + it("should render subtitle text", () => { + render(); + const text = screen.getByTestId("typography-text"); + expect(text).toBeInTheDocument(); + expect(text).toHaveAttribute("data-type", "secondary"); + expect(text).toHaveTextContent("Manage Single Sign-On authentication settings"); + }); + + it("should render two skeleton buttons with correct styles", () => { + render(); + const buttons = screen.getAllByTestId("skeleton-button"); + expect(buttons).toHaveLength(2); + + // First button + expect(buttons[0]).toHaveAttribute("data-active", "true"); + expect(buttons[0]).toHaveAttribute("data-size", "default"); + expect(buttons[0]).toHaveAttribute("data-style", JSON.stringify({ width: 170, height: 32 })); + + // Second button + expect(buttons[1]).toHaveAttribute("data-active", "true"); + expect(buttons[1]).toHaveAttribute("data-size", "default"); + expect(buttons[1]).toHaveAttribute("data-style", JSON.stringify({ width: 190, height: 32 })); + }); + }); + + describe("Descriptions Table", () => { + it("should render Descriptions component with bordered prop", () => { + render(); + const descriptions = screen.getByTestId("descriptions"); + expect(descriptions).toBeInTheDocument(); + expect(descriptions).toHaveAttribute("data-bordered", "true"); + }); + + it("should apply correct column configuration", () => { + render(); + const descriptions = screen.getByTestId("descriptions"); + const expectedColumn = { + xxl: 1, + xl: 1, + lg: 1, + md: 1, + sm: 1, + xs: 1, + }; + expect(descriptions).toHaveAttribute("data-column", JSON.stringify(expectedColumn)); + }); + + it("should render exactly 5 description items", () => { + render(); + const items = screen.getAllByTestId("descriptions-item"); + expect(items).toHaveLength(5); + }); + + describe("Description Items Structure", () => { + it("should render exactly 10 skeleton nodes total", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + expect(skeletonNodes).toHaveLength(10); + }); + + it("should render 5 skeleton nodes for labels with width 80", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + + const labelNodes = skeletonNodes.filter( + (node) => node.getAttribute("data-style") === JSON.stringify({ width: 80, height: 16 }), + ); + expect(labelNodes).toHaveLength(5); + + labelNodes.forEach((node) => { + expect(node).toHaveAttribute("data-active", "true"); + }); + }); + + it("should render skeleton nodes for content with correct widths", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + + // Expected content widths: [100, 200, 250, 180, 220] + const expectedWidths = [100, 200, 250, 180, 220]; + expectedWidths.forEach((width) => { + const contentNode = skeletonNodes.find( + (node) => node.getAttribute("data-style") === JSON.stringify({ width, height: 16 }), + ); + expect(contentNode).toBeInTheDocument(); + expect(contentNode).toHaveAttribute("data-active", "true"); + }); + }); + }); + }); + + describe("Accessibility and Structure", () => { + it("should have proper semantic structure", () => { + render(); + // Card contains Space + const card = screen.getByTestId("card"); + const space = screen.getByTestId("space"); + expect(card).toContainElement(space); + + // Space contains header section and descriptions + const descriptions = screen.getByTestId("descriptions"); + expect(space).toContainElement(descriptions); + }); + + it("should render all skeleton elements as active", () => { + render(); + const skeletonNodes = screen.getAllByTestId("skeleton-node"); + const skeletonButtons = screen.getAllByTestId("skeleton-button"); + + skeletonNodes.forEach((node) => { + expect(node).toHaveAttribute("data-active", "true"); + }); + + skeletonButtons.forEach((button) => { + expect(button).toHaveAttribute("data-active", "true"); + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx new file mode 100644 index 00000000000..59e34f255e3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/SSOSettingsLoadingSkeleton.tsx @@ -0,0 +1,66 @@ +"use client"; + +import { Card, Descriptions, Skeleton, Space, Typography } from "antd"; +import { Shield } from "lucide-react"; + +const { Title, Text } = Typography; +export default function SSOSettingsLoadingSkeleton() { + const descriptionsConfig = { + column: { + xxl: 1, + xl: 1, + lg: 1, + md: 1, + sm: 1, + xs: 1, + }, + }; + + return ( + + + {/* Header Section */} +
+
+ +
+ SSO Configuration + Manage Single Sign-On authentication settings +
+
+ +
+ + +
+
+ + {/* Descriptions Table Skeleton */} + + {/* Provider Row */} + }> +
+ +
+
+ + }> + + + + }> + + + + }> + + + + }> + + +
+
+
+ ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts new file mode 100644 index 00000000000..595a961401a --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/constants.ts @@ -0,0 +1,15 @@ +// SSO Provider logos +export const ssoProviderLogoMap: Record = { + google: "https://artificialanalysis.ai/img/logos/google_small.svg", + microsoft: "https://upload.wikimedia.org/wikipedia/commons/a/a8/Microsoft_Azure_Logo.svg", + okta: "https://www.okta.com/sites/default/files/Okta_Logo_BrightBlue_Medium.png", + generic: "", +}; + +// SSO Provider display names (consistent between select dropdown and table) +export const ssoProviderDisplayNames: Record = { + google: "Google SSO", + microsoft: "Microsoft SSO", + okta: "Okta / Auth0 SSO", + generic: "Generic SSO", +}; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts new file mode 100644 index 00000000000..1c878d7b7b3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.test.ts @@ -0,0 +1,274 @@ +import { processSSOSettingsPayload } from "./utils"; +import { describe, it, expect } from "vitest"; + +describe("processSSOSettingsPayload", () => { + describe("without role mappings", () => { + it("should return all fields except role mapping fields when use_role_mappings is false", () => { + const formValues = { + proxy_admin_teams: "team1, team2", + admin_viewer_teams: "viewer1", + internal_user_teams: "user1", + internal_viewer_teams: "viewer1", + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: false, + other_field: "value", + another_field: 123, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + other_field: "value", + another_field: 123, + }); + expect(result.role_mappings).toBeUndefined(); + }); + + it("should return all fields except role mapping fields when use_role_mappings is not present", () => { + const formValues = { + proxy_admin_teams: "team1", + admin_viewer_teams: "viewer1", + internal_user_teams: "user1", + internal_viewer_teams: "viewer1", + default_role: "proxy_admin", + group_claim: "groups", + other_field: "value", + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + other_field: "value", + }); + expect(result.role_mappings).toBeUndefined(); + }); + }); + + describe("with role mappings enabled", () => { + it("should create role mappings with all team types populated", () => { + const formValues = { + proxy_admin_teams: "admin1, admin2", + admin_viewer_teams: "viewer1, viewer2, viewer3", + internal_user_teams: "user1", + internal_viewer_teams: "internal_viewer1, internal_viewer2", + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: true, + other_field: "value", + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.other_field).toBe("value"); + expect(result.role_mappings).toEqual({ + provider: "generic", + group_claim: "groups", + default_role: "proxy_admin", + roles: { + proxy_admin: ["admin1", "admin2"], + proxy_admin_viewer: ["viewer1", "viewer2", "viewer3"], + internal_user: ["user1"], + internal_user_viewer: ["internal_viewer1", "internal_viewer2"], + }, + }); + }); + + it("should handle empty team strings", () => { + const formValues = { + proxy_admin_teams: "", + admin_viewer_teams: "", + internal_user_teams: "", + internal_viewer_teams: "", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: [], + internal_user_viewer: [], + }); + }); + + it("should handle undefined team fields", () => { + const formValues = { + default_role: "internal_user_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: [], + internal_user_viewer: [], + }); + }); + + it("should handle whitespace-only team strings", () => { + const formValues = { + proxy_admin_teams: " ", + admin_viewer_teams: ", , ,", + internal_user_teams: "user1, , user2", + internal_viewer_teams: "viewer1, ,viewer2", + default_role: "proxy_admin_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: [], + proxy_admin_viewer: [], + internal_user: ["user1", "user2"], + internal_user_viewer: ["viewer1", "viewer2"], + }); + }); + + it("should trim whitespace from team names", () => { + const formValues = { + proxy_admin_teams: " admin1 , admin2 ", + admin_viewer_teams: " viewer1 ", + internal_user_teams: " user1 , user2 ", + internal_viewer_teams: "viewer1,viewer2", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles).toEqual({ + proxy_admin: ["admin1", "admin2"], + proxy_admin_viewer: ["viewer1"], + internal_user: ["user1", "user2"], + internal_user_viewer: ["viewer1", "viewer2"], + }); + }); + + it("should filter out empty strings after trimming", () => { + const formValues = { + proxy_admin_teams: "admin1,,admin2, , admin3", + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.roles.proxy_admin).toEqual(["admin1", "admin2", "admin3"]); + }); + }); + + describe("default role mapping", () => { + it("should map internal_user_viewer correctly", () => { + const formValues = { + default_role: "internal_user_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user_viewer"); + }); + + it("should map internal_user correctly", () => { + const formValues = { + default_role: "internal_user", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + + it("should map proxy_admin_viewer correctly", () => { + const formValues = { + default_role: "proxy_admin_viewer", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("proxy_admin_viewer"); + }); + + it("should map proxy_admin correctly", () => { + const formValues = { + default_role: "proxy_admin", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("proxy_admin"); + }); + + it("should default to internal_user for unknown roles", () => { + const formValues = { + default_role: "unknown_role", + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + + it("should default to internal_user for undefined default_role", () => { + const formValues = { + group_claim: "groups", + use_role_mappings: true, + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result.role_mappings.default_role).toBe("internal_user"); + }); + }); + + describe("edge cases", () => { + it("should handle empty form values", () => { + const result = processSSOSettingsPayload({}); + + expect(result).toEqual({}); + }); + + it("should preserve other fields in the payload", () => { + const formValues = { + use_role_mappings: false, + sso_provider: "google", + client_id: "123", + client_secret: "secret", + redirect_url: "http://example.com", + custom_field: { nested: "value" }, + array_field: [1, 2, 3], + }; + + const result = processSSOSettingsPayload(formValues); + + expect(result).toEqual({ + sso_provider: "google", + client_id: "123", + client_secret: "secret", + redirect_url: "http://example.com", + custom_field: { nested: "value" }, + array_field: [1, 2, 3], + }); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts new file mode 100644 index 00000000000..3533e1226c2 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/SSOSettings/utils.ts @@ -0,0 +1,54 @@ +/** + * Processes SSO settings form values and transforms them into the payload format expected by the API + * Handles role mappings transformation and field extraction + */ +export const processSSOSettingsPayload = (formValues: Record): Record => { + const { + proxy_admin_teams, + admin_viewer_teams, + internal_user_teams, + internal_viewer_teams, + default_role, + group_claim, + use_role_mappings, + ...rest + } = formValues; + + const payload: any = { + ...rest, + }; + + // Add role mappings if use_role_mappings is checked + if (use_role_mappings) { + // Helper function to split comma-separated string into array + const splitTeams = (teams: string | undefined): string[] => { + if (!teams || teams.trim() === "") return []; + return teams + .split(",") + .map((team) => team.trim()) + .filter((team) => team.length > 0); + }; + + // Map default role display values to backend values + const defaultRoleMapping: Record = { + internal_user_viewer: "internal_user_viewer", + internal_user: "internal_user", + proxy_admin_viewer: "proxy_admin_viewer", + proxy_admin: "proxy_admin", + }; + + payload.role_mappings = { + provider: "generic", + group_claim, + default_role: defaultRoleMapping[default_role] || "internal_user", + roles: { + proxy_admin: splitTeams(proxy_admin_teams), + proxy_admin_viewer: splitTeams(admin_viewer_teams), + internal_user: splitTeams(internal_user_teams), + internal_user_viewer: splitTeams(internal_viewer_teams), + }, + }; + } + + return payload; +}; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 7680232ef0b..c9078383489 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -8,7 +8,7 @@ import { Alert, Card, Skeleton, Space, Switch, Typography } from "antd"; export default function UISettings() { const { accessToken } = useAuthorized(); - const { data, isLoading, isError, error } = useUISettings(accessToken); + const { data, isLoading, isError, error } = useUISettings(); const { mutate: updateSettings, isPending: isUpdating, error: updateError } = useUpdateUISettings(accessToken); const schema = data?.field_schema; diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx index 8f6bc411630..db268286007 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/EntityUsage/TopKeyView.tsx @@ -268,16 +268,7 @@ const TopKeyView: React.FC = ({ topKeys, teams, showTags = fals {/* Content */}
- +
diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx index d72515b8e4f..09415e87443 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.test.tsx @@ -1,8 +1,9 @@ import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; import { useCustomers } from "@/app/(dashboard)/hooks/customers/useCustomers"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; -import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { act, fireEvent, screen, waitFor } from "@testing-library/react"; import { beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { renderWithProviders } from "../../../../tests/test-utils"; import type { Organization } from "../../networking"; import * as networking from "../../networking"; import NewUsagePage from "./UsagePageView"; @@ -332,7 +333,7 @@ describe("NewUsage", () => { }); it("should render and fetch usage data on mount", async () => { - render(); + renderWithProviders(); // Wait for data to be fetched await waitFor(() => { @@ -349,7 +350,7 @@ describe("NewUsage", () => { }); it("should display usage metrics and charts", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); @@ -367,7 +368,7 @@ describe("NewUsage", () => { }); it("should switch between usage views correctly", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); @@ -401,7 +402,7 @@ describe("NewUsage", () => { }); it("should show organization usage banner and view for admins", async () => { - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); @@ -426,7 +427,7 @@ describe("NewUsage", () => { error: null, } as any); - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); @@ -450,7 +451,7 @@ describe("NewUsage", () => { error: null, } as any); - render(); + renderWithProviders(); await waitFor(() => { expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalled(); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx index 920955138b9..a9c86c064ef 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/components/UsagePage/components/UsagePageView.tsx @@ -27,7 +27,7 @@ import { Text, Title, } from "@tremor/react"; -import { Alert, Badge } from "antd"; +import { Alert } from "antd"; import React, { useCallback, useEffect, useMemo, useState } from "react"; import { useAgents } from "@/app/(dashboard)/hooks/agents/useAgents"; @@ -52,6 +52,7 @@ import { valueFormatterSpend } from "../utils/value_formatters"; import EntityUsage, { EntityList } from "./EntityUsage/EntityUsage"; import TopKeyView from "./EntityUsage/TopKeyView"; import { UsageOption, UsageViewSelect } from "./UsageViewSelect/UsageViewSelect"; +import { useCurrentUser } from "@/app/(dashboard)/hooks/users/useCurrentUser"; interface UsagePageProps { teams: Team[]; @@ -82,6 +83,9 @@ const UsagePage: React.FC = ({ teams, organizations }) => { const [allTags, setAllTags] = useState([]); const { data: customers = [] } = useCustomers(); const { data: agentsResponse } = useAgents(); + const { data: currentUser } = useCurrentUser(); + console.log(`currentUser: ${JSON.stringify(currentUser)}`); + console.log(`currentUser max budget: ${currentUser?.max_budget}`); const [modelViewType, setModelViewType] = useState<"groups" | "individual">("groups"); const [isCloudZeroModalOpen, setIsCloudZeroModalOpen] = useState(false); const [isGlobalExportModalOpen, setIsGlobalExportModalOpen] = useState(false); @@ -419,13 +423,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => {
- - setUsageView(value)} - isAdmin={all_admin_roles.includes(userRole || "")} - /> - + setUsageView(value)} + isAdmin={all_admin_roles.includes(userRole || "")} + />
{/* Your Usage Panel */} @@ -479,7 +481,11 @@ const UsagePage: React.FC = ({ teams, organizations }) => { )} - + diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx new file mode 100644 index 00000000000..cbd3d2c7320 --- /dev/null +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -0,0 +1,396 @@ +import { screen, waitFor, fireEvent } from "@testing-library/react"; +import { vi, it, expect, beforeEach, MockedFunction } from "vitest"; +import { renderWithProviders } from "../../../tests/test-utils"; +import { VirtualKeysTable } from "./VirtualKeysTable"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; +import { Organization } from "../networking"; +import { KeysResponse, useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; +import { useFilterLogic } from "../key_team_helpers/filter_logic"; + +// Mock network calls +vi.mock("./networking", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + userListCall: vi.fn().mockResolvedValue({ + users: [ + { + user_id: "user-1", + user_email: "user@example.com", + user_role: "user", + }, + ], + }), + }; +}); + +// Mock filter helpers +vi.mock("./key_team_helpers/filter_helpers", () => ({ + fetchAllKeyAliases: vi.fn().mockResolvedValue(["test-key-alias"]), + fetchAllTeams: vi.fn().mockResolvedValue([ + { + team_id: "team-1", + team_alias: "Test Team", + }, + ]), + fetchAllOrganizations: vi.fn().mockResolvedValue([ + { + organization_id: "org-1", + organization_alias: "Test Organization", + }, + ]), +})); + +// Mock useKeys hook +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ + useKeys: vi.fn(), +})); + +// Mock useFilterLogic hook +vi.mock("../key_team_helpers/filter_logic", () => ({ + useFilterLogic: vi.fn(), +})); + +const mockKey: KeyResponse = { + token: "sk-1234567890abcdef", + token_id: "key-1", + key_name: "test-key", + key_alias: "Test Key Alias", + spend: 5.5, + max_budget: 100, + expires: "2024-12-31T23:59:59Z", + models: ["gpt-3.5-turbo", "gpt-4"], + aliases: {}, + config: {}, + user_id: "user-1", + team_id: "team-1", + max_parallel_requests: 10, + metadata: {}, + tpm_limit: 1000, + rpm_limit: 100, + duration: "30d", + budget_duration: "1m", + budget_reset_at: "2024-12-01T00:00:00Z", + allowed_cache_controls: [], + allowed_routes: [], + permissions: {}, + model_spend: { "gpt-3.5-turbo": 2.5, "gpt-4": 3.0 }, + model_max_budget: { "gpt-3.5-turbo": 50, "gpt-4": 50 }, + soft_budget_cooldown: false, + blocked: false, + litellm_budget_table: {}, + organization_id: "org-1", + created_at: "2024-11-01T10:00:00Z", + updated_at: "2024-11-15T10:00:00Z", + team_spend: 5.5, + team_alias: "Test Team", + team_tpm_limit: 5000, + team_rpm_limit: 500, + team_max_budget: 500, + team_models: ["gpt-3.5-turbo", "gpt-4"], + team_blocked: false, + soft_budget: 50, + team_model_aliases: {}, + team_member_spend: 0, + team_metadata: {}, + end_user_id: "end-user-1", + end_user_tpm_limit: 100, + end_user_rpm_limit: 10, + end_user_max_budget: 10, + last_refreshed_at: Date.now(), + api_key: "sk-1234567890abcdef", + user_role: "user", + rpm_limit_per_model: {}, + tpm_limit_per_model: {}, + user_tpm_limit: 1000, + user_rpm_limit: 100, + user_email: "user@example.com", + user: { + user_email: "user@example.com", + user_id: "user-1", + }, +}; + +const mockTeam: Team = { + team_id: "team-1", + team_alias: "Test Team", + models: ["gpt-3.5-turbo", "gpt-4"], + max_budget: 500, + budget_duration: "1m", + tpm_limit: 5000, + rpm_limit: 500, + organization_id: "org-1", + created_at: "2024-10-01T10:00:00Z", + keys: [], + members_with_roles: [], +}; + +const mockOrganization: Organization = { + organization_id: "org-1", + organization_alias: "Test Organization", + budget_id: "budget-1", + metadata: {}, + models: ["gpt-3.5-turbo", "gpt-4"], + spend: 100, + model_spend: { "gpt-3.5-turbo": 50, "gpt-4": 50 }, + created_at: "2024-10-01T10:00:00Z", + created_by: "user-1", + updated_at: "2024-11-01T10:00:00Z", + updated_by: "user-1", + litellm_budget_table: {}, + teams: [], + users: [], + members: [], +}; + +// Mock hook implementations +const mockUseKeys = useKeys as MockedFunction; +const mockUseFilterLogic = useFilterLogic as MockedFunction; + +beforeEach(() => { + // Reset mocks before each test + vi.clearAllMocks(); + + // Setup default mock implementations + mockUseKeys.mockReturnValue({ + data: { + keys: [mockKey], + total_count: 1, + current_page: 1, + total_pages: 1, + } as KeysResponse, + isPending: false, + refetch: vi.fn(), + } as any); + + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "team-1", + "Organization ID": "org-1", + "Key Alias": "Test Key Alias", + "User ID": "user-1", + "User Email": "user@example.com", + "User Role": "user", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [mockKey], + allKeyAliases: ["test-key-alias"], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); +}); + +it("should render VirtualKeysTable component", () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); +}); + +it("should display key information correctly", async () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); + expect(screen.getByText("Test Team")).toBeInTheDocument(); + expect(screen.getByText("5.5000")).toBeInTheDocument(); + }); +}); + +it("should display user email correctly", async () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + await waitFor(() => { + expect(screen.getByText("user@example.com")).toBeInTheDocument(); + }); +}); + +it("should show skeleton loaders when isLoading is true", () => { + // Mock loading state + mockUseKeys.mockReturnValue({ + data: null, + isPending: true, + refetch: vi.fn(), + } as any); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // Check that loading message is shown + expect(screen.getByText("🚅 Loading keys...")).toBeInTheDocument(); + + // Check that actual key data is not shown + expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); + expect(screen.queryByText("Test Team")).not.toBeInTheDocument(); +}); + +it("should show 'No keys found' message when filteredKeys is empty", () => { + // Mock empty filteredKeys + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [], + allKeyAliases: [], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + expect(screen.getByText("No keys found")).toBeInTheDocument(); +}); + +it("should handle models with more than 3 entries to trigger expansion UI", () => { + const keyWithManyModels = { + ...mockKey, + models: ["gpt-3.5-turbo", "gpt-4", "gpt-4-turbo", "claude-3", "claude-3-5-sonnet"], + }; + + mockUseFilterLogic.mockReturnValue({ + filters: { + "Team ID": "", + "Organization ID": "", + "Key Alias": "", + "User ID": "", + "Sort By": "created_at", + "Sort Order": "desc", + }, + filteredKeys: [keyWithManyModels], + allKeyAliases: ["test-key-alias"], + allTeams: [mockTeam], + allOrganizations: [mockOrganization], + handleFilterChange: vi.fn(), + handleFilterReset: vi.fn(), + }); + + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // This test ensures the ChevronDownIcon import (line 6) is used + // by having a key with > 3 models which triggers the expansion logic + // that uses ChevronDownIcon and ChevronRightIcon + expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); +}); + +it("should render table headers correctly", () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // Check that main headers are rendered (testing the header.isPlaceholder condition path) + expect(screen.getByText("Key ID")).toBeInTheDocument(); + expect(screen.getByText("Key Alias")).toBeInTheDocument(); + expect(screen.getByText("Team Alias")).toBeInTheDocument(); + expect(screen.getByText("Models")).toBeInTheDocument(); + expect(screen.getByText("Spend (USD)")).toBeInTheDocument(); +}); + +it("should handle column resizing hover events", () => { + const mockProps = { + teams: [mockTeam], + organizations: [mockOrganization], + onSortChange: vi.fn(), + currentSort: { + sortBy: "created_at", + sortOrder: "desc" as const, + }, + }; + + renderWithProviders(); + + // Find a header cell with data-header-id attribute + const headerCell = document.querySelector("[data-header-id]") as HTMLElement; + + expect(headerCell).toBeInTheDocument(); + + // Check that the resizer element exists within the header + const resizer = headerCell?.querySelector(".resizer") as HTMLElement; + expect(resizer).toBeInTheDocument(); + + // Initially, resizer should have opacity 0 + expect(resizer.style.opacity).toBe("0"); + + // Simulate mouse enter using fireEvent - should set opacity to 0.5 (lines 612-616) + fireEvent.mouseEnter(headerCell); + expect(resizer.style.opacity).toBe("0.5"); + + // Simulate mouse leave using fireEvent - should set opacity back to 0 (lines 618-622) + fireEvent.mouseLeave(headerCell); + expect(resizer.style.opacity).toBe("0"); +}); diff --git a/ui/litellm-dashboard/src/components/all_keys_table.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx similarity index 74% rename from ui/litellm-dashboard/src/components/all_keys_table.tsx rename to ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index 210ce09fa34..3bda8ee2f02 100644 --- a/ui/litellm-dashboard/src/components/all_keys_table.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -1,132 +1,55 @@ "use client"; -import React, { useEffect, useState } from "react"; -import { ColumnDef, ColumnResizeMode, ColumnResizeDirection } from "@tanstack/react-table"; -import { Select, SelectItem } from "@tremor/react"; -import { Button } from "@tremor/react"; -import KeyInfoView from "./templates/key_info_view"; -import { Tooltip } from "antd"; -import { Team, KeyResponse } from "./key_team_helpers/key_list"; -import FilterComponent from "./molecules/filter"; -import { FilterOption } from "./molecules/filter"; -import { Organization, userListCall } from "./networking"; -import { useFilterLogic } from "./key_team_helpers/filter_logic"; -import { Setter } from "@/types"; -import { updateExistingKeys } from "@/utils/dataUtils"; -import { flexRender, getCoreRowModel, getSortedRowModel, SortingState, useReactTable } from "@tanstack/react-table"; -import { Table, TableHead, TableHeaderCell, TableBody, TableRow, TableCell, Icon } from "@tremor/react"; -import { SwitchVerticalIcon, ChevronUpIcon, ChevronDownIcon, ChevronRightIcon } from "@heroicons/react/outline"; -import { Badge, Text } from "@tremor/react"; -import { getModelDisplayName } from "./key_team_helpers/fetch_available_models_team_key"; +import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { formatNumberWithCommas } from "@/utils/dataUtils"; +import { ChevronDownIcon, ChevronRightIcon, ChevronUpIcon, SwitchVerticalIcon } from "@heroicons/react/outline"; +import { + ColumnDef, + flexRender, + getCoreRowModel, + getPaginationRowModel, + getSortedRowModel, + PaginationState, + SortingState, + useReactTable, +} from "@tanstack/react-table"; +import { + Badge, + Button, + Icon, + Table, + TableBody, + TableCell, + TableHead, + TableHeaderCell, + TableRow, + Text, +} from "@tremor/react"; +import { Skeleton, Tooltip } from "antd"; +import React, { useEffect, useState } from "react"; +import { getModelDisplayName } from "../key_team_helpers/fetch_available_models_team_key"; +import { useFilterLogic } from "../key_team_helpers/filter_logic"; +import { KeyResponse, Team } from "../key_team_helpers/key_list"; +import FilterComponent, { FilterOption } from "../molecules/filter"; +import { Organization } from "../networking"; +import KeyInfoView from "../templates/key_info_view"; -interface AllKeysTableProps { - keys: KeyResponse[]; - setKeys: (keys: KeyResponse[] | ((prev: KeyResponse[]) => KeyResponse[])) => void; - isLoading?: boolean; - pagination: { - currentPage: number; - totalPages: number; - totalCount: number; - }; - onPageChange: (page: number) => void; - pageSize?: number; +interface VirtualKeysTableProps { teams: Team[] | null; - selectedTeam: Team | null; - setSelectedTeam: (team: Team | null) => void; - selectedKeyAlias: string | null; - setSelectedKeyAlias: Setter; - accessToken: string | null; - userID: string | null; - userRole: string | null; organizations: Organization[] | null; - setCurrentOrg: React.Dispatch>; - refresh?: () => void; onSortChange?: (sortBy: string, sortOrder: "asc" | "desc") => void; currentSort?: { sortBy: string; sortOrder: "asc" | "desc"; }; - premiumUser: boolean; - setAccessToken?: (token: string) => void; } -// Define columns similar to our logs table - -interface UserResponse { - user_id: string; - user_email: string; - user_role: string; -} - -const TeamFilter = ({ - teams, - selectedTeam, - setSelectedTeam, -}: { - teams: Team[] | null; - selectedTeam: Team | null; - setSelectedTeam: (team: Team | null) => void; -}) => { - const handleTeamChange = (value: string) => { - const team = teams?.find((t) => t.team_id === value); - setSelectedTeam(team || null); - }; - - return ( -
-
- Where Team is - -
-
- ); -}; - /** - * AllKeysTable – a new table for keys that mimics the table styling used in view_logs. + * VirtualKeysTable – a new table for keys that mimics the table styling used in view_logs. * The team selector and filtering have been removed so that all keys are shown. */ -export function AllKeysTable({ - keys, - setKeys, - isLoading = false, - pagination, - onPageChange, - pageSize = 50, - teams, - selectedTeam, - setSelectedTeam, - selectedKeyAlias, - setSelectedKeyAlias, - accessToken, - userID, - userRole, - organizations, - setCurrentOrg, - refresh, - onSortChange, - currentSort, - premiumUser, - setAccessToken, -}: AllKeysTableProps) { - const [selectedKeyId, setSelectedKeyId] = useState(null); - const [userList, setUserList] = useState([]); - const [columnResizeMode, setColumnResizeMode] = React.useState("onChange"); - const [columnResizeDirection, setColumnResizeDirection] = React.useState("ltr"); +export function VirtualKeysTable({ teams, organizations, onSortChange, currentSort }: VirtualKeysTableProps) { + const [selectedKey, setSelectedKey] = useState(null); const [sorting, setSorting] = React.useState(() => { if (currentSort) { return [ @@ -143,34 +66,34 @@ export function AllKeysTable({ }, ]; }); + const [tablePagination, setTablePagination] = React.useState({ + pageIndex: 0, + pageSize: 50, + }); + + const { + data: keys, + isPending: isLoading, + isFetching, + refetch, + } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize); + const totalCount = keys?.total_count || 0; const [expandedAccordions, setExpandedAccordions] = useState>({}); // Use the filter logic hook const { filters, filteredKeys, allKeyAliases, allTeams, allOrganizations, handleFilterChange, handleFilterReset } = useFilterLogic({ - keys, + keys: keys?.keys || [], teams, organizations, - accessToken, }); - useEffect(() => { - if (accessToken) { - const user_IDs = keys.map((key) => key.user_id).filter((id) => id !== null); - const fetchUserList = async () => { - const userListData = await userListCall(accessToken, user_IDs, 1, 100); - setUserList(userListData.users); - }; - fetchUserList(); - } - }, [accessToken, keys]); - // Add a useEffect to call refresh when a key is created useEffect(() => { - if (refresh) { + if (refetch) { const handleStorageChange = () => { - refresh(); + refetch(); }; // Listen for storage events that might indicate a key was created @@ -180,7 +103,7 @@ export function AllKeysTable({ window.removeEventListener("storage", handleStorageChange); }; } - }, [refresh]); + }, [refetch]); const columns: ColumnDef[] = [ { @@ -206,7 +129,7 @@ export function AllKeysTable({ size="xs" variant="light" className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]" - onClick={() => setSelectedKeyId(info.getValue() as string)} + onClick={() => setSelectedKey(info.row.original)} > {info.getValue() ? `${(info.getValue() as string).slice(0, 7)}...` : "-"} @@ -269,18 +192,19 @@ export function AllKeysTable({ }, { id: "user_email", - accessorKey: "user_id", + accessorKey: "user", header: "User Email", size: 160, cell: (info) => { - const userId = info.getValue() as string; - const user = userList.find((u) => u.user_id === userId); - return user?.user_email ? ( - - {user?.user_email.slice(0, 20)}... + const user = info.getValue() as any; + const value = user?.user_email; + const width = info.cell.column.getSize(); + return ( + + + {value ?? "-"} + - ) : ( - "-" ); }, }, @@ -552,10 +476,11 @@ export function AllKeysTable({ const table = useReactTable({ data: filteredKeys, columns: columns.filter((col) => col.id !== "expander"), - columnResizeMode, - columnResizeDirection, + columnResizeMode: "onChange", + columnResizeDirection: "ltr", state: { sorting, + pagination: tablePagination, }, onSortingChange: (updaterOrValue) => { const newSorting = typeof updaterOrValue === "function" ? updaterOrValue(sorting) : updaterOrValue; @@ -574,10 +499,14 @@ export function AllKeysTable({ onSortChange?.(sortBy, sortOrder); } }, + onPaginationChange: setTablePagination, getCoreRowModel: getCoreRowModel(), getSortedRowModel: getSortedRowModel(), + getPaginationRowModel: getPaginationRowModel(), enableSorting: true, manualSorting: false, + manualPagination: true, + pageCount: Math.ceil(totalCount / tablePagination.pageSize), }); // Update local sorting state when currentSort prop changes @@ -592,34 +521,18 @@ export function AllKeysTable({ } }, [currentSort]); + const { pageIndex, pageSize } = table.getState().pagination; + const start = pageIndex * pageSize + 1; + const end = Math.min((pageIndex + 1) * pageSize, totalCount); + const rangeLabel = `${start} - ${end}`; return (
- {selectedKeyId ? ( + {selectedKey ? ( setSelectedKeyId(null)} - keyData={filteredKeys.find((k) => k.token === selectedKeyId)} - onKeyDataUpdate={(updatedKeyData) => { - setKeys((keys) => - keys.map((key) => { - if (key.token === updatedKeyData.token) { - return updateExistingKeys(key, updatedKeyData); - } - return key; - }), - ); - if (refresh) refresh(); // Minimal fix: refresh the full key list after an update - }} - onDelete={() => { - setKeys((keys) => keys.filter((key) => key.token !== selectedKeyId)); - if (refresh) refresh(); // Minimal fix: refresh the full key list after a delete - }} - accessToken={accessToken} - userID={userID} - userRole={userRole} + keyId={selectedKey.token} + onClose={() => setSelectedKey(null)} + keyData={selectedKey} teams={allTeams} - premiumUser={premiumUser} - setAccessToken={setAccessToken} /> ) : (
@@ -633,34 +546,46 @@ export function AllKeysTable({
- - Showing{" "} - {isLoading - ? "..." - : `${(pagination.currentPage - 1) * pageSize + 1} - ${Math.min(pagination.currentPage * pageSize, pagination.totalCount)}`}{" "} - of {isLoading ? "..." : pagination.totalCount} results - + {isLoading || isFetching ? ( + + ) : ( + + Showing {rangeLabel} of {totalCount} results + + )}
- - Page {isLoading ? "..." : pagination.currentPage} of {isLoading ? "..." : pagination.totalPages} - + {isLoading || isFetching ? ( + + ) : ( + + Page {pageIndex + 1} of {table.getPageCount()} + + )} - + {isLoading || isFetching ? ( + + ) : ( + + )} - + {isLoading || isFetching ? ( + + ) : ( + + )}
@@ -731,10 +656,6 @@ export function AllKeysTable({ userSelect: "none", touchAction: "none", opacity: header.column.getIsResizing() ? 1 : 0, - transform: - columnResizeMode === "onEnd" && header.column.getIsResizing() - ? `translateX(${(table.options.columnResizeDirection === "rtl" ? -1 : 1) * (table.getState().columnSizingInfo.deltaOffset ?? 0)}px)` - : "", }} />
@@ -744,7 +665,7 @@ export function AllKeysTable({ ))} - {isLoading ? ( + {isLoading || isFetching ? (
diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx new file mode 100644 index 00000000000..6f63fcfa549 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.test.tsx @@ -0,0 +1,276 @@ +import { renderHook, screen, waitFor, renderWithProviders } from "../../../tests/test-utils"; +import userEvent from "@testing-library/user-event"; +import { Form } from "antd"; +import type { UploadProps } from "antd/es/upload"; +import { describe, expect, it, vi } from "vitest"; +import type { Team } from "../key_team_helpers/key_list"; +import type { CredentialItem } from "../networking"; +import { Providers } from "../provider_info_helpers"; +import AddModelForm from "./AddModelForm"; + +vi.mock("../molecules/models/ProviderLogo", () => ({ + ProviderLogo: ({ provider, className }: { provider: string; className?: string }) => ( +
+ {provider} +
+ ), +})); + +vi.mock("../networking", async () => { + const actual = await vi.importActual("../networking"); + return { + ...actual, + getGuardrailsList: vi.fn().mockResolvedValue({ + guardrails: [{ guardrail_name: "test-guardrail-1" }, { guardrail_name: "test-guardrail-2" }], + }), + tagListCall: vi.fn().mockResolvedValue({}), + modelAvailableCall: vi.fn().mockResolvedValue({ + data: [{ id: "model-group-1" }, { id: "model-group-2" }], + }), + modelHubCall: vi.fn().mockResolvedValue({ + data: [ + { model_group: "gpt-4", mode: "chat" }, + { model_group: "gpt-3.5-turbo", mode: "chat" }, + ], + }), + getProviderCreateMetadata: vi.fn().mockResolvedValue([ + { + provider: "OpenAI", + provider_display_name: "OpenAI", + litellm_provider: "openai", + default_model_placeholder: "gpt-3.5-turbo", + credential_fields: [], + }, + ]), + }; +}); + +vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({ + useProviderFields: vi.fn().mockReturnValue({ + data: [ + { + provider: "OpenAI", + provider_display_name: "OpenAI", + litellm_provider: "openai", + default_model_placeholder: "gpt-3.5-turbo", + credential_fields: [], + }, + ], + isLoading: false, + error: null, + }), +})); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn(), +})); + +vi.mock("@/app/(dashboard)/hooks/guardrails/useGuardrails", () => ({ + useGuardrails: vi.fn().mockReturnValue({ + data: [{ guardrail_name: "test-guardrail" }], + isLoading: false, + error: null, + }), +})); + +vi.mock("@/app/(dashboard)/hooks/tags/useTags", () => ({ + useTags: vi.fn().mockReturnValue({ + data: { tag1: ["model1", "model2"] }, + isLoading: false, + error: null, + }), +})); + +const mockAuthorizedUser = (userRole: string, userId: string, premiumUser: boolean) => ({ + token: "test-token", + accessToken: "test-access-token", + userId, + userEmail: "test@example.com", + userRole, + premiumUser, + disabledPersonalKeyCreation: false, + showSSOBanner: false, +}); + +const testTeam: Team = { + team_id: "team-1", + team_alias: "Test Team", + models: ["gpt-4"], + max_budget: 100, + budget_duration: "monthly", + tpm_limit: null, + rpm_limit: null, + organization_id: "org-1", + created_at: "2024-01-01T00:00:00Z", + keys: [], + members_with_roles: [], +}; + +const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmin = false) => { + const { result } = renderHook(() => Form.useForm()); + const [form] = result.current; + + const teams = [ + { + ...testTeam, + members_with_roles: isTeamAdmin ? [{ user_id: userId, role: "admin" }] : [], + }, + ]; + + const credentials: CredentialItem[] = [ + { + credential_name: "test-credential", + credential_values: {}, + credential_info: { + custom_llm_provider: "openai", + description: "Test credential", + }, + }, + ]; + + const uploadProps: UploadProps = { + beforeUpload: () => false, + showUploadList: false, + }; + + return { + form, + handleOk: vi.fn(), + setSelectedProvider: vi.fn(), + setProviderModelsFn: vi.fn(), + getPlaceholder: vi.fn((provider: Providers) => `Enter ${provider} model name`), + setShowAdvancedSettings: vi.fn(), + selectedProvider: Providers.OpenAI, + providerModels: ["gpt-4", "gpt-3.5-turbo"], + showAdvancedSettings: false, + teams, + credentials, + uploadProps, + userRole, + userId, + }; +}; + +describe("AddModelForm", () => { + it("should render", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + + const props = createTestProps(); + + renderWithProviders(); + + expect(await screen.findByRole("heading", { name: "Add Model" })).toBeInTheDocument(); + }); + + it("should show proxy admin only (not team admin) - should not see Select Team dropdown unless switch is toggled", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + + const props = createTestProps("proxy_admin", "user-1", false); + + renderWithProviders(); + + await screen.findByText("Provider"); + + expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument(); + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + const teamSwitch = screen.getByRole("switch"); + expect(teamSwitch).toBeInTheDocument(); + + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + await userEvent.click(teamSwitch); + + expect(await screen.findByText("Select Team")).toBeInTheDocument(); + }); + + it("should show proxy admin who is also team admin - should not see Select Team dropdown unless switch is toggled", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + + const props = createTestProps("proxy_admin", "user-1", true); + + renderWithProviders(); + + await screen.findByText("Provider"); + + expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument(); + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + const teamSwitch = screen.getByRole("switch"); + expect(teamSwitch).toBeInTheDocument(); + + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + await userEvent.click(teamSwitch); + + expect(await screen.findByText("Select Team")).toBeInTheDocument(); + }); + + it("should show team admin (not proxy admin) - should see alert and team select, must select team before seeing remaining fields", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("team_member", "user-1", true)); + + const props = createTestProps("team_member", "user-1", true); + + renderWithProviders(); + + await screen.findByRole("heading", { name: "Add Model" }); + + expect(screen.getByText("Team Selection Required")).toBeInTheDocument(); + + expect(screen.getByText("Select Team")).toBeInTheDocument(); + + expect(screen.queryByText("Provider")).not.toBeInTheDocument(); + + const teamSelect = screen.getByRole("combobox"); + await userEvent.click(teamSelect); + await userEvent.click(screen.getByText("Test Team")); + + await waitFor(() => { + expect(screen.getByText("Provider")).toBeInTheDocument(); + }); + }); + + it("should show team admin (not proxy admin) - should not see team-BYOK switch", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("team_member", "user-1", true)); + + const props = createTestProps("team_member", "user-1", true); + + renderWithProviders(); + + await screen.findByText("Select Team"); + + const teamSelect = screen.getByRole("combobox"); + await userEvent.click(teamSelect); + await userEvent.click(screen.getByText("Test Team")); + + await waitFor(() => { + expect(screen.getByText("Provider")).toBeInTheDocument(); + }); + + expect(screen.queryByRole("switch")).not.toBeInTheDocument(); + }); + + it("should handle non-admin, non-team-admin users - should not see team selection or switch", async () => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("user", "user-1", false)); + + const props = createTestProps("user", "user-1", false); + + renderWithProviders(); + + await screen.findByRole("heading", { name: "Add Model" }); + + expect(screen.queryByText("Team Selection Required")).not.toBeInTheDocument(); + + expect(screen.queryByText("Select Team")).not.toBeInTheDocument(); + + expect(screen.queryByText("Provider")).not.toBeInTheDocument(); + + expect(screen.queryByRole("switch")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx new file mode 100644 index 00000000000..59ac63cffe6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -0,0 +1,421 @@ +import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; +import { useGuardrails } from "@/app/(dashboard)/hooks/guardrails/useGuardrails"; +import { useTags } from "@/app/(dashboard)/hooks/tags/useTags"; +import { all_admin_roles, isUserTeamAdminForAnyTeam } from "@/utils/roles"; +import { Switch, Text } from "@tremor/react"; +import type { FormInstance } from "antd"; +import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography, Alert } from "antd"; +import type { UploadProps } from "antd/es/upload"; +import React, { useEffect, useMemo, useState } from "react"; +import TeamDropdown from "../common_components/team_dropdown"; +import type { Team } from "../key_team_helpers/key_list"; +import { type CredentialItem, type ProviderCreateInfo, modelAvailableCall } from "../networking"; +import { Providers, providerLogoMap } from "../provider_info_helpers"; +import { ProviderLogo } from "../molecules/models/ProviderLogo"; +import AdvancedSettings from "./advanced_settings"; +import ConditionalPublicModelName from "./conditional_public_model_name"; +import LiteLLMModelNameField from "./litellm_model_name"; +import ConnectionErrorDisplay from "./model_connection_test"; +import ProviderSpecificFields from "./provider_specific_fields"; +import { TEST_MODES } from "./add_model_modes"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +interface AddModelFormProps { + form: FormInstance; // For the Add Model tab + handleOk: () => Promise; + selectedProvider: Providers; + setSelectedProvider: (provider: Providers) => void; + providerModels: string[]; + setProviderModelsFn: (provider: Providers) => void; + getPlaceholder: (provider: Providers) => string; + uploadProps: UploadProps; + showAdvancedSettings: boolean; + setShowAdvancedSettings: (show: boolean) => void; + teams: Team[] | null; + credentials: CredentialItem[]; +} + +const { Title, Link } = Typography; + +const AddModelForm: React.FC = ({ + form, + handleOk, + selectedProvider, + setSelectedProvider, + providerModels, + setProviderModelsFn, + getPlaceholder, + uploadProps, + showAdvancedSettings, + setShowAdvancedSettings, + teams, + credentials, +}) => { + const [testMode, setTestMode] = useState("chat"); + const [isResultModalVisible, setIsResultModalVisible] = useState(false); + const [isTestingConnection, setIsTestingConnection] = useState(false); + // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test + const [connectionTestId, setConnectionTestId] = useState(""); + + const { accessToken, userRole, premiumUser, userId } = useAuthorized(); + const { + data: providerMetadata, + isLoading: isProviderMetadataLoading, + error: providerMetadataError, + } = useProviderFields(); + const { data: guardrailsList, isLoading: isGuardrailsLoading, error: guardrailsError } = useGuardrails(); + const { data: tagsList, isLoading: isTagsLoading, error: tagsError } = useTags(); + + const handleTestConnection = async () => { + setIsTestingConnection(true); + setConnectionTestId(`test-${Date.now()}`); + setIsResultModalVisible(true); + }; + + const [isTeamOnly, setIsTeamOnly] = useState(false); + const [modelAccessGroups, setModelAccessGroups] = useState([]); + // Team admin specific state + const [teamAdminSelectedTeam, setTeamAdminSelectedTeam] = useState(null); + + useEffect(() => { + const fetchModelAccessGroups = async () => { + const response = await modelAvailableCall(accessToken, "", "", false, null, true, true); + setModelAccessGroups(response["data"].map((model: any) => model["id"])); + }; + fetchModelAccessGroups(); + }, [accessToken]); + + const sortedProviderMetadata: ProviderCreateInfo[] = useMemo(() => { + if (!providerMetadata) { + return []; + } + return [...providerMetadata].sort((a, b) => a.provider_display_name.localeCompare(b.provider_display_name)); + }, [providerMetadata]); + + const providerMetadataErrorText = providerMetadataError + ? providerMetadataError instanceof Error + ? providerMetadataError.message + : "Failed to load providers" + : null; + + const isAdmin = all_admin_roles.includes(userRole); + const isTeamAdmin = isUserTeamAdminForAnyTeam(teams, userId); + + return ( + <> + Add Model + + +
{ + console.log("🔥 Form onFinish triggered with values:", values); + await handleOk().then(() => { + setTeamAdminSelectedTeam(null); + }); + }} + onFinishFailed={(errorInfo) => { + console.log("💥 Form onFinishFailed triggered:", errorInfo); + }} + labelCol={{ span: 10 }} + wrapperCol={{ span: 16 }} + labelAlign="left" + > + <> + {isTeamAdmin && !isAdmin && ( + <> + + { + setTeamAdminSelectedTeam(value); + }} + /> + + {!teamAdminSelectedTeam && ( + + )} + + )} + {(isAdmin || (isTeamAdmin && teamAdminSelectedTeam)) && ( + <> + + { + setSelectedProvider(value as Providers); + setProviderModelsFn(value as Providers); + form.setFieldsValue({ + custom_llm_provider: value, + }); + form.setFieldsValue({ + model: [], + model_name: undefined, + }); + }} + > + {providerMetadataErrorText && sortedProviderMetadata.length === 0 && ( + + {providerMetadataErrorText} + + )} + {sortedProviderMetadata.map((providerInfo) => { + const displayName = providerInfo.provider_display_name; + const providerKey = providerInfo.provider; + const logoSrc = providerLogoMap[displayName] ?? ""; + + return ( + +
+ + {displayName} +
+
+ ); + })} +
+
+ + + {/* Conditionally Render "Public Model Name" */} + + + {/* Select Mode */} + + setTestMode(value)} + options={TEST_MODES} + /> + + + + + + Optional - LiteLLM endpoint to use when health checking this model{" "} + + Learn more + + + + + + {/* Credentials */} +
+ + Either select existing credentials OR enter new provider credentials below + +
+ + + (option?.label ?? "").toLowerCase().includes(input.toLowerCase())} + options={[ + { value: null, label: "None" }, + ...credentials.map((credential) => ({ + value: credential.credential_name, + label: credential.credential_name, + })), + ]} + allowClear + /> + + + + prevValues.litellm_credential_name !== currentValues.litellm_credential_name || + prevValues.provider !== currentValues.provider + } + > + {({ getFieldValue }) => { + const credentialName = getFieldValue("litellm_credential_name"); + console.log("🔑 Credential Name Changed:", credentialName); + // Only show provider specific fields if no credentials selected + if (!credentialName) { + return ( + <> +
+
+ OR +
+
+ + + ); + } + return null; + }} +
+
+
+ Additional Model Info Settings +
+
+ {/* Team-only Model Switch - Only show for proxy admins, not team admins */} + {(isAdmin || !isTeamAdmin) && ( + + + { + setIsTeamOnly(checked); + if (!checked) { + form.setFieldValue("team_id", undefined); + } + }} + disabled={!premiumUser} + /> + + + )} + + {/* Conditional Team Selection */} + {isTeamOnly && (isAdmin || !isTeamAdmin) && ( + + + + )} + {isAdmin && ( + <> + + ({ + value: group, + label: group, + }))} + maxTagCount="responsive" + allowClear + /> + + + )} + + + )} +
+ + Need Help? + +
+ + +
+
+ + +
+ + {/* Test Connection Results Modal */} + { + setIsResultModalVisible(false); + setIsTestingConnection(false); + }} + footer={[ + , + ]} + width={700} + > + {/* Only render the ConnectionErrorDisplay when modal is visible and we have a test ID */} + {isResultModalVisible && ( + { + setIsResultModalVisible(false); + setIsTestingConnection(false); + }} + onTestComplete={() => setIsTestingConnection(false)} + /> + )} + + + ); +}; + +export default AddModelForm; diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx index 9362cbb3260..197bcd6569f 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.test.tsx @@ -62,6 +62,14 @@ vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({ }), })); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: vi.fn().mockReturnValue({ + accessToken: "test-access-token", + userRole: "Admin", + premiumUser: true, + }), +})); + const createQueryClient = () => new QueryClient({ defaultOptions: { @@ -137,7 +145,6 @@ const createTestProps = () => { uploadProps, accessToken: "test-access-token", userRole: "Admin", - premiumUser: true, }; }; @@ -163,7 +170,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); @@ -192,7 +198,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); @@ -222,7 +227,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); @@ -251,7 +255,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); @@ -289,7 +292,6 @@ describe("Add Model Tab", () => { credentials={props.credentials} accessToken={props.accessToken} userRole={props.userRole} - premiumUser={props.premiumUser} /> , ); diff --git a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx index cfe00af5071..f9b6533ac60 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_model_tab.tsx @@ -1,34 +1,18 @@ -import { useProviderFields } from "@/app/(dashboard)/hooks/providers/useProviderFields"; -import { all_admin_roles } from "@/utils/roles"; -import { Switch, Tab, TabGroup, TabList, TabPanel, TabPanels, Text } from "@tremor/react"; +import { Tab, TabGroup, TabList, TabPanel, TabPanels } from "@tremor/react"; import type { FormInstance } from "antd"; -import { Select as AntdSelect, Button, Card, Col, Form, Modal, Row, Tooltip, Typography } from "antd"; +import { Form } from "antd"; import type { UploadProps } from "antd/es/upload"; -import React, { useEffect, useMemo, useState } from "react"; -import TeamDropdown from "../common_components/team_dropdown"; +import React from "react"; import type { Team } from "../key_team_helpers/key_list"; -import { - type CredentialItem, - type ProviderCreateInfo, - getGuardrailsList, - modelAvailableCall, - tagListCall, -} from "../networking"; -import { Providers, providerLogoMap } from "../provider_info_helpers"; -import { ProviderLogo } from "../molecules/models/ProviderLogo"; -import { Tag } from "../tag_management/types"; +import { type CredentialItem } from "../networking"; +import { Providers } from "../provider_info_helpers"; import AddAutoRouterTab from "./add_auto_router_tab"; -import { TEST_MODES } from "./add_model_modes"; -import AdvancedSettings from "./advanced_settings"; -import ConditionalPublicModelName from "./conditional_public_model_name"; +import AddModelForm from "./AddModelForm"; import { handleAddAutoRouterSubmit } from "./handle_add_auto_router_submit"; -import LiteLLMModelNameField from "./litellm_model_name"; -import ConnectionErrorDisplay from "./model_connection_test"; -import ProviderSpecificFields from "./provider_specific_fields"; interface AddModelTabProps { form: FormInstance; // For the Add Model tab - handleOk: () => void; + handleOk: (values?: any) => Promise; selectedProvider: Providers; setSelectedProvider: (provider: Providers) => void; providerModels: string[]; @@ -41,11 +25,8 @@ interface AddModelTabProps { credentials: CredentialItem[]; accessToken: string; userRole: string; - premiumUser: boolean; } -const { Title, Link } = Typography; - const AddModelTab: React.FC = ({ form, handleOk, @@ -61,90 +42,9 @@ const AddModelTab: React.FC = ({ credentials, accessToken, userRole, - premiumUser, }) => { // Create separate form instance for auto router const [autoRouterForm] = Form.useForm(); - // State for test mode and connection testing - const [testMode, setTestMode] = useState("chat"); - const [isResultModalVisible, setIsResultModalVisible] = useState(false); - const [isTestingConnection, setIsTestingConnection] = useState(false); - const [guardrailsList, setGuardrailsList] = useState([]); - const [tagsList, setTagsList] = useState>({}); - // Using a unique ID to force the ConnectionErrorDisplay to remount and run a fresh test - const [connectionTestId, setConnectionTestId] = useState(""); - - // Provider metadata for driving the provider select from backend config - const { - data: providerMetadata, - isLoading: isProviderMetadataLoading, - error: providerMetadataError, - } = useProviderFields(); - - useEffect(() => { - const fetchGuardrails = async () => { - try { - const response = await getGuardrailsList(accessToken); - const guardrailNames = response.guardrails.map((g: { guardrail_name: string }) => g.guardrail_name); - setGuardrailsList(guardrailNames); - } catch (error) { - console.error("Failed to fetch guardrails:", error); - } - }; - - fetchGuardrails(); - }, [accessToken]); - - useEffect(() => { - const fetchTags = async () => { - try { - const response = await tagListCall(accessToken); - setTagsList(response); - } catch (error) { - console.error("Failed to fetch tags:", error); - } - }; - - fetchTags(); - }, [accessToken]); - - // Test connection when button is clicked - const handleTestConnection = async () => { - setIsTestingConnection(true); - // Generate a new test ID (using timestamp for uniqueness) - // This forces React to create a new instance of ConnectionErrorDisplay - setConnectionTestId(`test-${Date.now()}`); - // Show the modal with the fresh test - setIsResultModalVisible(true); - }; - - // State for team-only switch - const [isTeamOnly, setIsTeamOnly] = useState(false); - - const [modelAccessGroups, setModelAccessGroups] = useState([]); - - useEffect(() => { - const fetchModelAccessGroups = async () => { - const response = await modelAvailableCall(accessToken, "", "", false, null, true, true); - setModelAccessGroups(response["data"].map((model: any) => model["id"])); - }; - fetchModelAccessGroups(); - }, [accessToken]); - - const sortedProviderMetadata: ProviderCreateInfo[] = useMemo(() => { - if (!providerMetadata) { - return []; - } - return [...providerMetadata].sort((a, b) => a.provider_display_name.localeCompare(b.provider_display_name)); - }, [providerMetadata]); - - const providerMetadataErrorText = providerMetadataError - ? providerMetadataError instanceof Error - ? providerMetadataError.message - : "Failed to load providers" - : null; - - const isAdmin = all_admin_roles.includes(userRole); const handleAutoRouterOk = () => { autoRouterForm @@ -166,247 +66,20 @@ const AddModelTab: React.FC = ({ - Add Model - -
{ - console.log("🔥 Form onFinish triggered with values:", values); - handleOk(); - }} - onFinishFailed={(errorInfo) => { - console.log("💥 Form onFinishFailed triggered:", errorInfo); - }} - labelCol={{ span: 10 }} - wrapperCol={{ span: 16 }} - labelAlign="left" - > - <> - {/* Provider Selection */} - - { - setSelectedProvider(value as Providers); - setProviderModelsFn(value as Providers); - form.setFieldsValue({ - custom_llm_provider: value, - }); - form.setFieldsValue({ - model: [], - model_name: undefined, - }); - }} - > - {providerMetadataErrorText && sortedProviderMetadata.length === 0 && ( - - {providerMetadataErrorText} - - )} - {sortedProviderMetadata.map((providerInfo) => { - const displayName = providerInfo.provider_display_name; - const providerKey = providerInfo.provider; - const logoSrc = providerLogoMap[displayName] ?? ""; - - return ( - -
- - {displayName} -
-
- ); - })} -
-
- - - {/* Conditionally Render "Public Model Name" */} - - - {/* Select Mode */} - - setTestMode(value)} - options={TEST_MODES} - /> - - - - - - Optional - LiteLLM endpoint to use when health checking this model{" "} - - Learn more - - - - - - {/* Credentials */} -
- - Either select existing credentials OR enter new provider credentials below - -
- - - - (option?.label ?? "").toLowerCase().includes(input.toLowerCase()) - } - options={[ - { value: null, label: "None" }, - ...credentials.map((credential) => ({ - value: credential.credential_name, - label: credential.credential_name, - })), - ]} - allowClear - /> - - - - prevValues.litellm_credential_name !== currentValues.litellm_credential_name || - prevValues.provider !== currentValues.provider - } - > - {({ getFieldValue }) => { - const credentialName = getFieldValue("litellm_credential_name"); - console.log("🔑 Credential Name Changed:", credentialName); - // Only show provider specific fields if no credentials selected - if (!credentialName) { - return ( - <> -
-
- OR -
-
- - - ); - } - return null; - }} -
-
-
- Additional Model Info Settings -
-
- {/* Team-only Model Switch */} - - - { - setIsTeamOnly(checked); - if (!checked) { - form.setFieldValue("team_id", undefined); - } - }} - disabled={!premiumUser} - /> - - - - {/* Conditional Team Selection */} - {isTeamOnly && ( - - - - )} - {isAdmin && ( - <> - - ({ - value: group, - label: group, - }))} - maxTagCount="responsive" - allowClear - /> - - - )} - - -
- - Need Help? - -
- - -
-
- - -
+
= ({
- - {/* Test Connection Results Modal */} - { - setIsResultModalVisible(false); - setIsTestingConnection(false); - }} - footer={[ - , - ]} - width={700} - > - {/* Only render the ConnectionErrorDisplay when modal is visible and we have a test ID */} - {isResultModalVisible && ( - { - setIsResultModalVisible(false); - setIsTestingConnection(false); - }} - onTestComplete={() => setIsTestingConnection(false)} - /> - )} - ); }; diff --git a/ui/litellm-dashboard/src/components/admins.tsx b/ui/litellm-dashboard/src/components/admins.tsx index 4ddd5cd5d1f..9de971bcd62 100644 --- a/ui/litellm-dashboard/src/components/admins.tsx +++ b/ui/litellm-dashboard/src/components/admins.tsx @@ -3,7 +3,7 @@ * Use this to avoid sharing master key with others */ import React, { useState, useEffect } from "react"; -import { Typography } from "antd"; +import { Alert, Typography } from "antd"; import { useRouter } from "next/navigation"; import { Button as Button2, Modal, Form, Input } from "antd"; import { Select, SelectItem } from "@tremor/react"; @@ -55,6 +55,7 @@ import { getSSOSettings, } from "./networking"; import UISettings from "./Settings/AdminSettings/UISettings/UISettings"; +import SSOSettings from "./Settings/AdminSettings/SSOSettings/SSOSettings"; const AdminPanel: React.FC = ({ searchParams, @@ -496,14 +497,24 @@ const AdminPanel: React.FC = ({ Go to 'Internal Users' page to add other admins. + SSO Settings Security Settings SCIM UI Settings + + + ✨ Security Settings +
{ - const actual = await importOriginal(); - return { - ...actual, - userListCall: vi.fn().mockResolvedValue({ - users: [ - { - user_id: "user-1", - user_email: "user@example.com", - user_role: "user", - }, - ], - }), - }; -}); - -// Mock filter helpers -vi.mock("./key_team_helpers/filter_helpers", () => ({ - fetchAllKeyAliases: vi.fn().mockResolvedValue(["test-key-alias"]), - fetchAllTeams: vi.fn().mockResolvedValue([ - { - team_id: "team-1", - team_alias: "Test Team", - }, - ]), - fetchAllOrganizations: vi.fn().mockResolvedValue([ - { - organization_id: "org-1", - organization_alias: "Test Organization", - }, - ]), -})); - -const mockKey: KeyResponse = { - token: "sk-1234567890abcdef", - token_id: "key-1", - key_name: "test-key", - key_alias: "Test Key Alias", - spend: 5.5, - max_budget: 100, - expires: "2024-12-31T23:59:59Z", - models: ["gpt-3.5-turbo", "gpt-4"], - aliases: {}, - config: {}, - user_id: "user-1", - team_id: "team-1", - max_parallel_requests: 10, - metadata: {}, - tpm_limit: 1000, - rpm_limit: 100, - duration: "30d", - budget_duration: "1m", - budget_reset_at: "2024-12-01T00:00:00Z", - allowed_cache_controls: [], - allowed_routes: [], - permissions: {}, - model_spend: { "gpt-3.5-turbo": 2.5, "gpt-4": 3.0 }, - model_max_budget: { "gpt-3.5-turbo": 50, "gpt-4": 50 }, - soft_budget_cooldown: false, - blocked: false, - litellm_budget_table: {}, - organization_id: "org-1", - created_at: "2024-11-01T10:00:00Z", - updated_at: "2024-11-15T10:00:00Z", - team_spend: 5.5, - team_alias: "Test Team", - team_tpm_limit: 5000, - team_rpm_limit: 500, - team_max_budget: 500, - team_models: ["gpt-3.5-turbo", "gpt-4"], - team_blocked: false, - soft_budget: 50, - team_model_aliases: {}, - team_member_spend: 0, - team_metadata: {}, - end_user_id: "end-user-1", - end_user_tpm_limit: 100, - end_user_rpm_limit: 10, - end_user_max_budget: 10, - last_refreshed_at: Date.now(), - api_key: "sk-1234567890abcdef", - user_role: "user", - rpm_limit_per_model: {}, - tpm_limit_per_model: {}, - user_tpm_limit: 1000, - user_rpm_limit: 100, - user_email: "user@example.com", -}; - -const mockTeam: Team = { - team_id: "team-1", - team_alias: "Test Team", - models: ["gpt-3.5-turbo", "gpt-4"], - max_budget: 500, - budget_duration: "1m", - tpm_limit: 5000, - rpm_limit: 500, - organization_id: "org-1", - created_at: "2024-10-01T10:00:00Z", - keys: [], - members_with_roles: [], -}; - -const mockOrganization: Organization = { - organization_id: "org-1", - organization_alias: "Test Organization", - budget_id: "budget-1", - metadata: {}, - models: ["gpt-3.5-turbo", "gpt-4"], - spend: 100, - model_spend: { "gpt-3.5-turbo": 50, "gpt-4": 50 }, - created_at: "2024-10-01T10:00:00Z", - created_by: "user-1", - updated_at: "2024-11-01T10:00:00Z", - updated_by: "user-1", - litellm_budget_table: {}, - teams: [], - users: [], - members: [], -}; - -it("should render AllKeysTable component", () => { - const mockProps = { - keys: [mockKey], - setKeys: vi.fn(), - isLoading: false, - pagination: { - currentPage: 1, - totalPages: 1, - totalCount: 1, - }, - onPageChange: vi.fn(), - pageSize: 50, - teams: [mockTeam], - selectedTeam: null, - setSelectedTeam: vi.fn(), - selectedKeyAlias: null, - setSelectedKeyAlias: vi.fn(), - accessToken: "test-token", - userID: "user-1", - userRole: "admin", - organizations: [mockOrganization], - setCurrentOrg: vi.fn(), - premiumUser: false, - }; - - renderWithProviders(); - - expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); -}); - -it("should display key information correctly", async () => { - const mockProps = { - keys: [mockKey], - setKeys: vi.fn(), - isLoading: false, - pagination: { - currentPage: 1, - totalPages: 1, - totalCount: 1, - }, - onPageChange: vi.fn(), - pageSize: 50, - teams: [mockTeam], - selectedTeam: null, - setSelectedTeam: vi.fn(), - selectedKeyAlias: null, - setSelectedKeyAlias: vi.fn(), - accessToken: "test-token", - userID: "user-1", - userRole: "admin", - organizations: [mockOrganization], - setCurrentOrg: vi.fn(), - premiumUser: false, - }; - - renderWithProviders(); - - await waitFor(() => { - expect(screen.getByText("Test Key Alias")).toBeInTheDocument(); - expect(screen.getByText("Test Team")).toBeInTheDocument(); - expect(screen.getByText("5.5000")).toBeInTheDocument(); - }); -}); diff --git a/ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx b/ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx index d6412025054..534693d3984 100644 --- a/ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx +++ b/ui/litellm-dashboard/src/components/budgets/budget_panel.test.tsx @@ -1,5 +1,6 @@ import * as networking from "../networking"; import { fireEvent, render, waitFor, screen } from "@testing-library/react"; +import { act } from "@testing-library/react"; import { afterEach, describe, expect, it, vi } from "vitest"; import BudgetPanel from "./budget_panel"; @@ -24,11 +25,11 @@ describe("Budget Panel", () => { }, ]); - const { getByText } = render(); + render(); await waitFor(() => { - expect(getByText("Create a budget to assign to customers.")).toBeInTheDocument(); - expect(getByText("budget-1")).toBeInTheDocument(); + expect(screen.getByText("Create a budget to assign to customers.")).toBeInTheDocument(); + expect(screen.getByText("budget-1")).toBeInTheDocument(); }); }); @@ -43,23 +44,102 @@ describe("Budget Panel", () => { }, ]); - const { getByText, container } = render(); + render(); await waitFor(() => { - expect(getByText("budget-to-delete")).toBeInTheDocument(); + expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); }); - // Find the first table row in tbody and click the second icon (trash/delete) - const bodyRows = container.querySelectorAll("tbody tr"); - expect(bodyRows.length).toBeGreaterThan(0); - const firstRow = bodyRows[0]; - const rowClickableIcons = firstRow.querySelectorAll(".cursor-pointer"); - expect(rowClickableIcons.length).toBeGreaterThan(1); + const deleteButton = screen.getByTestId("delete-budget-button"); - fireEvent.click(rowClickableIcons[1]); + act(() => { + fireEvent.click(deleteButton); + }); await waitFor(() => { expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); }); }); + + it("should successfully delete a budget", async () => { + vi.mocked(networking.getBudgetList).mockResolvedValue([ + { + budget_id: "budget-to-delete", + max_budget: "200", + rpm_limit: 20, + tpm_limit: 2000, + updated_at: "2024-01-02T00:00:00Z", + }, + ]); + vi.mocked(networking.budgetDeleteCall).mockResolvedValue(undefined); + + render(); + + await waitFor(() => { + expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); + }); + + // Open delete modal + const deleteButton = screen.getByTestId("delete-budget-button"); + act(() => { + fireEvent.click(deleteButton); + }); + + await waitFor(() => { + expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); + }); + + // Confirm delete + const confirmButton = screen.getByRole("button", { name: /delete/i }); + act(() => { + fireEvent.click(confirmButton); + }); + + await waitFor(() => { + expect(networking.budgetDeleteCall).toHaveBeenCalledWith("token-123", "budget-to-delete"); + expect(networking.getBudgetList).toHaveBeenCalledTimes(2); // Initial load + refresh after delete + }); + }); + + it("should handle delete error", async () => { + vi.mocked(networking.getBudgetList).mockResolvedValue([ + { + budget_id: "budget-to-delete", + max_budget: "200", + rpm_limit: 20, + tpm_limit: 2000, + updated_at: "2024-01-02T00:00:00Z", + }, + ]); + vi.mocked(networking.budgetDeleteCall).mockRejectedValue(new Error("Delete failed")); + + render(); + + await waitFor(() => { + expect(screen.getByText("budget-to-delete")).toBeInTheDocument(); + }); + + // Open delete modal + const deleteButton = screen.getByTestId("delete-budget-button"); + act(() => { + fireEvent.click(deleteButton); + }); + + await waitFor(() => { + expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); + }); + + // Confirm delete + const confirmButton = screen.getByRole("button", { name: /delete/i }); + act(() => { + fireEvent.click(confirmButton); + }); + + await waitFor(() => { + expect(networking.budgetDeleteCall).toHaveBeenCalledWith("token-123", "budget-to-delete"); + }); + + // Modal should still be open (error handling) + expect(screen.getByText("Delete Budget?")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/components/budgets/budget_panel.tsx b/ui/litellm-dashboard/src/components/budgets/budget_panel.tsx index 252287191b7..b52ef5ab947 100644 --- a/ui/litellm-dashboard/src/components/budgets/budget_panel.tsx +++ b/ui/litellm-dashboard/src/components/budgets/budget_panel.tsx @@ -27,6 +27,7 @@ import NotificationsManager from "../molecules/notifications_manager"; import { budgetDeleteCall, getBudgetList } from "../networking"; import BudgetModal from "./budget_modal"; import EditBudgetModal from "./edit_budget_modal"; +import { CREATE_END_USER_CURL_COMMAND, CHAT_COMPLETIONS_CURL_COMMAND, OPENAI_SDK_PYTHON_CODE } from "./constants"; interface BudgetSettingsPageProps { accessToken: string | null; @@ -110,139 +111,111 @@ const BudgetPanel: React.FC = ({ accessToken }) => { - - {selectedBudget && ( - - )} - - Create a budget to assign to customers. - - - - Budget ID - Max Budget - TPM - RPM - - + + + Budgets + Examples + + + +
+ + {selectedBudget && ( + + )} + + Create a budget to assign to customers. +
+ + + Budget ID + Max Budget + TPM + RPM + + - - {budgetList - .slice() // Creates a shallow copy to avoid mutating the original array - .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) // Sort by updated_at in descending order - .map((value: budgetItem, index: number) => ( - - {value.budget_id} - {value.max_budget ? value.max_budget : "n/a"} - {value.tpm_limit ? value.tpm_limit : "n/a"} - {value.rpm_limit ? value.rpm_limit : "n/a"} - handleEditCall(value)} - /> - handleDeleteClick(value)} - /> - - ))} - -
-
- -
- How to use budget id - - - Assign Budget to Customer - Test it (Curl) - - Test it (OpenAI SDK) - - - - - {` -curl -X POST --location '/end_user/new' \ - --H 'Authorization: Bearer ' \ - --H 'Content-Type: application/json' \ - --d '{"user_id": "my-customer-id', "budget_id": ""}' # 👈 KEY CHANGE - - `} - - - - - {` -curl -X POST --location '/chat/completions' \ - --H 'Authorization: Bearer ' \ - --H 'Content-Type: application/json' \ - --d '{ - "model": "gpt-3.5-turbo', - "messages":[{"role": "user", "content": "Hey, how's it going?"}], - "user": "my-customer-id" -}' # 👈 KEY CHANGE - - `} - - - - - {`from openai import OpenAI -client = OpenAI( - base_url="", - api_key="" -) - -completion = client.chat.completions.create( - model="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Hello!"} - ], - user="my-customer-id" -) - -print(completion.choices[0].message)`} - - - - -
+ + {budgetList + .slice() // Creates a shallow copy to avoid mutating the original array + .sort((a, b) => new Date(b.updated_at).getTime() - new Date(a.updated_at).getTime()) // Sort by updated_at in descending order + .map((value: budgetItem, index: number) => ( + + {value.budget_id} + {value.max_budget ? value.max_budget : "n/a"} + {value.tpm_limit ? value.tpm_limit : "n/a"} + {value.rpm_limit ? value.rpm_limit : "n/a"} + handleEditCall(value)} + dataTestId="edit-budget-button" + /> + handleDeleteClick(value)} + dataTestId="delete-budget-button" + /> + + ))} + + + + +
+
+ +
+ How to use budget id + + + Assign Budget to Customer + Test it (Curl) + Test it (OpenAI SDK) + + + + {CREATE_END_USER_CURL_COMMAND} + + + {CHAT_COMPLETIONS_CURL_COMMAND} + + + {OPENAI_SDK_PYTHON_CODE} + + + +
+
+
+
); }; diff --git a/ui/litellm-dashboard/src/components/budgets/constants.ts b/ui/litellm-dashboard/src/components/budgets/constants.ts new file mode 100644 index 00000000000..9d6736db1f1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/budgets/constants.ts @@ -0,0 +1,42 @@ +export const CREATE_END_USER_CURL_COMMAND = ` +curl -X POST --location '/end_user/new' \\ + +-H 'Authorization: Bearer ' \\ + +-H 'Content-Type: application/json' \\ + +-d '{"user_id": "my-customer-id', "budget_id": ""}' # 👈 KEY CHANGE + +`; + +export const CHAT_COMPLETIONS_CURL_COMMAND = ` +curl -X POST --location '/chat/completions' \\ + +-H 'Authorization: Bearer ' \\ + +-H 'Content-Type: application/json' \\ + +-d '{ + "model": "gpt-3.5-turbo', + "messages":[{"role": "user", "content": "Hey, how's it going?"}], + "user": "my-customer-id" +}' # 👈 KEY CHANGE + +`; + +export const OPENAI_SDK_PYTHON_CODE = `from openai import OpenAI +client = OpenAI( + base_url="", + api_key="" +) + +completion = client.chat.completions.create( + model="gpt-3.5-turbo", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"} + ], + user="my-customer-id" +) + +print(completion.choices[0].message)`; diff --git a/ui/litellm-dashboard/src/components/common_components/NewBadge.test.tsx b/ui/litellm-dashboard/src/components/common_components/NewBadge.test.tsx new file mode 100644 index 00000000000..a24b52db6b8 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/NewBadge.test.tsx @@ -0,0 +1,52 @@ +import { render, screen } from "@testing-library/react"; +import { describe, it, expect, vi, beforeEach } from "vitest"; +import NewBadge from "./NewBadge"; + +// Mock the hook directly +vi.mock("@/app/(dashboard)/hooks/useDisableShowNewBadge", () => ({ + useDisableShowNewBadge: vi.fn(), +})); + +import { useDisableShowNewBadge } from "@/app/(dashboard)/hooks/useDisableShowNewBadge"; + +const mockUseDisableShowNewBadge = vi.mocked(useDisableShowNewBadge); + +describe("NewBadge", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("should render the badge when disableShowNewBadge is false", () => { + mockUseDisableShowNewBadge.mockReturnValue(false); + + render(Test Content); + + expect(screen.getByText("New")).toBeInTheDocument(); + expect(screen.getByText("Test Content")).toBeInTheDocument(); + }); + + it("should render the badge when disableShowNewBadge is not set", () => { + mockUseDisableShowNewBadge.mockReturnValue(false); + + render(); + + expect(screen.getByText("New")).toBeInTheDocument(); + }); + + it("should render only children when disableShowNewBadge is true", () => { + mockUseDisableShowNewBadge.mockReturnValue(true); + + render(Test Content); + + expect(screen.queryByText("New")).not.toBeInTheDocument(); + expect(screen.getByText("Test Content")).toBeInTheDocument(); + }); + + it("should render nothing when disableShowNewBadge is true and no children", () => { + mockUseDisableShowNewBadge.mockReturnValue(true); + + const { container } = render(); + + expect(container.firstChild).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx b/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx new file mode 100644 index 00000000000..97cdea8cfbb --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/NewBadge.tsx @@ -0,0 +1,18 @@ +import { Badge } from "antd"; +import { useDisableShowNewBadge } from "@/app/(dashboard)/hooks/useDisableShowNewBadge"; + +export default function NewBadge({ children }: { children?: React.ReactNode }) { + const disableShowNewBadge = useDisableShowNewBadge(); + + if (disableShowNewBadge) { + return children ? <>{children} : null; + } + + return children ? ( + + {children} + + ) : ( + + ); +} diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx index a074a5484f6..5428efb28aa 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/filter_logic.tsx @@ -6,6 +6,7 @@ import { useQuery } from "@tanstack/react-query"; import { fetchAllKeyAliases, fetchAllOrganizations, fetchAllTeams } from "./filter_helpers"; import { debounce } from "lodash"; import { defaultPageSize } from "../constants"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; export interface FilterState { "Team ID": string; @@ -21,12 +22,10 @@ export function useFilterLogic({ keys, teams, organizations, - accessToken, }: { keys: KeyResponse[]; teams: Team[] | null; organizations: Organization[] | null; - accessToken: string | null; }) { const defaultFilters: FilterState = { "Team ID": "", @@ -36,6 +35,7 @@ export function useFilterLogic({ "Sort By": "created_at", "Sort Order": "desc", }; + const { accessToken } = useAuthorized(); const [filters, setFilters] = useState(defaultFilters); const [allTeams, setAllTeams] = useState(teams || []); const [allOrganizations, setAllOrganizations] = useState(organizations || []); diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 6bb014e6187..a04fbf3943d 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -1,6 +1,6 @@ -import { useState, useEffect } from "react"; -import { keyListCall, Member, Organization } from "../networking"; import { Setter } from "@/types"; +import { useEffect, useState } from "react"; +import { keyListCall, Member, Organization } from "../networking"; export interface Team { team_id: string; @@ -91,6 +91,10 @@ export interface KeyResponse { last_rotation_at?: string; key_rotation_at?: string; next_rotation_at?: string; + user?: { + user_id: string; + user_email: string; + }; } interface KeyListResponse { @@ -106,6 +110,7 @@ interface UseKeyListProps { selectedKeyAlias: string | null; accessToken: string; createClicked: boolean; + expand?: string[]; } interface PaginationData { @@ -129,6 +134,7 @@ const useKeyList = ({ selectedKeyAlias, accessToken, createClicked, + expand = [], }: UseKeyListProps): UseKeyListReturn => { const [keyData, setKeyData] = useState({ keys: [], @@ -151,7 +157,19 @@ const useKeyList = ({ const page = typeof params.page === "number" ? params.page : 1; const pageSize = typeof params.pageSize === "number" ? params.pageSize : 100; - const data = await keyListCall(accessToken, null, null, null, null, null, page, pageSize); + const data = await keyListCall( + accessToken, + null, + null, + null, + null, + null, + page, + pageSize, + null, + null, + expand.join(","), + ); console.log("data", data); setKeyData(data); setError(null); diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 11e957ae6d4..fc248ee049a 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -24,11 +24,12 @@ import { UserOutlined, } from "@ant-design/icons"; import type { MenuProps } from "antd"; -import { Badge, ConfigProvider, Layout, Menu } from "antd"; +import { ConfigProvider, Layout, Menu } from "antd"; import { useMemo } from "react"; import { all_admin_roles, internalUserRoles, isAdminRole, rolesWithWriteAccess } from "../utils/roles"; import type { Organization } from "./networking"; import UsageIndicator from "./usage_indicator"; +import NewBadge from "./common_components/NewBadge"; const { Sider } = Layout; // Define the props type @@ -103,11 +104,7 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse { key: "agents", page: "agents", - label: ( - - Agents - - ), + label: Agents, icon: , roles: rolesWithWriteAccess, }, @@ -155,11 +152,7 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse page: "new_usage", icon: , roles: [...all_admin_roles, ...internalUserRoles], - label: ( - - Usage - - ), + label: Usage, }, { key: "logs", @@ -267,7 +260,11 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse { key: "settings", page: "settings", - label: "Settings", + label: ( + + Settings + + ), icon: , roles: all_admin_roles, children: [ diff --git a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx index ba6739d07ce..e9f9636bf2c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/create_mcp_server.tsx @@ -111,6 +111,9 @@ const CreateMCPServer: React.FC = ({ transport, auth_type: AUTH_TYPE.OAUTH2, credentials: values.credentials, + authorization_url: values.authorization_url, + token_url: values.token_url, + registration_url: values.registration_url, mcp_access_groups: values.mcp_access_groups, static_headers: staticHeaders, command: values.command, @@ -608,6 +611,54 @@ const CreateMCPServer: React.FC = ({ size="large" /> + + Authorization URL Override (optional) + + + + + } + name="authorization_url" + > + + + + Token URL Override (optional) + + + + + } + name="token_url" + > + + + + Registration URL Override (optional) + + + + + } + name="registration_url" + > + +

Complete the OAuth authorization flow to fetch an access token and store it as the authentication value. diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx index 1bf719ef904..f6a5d6622d4 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_columns.tsx @@ -10,6 +10,7 @@ export const mcpServerColumns = ( onView: (serverId: string) => void, onEdit: (serverId: string) => void, onDelete: (serverId: string) => void, + isLoadingHealth?: boolean, ): ColumnDef[] => [ { accessorKey: "server_id", @@ -58,6 +59,19 @@ export const mcpServerColumns = ( const lastCheck = server.last_health_check; const error = server.health_check_error; + // Show loading spinner if health check is in progress + if (isLoadingHealth) { + return ( +

+ + + + + Loading... +
+ ); + } + const getStatusColor = (status: string) => { switch (status) { case "healthy": diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx index f27e58e4a93..b4486bdfb23 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_server_edit.tsx @@ -229,6 +229,9 @@ const MCPServerEdit: React.FC = ({ transport: mcpServer.transport, auth_type: mcpServer.auth_type, mcp_info: mcpServer.mcp_info, + authorization_url: mcpServer.authorization_url, + token_url: mcpServer.token_url, + registration_url: mcpServer.registration_url, }; const toolsResponse = await testMCPToolsListRequest(accessToken, mcpServerConfig, oauthAccessToken); @@ -495,6 +498,54 @@ const MCPServerEdit: React.FC = ({ size="large" /> + + Authorization URL Override (optional) + + + + + } + name="authorization_url" + > + + + + Token URL Override (optional) + + + + + } + name="token_url" + > + + + + Registration URL Override (optional) + + + + + } + name="registration_url" + > + +

Use OAuth to fetch a fresh access token and save it as the authentication value.

+
e.stopPropagation()} + > + Hide New Feature Indicators + { + setDisableShowNewBadge(checked); + if (checked) { + setLocalStorageItem("disableShowNewBadge", "true"); + emitLocalStorageChange("disableShowNewBadge"); + } else { + removeLocalStorageItem("disableShowNewBadge"); + emitLocalStorageChange("disableShowNewBadge"); + } + }} + aria-label="Toggle hide new feature indicators" + /> +
), @@ -148,11 +169,7 @@ const Navbar: React.FC = ({