diff --git a/.circleci/config.yml b/.circleci/config.yml index a518628afb9..0adfd5be529 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -3496,8 +3496,13 @@ jobs: command: | npx playwright test e2e_ui_tests/ --reporter=html --output=test-results no_output_timeout: 120m - - store_test_results: + - store_artifacts: path: test-results + destination: playwright-results + + - store_artifacts: + path: playwright-report + destination: playwright-report test_nonroot_image: machine: diff --git a/Dockerfile b/Dockerfile index f75706805e0..d8397ec4811 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,8 +1,8 @@ # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/python:latest-dev +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/python:latest-dev +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base # Builder stage FROM $LITELLM_BUILD_IMAGE AS builder @@ -12,11 +12,9 @@ WORKDIR /app USER root # Install build dependencies -RUN apk add --no-cache gcc python3-dev openssl openssl-dev +RUN apk add --no-cache bash gcc py3-pip python3 python3-dev openssl openssl-dev - -RUN pip install --upgrade pip>=24.3.1 && \ - pip install build +RUN python -m pip install build # Copy the current directory contents into the container at /app COPY . . @@ -48,10 +46,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root # Install runtime dependencies -RUN apk add --no-cache openssl tzdata nodejs npm - -# Upgrade pip to fix CVE-2025-8869 -RUN pip install --upgrade pip>=24.3.1 +RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip WORKDIR /app # Copy the current directory contents into the container at /app diff --git a/README.md b/README.md index 20483ce70c0..4f1a86a648f 100644 --- a/README.md +++ b/README.md @@ -348,7 +348,7 @@ curl 'http://0.0.0.0:4000/key/generate' \ | [Fireworks AI (`fireworks_ai`)](https://docs.litellm.ai/docs/providers/fireworks_ai) | ✅ | ✅ | ✅ | | | | | | | | | [FriendliAI (`friendliai`)](https://docs.litellm.ai/docs/providers/friendliai) | ✅ | ✅ | ✅ | | | | | | | | | [Galadriel (`galadriel`)](https://docs.litellm.ai/docs/providers/galadriel) | ✅ | ✅ | ✅ | | | | | | | | -| [GitHub Copilot (`github_copilot`)](https://docs.litellm.ai/docs/providers/github_copilot) | ✅ | ✅ | ✅ | | | | | | | | +| [GitHub Copilot (`github_copilot`)](https://docs.litellm.ai/docs/providers/github_copilot) | ✅ | ✅ | ✅ | ✅ | | | | | | | | [GitHub Models (`github`)](https://docs.litellm.ai/docs/providers/github) | ✅ | ✅ | ✅ | | | | | | | | | [Google - PaLM](https://docs.litellm.ai/docs/providers/palm) | ✅ | ✅ | ✅ | | | | | | | | | [Google - Vertex AI (`vertex_ai`)](https://docs.litellm.ai/docs/providers/vertex) | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | diff --git a/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py b/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py index f75a53e8798..fd53ece6604 100644 --- a/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py +++ b/cookbook/mock_guardrail_server/mock_bedrock_guardrail_server.py @@ -426,9 +426,12 @@ This is a beta API. Please help us improve it. class LitellmBasicGuardrailRequest(BaseModel): texts: List[str] images: Optional[List[str]] = None + tools: Optional[List[dict]] = None request_data: Dict[str, Any] = Field(default_factory=dict) additional_provider_specific_params: Dict[str, Any] = Field(default_factory=dict) input_type: Literal["request", "response"] + litellm_call_id: Optional[str] = None + litellm_trace_id: Optional[str] = None class LitellmBasicGuardrailResponse(BaseModel): diff --git a/deploy/charts/litellm-helm/Chart.yaml b/deploy/charts/litellm-helm/Chart.yaml index eedadebaa8e..7f14af7db5d 100644 --- a/deploy/charts/litellm-helm/Chart.yaml +++ b/deploy/charts/litellm-helm/Chart.yaml @@ -18,7 +18,7 @@ type: application # This is the chart version. This version number should be incremented each time you make changes # to the chart and its templates, including the app version. # Versions are expected to follow Semantic Versioning (https://semver.org/) -version: 0.4.8 +version: 0.4.9 # This is the version number of the application being deployed. This version number should be # incremented each time you make changes to the application. Versions are not expected to @@ -33,5 +33,5 @@ dependencies: condition: db.deployStandalone - name: redis version: ">=18.0.0" - repository: oci://registry-1.docker.io/bitnamicharts + repository: oci://registry-1.docker.io/bitnamicharts condition: redis.enabled diff --git a/deploy/charts/litellm-helm/README.md b/deploy/charts/litellm-helm/README.md index 352c3e9ddff..6fdc423a177 100644 --- a/deploy/charts/litellm-helm/README.md +++ b/deploy/charts/litellm-helm/README.md @@ -10,46 +10,48 @@ - Helm 3.8.0+ If `db.deployStandalone` is used: + - PV provisioner support in the underlying infrastructure If `db.useStackgresOperator` is used (not yet implemented): -- The Stackgres Operator must already be installed in the Kubernetes Cluster. This chart will **not** install the operator if it is missing. + +- The Stackgres Operator must already be installed in the Kubernetes Cluster. This chart will **not** install the operator if it is missing. ## Parameters ### LiteLLM Proxy Deployment Settings -| Name | Description | Value | -| ---------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----- | -| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` | -| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A | -| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A | -| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A | -| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | -| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | -| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` | -| `image.pullPolicy` | LiteLLM Proxy image pull policy | `IfNotPresent` | -| `image.tag` | Overrides the image tag whose default the latest version of LiteLLM at the time this chart was published. | `""` | -| `imagePullSecrets` | Registry credentials for the LiteLLM and initContainer images. | `[]` | -| `serviceAccount.create` | Whether or not to create a Kubernetes Service Account for this deployment. The default is `false` because LiteLLM has no need to access the Kubernetes API. | `false` | -| `service.type` | Kubernetes Service type (e.g. `LoadBalancer`, `ClusterIP`, etc.) | `ClusterIP` | -| `service.port` | TCP port that the Kubernetes Service will listen on. Also the TCP port within the Pod that the proxy will listen on. | `4000` | -| `service.loadBalancerClass` | Optional LoadBalancer implementation class (only used when `service.type` is `LoadBalancer`) | `""` | -| `ingress.*` | See [values.yaml](./values.yaml) for example settings | N/A | -| `proxyConfigMap.create` | When `true`, render a ConfigMap from `.Values.proxy_config` and mount it. | `true` | -| `proxyConfigMap.name` | When `create=false`, name of the existing ConfigMap to mount. | `""` | -| `proxyConfigMap.key` | Key in the ConfigMap that contains the proxy config file. | `"config.yaml"` | -| `proxy_config.*` | See [values.yaml](./values.yaml) for default settings. Rendered into the ConfigMap’s `config.yaml` only when `proxyConfigMap.create=true`. See [example_config_yaml](../../../litellm/proxy/example_config_yaml/) for configuration examples. | `N/A` | -| `extraContainers[]` | An array of additional containers to be deployed as sidecars alongside the LiteLLM Proxy. -| `pdb.enabled` | Enable a PodDisruptionBudget for the LiteLLM proxy Deployment | `false` | -| `pdb.minAvailable` | Minimum number/percentage of pods that must be available during **voluntary** disruptions (choose **one** of minAvailable/maxUnavailable) | `null` | -| `pdb.maxUnavailable` | Maximum number/percentage of pods that can be unavailable during **voluntary** disruptions (choose **one** of minAvailable/maxUnavailable) | `null` | -| `pdb.annotations` | Extra metadata annotations to add to the PDB | `{}` | -| `pdb.labels` | Extra metadata labels to add to the PDB | `{}` | +| Name | Description | Value | +| --------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------- | +| `replicaCount` | The number of LiteLLM Proxy pods to be deployed | `1` | +| `masterkeySecretName` | The name of the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use the generated secret name. | N/A | +| `masterkeySecretKey` | The key within the Kubernetes Secret that contains the Master API Key for LiteLLM. If not specified, use `masterkey` as the key. | N/A | +| `masterkey` | The Master API Key for LiteLLM. If not specified, a random key in the `sk-...` format is generated. | N/A | +| `environmentSecrets` | An optional array of Secret object names. The keys and values in these secrets will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | +| `environmentConfigMaps` | An optional array of ConfigMap object names. The keys and values in these configmaps will be presented to the LiteLLM proxy pod as environment variables. See below for an example Secret object. | `[]` | +| `image.repository` | LiteLLM Proxy image repository | `ghcr.io/berriai/litellm` | +| `image.pullPolicy` | LiteLLM Proxy image pull policy | `IfNotPresent` | +| `image.tag` | Overrides the image tag whose default the latest version of LiteLLM at the time this chart was published. | `""` | +| `imagePullSecrets` | Registry credentials for the LiteLLM and initContainer images. | `[]` | +| `serviceAccount.create` | Whether or not to create a Kubernetes Service Account for this deployment. The default is `false` because LiteLLM has no need to access the Kubernetes API. | `false` | +| `service.type` | Kubernetes Service type (e.g. `LoadBalancer`, `ClusterIP`, etc.) | `ClusterIP` | +| `service.port` | TCP port that the Kubernetes Service will listen on. Also the TCP port within the Pod that the proxy will listen on. | `4000` | +| `service.loadBalancerClass` | Optional LoadBalancer implementation class (only used when `service.type` is `LoadBalancer`) | `""` | +| `ingress.labels` | Additional labels for the Ingress resource | `{}` | +| `ingress.*` | See [values.yaml](./values.yaml) for example settings | N/A | +| `proxyConfigMap.create` | When `true`, render a ConfigMap from `.Values.proxy_config` and mount it. | `true` | +| `proxyConfigMap.name` | When `create=false`, name of the existing ConfigMap to mount. | `""` | +| `proxyConfigMap.key` | Key in the ConfigMap that contains the proxy config file. | `"config.yaml"` | +| `proxy_config.*` | See [values.yaml](./values.yaml) for default settings. Rendered into the ConfigMap’s `config.yaml` only when `proxyConfigMap.create=true`. See [example_config_yaml](../../../litellm/proxy/example_config_yaml/) for configuration examples. | `N/A` | +| `extraContainers[]` | An array of additional containers to be deployed as sidecars alongside the LiteLLM Proxy. | +| `pdb.enabled` | Enable a PodDisruptionBudget for the LiteLLM proxy Deployment | `false` | +| `pdb.minAvailable` | Minimum number/percentage of pods that must be available during **voluntary** disruptions (choose **one** of minAvailable/maxUnavailable) | `null` | +| `pdb.maxUnavailable` | Maximum number/percentage of pods that can be unavailable during **voluntary** disruptions (choose **one** of minAvailable/maxUnavailable) | `null` | +| `pdb.annotations` | Extra metadata annotations to add to the PDB | `{}` | +| `pdb.labels` | Extra metadata labels to add to the PDB | `{}` | #### Example `proxy_config` ConfigMap from values (default): - ``` proxyConfigMap: create: true @@ -67,7 +69,6 @@ proxy_config: #### Example using existing `proxyConfigMap` instead of creating it: - ``` proxyConfigMap: create: false @@ -77,8 +78,7 @@ proxyConfigMap: # proxy_config is ignored in this mode ``` -#### Example `environmentSecrets` Secret - +#### Example `environmentSecrets` Secret ``` apiVersion: v1 @@ -91,21 +91,23 @@ type: Opaque ``` ### Database Settings -| Name | Description | Value | -| ---------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----- | -| `db.useExisting` | Use an existing Postgres database. A Kubernetes Secret object must exist that contains credentials for connecting to the database. An example secret object definition is provided below. | `false` | -| `db.endpoint` | If `db.useExisting` is `true`, this is the IP, Hostname or Service Name of the Postgres server to connect to. | `localhost` | -| `db.database` | If `db.useExisting` is `true`, the name of the existing database to connect to. | `litellm` | -| `db.url` | If `db.useExisting` is `true`, the connection url of the existing database to connect to can be overwritten with this value. | `postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)` | -| `db.secret.name` | If `db.useExisting` is `true`, the name of the Kubernetes Secret that contains credentials. | `postgres` | -| `db.secret.usernameKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the username for authenticating with the Postgres instance. | `username` | -| `db.secret.passwordKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the password associates with the above user. | `password` | -| `db.useStackgresOperator` | Not yet implemented. | `false` | -| `db.deployStandalone` | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups. | `true` | -| `postgresql.*` | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml) | -| `postgresql.auth.*` | If `db.deployStandalone` is `true`, care should be taken to ensure the default `password` and `postgres-password` values are **NOT** used. | `NoTaGrEaTpAsSwOrD` | + +| Name | Description | Value | +| ------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------ | +| `db.useExisting` | Use an existing Postgres database. A Kubernetes Secret object must exist that contains credentials for connecting to the database. An example secret object definition is provided below. | `false` | +| `db.endpoint` | If `db.useExisting` is `true`, this is the IP, Hostname or Service Name of the Postgres server to connect to. | `localhost` | +| `db.database` | If `db.useExisting` is `true`, the name of the existing database to connect to. | `litellm` | +| `db.url` | If `db.useExisting` is `true`, the connection url of the existing database to connect to can be overwritten with this value. | `postgresql://$(DATABASE_USERNAME):$(DATABASE_PASSWORD)@$(DATABASE_HOST)/$(DATABASE_NAME)` | +| `db.secret.name` | If `db.useExisting` is `true`, the name of the Kubernetes Secret that contains credentials. | `postgres` | +| `db.secret.usernameKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the username for authenticating with the Postgres instance. | `username` | +| `db.secret.passwordKey` | If `db.useExisting` is `true`, the name of the key within the Kubernetes Secret that holds the password associates with the above user. | `password` | +| `db.useStackgresOperator` | Not yet implemented. | `false` | +| `db.deployStandalone` | Deploy a standalone, single instance deployment of Postgres, using the Bitnami postgresql chart. This is useful for getting started but doesn't provide HA or (by default) data backups. | `true` | +| `postgresql.*` | If `db.deployStandalone` is `true`, configuration passed to the Bitnami postgresql chart. See the [Bitnami Documentation](https://github.com/bitnami/charts/tree/main/bitnami/postgresql) for full configuration details. See [values.yaml](./values.yaml) for the default configuration. | See [values.yaml](./values.yaml) | +| `postgresql.auth.*` | If `db.deployStandalone` is `true`, care should be taken to ensure the default `password` and `postgres-password` values are **NOT** used. | `NoTaGrEaTpAsSwOrD` | #### Example Postgres `db.useExisting` Secret + ```yaml apiVersion: v1 kind: Secret @@ -143,7 +145,7 @@ metadata: name: litellm-env-secret type: Opaque data: - SOME_PASSWORD: cDZbUGVXeU5e0ZW # base64 encoded + SOME_PASSWORD: cDZbUGVXeU5e0ZW # base64 encoded ANOTHER_PASSWORD: AAZbUGVXeU5e0ZB # base64 encoded ``` @@ -153,23 +155,23 @@ Source: [GitHub Gist from troyharvey](https://gist.github.com/troyharvey/4506472 The migration job supports both ArgoCD and Helm hooks to ensure database migrations run at the appropriate time during deployments. -| Name | Description | Value | -| ---------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----- | -| `migrationJob.enabled` | Enable or disable the schema migration Job | `true` | -| `migrationJob.backoffLimit` | Backoff limit for Job restarts | `4` | -| `migrationJob.ttlSecondsAfterFinished` | TTL for completed migration jobs | `120` | -| `migrationJob.annotations` | Additional annotations for the migration job pod | `{}` | -| `migrationJob.extraContainers` | Additional containers to run alongside the migration job | `[]` | -| `migrationJob.hooks.argocd.enabled` | Enable ArgoCD hooks for the migration job (uses PreSync hook with BeforeHookCreation delete policy) | `true` | -| `migrationJob.hooks.helm.enabled` | Enable Helm hooks for the migration job (uses pre-install,pre-upgrade hooks with before-hook-creation delete policy) | `false` | -| `migrationJob.hooks.helm.weight` | Helm hook execution order (lower weights executed first). Optional - defaults to "1" if not specified. | N/A | - +| Name | Description | Value | +| -------------------------------------- | -------------------------------------------------------------------------------------------------------------------- | ------- | +| `migrationJob.enabled` | Enable or disable the schema migration Job | `true` | +| `migrationJob.backoffLimit` | Backoff limit for Job restarts | `4` | +| `migrationJob.ttlSecondsAfterFinished` | TTL for completed migration jobs | `120` | +| `migrationJob.annotations` | Additional annotations for the migration job pod | `{}` | +| `migrationJob.extraContainers` | Additional containers to run alongside the migration job | `[]` | +| `migrationJob.hooks.argocd.enabled` | Enable ArgoCD hooks for the migration job (uses PreSync hook with BeforeHookCreation delete policy) | `true` | +| `migrationJob.hooks.helm.enabled` | Enable Helm hooks for the migration job (uses pre-install,pre-upgrade hooks with before-hook-creation delete policy) | `false` | +| `migrationJob.hooks.helm.weight` | Helm hook execution order (lower weights executed first). Optional - defaults to "1" if not specified. | N/A | ## Accessing the Admin UI + When browsing to the URL published per the settings in `ingress.*`, you will -be prompted for **Admin Configuration**. The **Proxy Endpoint** is the internal +be prompted for **Admin Configuration**. The **Proxy Endpoint** is the internal (from the `litellm` pod's perspective) URL published by the `-litellm` -Kubernetes Service. If the deployment uses the default settings for this +Kubernetes Service. If the deployment uses the default settings for this service, the **Proxy Endpoint** should be set to `http://-litellm:4000`. The **Proxy Key** is the value specified for `masterkey` or, if a `masterkey` @@ -181,7 +183,8 @@ kubectl -n litellm get secret -litellm-masterkey -o jsonpath="{.data.ma ``` ## Admin UI Limitations -At the time of writing, the Admin UI is unable to add models. This is because + +At the time of writing, the Admin UI is unable to add models. This is because it would need to update the `config.yaml` file which is a exposed ConfigMap, and -therefore, read-only. This is a limitation of this helm chart, not the Admin UI +therefore, read-only. This is a limitation of this helm chart, not the Admin UI itself. diff --git a/deploy/charts/litellm-helm/templates/ingress.yaml b/deploy/charts/litellm-helm/templates/ingress.yaml index 09e8d715ab8..ea9ffcbb54c 100644 --- a/deploy/charts/litellm-helm/templates/ingress.yaml +++ b/deploy/charts/litellm-helm/templates/ingress.yaml @@ -18,6 +18,9 @@ metadata: name: {{ $fullName }} labels: {{- include "litellm.labels" . | nindent 4 }} + {{- with .Values.ingress.labels }} + {{- toYaml . | nindent 4 }} + {{- end }} {{- with .Values.ingress.annotations }} annotations: {{- toYaml . | nindent 4 }} diff --git a/deploy/charts/litellm-helm/tests/ingress_tests.yaml b/deploy/charts/litellm-helm/tests/ingress_tests.yaml new file mode 100644 index 00000000000..aad6ecfcee8 --- /dev/null +++ b/deploy/charts/litellm-helm/tests/ingress_tests.yaml @@ -0,0 +1,45 @@ +suite: Ingress Configuration Tests +templates: + - ingress.yaml +tests: + - it: should not create Ingress by default + asserts: + - hasDocuments: + count: 0 + + - it: should create Ingress when enabled + set: + ingress.enabled: true + asserts: + - hasDocuments: + count: 1 + - isKind: + of: Ingress + + - it: should add custom labels + set: + ingress.enabled: true + ingress.labels: + custom-label: "true" + another-label: "value" + asserts: + - isKind: + of: Ingress + - equal: + path: metadata.labels.custom-label + value: "true" + - equal: + path: metadata.labels.another-label + value: "value" + + - it: should add annotations + set: + ingress.enabled: true + ingress.annotations: + kubernetes.io/ingress.class: "nginx" + asserts: + - isKind: + of: Ingress + - equal: + path: metadata.annotations["kubernetes.io/ingress.class"] + value: "nginx" diff --git a/deploy/charts/litellm-helm/values.yaml b/deploy/charts/litellm-helm/values.yaml index acb8c9ca32f..35021157826 100644 --- a/deploy/charts/litellm-helm/values.yaml +++ b/deploy/charts/litellm-helm/values.yaml @@ -35,7 +35,8 @@ podAnnotations: {} podLabels: {} terminationGracePeriodSeconds: 90 -topologySpreadConstraints: [] +topologySpreadConstraints: + [] # - maxSkew: 1 # topologyKey: kubernetes.io/hostname # whenUnsatisfiable: DoNotSchedule @@ -46,7 +47,8 @@ topologySpreadConstraints: [] # At the time of writing, the litellm docker image requires write access to the # filesystem on startup so that prisma can install some dependencies. podSecurityContext: {} -securityContext: {} +securityContext: + {} # capabilities: # drop: # - ALL @@ -57,13 +59,15 @@ securityContext: {} # A list of Kubernetes Secret objects that will be exported to the LiteLLM proxy # pod as environment variables. These secrets can then be referenced in the # configuration file (or "litellm" ConfigMap) with `os.environ/` -environmentSecrets: [] +environmentSecrets: + [] # - litellm-env-secret # A list of Kubernetes ConfigMap objects that will be exported to the LiteLLM proxy # pod as environment variables. The ConfigMap kv-pairs can then be referenced in the # configuration file (or "litellm" ConfigMap) with `os.environ/` -environmentConfigMaps: [] +environmentConfigMaps: + [] # - litellm-env-configmap service: @@ -82,7 +86,9 @@ separateHealthPort: 8081 ingress: enabled: false className: "nginx" - annotations: {} + labels: {} + annotations: + {} # kubernetes.io/ingress.class: nginx # kubernetes.io/tls-acme: "true" hosts: @@ -129,7 +135,8 @@ proxy_config: general_settings: master_key: os.environ/PROXY_MASTER_KEY -resources: {} +resources: + {} # We usually recommend not to specify default resources and to leave this as a conscious # choice for the user. This also increases chances charts run on environments with little # resources, such as Minikube. If you do want to specify resources, uncomment the following @@ -231,7 +238,7 @@ migrationJob: # cpu: 100m # memory: 100Mi extraContainers: [] - + # Hook configuration hooks: argocd: @@ -240,30 +247,30 @@ migrationJob: enabled: false # Additional environment variables to be added to the deployment as a map of key-value pairs -envVars: { - # USE_DDTRACE: "true" -} +envVars: {} +# USE_DDTRACE: "true" # Additional environment variables to be added to the deployment as a list of k8s env vars -extraEnvVars: { - # - name: EXTRA_ENV_VAR - # value: EXTRA_ENV_VAR_VALUE -} +extraEnvVars: {} +# - name: EXTRA_ENV_VAR +# value: EXTRA_ENV_VAR_VALUE # Pod Disruption Budget pdb: enabled: false # Set exactly one of the following. If both are set, minAvailable takes precedence. - minAvailable: null # e.g. "50%" or 1 - maxUnavailable: null # e.g. 1 or "20%" + minAvailable: null # e.g. "50%" or 1 + maxUnavailable: null # e.g. 1 or "20%" annotations: {} labels: {} serviceMonitor: enabled: false - labels: {} + labels: + {} # test: test - annotations: {} + annotations: + {} # kubernetes.io/test: test interval: 15s scrapeTimeout: 10s @@ -273,4 +280,4 @@ serviceMonitor: # action: replace namespaceSelector: matchNames: [] - # - test-namespace \ No newline at end of file + # - test-namespace diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 09b5265191b..0e804cbfd12 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,8 +1,8 @@ # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/python:latest-dev +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/python:latest-dev +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base # Builder stage FROM $LITELLM_BUILD_IMAGE AS builder @@ -13,13 +13,15 @@ USER root # Install build dependencies RUN apk add --no-cache \ - build-base \ + bash \ + gcc \ + py3-pip \ + python3 \ python3-dev \ + openssl \ openssl-dev - -RUN pip install --upgrade pip && \ - pip install build +RUN python -m pip install build # Copy the current directory contents into the container at /app COPY . . @@ -46,7 +48,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root # Install runtime dependencies -RUN apk add --no-cache openssl +RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip WORKDIR /app # Copy the current directory contents into the container at /app diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 2dcb7cb4787..cd1633e319c 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,6 +1,6 @@ # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/python:latest-dev -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/python:latest-dev +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base # ----------------- # Builder Stage @@ -10,7 +10,18 @@ WORKDIR /app # Install build dependencies including Node.js for UI build USER root -RUN apk add --no-cache build-base bash nodejs npm \ +RUN apk add --no-cache \ + python3 \ + py3-pip \ + clang \ + llvm \ + lld \ + gcc \ + linux-headers \ + build-base \ + bash \ + nodejs \ + npm \ && pip install --no-cache-dir --upgrade pip build # Copy project files @@ -62,7 +73,7 @@ WORKDIR /app # Install runtime dependencies USER root RUN apk upgrade --no-cache && \ - apk add --no-cache bash libstdc++ ca-certificates openssl supervisor + apk add --no-cache python3 py3-pip bash openssl tzdata nodejs npm supervisor # Copy only necessary artifacts from builder stage for runtime COPY . . diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md new file mode 100644 index 00000000000..ef19e22dab3 --- /dev/null +++ b/docs/my-website/docs/a2a.md @@ -0,0 +1,204 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; +import Image from '@theme/IdealImage'; + +# /a2a - Agent Gateway (A2A Protocol) + +| Feature | Supported | +|---------|-----------| +| Logging | ✅ | +| Load Balancing | ✅ | +| Streaming | ✅ | + +:::tip + +LiteLLM follows the [A2A (Agent-to-Agent) Protocol](https://github.com/google/A2A) for invoking agents. + +::: + +## Adding your Agent + +You can add A2A-compatible agents through the LiteLLM Admin UI. + +1. Navigate to the **Agents** tab +2. Click **Add Agent** +3. Enter the agent name (e.g., `ij-local`) and the URL of your A2A agent + + + +The URL should be the invocation URL for your A2A agent (e.g., `http://localhost:10001`). + +## Invoking your Agents + +Use the [A2A Python SDK](https://pypi.org/project/a2a/) to invoke agents through LiteLLM: + +- `base_url`: Your LiteLLM proxy URL + `/a2a/{agent_name}` +- `headers`: Include your LiteLLM Virtual Key for authentication + +```python showLineNumbers title="invoke_a2a_agent.py" +from uuid import uuid4 +import httpx +import asyncio +from a2a.client import A2ACardResolver, A2AClient +from a2a.types import MessageSendParams, SendMessageRequest + +# === CONFIGURE THESE === +LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL +LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key +LITELLM_AGENT_NAME = "ij-local" # Agent name registered in LiteLLM +# ======================= + +async def main(): + base_url = f"{LITELLM_BASE_URL}/a2a/{LITELLM_AGENT_NAME}" + headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"} + + async with httpx.AsyncClient(headers=headers) as httpx_client: + # Resolve agent card and create client + resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url) + agent_card = await resolver.get_agent_card() + client = A2AClient(httpx_client=httpx_client, agent_card=agent_card) + + # Send a message + request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": "Hello, what can you do?"}], + "messageId": uuid4().hex, + } + ), + ) + response = await client.send_message(request) + print(response.model_dump(mode="json", exclude_none=True)) + +if __name__ == "__main__": + asyncio.run(main()) +``` + +### Streaming Responses + +For streaming responses, use `send_message_streaming`: + +```python showLineNumbers title="invoke_a2a_agent_streaming.py" +from uuid import uuid4 +import httpx +import asyncio +from a2a.client import A2ACardResolver, A2AClient +from a2a.types import MessageSendParams, SendStreamingMessageRequest + +# === CONFIGURE THESE === +LITELLM_BASE_URL = "http://localhost:4000" # Your LiteLLM proxy URL +LITELLM_VIRTUAL_KEY = "sk-1234" # Your LiteLLM Virtual Key +LITELLM_AGENT_NAME = "ij-local" # Agent name registered in LiteLLM +# ======================= + +async def main(): + base_url = f"{LITELLM_BASE_URL}/a2a/{LITELLM_AGENT_NAME}" + headers = {"Authorization": f"Bearer {LITELLM_VIRTUAL_KEY}"} + + async with httpx.AsyncClient(headers=headers) as httpx_client: + # Resolve agent card and create client + resolver = A2ACardResolver(httpx_client=httpx_client, base_url=base_url) + agent_card = await resolver.get_agent_card() + client = A2AClient(httpx_client=httpx_client, agent_card=agent_card) + + # Send a streaming message + request = SendStreamingMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": "Hello, what can you do?"}], + "messageId": uuid4().hex, + } + ), + ) + + # Stream the response + async for chunk in client.send_message_streaming(request): + print(chunk.model_dump(mode="json", exclude_none=True)) + +if __name__ == "__main__": + asyncio.run(main()) +``` + +## Tracking Agent Logs + +After invoking an agent, you can view the request logs in the LiteLLM **Logs** tab. + +The logs show: +- **Request/Response content** sent to and received from the agent +- **User, Key, Team** information for tracking who made the request +- **Latency and cost** metrics + + + +## API Reference + +### Endpoint + +``` +POST /a2a/{agent_name}/message/send +``` + +### Authentication + +Include your LiteLLM Virtual Key in the `Authorization` header: + +``` +Authorization: Bearer sk-your-litellm-key +``` + +### Request Format + +LiteLLM follows the [A2A JSON-RPC 2.0 specification](https://github.com/google/A2A): + +```json title="Request Body" +{ + "jsonrpc": "2.0", + "id": "unique-request-id", + "method": "message/send", + "params": { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Your message here"}], + "messageId": "unique-message-id" + } + } +} +``` + +### Response Format + +```json title="Response" +{ + "jsonrpc": "2.0", + "id": "unique-request-id", + "result": { + "kind": "task", + "id": "task-id", + "contextId": "context-id", + "status": {"state": "completed", "timestamp": "2025-01-01T00:00:00Z"}, + "artifacts": [ + { + "artifactId": "artifact-id", + "name": "response", + "parts": [{"kind": "text", "text": "Agent response here"}] + } + ] + } +} +``` + +## Agent Registry + +Want to create a central registry so your team can discover what agents are available within your company? + +Use the [AI Hub](./proxy/ai_hub) to make agents public and discoverable across your organization. This allows developers to browse available agents without needing to rebuild them. diff --git a/docs/my-website/docs/adding_provider/generic_guardrail_api.md b/docs/my-website/docs/adding_provider/generic_guardrail_api.md index 9392deeb5db..f8c07b25f9d 100644 --- a/docs/my-website/docs/adding_provider/generic_guardrail_api.md +++ b/docs/my-website/docs/adding_provider/generic_guardrail_api.md @@ -21,6 +21,20 @@ The **Generic Guardrail API** lets you integrate with LiteLLM **instantly** by i 5. **Custom Parameters** - Pass provider-specific params via config 6. **Full Control** - You own and maintain your guardrail API +## Supported Endpoints + +The Generic Guardrail API works with the following LiteLLM endpoints: + +- `/v1/chat/completions` - OpenAI Chat Completions +- `/v1/completions` - OpenAI Text Completions +- `/v1/responses` - OpenAI Responses API +- `/v1/images/generations` - OpenAI Image Generation +- `/v1/audio/transcriptions` - OpenAI Audio Transcriptions +- `/v1/audio/speech` - OpenAI Text-to-Speech +- `/v1/messages` - Anthropic Messages +- `/v1/rerank` - Cohere Rerank +- Pass-through endpoints + ## How It Works 1. LiteLLM extracts text and images from any request (chat messages, embeddings, image prompts, etc.) @@ -40,6 +54,21 @@ Implement `POST /beta/litellm_basic_guardrail_api` { "texts": ["extracted text from the request"], // array of text strings "images": ["base64_encoded_image_data"], // optional array of images + "tools": [ // optional array of tools (OpenAI ChatCompletionToolParam format) + { + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"} + } + } + } + } + ], "request_data": { "user_api_key_hash": "hash of the litellm virtual key used", "user_api_key_alias": "alias of the litellm virtual key used", @@ -75,6 +104,49 @@ Implement `POST /beta/litellm_basic_guardrail_api` - `NONE` - Request proceeds unchanged - `GUARDRAIL_INTERVENED` - Request proceeds with modified texts/images (provide `texts` and/or `images` fields) +## Parameters + +### `tools` Parameter + +The `tools` parameter provides information about available function/tool definitions in the request. + +**Format:** OpenAI `ChatCompletionToolParam` format (see [OpenAI API reference](https://platform.openai.com/docs/api-reference/chat/create#chat-create-tools)) + +**Example:** +```json +{ + "type": "function", + "function": { + "name": "get_weather", + "description": "Get the current weather in a location", + "parameters": { + "type": "object", + "properties": { + "location": { + "type": "string", + "description": "City and state, e.g. San Francisco, CA" + }, + "unit": { + "type": "string", + "enum": ["celsius", "fahrenheit"] + } + }, + "required": ["location"] + } + } +} +``` + +**Limitations:** +- **Input only:** Tools are only passed for `input_type="request"` (pre-call guardrails). Output/response guardrails do not currently receive tool information. +- **Supported endpoints:** The `tools` parameter is supported on: `/v1/chat/completions`, `/v1/responses`, and `/v1/messages`. Other endpoints do not have tool support. + +**Use cases:** +- Enforce tool permission policies (e.g., only allow certain users/teams to access specific tools) +- Validate tool schemas before sending to LLM +- Log tool usage for audit purposes +- Block sensitive tools based on user context + ## LiteLLM Configuration Add to `config.yaml`: @@ -138,6 +210,7 @@ app = FastAPI() class GuardrailRequest(BaseModel): texts: List[str] images: Optional[List[str]] = None + tools: Optional[List[Dict[str, Any]]] = None # OpenAI ChatCompletionToolParam format request_data: Dict[str, Any] input_type: str # "request" or "response" litellm_call_id: Optional[str] = None @@ -153,6 +226,8 @@ class GuardrailResponse(BaseModel): @app.post("/beta/litellm_basic_guardrail_api") async def apply_guardrail(request: GuardrailRequest): # Your guardrail logic here + + # Example: Check text content for text in request.texts: if "badword" in text.lower(): return GuardrailResponse( @@ -160,6 +235,18 @@ async def apply_guardrail(request: GuardrailRequest): blocked_reason="Content contains prohibited terms" ) + # Example: Check tools (if present in request) + if request.tools: + for tool in request.tools: + if tool.get("type") == "function": + function_name = tool.get("function", {}).get("name", "") + # Block sensitive tools + if function_name in ["delete_data", "access_admin_panel"]: + return GuardrailResponse( + action="BLOCKED", + blocked_reason=f"Tool '{function_name}' is not allowed" + ) + return GuardrailResponse(action="NONE") ``` diff --git a/docs/my-website/docs/completion/knowledgebase.md b/docs/my-website/docs/completion/knowledgebase.md index 759f7912d87..fd6ef7a9982 100644 --- a/docs/my-website/docs/completion/knowledgebase.md +++ b/docs/my-website/docs/completion/knowledgebase.md @@ -21,6 +21,7 @@ LiteLLM integrates with vector stores, allowing your models to access your organ - [Azure Vector Stores](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/file-search?tabs=python#vector-stores) (Cannot be directly queried. Only available for calling in Assistants messages. We will be adding Azure AI Search Vector Store API support soon.) - [Vertex AI RAG API](https://cloud.google.com/vertex-ai/generative-ai/docs/rag-overview) - [Gemini File Search](https://ai.google.dev/gemini-api/docs/file-search) +- [RAGFlow Datasets](/docs/providers/ragflow_vector_store.md) (Dataset management only, search not supported) ## Quick Start diff --git a/docs/my-website/docs/files_endpoints.md b/docs/my-website/docs/files_endpoints.md index fc0484e9219..30677c748a9 100644 --- a/docs/my-website/docs/files_endpoints.md +++ b/docs/my-website/docs/files_endpoints.md @@ -301,6 +301,17 @@ content = await litellm.afile_content( print("file content=", content) ``` +**Get File Content (Bedrock)** +```python +# For Bedrock batch output files stored in S3 +content = await litellm.afile_content( + file_id="s3://bucket-name/path/to/file.jsonl", # S3 URI or unified file ID + custom_llm_provider="bedrock", + aws_region_name="us-west-2" +) +print("file content=", content.text) +``` + @@ -313,4 +324,6 @@ print("file content=", content) ### [Vertex AI](./providers/vertex#batch-apis) +### [Bedrock](./providers/bedrock_batches#4-retrieve-batch-results) + ## [Swagger API Reference](https://litellm-api.up.railway.app/#/files) diff --git a/docs/my-website/docs/getting_started.md b/docs/my-website/docs/getting_started.md deleted file mode 100644 index 6b2c1fd531e..00000000000 --- a/docs/my-website/docs/getting_started.md +++ /dev/null @@ -1,108 +0,0 @@ -# Getting Started - -import QuickStart from '../src/components/QuickStart.js' - -LiteLLM simplifies LLM API calls by mapping them all to the [OpenAI ChatCompletion format](https://platform.openai.com/docs/api-reference/chat). - -## basic usage - -By default we provide a free $10 community-key to try all providers supported on LiteLLM. - -```python -from litellm import completion - -## set ENV variables -os.environ["OPENAI_API_KEY"] = "your-api-key" -os.environ["COHERE_API_KEY"] = "your-api-key" - -messages = [{ "content": "Hello, how are you?","role": "user"}] - -# openai call -response = completion(model="gpt-3.5-turbo", messages=messages) - -# cohere call -response = completion("command-nightly", messages) -``` - -**Need a dedicated key?** -Email us @ krrish@berri.ai - -Next Steps 👉 [Call all supported models - e.g. Claude-2, Llama2-70b, etc.](./proxy_api.md#supported-models) - -More details 👉 - -- [Completion() function details](./completion/) -- [Overview of supported models / providers on LiteLLM](./providers/) -- [Search all models / providers](https://models.litellm.ai/) -- [Build your own OpenAI proxy](https://github.com/BerriAI/liteLLM-proxy/tree/main) - -## streaming - -Same example from before. Just pass in `stream=True` in the completion args. - -```python -from litellm import completion - -## set ENV variables -os.environ["OPENAI_API_KEY"] = "openai key" -os.environ["COHERE_API_KEY"] = "cohere key" - -messages = [{ "content": "Hello, how are you?","role": "user"}] - -# openai call -response = completion(model="gpt-3.5-turbo", messages=messages, stream=True) - -# cohere call -response = completion("command-nightly", messages, stream=True) - -print(response) -``` - -More details 👉 - -- [streaming + async](./completion/stream.md) -- [tutorial for streaming Llama2 on TogetherAI](./tutorials/TogetherAI_liteLLM.md) - -## exception handling - -LiteLLM maps exceptions across all supported providers to the OpenAI exceptions. All our exceptions inherit from OpenAI's exception types, so any error-handling you have for that, should work out of the box with LiteLLM. - -```python -from openai.error import OpenAIError -from litellm import completion - -os.environ["ANTHROPIC_API_KEY"] = "bad-key" -try: - # some code - completion(model="claude-instant-1", messages=[{"role": "user", "content": "Hey, how's it going?"}]) -except OpenAIError as e: - print(e) -``` - -## Logging Observability - Log LLM Input/Output ([Docs](https://docs.litellm.ai/docs/observability/callbacks)) - -LiteLLM exposes pre defined callbacks to send data to MLflow, Lunary, Langfuse, Helicone, Promptlayer, Traceloop, Slack - -```python -from litellm import completion - -## set env variables for logging tools (API key set up is not required when using MLflow) -os.environ["LUNARY_PUBLIC_KEY"] = "your-lunary-public-key" # get your public key at https://app.lunary.ai/settings -os.environ["HELICONE_API_KEY"] = "your-helicone-key" -os.environ["LANGFUSE_PUBLIC_KEY"] = "" -os.environ["LANGFUSE_SECRET_KEY"] = "" - -os.environ["OPENAI_API_KEY"] - -# set callbacks -litellm.success_callback = ["lunary", "mlflow", "langfuse", "helicone"] # log input/output to MLflow, langfuse, lunary, helicone - -#openai call -response = completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}]) -``` - -More details 👉 - -- [exception mapping](./exception_mapping.md) -- [retries + model fallbacks for completion()](./completion/reliable_completions.md) -- [tutorial for model fallbacks with completion()](./tutorials/fallbacks.md) diff --git a/docs/my-website/docs/observability/datadog.md b/docs/my-website/docs/observability/datadog.md index 5cb5ab3af2d..b2901650ea6 100644 --- a/docs/my-website/docs/observability/datadog.md +++ b/docs/my-website/docs/observability/datadog.md @@ -71,17 +71,19 @@ DD_SOURCE="litellm_dev" # [OPTIONAL] your datadog source. use to different Send logs through a local DataDog agent (useful for containerized environments): ```shell -DD_AGENT_HOST="localhost" # hostname or IP of DataDog agent -DD_AGENT_PORT="10518" # [OPTIONAL] port of DataDog agent (default: 10518) -DD_API_KEY="5f2d0f310***********" # [OPTIONAL] your datadog API Key (agent handles auth) -DD_SOURCE="litellm_dev" # [OPTIONAL] your datadog source +LITELLM_DD_AGENT_HOST="localhost" # hostname or IP of DataDog agent +LITELLM_DD_AGENT_PORT="10518" # [OPTIONAL] port of DataDog agent (default: 10518) +DD_API_KEY="5f2d0f310***********" # [OPTIONAL] your datadog API Key (agent handles auth) +DD_SOURCE="litellm_dev" # [OPTIONAL] your datadog source ``` -When `DD_AGENT_HOST` is set, logs are sent to the agent instead of directly to DataDog API. This is useful for: +When `LITELLM_DD_AGENT_HOST` is set, logs are sent to the agent instead of directly to DataDog API. This is useful for: - Centralized log shipping in containerized environments - Reducing direct API calls from multiple services - Leveraging agent-side processing and filtering +**Note:** We use `LITELLM_DD_AGENT_HOST` instead of `DD_AGENT_HOST` to avoid conflicts with `ddtrace` which automatically sets `DD_AGENT_HOST` for APM tracing. + **Step 3**: Start the proxy, make a test request Start proxy @@ -191,8 +193,8 @@ LiteLLM supports customizing the following Datadog environment variables |---------------------|-------------|---------------|----------| | `DD_API_KEY` | Your Datadog API key for authentication (required for direct API, optional for agent) | None | Conditional* | | `DD_SITE` | Your Datadog site (e.g., "us5.datadoghq.com") (required for direct API) | None | Conditional* | -| `DD_AGENT_HOST` | Hostname or IP of DataDog agent (e.g., "localhost"). When set, logs are sent to agent instead of direct API | None | ❌ No | -| `DD_AGENT_PORT` | Port of DataDog agent for log intake | "10518" | ❌ No | +| `LITELLM_DD_AGENT_HOST` | Hostname or IP of DataDog agent (e.g., "localhost"). When set, logs are sent to agent instead of direct API | None | ❌ No | +| `LITELLM_DD_AGENT_PORT` | Port of DataDog agent for log intake | "10518" | ❌ No | | `DD_ENV` | Environment tag for your logs (e.g., "production", "staging") | "unknown" | ❌ No | | `DD_SERVICE` | Service name for your logs | "litellm-server" | ❌ No | | `DD_SOURCE` | Source name for your logs | "litellm" | ❌ No | @@ -201,5 +203,5 @@ LiteLLM supports customizing the following Datadog environment variables | `POD_NAME` | Pod name tag (useful for Kubernetes deployments) | "unknown" | ❌ No | \* **Required when using Direct API** (default): `DD_API_KEY` and `DD_SITE` are required -\* **Optional when using DataDog Agent**: Set `DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required +\* **Optional when using DataDog Agent**: Set `LITELLM_DD_AGENT_HOST` to use agent mode; `DD_API_KEY` and `DD_SITE` are not required diff --git a/docs/my-website/docs/observability/phoenix_integration.md b/docs/my-website/docs/observability/phoenix_integration.md index ad337439934..898d780668d 100644 --- a/docs/my-website/docs/observability/phoenix_integration.md +++ b/docs/my-website/docs/observability/phoenix_integration.md @@ -6,7 +6,7 @@ Open source tracing and evaluation platform :::tip -This is community maintained, Please make an issue if you run into a bug +This is community maintained. Please make an issue if you run into a bug: https://github.com/BerriAI/litellm ::: @@ -31,19 +31,16 @@ litellm.callbacks = ["arize_phoenix"] import litellm import os -os.environ["PHOENIX_API_KEY"] = "" # Necessary only using Phoenix Cloud -os.environ["PHOENIX_COLLECTOR_HTTP_ENDPOINT"] = "" # The URL of your Phoenix OSS instance e.g. http://localhost:6006/v1/traces -os.environ["PHOENIX_PROJECT_NAME"]="litellm" # OPTIONAL: you can configure project names, otherwise traces would go to "default" project +# Set env variables +os.environ["PHOENIX_API_KEY"] = "d0*****" # Set the Phoenix API key here. It is necessary only when using Phoenix Cloud. +os.environ["PHOENIX_COLLECTOR_HTTP_ENDPOINT"] = "https://app.phoenix.arize.com/s//v1/traces" # Set the URL of your Phoenix OSS instance, otherwise tracer would use https://app.phoenix.arize.com/v1/traces for Phoenix Cloud. +os.environ["PHOENIX_PROJECT_NAME"] = "litellm" # Configure the project name, otherwise traces would go to "default" project. +os.environ['OPENAI_API_KEY'] = "fake-key" # Set the OpenAI API key here. -# This defaults to https://app.phoenix.arize.com/v1/traces for Phoenix Cloud - -# LLM API Keys -os.environ['OPENAI_API_KEY']="" - -# set arize as a callback, litellm will send the data to arize +# Set arize_phoenix as a callback & LiteLLM will send the data to Phoenix. litellm.callbacks = ["arize_phoenix"] - -# openai call + +# OpenAI call response = litellm.completion( model="gpt-3.5-turbo", messages=[ @@ -52,8 +49,9 @@ response = litellm.completion( ) ``` -### Using with LiteLLM Proxy +## Using with LiteLLM Proxy +1. Setup config.yaml ```yaml model_list: @@ -66,12 +64,63 @@ model_list: litellm_settings: callbacks: ["arize_phoenix"] +general_settings: + master_key: "sk-1234" + environment_variables: PHOENIX_API_KEY: "d0*****" - PHOENIX_COLLECTOR_ENDPOINT: "https://app.phoenix.arize.com/v1/traces" # OPTIONAL, for setting the GRPC endpoint - PHOENIX_COLLECTOR_HTTP_ENDPOINT: "https://app.phoenix.arize.com/v1/traces" # OPTIONAL, for setting the HTTP endpoint + PHOENIX_COLLECTOR_ENDPOINT: "https://app.phoenix.arize.com/s//v1/traces" # OPTIONAL - For setting the gRPC endpoint + PHOENIX_COLLECTOR_HTTP_ENDPOINT: "https://app.phoenix.arize.com/s//v1/traces" # OPTIONAL - For setting the HTTP endpoint ``` +2. Start the proxy + +```bash +litellm --config config.yaml +``` + +3. Test it! + +```bash +curl -X POST 'http://0.0.0.0:4000/chat/completions' \ +-H 'Content-Type: application/json' \ +-H 'Authorization: Bearer sk-1234' \ +-d '{ "model": "gpt-4o", "messages": [{"role": "user", "content": "Hi 👋 - i'm openai"}]}' +``` + +## Supported Phoenix Endpoints +Phoenix now supports multiple deployment types. The correct endpoint depends on which version of Phoenix Cloud you are using. + +**Phoenix Cloud (With Spaces - New Version)** +Use this if your Phoenix URL contains `/s/` path. + +```bash +https://app.phoenix.arize.com/s//v1/traces +``` + +**Phoenix Cloud (Legacy - Deprecated)** +Use this only if your deployment still shows the `/legacy` pattern. + +```bash +https://app.phoenix.arize.com/legacy/v1/traces +``` + +**Phoenix Cloud (Without Spaces - Old Version)** +Use this if your Phoenix Cloud URL does not contain `/s/` or `/legacy` path. + +```bash +https://app.phoenix.arize.com/v1/traces +``` + +**Self-Hosted Phoenix (Local Instance)** +Use this when running Phoenix on your machine or a private server. + +```bash +http://localhost:6006/v1/traces +``` + +Depending on which Phoenix Cloud version or deployment you are using, you should set the corresponding endpoint in `PHOENIX_COLLECTOR_HTTP_ENDPOINT` or `PHOENIX_COLLECTOR_ENDPOINT`. + ## Support & Talk to Founders - [Schedule Demo 👋](https://calendly.com/d/4mp-gd3-k5k/berriai-1-1-onboarding-litellm-hosted-version) diff --git a/docs/my-website/docs/projects/Agent Lightning.md b/docs/my-website/docs/projects/Agent Lightning.md new file mode 100644 index 00000000000..28e5546e398 --- /dev/null +++ b/docs/my-website/docs/projects/Agent Lightning.md @@ -0,0 +1,10 @@ + +# Agent Lightning + +[Agent Lightning](https://github.com/microsoft/agent-lightning) is Microsoft's open-source framework for training and optimizing AI agents with Reinforcement Learning, Automatic Prompt Optimization, and Supervised Fine-tuning — with almost zero code changes. + +It works with any agent framework including LangChain, OpenAI Agents SDK, AutoGen, and CrewAI. Agent Lightning uses LiteLLM Proxy under the hood to route LLM requests and collect traces that power its training algorithms. + +- [GitHub](https://github.com/microsoft/agent-lightning) +- [Docs](https://microsoft.github.io/agent-lightning/) +- [arXiv Paper](https://arxiv.org/abs/2508.03680) diff --git a/docs/my-website/docs/projects/Google ADK.md b/docs/my-website/docs/projects/Google ADK.md new file mode 100644 index 00000000000..25e910dcbad --- /dev/null +++ b/docs/my-website/docs/projects/Google ADK.md @@ -0,0 +1,21 @@ + +# Google ADK (Agent Development Kit) + +[Google ADK](https://github.com/google/adk-python) is an open-source, code-first Python framework for building, evaluating, and deploying sophisticated AI agents. While optimized for Gemini, ADK is model-agnostic and supports LiteLLM for using 100+ providers. + +```python +from google.adk.agents.llm_agent import Agent +from google.adk.models.lite_llm import LiteLlm + +root_agent = Agent( + model=LiteLlm(model="openai/gpt-4o"), # Or any LiteLLM-supported model + name="my_agent", + description="An agent using LiteLLM", + instruction="You are a helpful assistant.", + tools=[your_tools], +) +``` + +- [GitHub](https://github.com/google/adk-python) +- [Documentation](https://google.github.io/adk-docs) +- [LiteLLM Samples](https://github.com/google/adk-python/tree/main/contributing/samples/hello_world_litellm) diff --git a/docs/my-website/docs/projects/Harbor.md b/docs/my-website/docs/projects/Harbor.md new file mode 100644 index 00000000000..684dfa93720 --- /dev/null +++ b/docs/my-website/docs/projects/Harbor.md @@ -0,0 +1,24 @@ + +# Harbor + +[Harbor](https://github.com/laude-institute/harbor) is a framework from the creators of Terminal-Bench for evaluating and optimizing agents and language models. It uses LiteLLM to call 100+ LLM providers. + +```bash +# Install +pip install harbor + +# Run a benchmark with any LiteLLM-supported model +harbor run --dataset terminal-bench@2.0 \ + --agent claude-code \ + --model anthropic/claude-opus-4-1 \ + --n-concurrent 4 +``` + +Key features: +- Evaluate agents like Claude Code, OpenHands, Codex CLI +- Build and share benchmarks and environments +- Run experiments in parallel across cloud providers (Daytona, Modal) +- Generate rollouts for RL optimization + +- [GitHub](https://github.com/laude-institute/harbor) +- [Documentation](https://harborframework.com/docs) diff --git a/docs/my-website/docs/providers/anthropic.md b/docs/my-website/docs/providers/anthropic.md index d84c1c23048..f78af51bd90 100644 --- a/docs/my-website/docs/providers/anthropic.md +++ b/docs/my-website/docs/providers/anthropic.md @@ -201,6 +201,30 @@ Without `LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX`: With `LITELLM_ANTHROPIC_DISABLE_URL_SUFFIX=true`: - Base URL `https://my-proxy.com/custom/path` → `https://my-proxy.com/custom/path` (unchanged) +### Azure AI Foundry (Alternative Method) + +:::tip Recommended Method +For full Azure support including Azure AD authentication, use the dedicated [Azure Anthropic provider](./azure/azure_anthropic) with `azure_ai/` prefix. +::: + +As an alternative, you can use the `anthropic/` provider directly with your Azure endpoint since Azure exposes Claude using Anthropic's native API. + +```python +from litellm import completion + +response = completion( + model="anthropic/claude-sonnet-4-5", + api_base="https://.services.ai.azure.com/anthropic", + api_key="", + messages=[{"role": "user", "content": "Hello!"}], +) +print(response) +``` + +:::info +**Finding your Azure endpoint:** Go to Azure AI Foundry → Your deployment → Overview. Your base URL will be `https://.services.ai.azure.com/anthropic` +::: + ## Usage ```python diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index a9ac85a7571..17c0d38111d 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -7,7 +7,7 @@ ALL Bedrock models (Anthropic, Meta, Deepseek, Mistral, Amazon, etc.) are Suppor | Property | Details | |-------|-------| | Description | Amazon Bedrock is a fully managed service that offers a choice of high-performing foundation models (FMs). | -| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models), [`bedrock/openai/`](./bedrock_imported.md#openai-compatible-imported-models-qwen-25-vl-etc) | +| Provider Route on LiteLLM | `bedrock/`, [`bedrock/converse/`](#set-converse--invoke-route), [`bedrock/invoke/`](#set-invoke-route), [`bedrock/converse_like/`](#calling-via-internal-proxy), [`bedrock/llama/`](#deepseek-not-r1), [`bedrock/deepseek_r1/`](#deepseek-r1), [`bedrock/qwen3/`](#qwen3-imported-models), [`bedrock/qwen2/`](./bedrock_imported.md#qwen2-imported-models), [`bedrock/openai/`](./bedrock_imported.md#openai-compatible-imported-models-qwen-25-vl-etc) | | Provider Doc | [Amazon Bedrock ↗](https://docs.aws.amazon.com/bedrock/latest/userguide/what-is-bedrock.html) | | Supported OpenAI Endpoints | `/chat/completions`, `/completions`, `/embeddings`, `/images/generations` | | Rerank Endpoint | `/rerank` | @@ -43,6 +43,8 @@ export AWS_BEARER_TOKEN_BEDROCK="your-api-key" Option 2: use the api_key parameter to pass in API key for completion, embedding, image_generation API calls. + + ```python response = completion( model="bedrock/anthropic.claude-3-sonnet-20240229-v1:0", @@ -50,7 +52,17 @@ response = completion( api_key="your-api-key" ) ``` - + + +```yaml +model_list: + - model_name: bedrock-claude-3-sonnet + litellm_params: + model: bedrock/anthropic.claude-3-sonnet-20240229-v1:0 + api_key: os.environ/AWS_BEARER_TOKEN_BEDROCK +``` + + ## Usage diff --git a/docs/my-website/docs/providers/bedrock_batches.md b/docs/my-website/docs/providers/bedrock_batches.md index a1116f41076..19446fda837 100644 --- a/docs/my-website/docs/providers/bedrock_batches.md +++ b/docs/my-website/docs/providers/bedrock_batches.md @@ -172,6 +172,97 @@ curl http://localhost:4000/v1/batches \ +### 4. Retrieve batch results + +Once the batch job is completed, download the results from S3: + + + + +```python showLineNumbers title="bedrock_batch.py" +... +# Wait for batch completion (check status periodically) +batch_status = client.batches.retrieve(batch_id=batch.id) + +if batch_status.status == "completed": + # Download the output file + result = client.files.content( + file_id=batch_status.output_file_id, + extra_headers={"custom-llm-provider": "bedrock"} + ) + + # Save or process the results + with open("batch_output.jsonl", "wb") as f: + f.write(result.content) + + # Parse JSONL results + for line in result.text.strip().split('\n'): + record = json.loads(line) + print(f"Record ID: {record['recordId']}") + print(f"Output: {record.get('modelOutput', {})}") +``` + + + + +```bash showLineNumbers title="Download Batch Results" +# First retrieve batch to get output_file_id +curl http://localhost:4000/v1/batches/batch_abc123 \ + -H "Authorization: Bearer sk-1234" + +# Then download the output file +curl http://localhost:4000/v1/files/{output_file_id}/content \ + -H "Authorization: Bearer sk-1234" \ + -H "custom-llm-provider: bedrock" \ + -o batch_output.jsonl +``` + + + + +```python showLineNumbers title="bedrock_batch.py" +import litellm +from litellm import file_content + +# Download using litellm directly (bypasses proxy managed files) +result = file_content( + file_id=batch_status.output_file_id, # Can be S3 URI or unified file ID + custom_llm_provider="bedrock", + aws_region_name="us-west-2", +) + +# Process results +print(result.text) +``` + + + + +**Output Format:** + +The batch output file is in JSONL format with each line containing: + +```json +{ + "recordId": "request-1", + "modelInput": { + "messages": [...], + "max_tokens": 1000 + }, + "modelOutput": { + "content": [...], + "id": "msg_abc123", + "model": "claude-3-5-sonnet-20240620-v1:0", + "role": "assistant", + "stop_reason": "end_turn", + "usage": { + "input_tokens": 15, + "output_tokens": 10 + } + } +} +``` + ## FAQ ### Where are my files written? diff --git a/docs/my-website/docs/providers/bedrock_imported.md b/docs/my-website/docs/providers/bedrock_imported.md index 8b0dd721c3c..0784f716925 100644 --- a/docs/my-website/docs/providers/bedrock_imported.md +++ b/docs/my-website/docs/providers/bedrock_imported.md @@ -203,6 +203,71 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ +### Qwen2 Imported Models + +| Property | Details | +|----------|---------| +| Provider Route | `bedrock/qwen2/{model_arn}` | +| Provider Documentation | [Bedrock Imported Models](https://docs.aws.amazon.com/bedrock/latest/userguide/model-customization-import-model.html) | +| Note | Qwen2 and Qwen3 architectures are mostly similar. The main difference is in the response format: Qwen2 uses "text" field while Qwen3 uses "generation" field. | + + + + +```python +from litellm import completion +import os + +response = completion( + model="bedrock/qwen2/arn:aws:bedrock:us-east-1:086734376398:imported-model/your-qwen2-model", # bedrock/qwen2/{your-model-arn} + messages=[{"role": "user", "content": "Tell me a joke"}], + max_tokens=100, + temperature=0.7 +) +``` + + + + + +**1. Add to config** + +```yaml +model_list: + - model_name: Qwen2-72B + litellm_params: + model: bedrock/qwen2/arn:aws:bedrock:us-east-1:086734376398:imported-model/your-qwen2-model + +``` + +**2. Start proxy** + +```bash +litellm --config /path/to/config.yaml + +# RUNNING at http://0.0.0.0:4000 +``` + +**3. Test it!** + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data '{ + "model": "Qwen2-72B", # 👈 the 'model_name' in config + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ], + }' +``` + + + + ### OpenAI-Compatible Imported Models (Qwen 2.5 VL, etc.) Use this route for Bedrock imported models that follow the **OpenAI Chat Completions API spec**. This includes models like Qwen 2.5 VL that accept OpenAI-formatted messages with support for vision (images), tool calling, and other OpenAI features. diff --git a/docs/my-website/docs/providers/github_copilot.md b/docs/my-website/docs/providers/github_copilot.md index 2ebe6eacb1c..306c9f949ec 100644 --- a/docs/my-website/docs/providers/github_copilot.md +++ b/docs/my-website/docs/providers/github_copilot.md @@ -15,7 +15,7 @@ https://docs.github.com/en/copilot |-------|-------| | Description | GitHub Copilot Chat API provides access to GitHub's AI-powered coding assistant. | | Provider Route on LiteLLM | `github_copilot/` | -| Supported Endpoints | `/chat/completions` | +| Supported Endpoints | `/chat/completions`, `/embeddings` | | API Reference | [GitHub Copilot docs](https://docs.github.com/en/copilot) | ## Authentication @@ -62,6 +62,34 @@ for chunk in stream: print(chunk.choices[0].delta.content, end="") ``` +### Responses + +For GPT Codex models, only responses API is supported. + +```python showLineNumbers title="GitHub Copilot Responses" +import litellm + +response = await litellm.aresponses( + model="github_copilot/gpt-5.1-codex", + input="Write a Python hello world", + max_output_tokens=500 +) + +print(response) +``` + +### Embedding + +```python showLineNumbers title="GitHub Copilot Embedding" +import litellm + +response = litellm.embedding( + model="github_copilot/text-embedding-3-small", + input=["good morning from litellm"] +) +print(response) +``` + ## Usage - LiteLLM Proxy Add the following to your LiteLLM Proxy configuration file: @@ -71,6 +99,16 @@ model_list: - model_name: github_copilot/gpt-4 litellm_params: model: github_copilot/gpt-4 + - model_name: github_copilot/gpt-5.1-codex + model_info: + mode: responses + litellm_params: + model: github_copilot/gpt-5.1-codex + - model_name: github_copilot/text-embedding-ada-002 + model_info: + mode: embedding + litellm_params: + model: github_copilot/text-embedding-ada-002 ``` Start your LiteLLM Proxy server: @@ -180,7 +218,7 @@ extra_headers = { "editor-version": "vscode/1.85.1", # Editor version "editor-plugin-version": "copilot/1.155.0", # Plugin version "Copilot-Integration-Id": "vscode-chat", # Integration ID - "user-agent": "GithubCopilot/1.155.0" # User agent + "user-agent": "GithubCopilot/1.155.0" # User agent } ``` diff --git a/docs/my-website/docs/providers/ragflow.md b/docs/my-website/docs/providers/ragflow.md new file mode 100644 index 00000000000..73223bd07b5 --- /dev/null +++ b/docs/my-website/docs/providers/ragflow.md @@ -0,0 +1,244 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# RAGFlow + +Litellm supports Ragflow's chat completions APIs + +## Supported Features + +- ✅ Chat completions +- ✅ Streaming responses +- ✅ Both chat and agent endpoints +- ✅ Multiple credential sources (params, env vars, litellm_params) +- ✅ OpenAI-compatible API format + + +## API Key + +```python +# env variable +os.environ['RAGFLOW_API_KEY'] +``` + +## API Base + +```python +# env variable +os.environ['RAGFLOW_API_BASE'] +``` + +## Overview + +RAGFlow provides OpenAI-compatible APIs with unique path structures that include chat and agent IDs: + +- **Chat endpoint**: `/api/v1/chats_openai/{chat_id}/chat/completions` +- **Agent endpoint**: `/api/v1/agents_openai/{agent_id}/chat/completions` + +The model name format embeds the endpoint type and ID: +- Chat: `ragflow/chat/{chat_id}/{model_name}` +- Agent: `ragflow/agent/{agent_id}/{model_name}` + + +## Sample Usage - Chat Endpoint + +```python +from litellm import completion +import os + +os.environ['RAGFLOW_API_KEY'] = "your-ragflow-api-key" +os.environ['RAGFLOW_API_BASE'] = "http://localhost:9380" # or your hosted URL + +response = completion( + model="ragflow/chat/my-chat-id/gpt-4o-mini", + messages=[{"role": "user", "content": "How does the deep doc understanding work?"}] +) +print(response) +``` + +## Sample Usage - Agent Endpoint + +```python +from litellm import completion +import os + +os.environ['RAGFLOW_API_KEY'] = "your-ragflow-api-key" +os.environ['RAGFLOW_API_BASE'] = "http://localhost:9380" # or your hosted URL + +response = completion( + model="ragflow/agent/my-agent-id/gpt-4o-mini", + messages=[{"role": "user", "content": "What are the key features?"}] +) +print(response) +``` + +## Sample Usage - With Parameters + +You can also pass `api_key` and `api_base` directly as parameters: + +```python +from litellm import completion + +response = completion( + model="ragflow/chat/my-chat-id/gpt-4o-mini", + messages=[{"role": "user", "content": "Hello!"}], + api_key="your-ragflow-api-key", + api_base="http://localhost:9380" +) +print(response) +``` + +## Sample Usage - Streaming + +```python +from litellm import completion +import os + +os.environ['RAGFLOW_API_KEY'] = "your-ragflow-api-key" +os.environ['RAGFLOW_API_BASE'] = "http://localhost:9380" + +response = completion( + model="ragflow/agent/my-agent-id/gpt-4o-mini", + messages=[{"role": "user", "content": "Explain RAGFlow"}], + stream=True +) + +for chunk in response: + print(chunk) +``` + +## Model Name Format + +The model name must follow one of these formats: + +### Chat Endpoint +``` +ragflow/chat/{chat_id}/{model_name} +``` + +Example: `ragflow/chat/my-chat-id/gpt-4o-mini` + +### Agent Endpoint +``` +ragflow/agent/{agent_id}/{model_name} +``` + +Example: `ragflow/agent/my-agent-id/gpt-4o-mini` + +Where: +- `{chat_id}` or `{agent_id}` is the ID of your chat or agent in RAGFlow +- `{model_name}` is the actual model name (e.g., `gpt-4o-mini`, `gpt-4o`, etc.) + +## Configuration Sources + +LiteLLM supports multiple ways to provide credentials, checked in this order: + +1. **Function parameters**: `api_key="..."`, `api_base="..."` +2. **litellm_params**: `litellm_params={"api_key": "...", "api_base": "..."}` +3. **Environment variables**: `RAGFLOW_API_KEY`, `RAGFLOW_API_BASE` +4. **Global litellm settings**: `litellm.api_key`, `litellm.api_base` + +## Usage - LiteLLM Proxy Server + +### 1. Save key in your environment + +```bash +export RAGFLOW_API_KEY="your-ragflow-api-key" +export RAGFLOW_API_BASE="http://localhost:9380" +``` + +### 2. Start the proxy + + + + +```yaml +model_list: + - model_name: ragflow-chat-gpt4 + litellm_params: + model: ragflow/chat/my-chat-id/gpt-4o-mini + api_key: os.environ/RAGFLOW_API_KEY + api_base: os.environ/RAGFLOW_API_BASE + - model_name: ragflow-agent-gpt4 + litellm_params: + model: ragflow/agent/my-agent-id/gpt-4o-mini + api_key: os.environ/RAGFLOW_API_KEY + api_base: os.environ/RAGFLOW_API_BASE +``` + + + + +```bash +$ litellm --config /path/to/config.yaml + +# Server running on http://0.0.0.0:4000 +``` + + + + +### 3. Test it + + + + +```bash +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "ragflow-chat-gpt4", + "messages": [ + {"role": "user", "content": "How does RAGFlow work?"} + ] + }' +``` + + + + +```python +from openai import OpenAI + +client = OpenAI( + api_key="sk-1234", # Your LiteLLM proxy key + base_url="http://0.0.0.0:4000" +) + +response = client.chat.completions.create( + model="ragflow-chat-gpt4", + messages=[ + {"role": "user", "content": "How does RAGFlow work?"} + ] +) +print(response) +``` + + + + +## API Base URL Handling + +The `api_base` parameter can be provided with or without `/v1` suffix. LiteLLM will automatically handle it: + +- `http://localhost:9380` → `http://localhost:9380/api/v1/chats_openai/{chat_id}/chat/completions` +- `http://localhost:9380/v1` → `http://localhost:9380/api/v1/chats_openai/{chat_id}/chat/completions` +- `http://localhost:9380/api/v1` → `http://localhost:9380/api/v1/chats_openai/{chat_id}/chat/completions` + +All three formats will work correctly. + +## Error Handling + +If you encounter errors: + +1. **Invalid model format**: Ensure your model name follows `ragflow/{chat|agent}/{id}/{model_name}` format +2. **Missing api_base**: Provide `api_base` via parameter, environment variable, or litellm_params +3. **Connection errors**: Verify your RAGFlow server is running and accessible at the provided `api_base` + +:::info + +For more information about passing provider-specific parameters, [go here](../completion/provider_specific_params.md) + +::: + diff --git a/docs/my-website/docs/providers/ragflow_vector_store.md b/docs/my-website/docs/providers/ragflow_vector_store.md new file mode 100644 index 00000000000..bc014cacbe6 --- /dev/null +++ b/docs/my-website/docs/providers/ragflow_vector_store.md @@ -0,0 +1,349 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; +import Image from '@theme/IdealImage'; + +# RAGFlow Vector Stores + +Litellm support creation and management of datasets for document processing and knowledge base management in Ragflow. + +| Property | Details | +|----------|---------| +| Description | RAGFlow datasets enable document processing, chunking, and knowledge base management for RAG applications. | +| Provider Route on LiteLLM | `ragflow` in the litellm vector_store_registry | +| Provider Doc | [RAGFlow API Documentation ↗](https://ragflow.io/docs) | +| Supported Operations | Dataset Management (Create, List, Update, Delete) | +| Search/Retrieval | ❌ Not supported (management only) | + +## Quick Start + +### LiteLLM Python SDK + +```python showLineNumbers title="Example using LiteLLM Python SDK" +import os +import litellm + +# Set RAGFlow credentials +os.environ["RAGFLOW_API_KEY"] = "your-ragflow-api-key" +os.environ["RAGFLOW_API_BASE"] = "http://localhost:9380" # Optional, defaults to localhost:9380 + +# Create a RAGFlow dataset +response = litellm.vector_stores.create( + name="my-dataset", + custom_llm_provider="ragflow", + metadata={ + "description": "My knowledge base dataset", + "embedding_model": "BAAI/bge-large-zh-v1.5@BAAI", + "chunk_method": "naive" + } +) + +print(f"Created dataset ID: {response.id}") +print(f"Dataset name: {response.name}") +``` + +### LiteLLM Proxy + +#### 1. Configure your vector_store_registry + + + + +```yaml +model_list: + - model_name: gpt-4o-mini + litellm_params: + model: gpt-4o-mini + api_key: os.environ/OPENAI_API_KEY + +vector_store_registry: + - vector_store_name: "ragflow-knowledge-base" + litellm_params: + vector_store_id: "your-dataset-id" + custom_llm_provider: "ragflow" + api_key: os.environ/RAGFLOW_API_KEY + api_base: os.environ/RAGFLOW_API_BASE # Optional + vector_store_description: "RAGFlow dataset for knowledge base" + vector_store_metadata: + source: "Company documentation" +``` + + + + + +On the LiteLLM UI, Navigate to Experimental > Vector Stores > Create Vector Store. On this page you can create a vector store with a name, vector store id and credentials. + + + + + + +#### 2. Create a dataset via Proxy + + + + +```bash +curl http://localhost:4000/v1/vector_stores \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer $LITELLM_API_KEY" \ + -d '{ + "name": "my-ragflow-dataset", + "custom_llm_provider": "ragflow", + "metadata": { + "description": "Test dataset", + "chunk_method": "naive" + } + }' +``` + + + + + +```python +from openai import OpenAI + +# Initialize client with your LiteLLM proxy URL +client = OpenAI( + base_url="http://localhost:4000", + api_key="your-litellm-api-key" +) + +# Create a RAGFlow dataset +response = client.vector_stores.create( + name="my-ragflow-dataset", + custom_llm_provider="ragflow", + metadata={ + "description": "Test dataset", + "chunk_method": "naive" + } +) + +print(f"Created dataset: {response.id}") +``` + + + + +## Configuration + +### Environment Variables + +RAGFlow vector stores support configuration via environment variables: + +- `RAGFLOW_API_KEY` - Your RAGFlow API key (required) +- `RAGFLOW_API_BASE` - RAGFlow API base URL (optional, defaults to `http://localhost:9380`) + +### Parameters + +You can also pass these via `litellm_params`: + +- `api_key` - RAGFlow API key (overrides `RAGFLOW_API_KEY` env var) +- `api_base` - RAGFlow API base URL (overrides `RAGFLOW_API_BASE` env var) + +## Dataset Creation Options + +### Basic Dataset Creation + +```python +response = litellm.vector_stores.create( + name="basic-dataset", + custom_llm_provider="ragflow" +) +``` + +### Dataset with Chunk Method + +RAGFlow supports various chunk methods for different document types: + + + + +```python +response = litellm.vector_stores.create( + name="general-dataset", + custom_llm_provider="ragflow", + metadata={ + "chunk_method": "naive", + "parser_config": { + "chunk_token_num": 512, + "delimiter": "\n", + "html4excel": False, + "layout_recognize": "DeepDOC" + } + } +) +``` + + + + + +```python +response = litellm.vector_stores.create( + name="book-dataset", + custom_llm_provider="ragflow", + metadata={ + "chunk_method": "book", + "parser_config": { + "raptor": { + "use_raptor": False + } + } + } +) +``` + + + + + +```python +response = litellm.vector_stores.create( + name="qa-dataset", + custom_llm_provider="ragflow", + metadata={ + "chunk_method": "qa", + "parser_config": { + "raptor": { + "use_raptor": False + } + } + } +) +``` + + + + + +```python +response = litellm.vector_stores.create( + name="paper-dataset", + custom_llm_provider="ragflow", + metadata={ + "chunk_method": "paper", + "parser_config": { + "raptor": { + "use_raptor": False + } + } + } +) +``` + + + + +### Dataset with Ingestion Pipeline + +Instead of using a chunk method, you can use an ingestion pipeline: + +```python +response = litellm.vector_stores.create( + name="pipeline-dataset", + custom_llm_provider="ragflow", + metadata={ + "parse_type": 2, # Number of parsers in your pipeline + "pipeline_id": "d0bebe30ae2211f0970942010a8e0005" # 32-character hex ID + } +) +``` + +**Note**: `chunk_method` and `pipeline_id` are mutually exclusive. Use one or the other. + +### Advanced Parser Configuration + +```python +response = litellm.vector_stores.create( + name="advanced-dataset", + custom_llm_provider="ragflow", + metadata={ + "chunk_method": "naive", + "description": "Advanced dataset with custom parser config", + "embedding_model": "BAAI/bge-large-zh-v1.5@BAAI", + "permission": "me", # or "team" + "parser_config": { + "chunk_token_num": 1024, + "delimiter": "\n!?;。；！？", + "html4excel": True, + "layout_recognize": "DeepDOC", + "auto_keywords": 5, + "auto_questions": 3, + "task_page_size": 12, + "raptor": { + "use_raptor": True + }, + "graphrag": { + "use_graphrag": False + } + } + } +) +``` + +## Supported Chunk Methods + +RAGFlow supports the following chunk methods: + +- `naive` - General purpose (default) +- `book` - For book documents +- `email` - For email documents +- `laws` - For legal documents +- `manual` - Manual chunking +- `one` - Single chunk +- `paper` - For academic papers +- `picture` - For image documents +- `presentation` - For presentation documents +- `qa` - Q&A format +- `table` - For table documents +- `tag` - Tag-based chunking + +## RAGFlow-Specific Parameters + +All RAGFlow-specific parameters should be passed via the `metadata` field: + +| Parameter | Type | Description | +|-----------|------|-------------| +| `avatar` | string | Base64 encoding of the avatar (max 65535 chars) | +| `description` | string | Brief description of the dataset (max 65535 chars) | +| `embedding_model` | string | Embedding model name (e.g., "BAAI/bge-large-zh-v1.5@BAAI") | +| `permission` | string | Access permission: "me" (default) or "team" | +| `chunk_method` | string | Chunking method (see supported methods above) | +| `parser_config` | object | Parser configuration (varies by chunk_method) | +| `parse_type` | int | Number of parsers in pipeline (required with pipeline_id) | +| `pipeline_id` | string | 32-character hex pipeline ID (required with parse_type) | + +## Error Handling + +RAGFlow returns error responses in the following format: + +```json +{ + "code": 101, + "message": "Dataset name 'my-dataset' already exists" +} +``` + +LiteLLM automatically maps these to appropriate exceptions: + +- `code != 0` → Raises exception with the error message +- Missing required fields → Raises `ValueError` +- Mutually exclusive parameters → Raises `ValueError` + +## Limitations + +- **Search/Retrieval**: RAGFlow vector stores support dataset management only. Search operations are not supported and will raise `NotImplementedError`. +- **List/Update/Delete**: These operations are not yet implemented through the standard vector store API. Use RAGFlow's native API endpoints directly. + +## Further Reading + +Vector Stores: +- [Vector Store Creation](../vector_stores/create.md) +- [Using Vector Stores with Completions](../completion/knowledgebase.md) +- [Vector Store Registry](../completion/knowledgebase.md#vectorstoreregistry) + diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 70babea3814..da2997f6202 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -2550,355 +2550,6 @@ print(response) - -## **Gemini TTS (Text-to-Speech) Audio Output** - -:::info - -LiteLLM supports Gemini TTS models on Vertex AI that can generate audio responses using the OpenAI-compatible `audio` parameter format. - -::: - -### Supported Models - -LiteLLM supports Gemini TTS models with audio capabilities on Vertex AI (e.g. `vertex_ai/gemini-2.5-flash-preview-tts` and `vertex_ai/gemini-2.5-pro-preview-tts`). For the complete list of available TTS models and voices, see the [official Gemini TTS documentation](https://ai.google.dev/gemini-api/docs/speech-generation). - -### Limitations - -:::warning - -**Important Limitations**: -- Gemini TTS models only support the `pcm16` audio format -- **Streaming support has not been added** to TTS models yet -- The `modalities` parameter must be set to `['audio']` for TTS requests - -::: - -### Quick Start - - - - -```python -from litellm import completion -import json - -## GET CREDENTIALS -file_path = 'path/to/vertex_ai_service_account.json' - -# Load the JSON file -with open(file_path, 'r') as file: - vertex_credentials = json.load(file) - -# Convert to JSON string -vertex_credentials_json = json.dumps(vertex_credentials) - -response = completion( - model="vertex_ai/gemini-2.5-flash-preview-tts", - messages=[{"role": "user", "content": "Say hello in a friendly voice"}], - modalities=["audio"], # Required for TTS models - audio={ - "voice": "Kore", - "format": "pcm16" # Required: must be "pcm16" - }, - vertex_credentials=vertex_credentials_json -) - -print(response) -``` - - - - -1. Setup config.yaml - -```yaml -model_list: - - model_name: gemini-tts-flash - litellm_params: - model: vertex_ai/gemini-2.5-flash-preview-tts - vertex_project: "your-project-id" - vertex_location: "us-central1" - vertex_credentials: "/path/to/service_account.json" - - model_name: gemini-tts-pro - litellm_params: - model: vertex_ai/gemini-2.5-pro-preview-tts - vertex_project: "your-project-id" - vertex_location: "us-central1" - vertex_credentials: "/path/to/service_account.json" -``` - -2. Start proxy - -```bash -litellm --config /path/to/config.yaml -``` - -3. Make TTS request - -```bash -curl http://0.0.0.0:4000/v1/chat/completions \ - -H "Content-Type: application/json" \ - -H "Authorization: Bearer " \ - -d '{ - "model": "gemini-tts-flash", - "messages": [{"role": "user", "content": "Say hello in a friendly voice"}], - "modalities": ["audio"], - "audio": { - "voice": "Kore", - "format": "pcm16" - } - }' -``` - - - - -### Advanced Usage - -You can combine TTS with other Gemini features: - -```python -response = completion( - model="vertex_ai/gemini-2.5-pro-preview-tts", - messages=[ - {"role": "system", "content": "You are a helpful assistant that speaks clearly."}, - {"role": "user", "content": "Explain quantum computing in simple terms"} - ], - modalities=["audio"], - audio={ - "voice": "Charon", - "format": "pcm16" - }, - temperature=0.7, - max_tokens=150, - vertex_credentials=vertex_credentials_json -) -``` - -For more information about Gemini's TTS capabilities and available voices, see the [official Gemini TTS documentation](https://ai.google.dev/gemini-api/docs/speech-generation). - -## **Text to Speech APIs** - -:::info - -LiteLLM supports calling [Vertex AI Text to Speech API](https://console.cloud.google.com/vertex-ai/generative/speech/text-to-speech) in the OpenAI text to speech API format - -::: - - - -### Usage - Basic - - - - -Vertex AI does not support passing a `model` param - so passing `model=vertex_ai/` is the only required param - -**Sync Usage** - -```python -speech_file_path = Path(__file__).parent / "speech_vertex.mp3" -response = litellm.speech( - model="vertex_ai/", - input="hello what llm guardrail do you have", -) -response.stream_to_file(speech_file_path) -``` - -**Async Usage** -```python -speech_file_path = Path(__file__).parent / "speech_vertex.mp3" -response = litellm.aspeech( - model="vertex_ai/", - input="hello what llm guardrail do you have", -) -response.stream_to_file(speech_file_path) -``` - - - - -1. Add model to config.yaml -```yaml -model_list: - - model_name: vertex-tts - litellm_params: - model: vertex_ai/ # Vertex AI does not support passing a `model` param - so passing `model=vertex_ai/` is the only required param - vertex_project: "adroit-crow-413218" - vertex_location: "us-central1" - vertex_credentials: adroit-crow-413218-a956eef1a2a8.json - -litellm_settings: - drop_params: True -``` - -2. Start Proxy - -``` -$ litellm --config /path/to/config.yaml -``` - -3. Make Request use OpenAI Python SDK - - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -# see supported values for "voice" on vertex here: -# https://console.cloud.google.com/vertex-ai/generative/speech/text-to-speech -response = client.audio.speech.create( - model = "vertex-tts", - input="the quick brown fox jumped over the lazy dogs", - voice={'languageCode': 'en-US', 'name': 'en-US-Studio-O'} -) -print("response from proxy", response) -``` - - - - - -### Usage - `ssml` as input - -Pass your `ssml` as input to the `input` param, if it contains ``, it will be automatically detected and passed as `ssml` to the Vertex AI API - -If you need to force your `input` to be passed as `ssml`, set `use_ssml=True` - - - - -Vertex AI does not support passing a `model` param - so passing `model=vertex_ai/` is the only required param - - -```python -speech_file_path = Path(__file__).parent / "speech_vertex.mp3" - - -ssml = """ - -

Hello, world!

-

This is a test of the text-to-speech API.

- -""" - -response = litellm.speech( - input=ssml, - model="vertex_ai/test", - voice={ - "languageCode": "en-UK", - "name": "en-UK-Studio-O", - }, - audioConfig={ - "audioEncoding": "LINEAR22", - "speakingRate": "10", - }, -) -response.stream_to_file(speech_file_path) -``` - - - - - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -ssml = """ - -

Hello, world!

-

This is a test of the text-to-speech API.

- -""" - -# see supported values for "voice" on vertex here: -# https://console.cloud.google.com/vertex-ai/generative/speech/text-to-speech -response = client.audio.speech.create( - model = "vertex-tts", - input=ssml, - voice={'languageCode': 'en-US', 'name': 'en-US-Studio-O'}, -) -print("response from proxy", response) -``` - - - - - -### Forcing SSML Usage - -You can force the use of SSML by setting the `use_ssml` parameter to `True`. This is useful when you want to ensure that your input is treated as SSML, even if it doesn't contain the `` tags. - -Here are examples of how to force SSML usage: - - - - - -Vertex AI does not support passing a `model` param - so passing `model=vertex_ai/` is the only required param - - -```python -speech_file_path = Path(__file__).parent / "speech_vertex.mp3" - - -ssml = """ - -

Hello, world!

-

This is a test of the text-to-speech API.

- -""" - -response = litellm.speech( - input=ssml, - use_ssml=True, - model="vertex_ai/test", - voice={ - "languageCode": "en-UK", - "name": "en-UK-Studio-O", - }, - audioConfig={ - "audioEncoding": "LINEAR22", - "speakingRate": "10", - }, -) -response.stream_to_file(speech_file_path) -``` - - - - - -```python -import openai - -client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") - -ssml = """ - -

Hello, world!

-

This is a test of the text-to-speech API.

- -""" - -# see supported values for "voice" on vertex here: -# https://console.cloud.google.com/vertex-ai/generative/speech/text-to-speech -response = client.audio.speech.create( - model = "vertex-tts", - input=ssml, # pass as None since OpenAI SDK requires this param - voice={'languageCode': 'en-US', 'name': 'en-US-Studio-O'}, - extra_body={"use_ssml": True}, -) -print("response from proxy", response) -``` - - - - ## **Fine Tuning APIs** diff --git a/docs/my-website/docs/providers/vertex_speech.md b/docs/my-website/docs/providers/vertex_speech.md new file mode 100644 index 00000000000..d0acacb5aec --- /dev/null +++ b/docs/my-website/docs/providers/vertex_speech.md @@ -0,0 +1,423 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Vertex AI Text to Speech + +| Property | Details | +|-------|-------| +| Description | Google Cloud Text-to-Speech with Chirp3 HD voices and Gemini TTS | +| Provider Route on LiteLLM | `vertex_ai/chirp` (Chirp), `vertex_ai/gemini-*-tts` (Gemini) | + +## Chirp3 HD Voices + +Google Cloud Text-to-Speech API with high-quality Chirp3 HD voices. + +### Quick Start + +#### LiteLLM Python SDK + +```python showLineNumbers title="Chirp3 Quick Start" +from litellm import speech +from pathlib import Path + +speech_file_path = Path(__file__).parent / "speech.mp3" +response = speech( + model="vertex_ai/chirp", + voice="alloy", # OpenAI voice name - automatically mapped + input="Hello, this is Vertex AI Text to Speech", + vertex_project="your-project-id", + vertex_location="us-central1", +) +response.stream_to_file(speech_file_path) +``` + +#### LiteLLM AI Gateway + +**1. Setup config.yaml** + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: vertex-tts + litellm_params: + model: vertex_ai/chirp + vertex_project: "your-project-id" + vertex_location: "us-central1" + vertex_credentials: "/path/to/service_account.json" +``` + +**2. Start the proxy** + +```bash title="Start LiteLLM Proxy" +litellm --config /path/to/config.yaml +``` + +**3. Make requests** + + + + +```bash showLineNumbers title="Chirp3 Quick Start" +curl http://0.0.0.0:4000/v1/audio/speech \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "vertex-tts", + "voice": "alloy", + "input": "Hello, this is Vertex AI Text to Speech" + }' \ + --output speech.mp3 +``` + + + + +```python showLineNumbers title="Chirp3 Quick Start" +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.audio.speech.create( + model="vertex-tts", + voice="alloy", + input="Hello, this is Vertex AI Text to Speech", +) +response.stream_to_file("speech.mp3") +``` + + + + +### Voice Mapping + +LiteLLM maps OpenAI voice names to Google Cloud voices. You can use either OpenAI voices or Google Cloud voices directly. + +| OpenAI Voice | Google Cloud Voice | +|-------------|-------------------| +| `alloy` | en-US-Studio-O | +| `echo` | en-US-Studio-M | +| `fable` | en-GB-Studio-B | +| `onyx` | en-US-Wavenet-D | +| `nova` | en-US-Studio-O | +| `shimmer` | en-US-Wavenet-F | + +### Using Google Cloud Voices Directly + +#### LiteLLM Python SDK + +```python showLineNumbers title="Chirp3 HD Voice" +from litellm import speech + +# Pass Chirp3 HD voice name directly +response = speech( + model="vertex_ai/chirp", + voice="en-US-Chirp3-HD-Charon", + input="Hello with a Chirp3 HD voice", + vertex_project="your-project-id", +) +response.stream_to_file("speech.mp3") +``` + +```python showLineNumbers title="Voice as Dict (Multilingual)" +from litellm import speech + +# Pass as dict for full control over language and voice +response = speech( + model="vertex_ai/chirp", + voice={ + "languageCode": "de-DE", + "name": "de-DE-Chirp3-HD-Charon", + }, + input="Hallo, dies ist ein Test", + vertex_project="your-project-id", +) +response.stream_to_file("speech.mp3") +``` + +#### LiteLLM AI Gateway + + + + +```bash showLineNumbers title="Chirp3 HD Voice" +curl http://0.0.0.0:4000/v1/audio/speech \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "vertex-tts", + "voice": "en-US-Chirp3-HD-Charon", + "input": "Hello with a Chirp3 HD voice" + }' \ + --output speech.mp3 +``` + +```bash showLineNumbers title="Voice as Dict (Multilingual)" +curl http://0.0.0.0:4000/v1/audio/speech \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "vertex-tts", + "voice": {"languageCode": "de-DE", "name": "de-DE-Chirp3-HD-Charon"}, + "input": "Hallo, dies ist ein Test" + }' \ + --output speech.mp3 +``` + + + + +```python showLineNumbers title="Chirp3 HD Voice" +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.audio.speech.create( + model="vertex-tts", + voice="en-US-Chirp3-HD-Charon", + input="Hello with a Chirp3 HD voice", +) +response.stream_to_file("speech.mp3") +``` + +```python showLineNumbers title="Voice as Dict (Multilingual)" +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.audio.speech.create( + model="vertex-tts", + voice={"languageCode": "de-DE", "name": "de-DE-Chirp3-HD-Charon"}, + input="Hallo, dies ist ein Test", +) +response.stream_to_file("speech.mp3") +``` + + + + +Browse available voices: [Google Cloud Text-to-Speech Console](https://console.cloud.google.com/vertex-ai/generative/speech/text-to-speech) + +### Passing Raw SSML + +LiteLLM auto-detects SSML when your input contains `` tags and passes it through unchanged. + +#### LiteLLM Python SDK + +```python showLineNumbers title="SSML Input" +from litellm import speech + +ssml = """ + +

Hello, world!

+

This is a test of the text-to-speech API.

+ +""" + +response = speech( + model="vertex_ai/chirp", + voice="en-US-Studio-O", + input=ssml, # Auto-detected as SSML + vertex_project="your-project-id", +) +response.stream_to_file("speech.mp3") +``` + +```python showLineNumbers title="Force SSML Mode" +from litellm import speech + +# Force SSML mode with use_ssml=True +response = speech( + model="vertex_ai/chirp", + voice="en-US-Studio-O", + input="Speaking slowly", + use_ssml=True, + vertex_project="your-project-id", +) +response.stream_to_file("speech.mp3") +``` + +#### LiteLLM AI Gateway + + + + +```bash showLineNumbers title="SSML Input" +curl http://0.0.0.0:4000/v1/audio/speech \ + -H "Authorization: Bearer sk-1234" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "vertex-tts", + "voice": "en-US-Studio-O", + "input": "

Hello!

How are you?

" + }' \ + --output speech.mp3 +``` + + + + +```python showLineNumbers title="SSML Input" +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +ssml = """

Hello!

How are you?

""" + +response = client.audio.speech.create( + model="vertex-tts", + voice="en-US-Studio-O", + input=ssml, +) +response.stream_to_file("speech.mp3") +``` + + + + +### Supported Parameters + +| Parameter | Description | Values | +|-----------|-------------|--------| +| `voice` | Voice selection | OpenAI voice, Google Cloud voice name, or dict | +| `input` | Text to convert | Plain text or SSML | +| `speed` | Speaking rate | 0.25 to 4.0 (default: 1.0) | +| `response_format` | Audio format | `mp3`, `opus`, `wav`, `pcm`, `flac` | +| `use_ssml` | Force SSML mode | `True` / `False` | + +### Async Usage + +```python showLineNumbers title="Async Speech Generation" +import asyncio +from litellm import aspeech + +async def main(): + response = await aspeech( + model="vertex_ai/chirp", + voice="alloy", + input="Hello from async", + vertex_project="your-project-id", + ) + response.stream_to_file("speech.mp3") + +asyncio.run(main()) +``` + +--- + +## Gemini TTS + +Gemini models with audio output capabilities using the chat completions API. + +:::warning +**Limitations:** +- Only supports `pcm16` audio format +- Streaming not yet supported +- Must set `modalities: ["audio"]` +::: + +### Quick Start + +#### LiteLLM Python SDK + +```python showLineNumbers title="Gemini TTS Quick Start" +from litellm import completion +import json + +# Load credentials +with open('path/to/service_account.json', 'r') as file: + vertex_credentials = json.dumps(json.load(file)) + +response = completion( + model="vertex_ai/gemini-2.5-flash-preview-tts", + messages=[{"role": "user", "content": "Say hello in a friendly voice"}], + modalities=["audio"], + audio={ + "voice": "Kore", + "format": "pcm16" + }, + vertex_credentials=vertex_credentials +) +print(response) +``` + +#### LiteLLM AI Gateway + +**1. Setup config.yaml** + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: gemini-tts + litellm_params: + model: vertex_ai/gemini-2.5-flash-preview-tts + vertex_project: "your-project-id" + vertex_location: "us-central1" + vertex_credentials: "/path/to/service_account.json" +``` + +**2. Start the proxy** + +```bash title="Start LiteLLM Proxy" +litellm --config /path/to/config.yaml +``` + +**3. Make requests** + + + + +```bash showLineNumbers title="Gemini TTS Request" +curl http://0.0.0.0:4000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "gemini-tts", + "messages": [{"role": "user", "content": "Say hello in a friendly voice"}], + "modalities": ["audio"], + "audio": {"voice": "Kore", "format": "pcm16"} + }' +``` + + + + +```python showLineNumbers title="Gemini TTS Request" +import openai + +client = openai.OpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000") + +response = client.chat.completions.create( + model="gemini-tts", + messages=[{"role": "user", "content": "Say hello in a friendly voice"}], + modalities=["audio"], + audio={"voice": "Kore", "format": "pcm16"}, +) +print(response) +``` + + + + +### Supported Models + +- `vertex_ai/gemini-2.5-flash-preview-tts` +- `vertex_ai/gemini-2.5-pro-preview-tts` + +See [Gemini TTS documentation](https://ai.google.dev/gemini-api/docs/speech-generation) for available voices. + +### Advanced Usage + +```python showLineNumbers title="Gemini TTS with System Prompt" +from litellm import completion + +response = completion( + model="vertex_ai/gemini-2.5-pro-preview-tts", + messages=[ + {"role": "system", "content": "You are a helpful assistant that speaks clearly."}, + {"role": "user", "content": "Explain quantum computing in simple terms"} + ], + modalities=["audio"], + audio={"voice": "Charon", "format": "pcm16"}, + temperature=0.7, + max_tokens=150, + vertex_credentials=vertex_credentials +) +``` diff --git a/docs/my-website/docs/proxy/config_settings.md b/docs/my-website/docs/proxy/config_settings.md index 7140c99e6fb..b71e100e157 100644 --- a/docs/my-website/docs/proxy/config_settings.md +++ b/docs/my-website/docs/proxy/config_settings.md @@ -113,7 +113,7 @@ general_settings: # Database Settings database_url: string - database_connection_pool_limit: 0 # default 100 + database_connection_pool_limit: 0 # default 10 database_connection_timeout: 0 # default 60s allow_requests_on_db_unavailable: boolean # if true, will allow requests that can not connect to the DB to verify Virtual Key to still work @@ -234,7 +234,7 @@ router_settings: | max_response_size_mb | int | The maximum size for responses in MB. LLM Responses above this size will not be sent. | | proxy_budget_rescheduler_min_time | int | The minimum time (in seconds) to wait before checking db for budget resets. **Default is 597 seconds** | | proxy_budget_rescheduler_max_time | int | The maximum time (in seconds) to wait before checking db for budget resets. **Default is 605 seconds** | -| proxy_batch_write_at | int | Time (in seconds) to wait before batch writing spend logs to the db. **Default is 30 seconds** | +| proxy_batch_write_at | int | Time (in seconds) to wait before batch writing spend logs to the db. **Default is 10 seconds** | | proxy_batch_polling_interval | int | Time (in seconds) to wait before polling a batch, to check if it's completed. **Default is 6000 seconds (1 hour)** | | alerting_args | dict | Args for Slack Alerting [Doc on Slack Alerting](./alerting.md) | | custom_key_generate | str | Custom function for key generation [Doc on custom key generation](./virtual_keys.md#custom--key-generate) | @@ -377,6 +377,7 @@ router_settings: | ATHINA_API_KEY | API key for Athina service | ATHINA_BASE_URL | Base URL for Athina service (defaults to `https://log.athina.ai`) | AUTH_STRATEGY | Strategy used for authentication (e.g., OAuth, API key) +| AUTO_REDIRECT_UI_LOGIN_TO_SSO | Flag to enable automatic redirect of UI login page to SSO when SSO is configured. Default is **true** | ANTHROPIC_API_KEY | API key for Anthropic service | ANTHROPIC_API_BASE | Base URL for Anthropic API. Default is https://api.anthropic.com | AWS_ACCESS_KEY_ID | Access Key ID for AWS services @@ -763,7 +764,7 @@ router_settings: | PROMPTLAYER_API_KEY | API key for PromptLayer integration | PROXY_ADMIN_ID | Admin identifier for proxy server | PROXY_BASE_URL | Base URL for proxy service -| PROXY_BATCH_WRITE_AT | Time in seconds to wait before batch writing spend logs to the database. Default is 30 +| PROXY_BATCH_WRITE_AT | Time in seconds to wait before batch writing spend logs to the database. Default is 10 | PROXY_BATCH_POLLING_INTERVAL | Time in seconds to wait before polling a batch, to check if it's completed. Default is 6000s (1 hour) | PROXY_BUDGET_RESCHEDULER_MAX_TIME | Maximum time in seconds to wait before checking database for budget resets. Default is 605 | PROXY_BUDGET_RESCHEDULER_MIN_TIME | Minimum time in seconds to wait before checking database for budget resets. Default is 597 diff --git a/docs/my-website/docs/proxy/configs.md b/docs/my-website/docs/proxy/configs.md index 18177b7c4d2..77ab3158f74 100644 --- a/docs/my-website/docs/proxy/configs.md +++ b/docs/my-website/docs/proxy/configs.md @@ -576,7 +576,7 @@ custom_tokenizer: ```yaml general_settings: - database_connection_pool_limit: 100 # sets connection pool for prisma client to postgres db at 100 + database_connection_pool_limit: 10 # sets connection pool for prisma client to postgres db (default: 10, recommended: 10-20) database_connection_timeout: 60 # sets a 60s timeout for any connection call to the db ``` diff --git a/docs/my-website/docs/proxy/error_diagnosis.md b/docs/my-website/docs/proxy/error_diagnosis.md new file mode 100644 index 00000000000..9629fc52b0c --- /dev/null +++ b/docs/my-website/docs/proxy/error_diagnosis.md @@ -0,0 +1,90 @@ +# Diagnosing Errors - Provider vs Gateway + +Having trouble diagnosing if an error is from the **LLM Provider** (OpenAI, Anthropic, etc.) or from the **LiteLLM AI Gateway** itself? Here's how to tell. + +## Quick Rule + +**If the error contains `Exception`, it's from the provider.** + +| Error Contains | Error Source | +|----------------|--------------| +| `AnthropicException` | Anthropic | +| `OpenAIException` | OpenAI | +| `AzureException` | Azure | +| `BedrockException` | AWS Bedrock | +| `VertexAIException` | Google Vertex AI | +| No provider name | LiteLLM AI Gateway | + +## Examples + +### Provider Error (from AWS Bedrock) + +``` +{ + "error": { + "message": "litellm.BadRequestError: BedrockException - {\"message\":\"The model returned the following errors: messages.1.content.0.type: Expected `thinking` or `redacted_thinking`, but found `text`.\"}", + "type": "invalid_request_error", + "param": null, + "code": "400" + } +} +``` + +This error is from **AWS Bedrock** (notice `BedrockException`). The Bedrock API is rejecting the request due to invalid message format - this is not a LiteLLM issue. + +### Provider Error (from OpenAI) + +``` +{ + "error": { + "message": "litellm.AuthenticationError: OpenAIException - Incorrect API key provided: . You can find your API key at https://platform.openai.com/account/api-keys.", + "type": "invalid_request_error", + "param": null, + "code": "invalid_api_key" + } +} +``` + +This error is from **OpenAI** (notice `OpenAIException`). The OpenAI API key configured in LiteLLM is invalid. + +### Provider Error (from Anthropic) + +``` +{ + "error": { + "message": "litellm.InternalServerError: AnthropicException - Overloaded. Handle with `litellm.InternalServerError`.", + "type": "internal_server_error", + "param": null, + "code": "500" + } +} +``` + +This error is from **Anthropic** (notice `AnthropicException`). The Anthropic API is overloaded - this is not a LiteLLM issue. + +### Gateway Error (from LiteLLM) + +``` +{ + "error": { + "message": "Invalid API Key. Please check your LiteLLM API key.", + "type": "auth_error", + "param": null, + "code": "401" + } +} +``` + +This error is from the **LiteLLM AI Gateway** (no provider name). Your LiteLLM virtual key is invalid. + +## What to do? + +| Error Source | Action | +|--------------|--------| +| Provider Error | Check the provider's status page, adjust rate limits, or retry later | +| Gateway Error | Check your LiteLLM configuration, API keys, or [open an issue](https://github.com/BerriAI/litellm/issues) | + +## See Also + +- [Debugging](/docs/proxy/debugging) - Enable debug logs to see detailed request/response info +- [Exception Mapping](/docs/exception_mapping) - Full list of LiteLLM exception types diff --git a/docs/my-website/docs/proxy/pass_through.md b/docs/my-website/docs/proxy/pass_through.md index 7309cdeda26..03454004b8c 100644 --- a/docs/my-website/docs/proxy/pass_through.md +++ b/docs/my-website/docs/proxy/pass_through.md @@ -275,6 +275,20 @@ In this video, we'll add the Azure OpenAI Assistants API as a pass through endpo - Check LiteLLM proxy logs for error details - Verify the target API's expected request format +### Allowing Team JWTs to use pass-through routes + +If you are using pass-through provider routes (e.g., `/anthropic/*`) and want your JWT team tokens to access these routes, add `mapped_pass_through_routes` to the `team_allowed_routes` in `litellm_jwtauth` or explicitly add the relevant route(s). + +Example (`proxy_server_config.yaml`): + +```yaml +general_settings: + enable_jwt_auth: True + litellm_jwtauth: + team_ids_jwt_field: "team_ids" + team_allowed_routes: ["openai_routes","info_routes","mapped_pass_through_routes"] +``` + ### Getting Help [Schedule Demo 👋](https://calendly.com/d/4mp-gd3-k5k/berriai-1-1-onboarding-litellm-hosted-version) diff --git a/docs/my-website/docs/proxy/token_auth.md b/docs/my-website/docs/proxy/token_auth.md index c465c1022e4..1db1b2a8965 100644 --- a/docs/my-website/docs/proxy/token_auth.md +++ b/docs/my-website/docs/proxy/token_auth.md @@ -338,6 +338,58 @@ general_settings: team_allowed_routes: ["/v1/chat/completions"] # 👈 Set accepted routes ``` +### Allowing other provider routes for Teams + +To enable team JWT tokens to access Anthropic-style endpoints such as `/v1/messages`, update `team_allowed_routes` in your `litellm_jwtauth` configuration. `team_allowed_routes` supports the following values: + +- Named route groups from `LiteLLMRoutes` (e.g., `openai_routes`, `anthropic_routes`, `info_routes`, `mapped_pass_through_routes`). + +Below is a quick reference for the route groups you can use and example representative routes from each group. If you need the exhaustive list, see the `LiteLLMRoutes` enum in `litellm/proxy/_types.py` for the authoritative list. + +| Route Group | What it contains | Representative routes | +|-------------|------------------|-----------------------| +| `openai_routes` | OpenAI-compatible REST endpoints (chat, completion, embeddings, images, responses, models, etc.) | `/v1/chat/completions`, `/v1/completions`, `/v1/embeddings`, `/v1/images/generations`, `/v1/models` | +| `anthropic_routes` | Anthropic-style endpoints (`/v1/messages` and related) | `/v1/messages`, `/v1/messages/count_tokens`, `/v1/skills` | +| `mapped_pass_through_routes` | Provider-specific pass-through route prefixes (e.g., Anthropic when proxied via `/anthropic`). Use with `mapped_pass_through_routes` for provider wildcard mapping | `/anthropic/*`, `/vertex-ai/*`, `/bedrock/*` | +| `passthrough_routes_wildcard` | Wildcard mapping for providers (e.g., `/anthropic/*`) - precomputed wildcard list used by the proxy | `/anthropic/*`, `/vllm/*` | +| `google_routes` | Google-specific (e.g., Vertex / Batching endpoints) | `/v1beta/models/{model_name}:generateContent` | +| `mcp_routes` | Internal MCP management endpoints | `/mcp/tools`, `/mcp/tools/call` | +| `info_routes` | Read-only & info endpoints used by the UI | `/key/info`, `/team/info`, `/v1/models` | +| `management_routes` | Admin-only management endpoints (create/update/delete user/team/model) | `/team/new`, `/key/generate`, `/model/new` | +| `spend_tracking_routes` | Budget/spend related endpoints | `/spend/logs`, `/spend/keys` | +| `public_routes` | Public and unauthenticated endpoints | `/`, `/routes`, `/.well-known/litellm-ui-config` | + +Note: `llm_api_routes` is the union of OpenAI, Anthropic, Google, pass-through and other LLM routes (`openai_routes + anthropic_routes + google_routes + mapped_pass_through_routes + passthrough_routes_wildcard + apply_guardrail_routes + mcp_routes + litellm_native_routes`). + +Defaults (what the proxy uses if you don't override them in `litellm_jwtauth`): + +- `admin_jwt_scope`: `litellm_proxy_admin` +- `admin_allowed_routes` (default): `management_routes`, `spend_tracking_routes`, `global_spend_tracking_routes`, `info_routes` +- `team_allowed_routes` (default): `openai_routes`, `info_routes` +- `public_allowed_routes` (default): `public_routes` + + +Example: Allow team JWTs to call Anthropic `/v1/messages` (either by route group or by explicit route string): + +```yaml +general_settings: + enable_jwt_auth: True + litellm_jwtauth: + team_ids_jwt_field: "team_ids" + team_allowed_routes: ["openai_routes", "info_routes", "anthropic_routes"] +``` + +Or selectively allow the exact Anthropic message endpoint only: + +```yaml +general_settings: + enable_jwt_auth: True + litellm_jwtauth: + team_ids_jwt_field: "team_ids" + team_allowed_routes: ["/v1/messages", "info_routes"] +``` + + ### Caching Public Keys Control how long public keys are cached for (in seconds). diff --git a/docs/my-website/docs/secret_managers/cyberark.md b/docs/my-website/docs/secret_managers/cyberark.md index 37aa1086691..c33aa286703 100644 --- a/docs/my-website/docs/secret_managers/cyberark.md +++ b/docs/my-website/docs/secret_managers/cyberark.md @@ -41,6 +41,7 @@ CYBERARK_CLIENT_KEY="path/to/client.key" # OPTIONAL CYBERARK_REFRESH_INTERVAL="300" # defaults to 300 seconds (5 minutes), frequency of token refresh +CYBERARK_SSL_VERIFY="true" # defaults to true, set to "false" to disable SSL verification (for self-signed certificates) ``` **Step 2.** Add to proxy config.yaml @@ -172,6 +173,24 @@ If these commands work successfully against your CyberArk instance, then CyberAr - The `CYBERARK_API_BASE` URL is accessible from your LiteLLM instance - Your API key or certificates have the necessary permissions in CyberArk +### SSL Certificate Errors + +If you encounter SSL certificate verification errors like: + +``` +RuntimeError: Could not authenticate to CyberArk Conjur: [SSL: CERTIFICATE_VERIFY_FAILED] certificate verify failed: self-signed certificate in certificate chain +``` + +This typically occurs when your CyberArk Conjur instance uses a self-signed certificate. You can disable SSL verification by setting: + +```bash +CYBERARK_SSL_VERIFY="false" +``` + +:::warning +Disabling SSL verification is insecure and should only be used for testing or development environments with self-signed certificates. For production, configure your certificate chain properly or use certificate-based authentication with `CYBERARK_CLIENT_CERT` and `CYBERARK_CLIENT_KEY`. +::: + ## Video Walkthrough This video walks through using CyberArk Conjur as a secret manager with LiteLLM. We create a virtual key in the LiteLLM Admin UI and verify it exists in CyberArk. Then we rotate the secret key and verify it exists in CyberArk. diff --git a/docs/my-website/docs/vector_stores/create.md b/docs/my-website/docs/vector_stores/create.md index 19b4f39cd9e..7025c490a32 100644 --- a/docs/my-website/docs/vector_stores/create.md +++ b/docs/my-website/docs/vector_stores/create.md @@ -14,6 +14,7 @@ Create a vector store which can be used to store and search document chunks for | End-user Tracking | ✅ | | | Support LLM Providers (OpenAI `/vector_stores` API) | **OpenAI** | Full vector stores API support across providers | | Support LLM Providers (Passthrough API) | [**Azure AI**](/docs/providers/azure_ai/azure_ai_vector_stores_passthrough) | Full vector stores API support across providers | +| Support LLM Providers (Dataset Management) | [**RAGFlow**](/docs/providers/ragflow_vector_store.md) | Dataset creation and management (search not supported) | ## Usage diff --git a/docs/my-website/img/add_agent1.png b/docs/my-website/img/add_agent1.png new file mode 100644 index 00000000000..e69de29bb2d diff --git a/docs/my-website/img/add_agent_1.png b/docs/my-website/img/add_agent_1.png new file mode 100644 index 00000000000..e60435996a9 Binary files /dev/null and b/docs/my-website/img/add_agent_1.png differ diff --git a/docs/my-website/img/agent2.png b/docs/my-website/img/agent2.png new file mode 100644 index 00000000000..412047a6aa3 Binary files /dev/null and b/docs/my-website/img/agent2.png differ diff --git a/docs/my-website/release_notes/v1.80.5-stable/index.md b/docs/my-website/release_notes/v1.80.5-stable/index.md index 4324fdef776..7c3283ce349 100644 --- a/docs/my-website/release_notes/v1.80.5-stable/index.md +++ b/docs/my-website/release_notes/v1.80.5-stable/index.md @@ -1,5 +1,5 @@ --- -title: "[PREVIEW] v1.80.5.rc.2 - Gemini 3.0 Support" +title: "v1.80.5-stable - Gemini 3.0 Support" slug: "v1-80-5" date: 2025-11-22T10:00:00 authors: @@ -27,7 +27,7 @@ import TabItem from '@theme/TabItem'; docker run \ -e STORE_MODEL_IN_DB=True \ -p 4000:4000 \ -ghcr.io/berriai/litellm:v1.80.5.rc.2 +ghcr.io/berriai/litellm:v1.80.5-stable ``` diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index e467711b59d..cb8f3be0344 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -141,6 +141,7 @@ const sidebars = { "proxy/quick_start", "proxy/cli", "proxy/debugging", + "proxy/error_diagnosis", "proxy/deploy", "proxy/health", "proxy/master_key_rotations", @@ -315,6 +316,7 @@ const sidebars = { slug: "/supported_endpoints", }, items: [ + "a2a", "assistants", { type: "category", @@ -520,6 +522,7 @@ const sidebars = { "providers/vertex_partner", "providers/vertex_self_deployed", "providers/vertex_image", + "providers/vertex_speech", "providers/vertex_batch", "providers/vertex_ocr", ] @@ -625,6 +628,7 @@ const sidebars = { "providers/petals", "providers/publicai", "providers/predibase", + "providers/ragflow", "providers/recraft", "providers/replicate", { @@ -818,10 +822,13 @@ const sidebars = { "Learn how to deploy + call models from different providers on LiteLLM", slug: "/project", }, - items: [ + items: [ "projects/smolagents", "projects/mini-swe-agent", "projects/openai-agents", + "projects/Google ADK", + "projects/Agent Lightning", + "projects/Harbor", "projects/Docq.AI", "projects/PDL", "projects/OpenInterpreter", diff --git a/docs/my-website/src/pages/intro.md b/docs/my-website/src/pages/intro.md deleted file mode 100644 index 8a2e69d95f9..00000000000 --- a/docs/my-website/src/pages/intro.md +++ /dev/null @@ -1,47 +0,0 @@ ---- -sidebar_position: 1 ---- - -# Tutorial Intro - -Let's discover **Docusaurus in less than 5 minutes**. - -## Getting Started - -Get started by **creating a new site**. - -Or **try Docusaurus immediately** with **[docusaurus.new](https://docusaurus.new)**. - -### What you'll need - -- [Node.js](https://nodejs.org/en/download/) version 16.14 or above: - - When installing Node.js, you are recommended to check all checkboxes related to dependencies. - -## Generate a new site - -Generate a new Docusaurus site using the **classic template**. - -The classic template will automatically be added to your project after you run the command: - -```bash -npm init docusaurus@latest my-website classic -``` - -You can type this command into Command Prompt, Powershell, Terminal, or any other integrated terminal of your code editor. - -The command also installs all necessary dependencies you need to run Docusaurus. - -## Start your site - -Run the development server: - -```bash -cd my-website -npm run start -``` - -The `cd` command changes the directory you're working with. In order to work with your newly created Docusaurus site, you'll need to navigate the terminal there. - -The `npm run start` command builds your website locally and serves it through a development server, ready for you to view at http://localhost:3000/. - -Open `docs/intro.md` (this page) and edit some lines: the site **reloads automatically** and displays your changes. diff --git a/docs/my-website/src/pages/tutorial-basics/_category_.json b/docs/my-website/src/pages/tutorial-basics/_category_.json deleted file mode 100644 index 2e6db55b1eb..00000000000 --- a/docs/my-website/src/pages/tutorial-basics/_category_.json +++ /dev/null @@ -1,8 +0,0 @@ -{ - "label": "Tutorial - Basics", - "position": 2, - "link": { - "type": "generated-index", - "description": "5 minutes to learn the most important Docusaurus concepts." - } -} diff --git a/docs/my-website/src/pages/tutorial-basics/congratulations.md b/docs/my-website/src/pages/tutorial-basics/congratulations.md deleted file mode 100644 index 04771a00b72..00000000000 --- a/docs/my-website/src/pages/tutorial-basics/congratulations.md +++ /dev/null @@ -1,23 +0,0 @@ ---- -sidebar_position: 6 ---- - -# Congratulations! - -You have just learned the **basics of Docusaurus** and made some changes to the **initial template**. - -Docusaurus has **much more to offer**! - -Have **5 more minutes**? Take a look at **[versioning](../tutorial-extras/manage-docs-versions.md)** and **[i18n](../tutorial-extras/translate-your-site.md)**. - -Anything **unclear** or **buggy** in this tutorial? [Please report it!](https://github.com/facebook/docusaurus/discussions/4610) - -## What's next? - -- Read the [official documentation](https://docusaurus.io/) -- Modify your site configuration with [`docusaurus.config.js`](https://docusaurus.io/docs/api/docusaurus-config) -- Add navbar and footer items with [`themeConfig`](https://docusaurus.io/docs/api/themes/configuration) -- Add a custom [Design and Layout](https://docusaurus.io/docs/styling-layout) -- Add a [search bar](https://docusaurus.io/docs/search) -- Find inspirations in the [Docusaurus showcase](https://docusaurus.io/showcase) -- Get involved in the [Docusaurus Community](https://docusaurus.io/community/support) diff --git a/docs/my-website/src/pages/tutorial-basics/create-a-blog-post.md b/docs/my-website/src/pages/tutorial-basics/create-a-blog-post.md deleted file mode 100644 index ea472bbaf87..00000000000 --- a/docs/my-website/src/pages/tutorial-basics/create-a-blog-post.md +++ /dev/null @@ -1,34 +0,0 @@ ---- -sidebar_position: 3 ---- - -# Create a Blog Post - -Docusaurus creates a **page for each blog post**, but also a **blog index page**, a **tag system**, an **RSS** feed... - -## Create your first Post - -Create a file at `blog/2021-02-28-greetings.md`: - -```md title="blog/2021-02-28-greetings.md" ---- -slug: greetings -title: Greetings! -authors: - - name: Joel Marcey - title: Co-creator of Docusaurus 1 - url: https://github.com/JoelMarcey - image_url: https://github.com/JoelMarcey.png - - name: Sébastien Lorber - title: Docusaurus maintainer - url: https://sebastienlorber.com - image_url: https://github.com/slorber.png -tags: [greetings] ---- - -Congratulations, you have made your first post! - -Feel free to play around and edit this post as much you like. -``` - -A new blog post is now available at [http://localhost:3000/blog/greetings](http://localhost:3000/blog/greetings). diff --git a/docs/my-website/src/pages/tutorial-basics/create-a-document.md b/docs/my-website/src/pages/tutorial-basics/create-a-document.md deleted file mode 100644 index ffddfa8eb8a..00000000000 --- a/docs/my-website/src/pages/tutorial-basics/create-a-document.md +++ /dev/null @@ -1,57 +0,0 @@ ---- -sidebar_position: 2 ---- - -# Create a Document - -Documents are **groups of pages** connected through: - -- a **sidebar** -- **previous/next navigation** -- **versioning** - -## Create your first Doc - -Create a Markdown file at `docs/hello.md`: - -```md title="docs/hello.md" -# Hello - -This is my **first Docusaurus document**! -``` - -A new document is now available at [http://localhost:3000/docs/hello](http://localhost:3000/docs/hello). - -## Configure the Sidebar - -Docusaurus automatically **creates a sidebar** from the `docs` folder. - -Add metadata to customize the sidebar label and position: - -```md title="docs/hello.md" {1-4} ---- -sidebar_label: 'Hi!' -sidebar_position: 3 ---- - -# Hello - -This is my **first Docusaurus document**! -``` - -It is also possible to create your sidebar explicitly in `sidebars.js`: - -```js title="sidebars.js" -module.exports = { - tutorialSidebar: [ - 'intro', - // highlight-next-line - 'hello', - { - type: 'category', - label: 'Tutorial', - items: ['tutorial-basics/create-a-document'], - }, - ], -}; -``` diff --git a/docs/my-website/src/pages/tutorial-basics/create-a-page.md b/docs/my-website/src/pages/tutorial-basics/create-a-page.md deleted file mode 100644 index 20e2ac30055..00000000000 --- a/docs/my-website/src/pages/tutorial-basics/create-a-page.md +++ /dev/null @@ -1,43 +0,0 @@ ---- -sidebar_position: 1 ---- - -# Create a Page - -Add **Markdown or React** files to `src/pages` to create a **standalone page**: - -- `src/pages/index.js` → `localhost:3000/` -- `src/pages/foo.md` → `localhost:3000/foo` -- `src/pages/foo/bar.js` → `localhost:3000/foo/bar` - -## Create your first React Page - -Create a file at `src/pages/my-react-page.js`: - -```jsx title="src/pages/my-react-page.js" -import React from 'react'; -import Layout from '@theme/Layout'; - -export default function MyReactPage() { - return ( - -

My React page

-

This is a React page

- - ); -} -``` - -A new page is now available at [http://localhost:3000/my-react-page](http://localhost:3000/my-react-page). - -## Create your first Markdown Page - -Create a file at `src/pages/my-markdown-page.md`: - -```mdx title="src/pages/my-markdown-page.md" -# My Markdown page - -This is a Markdown page -``` - -A new page is now available at [http://localhost:3000/my-markdown-page](http://localhost:3000/my-markdown-page). diff --git a/docs/my-website/src/pages/tutorial-basics/deploy-your-site.md b/docs/my-website/src/pages/tutorial-basics/deploy-your-site.md deleted file mode 100644 index 1c50ee063ef..00000000000 --- a/docs/my-website/src/pages/tutorial-basics/deploy-your-site.md +++ /dev/null @@ -1,31 +0,0 @@ ---- -sidebar_position: 5 ---- - -# Deploy your site - -Docusaurus is a **static-site-generator** (also called **[Jamstack](https://jamstack.org/)**). - -It builds your site as simple **static HTML, JavaScript and CSS files**. - -## Build your site - -Build your site **for production**: - -```bash -npm run build -``` - -The static files are generated in the `build` folder. - -## Deploy your site - -Test your production build locally: - -```bash -npm run serve -``` - -The `build` folder is now served at [http://localhost:3000/](http://localhost:3000/). - -You can now deploy the `build` folder **almost anywhere** easily, **for free** or very small cost (read the **[Deployment Guide](https://docusaurus.io/docs/deployment)**). diff --git a/docs/my-website/src/pages/tutorial-basics/markdown-features.mdx b/docs/my-website/src/pages/tutorial-basics/markdown-features.mdx deleted file mode 100644 index 0337f34d6a5..00000000000 --- a/docs/my-website/src/pages/tutorial-basics/markdown-features.mdx +++ /dev/null @@ -1,150 +0,0 @@ ---- -sidebar_position: 4 ---- - -# Markdown Features - -Docusaurus supports **[Markdown](https://daringfireball.net/projects/markdown/syntax)** and a few **additional features**. - -## Front Matter - -Markdown documents have metadata at the top called [Front Matter](https://jekyllrb.com/docs/front-matter/): - -```text title="my-doc.md" -// highlight-start ---- -id: my-doc-id -title: My document title -description: My document description -slug: /my-custom-url ---- -// highlight-end - -## Markdown heading - -Markdown text with [links](./hello.md) -``` - -## Links - -Regular Markdown links are supported, using url paths or relative file paths. - -```md -Let's see how to [Create a page](/create-a-page). -``` - -```md -Let's see how to [Create a page](./create-a-page.md). -``` - -**Result:** Let's see how to [Create a page](./create-a-page.md). - -## Images - -Regular Markdown images are supported. - -You can use absolute paths to reference images in the static directory (`static/img/docusaurus.png`): - -```md -![Docusaurus logo](/img/docusaurus.png) -``` - -![Docusaurus logo](/img/docusaurus.png) - -You can reference images relative to the current file as well. This is particularly useful to colocate images close to the Markdown files using them: - -```md -![Docusaurus logo](./img/docusaurus.png) -``` - -## Code Blocks - -Markdown code blocks are supported with Syntax highlighting. - - ```jsx title="src/components/HelloDocusaurus.js" - function HelloDocusaurus() { - return ( -

Hello, Docusaurus!

- ) - } - ``` - -```jsx title="src/components/HelloDocusaurus.js" -function HelloDocusaurus() { - return

Hello, Docusaurus!

; -} -``` - -## Admonitions - -Docusaurus has a special syntax to create admonitions and callouts: - - :::tip My tip - - Use this awesome feature option - - ::: - - :::danger Take care - - This action is dangerous - - ::: - -:::tip My tip - -Use this awesome feature option - -::: - -:::danger Take care - -This action is dangerous - -::: - -## MDX and React Components - -[MDX](https://mdxjs.com/) can make your documentation more **interactive** and allows using any **React components inside Markdown**: - -```jsx -export const Highlight = ({children, color}) => ( - { - alert(`You clicked the color ${color} with label ${children}`) - }}> - {children} - -); - -This is Docusaurus green ! - -This is Facebook blue ! -``` - -export const Highlight = ({children, color}) => ( - { - alert(`You clicked the color ${color} with label ${children}`); - }}> - {children} - -); - -This is Docusaurus green ! - -This is Facebook blue ! diff --git a/docs/my-website/src/pages/tutorial-extras/_category_.json b/docs/my-website/src/pages/tutorial-extras/_category_.json deleted file mode 100644 index a8ffcc19300..00000000000 --- a/docs/my-website/src/pages/tutorial-extras/_category_.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "label": "Tutorial - Extras", - "position": 3, - "link": { - "type": "generated-index" - } -} diff --git a/docs/my-website/src/pages/tutorial-extras/img/docsVersionDropdown.png b/docs/my-website/src/pages/tutorial-extras/img/docsVersionDropdown.png deleted file mode 100644 index 97e4164618b..00000000000 Binary files a/docs/my-website/src/pages/tutorial-extras/img/docsVersionDropdown.png and /dev/null differ diff --git a/docs/my-website/src/pages/tutorial-extras/img/localeDropdown.png b/docs/my-website/src/pages/tutorial-extras/img/localeDropdown.png deleted file mode 100644 index e257edc1f93..00000000000 Binary files a/docs/my-website/src/pages/tutorial-extras/img/localeDropdown.png and /dev/null differ diff --git a/docs/my-website/src/pages/tutorial-extras/manage-docs-versions.md b/docs/my-website/src/pages/tutorial-extras/manage-docs-versions.md deleted file mode 100644 index e12c3f3444f..00000000000 --- a/docs/my-website/src/pages/tutorial-extras/manage-docs-versions.md +++ /dev/null @@ -1,55 +0,0 @@ ---- -sidebar_position: 1 ---- - -# Manage Docs Versions - -Docusaurus can manage multiple versions of your docs. - -## Create a docs version - -Release a version 1.0 of your project: - -```bash -npm run docusaurus docs:version 1.0 -``` - -The `docs` folder is copied into `versioned_docs/version-1.0` and `versions.json` is created. - -Your docs now have 2 versions: - -- `1.0` at `http://localhost:3000/docs/` for the version 1.0 docs -- `current` at `http://localhost:3000/docs/next/` for the **upcoming, unreleased docs** - -## Add a Version Dropdown - -To navigate seamlessly across versions, add a version dropdown. - -Modify the `docusaurus.config.js` file: - -```js title="docusaurus.config.js" -module.exports = { - themeConfig: { - navbar: { - items: [ - // highlight-start - { - type: 'docsVersionDropdown', - }, - // highlight-end - ], - }, - }, -}; -``` - -The docs version dropdown appears in your navbar: - -![Docs Version Dropdown](./img/docsVersionDropdown.png) - -## Update an existing version - -It is possible to edit versioned docs in their respective folder: - -- `versioned_docs/version-1.0/hello.md` updates `http://localhost:3000/docs/hello` -- `docs/hello.md` updates `http://localhost:3000/docs/next/hello` diff --git a/docs/my-website/src/pages/tutorial-extras/translate-your-site.md b/docs/my-website/src/pages/tutorial-extras/translate-your-site.md deleted file mode 100644 index caeaffb0554..00000000000 --- a/docs/my-website/src/pages/tutorial-extras/translate-your-site.md +++ /dev/null @@ -1,88 +0,0 @@ ---- -sidebar_position: 2 ---- - -# Translate your site - -Let's translate `docs/intro.md` to French. - -## Configure i18n - -Modify `docusaurus.config.js` to add support for the `fr` locale: - -```js title="docusaurus.config.js" -module.exports = { - i18n: { - defaultLocale: 'en', - locales: ['en', 'fr'], - }, -}; -``` - -## Translate a doc - -Copy the `docs/intro.md` file to the `i18n/fr` folder: - -```bash -mkdir -p i18n/fr/docusaurus-plugin-content-docs/current/ - -cp docs/intro.md i18n/fr/docusaurus-plugin-content-docs/current/intro.md -``` - -Translate `i18n/fr/docusaurus-plugin-content-docs/current/intro.md` in French. - -## Start your localized site - -Start your site on the French locale: - -```bash -npm run start -- --locale fr -``` - -Your localized site is accessible at [http://localhost:3000/fr/](http://localhost:3000/fr/) and the `Getting Started` page is translated. - -:::caution - -In development, you can only use one locale at a same time. - -::: - -## Add a Locale Dropdown - -To navigate seamlessly across languages, add a locale dropdown. - -Modify the `docusaurus.config.js` file: - -```js title="docusaurus.config.js" -module.exports = { - themeConfig: { - navbar: { - items: [ - // highlight-start - { - type: 'localeDropdown', - }, - // highlight-end - ], - }, - }, -}; -``` - -The locale dropdown now appears in your navbar: - -![Locale Dropdown](./img/localeDropdown.png) - -## Build your localized site - -Build your site for a specific locale: - -```bash -npm run build -- --locale fr -``` - -Or build your site to include all the locales at once: - -```bash -npm run build -``` diff --git a/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py new file mode 100644 index 00000000000..fdb1dba372f --- /dev/null +++ b/enterprise/litellm_enterprise/proxy/vector_stores/endpoints.py @@ -0,0 +1,337 @@ +""" +VECTOR STORE MANAGEMENT + +All /vector_store management endpoints + +/vector_store/new +/vector_store/delete +/vector_store/list +""" + +import copy +import json +from typing import List, Optional + +from fastapi import APIRouter, Depends, HTTPException + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.proxy._types import ( + LiteLLM_ManagedVectorStoresTable, + ResponseLiteLLM_ManagedVectorStore, + UserAPIKeyAuth, +) +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.vector_stores import ( + LiteLLM_ManagedVectorStore, + LiteLLM_ManagedVectorStoreListResponse, + VectorStoreDeleteRequest, + VectorStoreInfoRequest, + VectorStoreUpdateRequest, +) +from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + +router = APIRouter() + + +######################################################## +# Management Endpoints +######################################################## +@router.post( + "/vector_store/new", + tags=["vector store management"], + dependencies=[Depends(user_api_key_auth)], +) +async def new_vector_store( + vector_store: LiteLLM_ManagedVectorStore, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Create a new vector store. + + Parameters: + - vector_store_id: str - Unique identifier for the vector store + - custom_llm_provider: str - Provider of the vector store + - vector_store_name: Optional[str] - Name of the vector store + - vector_store_description: Optional[str] - Description of the vector store + - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store + """ + from litellm.proxy.proxy_server import prisma_client + from litellm.types.router import GenericLiteLLMParams + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + # Check if vector store already exists + existing_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": vector_store.get("vector_store_id")} + ) + ) + if existing_vector_store is not None: + raise HTTPException( + status_code=400, + detail=f"Vector store with ID {vector_store.get('vector_store_id')} already exists", + ) + + if vector_store.get("vector_store_metadata") is not None: + vector_store["vector_store_metadata"] = safe_dumps( + vector_store.get("vector_store_metadata") + ) + + # Safely handle JSON serialization of litellm_params + litellm_params_json: Optional[str] = None + _input_litellm_params: dict = vector_store.get("litellm_params", {}) or {} + if _input_litellm_params is not None: + litellm_params_dict = GenericLiteLLMParams( + **_input_litellm_params + ).model_dump(exclude_none=True) + litellm_params_json = safe_dumps(litellm_params_dict) + del vector_store["litellm_params"] + + _new_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.create( + data={ + **vector_store, + "litellm_params": litellm_params_json, + } + ) + ) + + new_vector_store: LiteLLM_ManagedVectorStore = LiteLLM_ManagedVectorStore( + **_new_vector_store.model_dump() + ) + + # Add vector store to registry + if litellm.vector_store_registry is not None: + litellm.vector_store_registry.add_vector_store_to_registry( + vector_store=new_vector_store + ) + + return { + "status": "success", + "message": f"Vector store {vector_store.get('vector_store_id')} created successfully", + "vector_store": new_vector_store, + } + except Exception as e: + verbose_proxy_logger.exception(f"Error creating vector store: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/vector_store/list", + tags=["vector store management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_ManagedVectorStoreListResponse, +) +async def list_vector_stores( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), + page: int = 1, + page_size: int = 100, +): + """ + List all available vector stores with optional filtering and pagination. + Combines both in-memory vector stores and those stored in the database. + + Parameters: + - page: int - Page number for pagination (default: 1) + - page_size: int - Number of items per page (default: 100) + """ + from litellm.proxy.proxy_server import prisma_client + + seen_vector_store_ids = set() + + try: + # Get in-memory vector stores + in_memory_vector_stores: List[LiteLLM_ManagedVectorStore] = [] + if litellm.vector_store_registry is not None: + in_memory_vector_stores = copy.deepcopy( + litellm.vector_store_registry.vector_stores + ) + + # Get vector stores from database + vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db( + prisma_client=prisma_client + ) + + # Combine in-memory and database vector stores + combined_vector_stores: List[LiteLLM_ManagedVectorStore] = [] + for vector_store in in_memory_vector_stores + vector_stores_from_db: + vector_store_id = vector_store.get("vector_store_id", None) + if vector_store_id not in seen_vector_store_ids: + combined_vector_stores.append(vector_store) + seen_vector_store_ids.add(vector_store_id) + + total_count = len(combined_vector_stores) + total_pages = (total_count + page_size - 1) // page_size + + # Format response using LiteLLM_ManagedVectorStoreListResponse + response = LiteLLM_ManagedVectorStoreListResponse( + object="list", + data=combined_vector_stores, + total_count=total_count, + current_page=page, + total_pages=total_pages, + ) + + return response + except Exception as e: + verbose_proxy_logger.exception(f"Error listing vector stores: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/vector_store/delete", + tags=["vector store management"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_vector_store( + data: VectorStoreDeleteRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Delete a vector store. + + Parameters: + - vector_store_id: str - ID of the vector store to delete + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + # Check if vector store exists + existing_vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": data.vector_store_id} + ) + ) + if existing_vector_store is None: + raise HTTPException( + status_code=404, + detail=f"Vector store with ID {data.vector_store_id} not found", + ) + + # Delete vector store + await prisma_client.db.litellm_managedvectorstorestable.delete( + where={"vector_store_id": data.vector_store_id} + ) + + # Delete vector store from registry + if litellm.vector_store_registry is not None: + litellm.vector_store_registry.delete_vector_store_from_registry( + vector_store_id=data.vector_store_id + ) + + return {"message": f"Vector store {data.vector_store_id} deleted successfully"} + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/vector_store/info", + tags=["vector store management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ResponseLiteLLM_ManagedVectorStore, +) +async def get_vector_store_info( + data: VectorStoreInfoRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Return a single vector store's details""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + if litellm.vector_store_registry is not None: + vector_store = litellm.vector_store_registry.get_litellm_managed_vector_store_from_registry( + vector_store_id=data.vector_store_id + ) + if vector_store is not None: + vector_store_metadata = vector_store.get("vector_store_metadata") + # Parse metadata if it's a JSON string + parsed_metadata: Optional[dict] = None + if isinstance(vector_store_metadata, str): + parsed_metadata = json.loads(vector_store_metadata) + elif isinstance(vector_store_metadata, dict): + parsed_metadata = vector_store_metadata + + vector_store_pydantic_obj = LiteLLM_ManagedVectorStoresTable( + vector_store_id=vector_store.get("vector_store_id") or "", + custom_llm_provider=vector_store.get("custom_llm_provider") or "", + vector_store_name=vector_store.get("vector_store_name") or None, + vector_store_description=vector_store.get( + "vector_store_description" + ) + or None, + vector_store_metadata=parsed_metadata, + created_at=vector_store.get("created_at") or None, + updated_at=vector_store.get("updated_at") or None, + litellm_credential_name=vector_store.get("litellm_credential_name"), + litellm_params=vector_store.get("litellm_params") or None, + ) + return {"vector_store": vector_store_pydantic_obj} + + vector_store = ( + await prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": data.vector_store_id} + ) + ) + if vector_store is None: + raise HTTPException( + status_code=404, + detail=f"Vector store with ID {data.vector_store_id} not found", + ) + + vector_store_dict = vector_store.model_dump() # type: ignore[attr-defined] + return {"vector_store": vector_store_dict} + except Exception as e: + verbose_proxy_logger.exception(f"Error getting vector store info: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/vector_store/update", + tags=["vector store management"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_vector_store( + data: VectorStoreUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Update vector store details""" + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException(status_code=500, detail="Database not connected") + + try: + update_data = data.model_dump(exclude_unset=True) + vector_store_id = update_data.pop("vector_store_id") + if update_data.get("vector_store_metadata") is not None: + update_data["vector_store_metadata"] = safe_dumps( + update_data["vector_store_metadata"] + ) + + updated = await prisma_client.db.litellm_managedvectorstorestable.update( + where={"vector_store_id": vector_store_id}, + data=update_data, + ) + + updated_vs = LiteLLM_ManagedVectorStore(**updated.model_dump()) + + if litellm.vector_store_registry is not None: + litellm.vector_store_registry.update_vector_store_in_registry( + vector_store_id=vector_store_id, + updated_data=updated_vs, + ) + + return {"vector_store": updated_vs} + except Exception as e: + verbose_proxy_logger.exception(f"Error updating vector store: {str(e)}") + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/__init__.py b/litellm/__init__.py index 007eff892c8..c9f0ddd1426 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -151,6 +151,7 @@ _custom_logger_compatible_callbacks_literal = Literal[ "mlflow", "langfuse", "langfuse_otel", + "weave_otel", "pagerduty", "humanloop", "gcs_pubsub", @@ -1056,57 +1057,10 @@ from .timeout import timeout 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 from litellm.litellm_core_utils.token_counter import get_modified_max_tokens -from .utils import ( - client, - exception_type, - get_optional_params, - get_response_string, - token_counter, - create_pretrained_tokenizer, - create_tokenizer, - supports_function_calling, - supports_web_search, - supports_url_context, - supports_response_schema, - supports_parallel_function_calling, - supports_vision, - supports_audio_input, - supports_audio_output, - supports_system_messages, - supports_reasoning, - get_litellm_params, - acreate, - get_max_tokens, - get_model_info, - register_prompt_template, - validate_environment, - check_valid_key, - register_model, - encode, - decode, - _calculate_retry_after, - _should_retry, - get_supported_openai_params, - get_api_base, - get_first_chars_messages, - ModelResponse, - ModelResponseStream, - EmbeddingResponse, - ImageResponse, - TranscriptionResponse, - TextCompletionResponse, - get_provider_fields, - ModelResponseListIterator, - get_valid_models, -) - -ALL_LITELLM_RESPONSE_TYPES = [ - ModelResponse, - EmbeddingResponse, - ImageResponse, - TranscriptionResponse, - TextCompletionResponse, -] +# client must be imported immediately as it's used as a decorator at function definition time +from .utils import client +# Note: Most other utils imports are lazy-loaded via __getattr__ to avoid loading utils.py +# (which imports tiktoken) at import time from .llms.bytez.chat.transformation import BytezChatConfig from .llms.custom_llm import CustomLLM @@ -1210,6 +1164,9 @@ from .llms.bedrock.chat.invoke_transformations.amazon_ai21_transformation import from .llms.bedrock.chat.invoke_transformations.amazon_nova_transformation import ( AmazonInvokeNovaConfig, ) +from .llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformation import ( + AmazonQwen2Config, +) from .llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( AmazonQwen3Config, ) @@ -1387,6 +1344,7 @@ from .llms.docker_model_runner.chat.transformation import DockerModelRunnerChatC from .llms.v0.chat.transformation import V0ChatConfig from .llms.oci.chat.transformation import OCIChatConfig from .llms.morph.chat.transformation import MorphChatConfig +from .llms.ragflow.chat.transformation import RAGFlowConfig from .llms.lambda_ai.chat.transformation import LambdaAIChatConfig from .llms.hyperbolic.chat.transformation import HyperbolicChatConfig from .llms.vercel_ai_gateway.chat.transformation import VercelAIGatewayConfig @@ -1537,56 +1495,6 @@ def set_global_gitlab_config(config: Dict[str, Any]) -> None: # Lazy loading system for heavy modules to reduce initial import time and memory usage -def _lazy_import_cost_calculator(name: str) -> Any: - """Lazy import for cost_calculator functions.""" - from .cost_calculator import ( - completion_cost as _completion_cost, - cost_per_token as _cost_per_token, - response_cost_calculator as _response_cost_calculator, - ) - - _cost_functions = { - "completion_cost": _completion_cost, - "cost_per_token": _cost_per_token, - "response_cost_calculator": _response_cost_calculator, - } - - func = _cost_functions[name] - globals()[name] = func - return func - - -def _lazy_import_litellm_logging(name: str) -> Any: - """Lazy import for litellm_logging module.""" - try: - from litellm.litellm_core_utils.litellm_logging import ( - Logging as _Logging, - modify_integration as _modify_integration, - ) - - _logging_objects = { - "Logging": _Logging, - "modify_integration": _modify_integration, - } - - obj = _logging_objects[name] - globals()[name] = obj - return obj - except Exception as e: - raise AttributeError( - f"module {__name__!r} has no attribute {name!r}. " - f"Lazy import failed: {e}" - ) from e - - -_LAZY_LOAD_REGISTRY: Dict[str, Callable[[str], Any]] = { - "completion_cost": _lazy_import_cost_calculator, - "cost_per_token": _lazy_import_cost_calculator, - "response_cost_calculator": _lazy_import_cost_calculator, - "Logging": _lazy_import_litellm_logging, - "modify_integration": _lazy_import_litellm_logging, -} - if TYPE_CHECKING: cost_per_token: Callable[..., Tuple[float, float]] @@ -1597,7 +1505,45 @@ if TYPE_CHECKING: def __getattr__(name: str) -> Any: """Lazy import handler for cost_calculator and litellm_logging functions.""" - if name in _LAZY_LOAD_REGISTRY: - return _LAZY_LOAD_REGISTRY[name](name) + # Lazy load cost_calculator functions + _cost_calculator_names = ( + "completion_cost", + "cost_per_token", + "response_cost_calculator", + ) + if name in _cost_calculator_names: + from ._lazy_imports import _lazy_import_cost_calculator + return _lazy_import_cost_calculator(name) + + # Lazy load litellm_logging functions + _litellm_logging_names = ( + "Logging", + "modify_integration", + ) + if name in _litellm_logging_names: + from ._lazy_imports import _lazy_import_litellm_logging + return _lazy_import_litellm_logging(name) + + # Lazy load utils functions + _utils_names = ( + "exception_type", "get_optional_params", "get_response_string", "token_counter", + "create_pretrained_tokenizer", "create_tokenizer", "supports_function_calling", + "supports_web_search", "supports_url_context", "supports_response_schema", + "supports_parallel_function_calling", "supports_vision", "supports_audio_input", + "supports_audio_output", "supports_system_messages", "supports_reasoning", + "get_litellm_params", "acreate", "get_max_tokens", "get_model_info", + "register_prompt_template", "validate_environment", "check_valid_key", + "register_model", "encode", "decode", "_calculate_retry_after", "_should_retry", + "get_supported_openai_params", "get_api_base", "get_first_chars_messages", + "ModelResponse", "ModelResponseStream", "EmbeddingResponse", "ImageResponse", + "TranscriptionResponse", "TextCompletionResponse", "get_provider_fields", + "ModelResponseListIterator", "get_valid_models", + ) + if name in _utils_names: + from ._lazy_imports import _lazy_import_utils + return _lazy_import_utils(name) raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +# ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py new file mode 100644 index 00000000000..91b16864de1 --- /dev/null +++ b/litellm/_lazy_imports.py @@ -0,0 +1,259 @@ +from typing import Any +import sys + +def _get_litellm_globals() -> dict: + """Helper to get the globals dictionary of the litellm module.""" + return sys.modules["litellm"].__dict__ + +# Lazy import for utils module - imports only the requested item by name. +# Note: PLR0915 (too many statements) is suppressed because the many if statements +# are intentional - each attribute is imported individually only when requested, +# ensuring true lazy imports rather than importing the entire utils module. +def _lazy_import_utils(name: str) -> Any: # noqa: PLR0915 + """Lazy import for utils module - imports only the requested item by name.""" + _globals = _get_litellm_globals() + if name == "exception_type": + from .utils import exception_type as _exception_type + _globals["exception_type"] = _exception_type + return _exception_type + + if name == "get_optional_params": + from .utils import get_optional_params as _get_optional_params + _globals["get_optional_params"] = _get_optional_params + return _get_optional_params + + if name == "get_response_string": + from .utils import get_response_string as _get_response_string + _globals["get_response_string"] = _get_response_string + return _get_response_string + + if name == "token_counter": + from .utils import token_counter as _token_counter + _globals["token_counter"] = _token_counter + return _token_counter + + if name == "create_pretrained_tokenizer": + from .utils import create_pretrained_tokenizer as _create_pretrained_tokenizer + _globals["create_pretrained_tokenizer"] = _create_pretrained_tokenizer + return _create_pretrained_tokenizer + + if name == "create_tokenizer": + from .utils import create_tokenizer as _create_tokenizer + _globals["create_tokenizer"] = _create_tokenizer + return _create_tokenizer + + if name == "supports_function_calling": + from .utils import supports_function_calling as _supports_function_calling + _globals["supports_function_calling"] = _supports_function_calling + return _supports_function_calling + + if name == "supports_web_search": + from .utils import supports_web_search as _supports_web_search + _globals["supports_web_search"] = _supports_web_search + return _supports_web_search + + if name == "supports_url_context": + from .utils import supports_url_context as _supports_url_context + _globals["supports_url_context"] = _supports_url_context + return _supports_url_context + + if name == "supports_response_schema": + from .utils import supports_response_schema as _supports_response_schema + _globals["supports_response_schema"] = _supports_response_schema + return _supports_response_schema + + if name == "supports_parallel_function_calling": + from .utils import supports_parallel_function_calling as _supports_parallel_function_calling + _globals["supports_parallel_function_calling"] = _supports_parallel_function_calling + return _supports_parallel_function_calling + + if name == "supports_vision": + from .utils import supports_vision as _supports_vision + _globals["supports_vision"] = _supports_vision + return _supports_vision + + if name == "supports_audio_input": + from .utils import supports_audio_input as _supports_audio_input + _globals["supports_audio_input"] = _supports_audio_input + return _supports_audio_input + + if name == "supports_audio_output": + from .utils import supports_audio_output as _supports_audio_output + _globals["supports_audio_output"] = _supports_audio_output + return _supports_audio_output + + if name == "supports_system_messages": + from .utils import supports_system_messages as _supports_system_messages + _globals["supports_system_messages"] = _supports_system_messages + return _supports_system_messages + + if name == "supports_reasoning": + from .utils import supports_reasoning as _supports_reasoning + _globals["supports_reasoning"] = _supports_reasoning + return _supports_reasoning + + if name == "get_litellm_params": + from .utils import get_litellm_params as _get_litellm_params + _globals["get_litellm_params"] = _get_litellm_params + return _get_litellm_params + + if name == "acreate": + from .utils import acreate as _acreate + _globals["acreate"] = _acreate + return _acreate + + if name == "get_max_tokens": + from .utils import get_max_tokens as _get_max_tokens + _globals["get_max_tokens"] = _get_max_tokens + return _get_max_tokens + + if name == "get_model_info": + from .utils import get_model_info as _get_model_info + _globals["get_model_info"] = _get_model_info + return _get_model_info + + if name == "register_prompt_template": + from .utils import register_prompt_template as _register_prompt_template + _globals["register_prompt_template"] = _register_prompt_template + return _register_prompt_template + + if name == "validate_environment": + from .utils import validate_environment as _validate_environment + _globals["validate_environment"] = _validate_environment + return _validate_environment + + if name == "check_valid_key": + from .utils import check_valid_key as _check_valid_key + _globals["check_valid_key"] = _check_valid_key + return _check_valid_key + + if name == "register_model": + from .utils import register_model as _register_model + _globals["register_model"] = _register_model + return _register_model + + if name == "encode": + from .utils import encode as _encode + _globals["encode"] = _encode + return _encode + + if name == "decode": + from .utils import decode as _decode + _globals["decode"] = _decode + return _decode + + if name == "_calculate_retry_after": + from .utils import _calculate_retry_after as __calculate_retry_after + _globals["_calculate_retry_after"] = __calculate_retry_after + return __calculate_retry_after + + if name == "_should_retry": + from .utils import _should_retry as __should_retry + _globals["_should_retry"] = __should_retry + return __should_retry + + if name == "get_supported_openai_params": + from .utils import get_supported_openai_params as _get_supported_openai_params + _globals["get_supported_openai_params"] = _get_supported_openai_params + return _get_supported_openai_params + + if name == "get_api_base": + from .utils import get_api_base as _get_api_base + _globals["get_api_base"] = _get_api_base + return _get_api_base + + if name == "get_first_chars_messages": + from .utils import get_first_chars_messages as _get_first_chars_messages + _globals["get_first_chars_messages"] = _get_first_chars_messages + return _get_first_chars_messages + + if name == "ModelResponse": + from .utils import ModelResponse as _ModelResponse + _globals["ModelResponse"] = _ModelResponse + return _ModelResponse + + if name == "ModelResponseStream": + from .utils import ModelResponseStream as _ModelResponseStream + _globals["ModelResponseStream"] = _ModelResponseStream + return _ModelResponseStream + + if name == "EmbeddingResponse": + from .utils import EmbeddingResponse as _EmbeddingResponse + _globals["EmbeddingResponse"] = _EmbeddingResponse + return _EmbeddingResponse + + if name == "ImageResponse": + from .utils import ImageResponse as _ImageResponse + _globals["ImageResponse"] = _ImageResponse + return _ImageResponse + + if name == "TranscriptionResponse": + from .utils import TranscriptionResponse as _TranscriptionResponse + _globals["TranscriptionResponse"] = _TranscriptionResponse + return _TranscriptionResponse + + if name == "TextCompletionResponse": + from .utils import TextCompletionResponse as _TextCompletionResponse + _globals["TextCompletionResponse"] = _TextCompletionResponse + return _TextCompletionResponse + + if name == "get_provider_fields": + from .utils import get_provider_fields as _get_provider_fields + _globals["get_provider_fields"] = _get_provider_fields + return _get_provider_fields + + if name == "ModelResponseListIterator": + from .utils import ModelResponseListIterator as _ModelResponseListIterator + _globals["ModelResponseListIterator"] = _ModelResponseListIterator + return _ModelResponseListIterator + + if name == "get_valid_models": + from .utils import get_valid_models as _get_valid_models + _globals["get_valid_models"] = _get_valid_models + return _get_valid_models + + raise AttributeError(f"Utils lazy import: unknown attribute {name!r}") + + +def _lazy_import_cost_calculator(name: str) -> Any: + """Lazy import for cost_calculator functions.""" + _globals = _get_litellm_globals() + from .cost_calculator import ( + completion_cost as _completion_cost, + cost_per_token as _cost_per_token, + response_cost_calculator as _response_cost_calculator, + ) + + _cost_functions = { + "completion_cost": _completion_cost, + "cost_per_token": _cost_per_token, + "response_cost_calculator": _response_cost_calculator, + } + + func = _cost_functions[name] + _globals[name] = func + return func + + +def _lazy_import_litellm_logging(name: str) -> Any: + """Lazy import for litellm_logging module.""" + _globals = _get_litellm_globals() + try: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as _Logging, + modify_integration as _modify_integration, + ) + + _logging_objects = { + "Logging": _Logging, + "modify_integration": _modify_integration, + } + + obj = _logging_objects[name] + _globals[name] = obj + return obj + except Exception as e: + raise AttributeError( + f"module 'litellm' has no attribute {name!r}. " + f"Lazy import failed: {e}" + ) from e \ No newline at end of file diff --git a/litellm/a2a_protocol/__init__.py b/litellm/a2a_protocol/__init__.py new file mode 100644 index 00000000000..d8d349bb98a --- /dev/null +++ b/litellm/a2a_protocol/__init__.py @@ -0,0 +1,59 @@ +""" +LiteLLM A2A - Wrapper for invoking A2A protocol agents. + +This module provides a thin wrapper around the official `a2a` SDK that: +- Handles httpx client creation and agent card resolution +- Adds LiteLLM logging via @client decorator +- Matches the A2A SDK interface (SendMessageRequest, SendMessageResponse, etc.) + +Example usage (standalone functions with @client decorator): + ```python + from litellm.a2a_protocol import asend_message + from a2a.types import SendMessageRequest, MessageSendParams + from uuid import uuid4 + + request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": "Hello!"}], + "messageId": uuid4().hex, + } + ) + ) + response = await asend_message( + base_url="http://localhost:10001", + request=request, + ) + print(response.model_dump(mode='json', exclude_none=True)) + ``` + +Example usage (class-based): + ```python + from litellm.a2a_protocol import A2AClient + + client = A2AClient(base_url="http://localhost:10001") + response = await client.send_message(request) + ``` +""" + +from litellm.a2a_protocol.client import A2AClient +from litellm.a2a_protocol.main import ( + aget_agent_card, + asend_message, + asend_message_streaming, + create_a2a_client, + send_message, +) +from litellm.types.agents import LiteLLMSendMessageResponse + +__all__ = [ + "A2AClient", + "asend_message", + "send_message", + "asend_message_streaming", + "aget_agent_card", + "create_a2a_client", + "LiteLLMSendMessageResponse", +] diff --git a/litellm/a2a_protocol/client.py b/litellm/a2a_protocol/client.py new file mode 100644 index 00000000000..31f7c3b6a90 --- /dev/null +++ b/litellm/a2a_protocol/client.py @@ -0,0 +1,107 @@ +""" +LiteLLM A2A Client class. + +Provides a class-based interface for A2A agent invocation. +""" + +from typing import TYPE_CHECKING, AsyncIterator, Dict, Optional + +from litellm.types.agents import LiteLLMSendMessageResponse + +if TYPE_CHECKING: + from a2a.client import A2AClient as A2AClientType + from a2a.types import ( + AgentCard, + SendMessageRequest, + SendStreamingMessageRequest, + SendStreamingMessageResponse, + ) + + +class A2AClient: + """ + LiteLLM wrapper for A2A agent invocation. + + Creates the underlying A2A client once on first use and reuses it. + + Example: + ```python + from litellm.a2a_protocol import A2AClient + from a2a.types import SendMessageRequest, MessageSendParams + from uuid import uuid4 + + client = A2AClient(base_url="http://localhost:10001") + + request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": "Hello!"}], + "messageId": uuid4().hex, + } + ) + ) + response = await client.send_message(request) + ``` + """ + + def __init__( + self, + base_url: str, + timeout: float = 60.0, + extra_headers: Optional[Dict[str, str]] = None, + ): + """ + Initialize the A2A client wrapper. + + Args: + base_url: The base URL of the A2A agent (e.g., "http://localhost:10001") + timeout: Request timeout in seconds (default: 60.0) + extra_headers: Optional additional headers to include in requests + """ + self.base_url = base_url + self.timeout = timeout + self.extra_headers = extra_headers + self._a2a_client: Optional["A2AClientType"] = None + + async def _get_client(self) -> "A2AClientType": + """Get or create the underlying A2A client.""" + if self._a2a_client is None: + from litellm.a2a_protocol.main import create_a2a_client + + self._a2a_client = await create_a2a_client( + base_url=self.base_url, + timeout=self.timeout, + extra_headers=self.extra_headers, + ) + return self._a2a_client + + async def get_agent_card(self) -> "AgentCard": + """Fetch the agent card from the server.""" + from litellm.a2a_protocol.main import aget_agent_card + + return await aget_agent_card( + base_url=self.base_url, + timeout=self.timeout, + extra_headers=self.extra_headers, + ) + + async def send_message( + self, request: "SendMessageRequest" + ) -> LiteLLMSendMessageResponse: + """Send a message to the A2A agent.""" + from litellm.a2a_protocol.main import asend_message + + a2a_client = await self._get_client() + return await asend_message(a2a_client=a2a_client, request=request) + + async def send_message_streaming( + self, request: "SendStreamingMessageRequest" + ) -> AsyncIterator["SendStreamingMessageResponse"]: + """Send a streaming message to the A2A agent.""" + from litellm.a2a_protocol.main import asend_message_streaming + + a2a_client = await self._get_client() + async for chunk in asend_message_streaming(a2a_client=a2a_client, request=request): + yield chunk diff --git a/litellm/a2a_protocol/cost_calculator.py b/litellm/a2a_protocol/cost_calculator.py new file mode 100644 index 00000000000..2d821bb9c3b --- /dev/null +++ b/litellm/a2a_protocol/cost_calculator.py @@ -0,0 +1,36 @@ +""" +Cost calculator for A2A (Agent-to-Agent) calls. +""" + +from typing import TYPE_CHECKING, Any, Optional + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LitellmLoggingObject, + ) +else: + LitellmLoggingObject = Any + + +class A2ACostCalculator: + @staticmethod + def calculate_a2a_cost( + litellm_logging_obj: Optional[LitellmLoggingObject], + ) -> float: + """ + Calculate the cost of an A2A send_message call. + + Default is 0.0. In the future, users can configure cost per agent call. + """ + if litellm_logging_obj is None: + return 0.0 + + # Check if user set a custom response cost + response_cost = litellm_logging_obj.model_call_details.get( + "response_cost", None + ) + if response_cost is not None: + return response_cost + + # Default to 0.0 for A2A calls + return 0.0 diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py new file mode 100644 index 00000000000..5cb238904f0 --- /dev/null +++ b/litellm/a2a_protocol/main.py @@ -0,0 +1,298 @@ +""" +LiteLLM A2A SDK functions. + +Provides standalone functions with @client decorator for LiteLLM logging integration. +""" + +import asyncio +from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union + +from litellm._logging import verbose_logger +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, + httpxSpecialProvider, +) +from litellm.types.agents import LiteLLMSendMessageResponse +from litellm.utils import client + +if TYPE_CHECKING: + from a2a.client import A2AClient as A2AClientType + from a2a.types import ( + AgentCard, + SendMessageRequest, + SendStreamingMessageRequest, + SendStreamingMessageResponse, + ) + +# Runtime imports with availability check +A2A_SDK_AVAILABLE = False +A2ACardResolver: Any = None +_A2AClient: Any = None + +try: + from a2a.client import A2ACardResolver # type: ignore[no-redef] + from a2a.client import A2AClient as _A2AClient # type: ignore[no-redef] + + A2A_SDK_AVAILABLE = True +except ImportError: + pass + + +def _get_a2a_model_info(a2a_client: Any, kwargs: Dict[str, Any]) -> str: + """ + Extract agent info and set model/custom_llm_provider for cost tracking. + + Sets model info on the litellm_logging_obj if available. + Returns the agent name for logging. + """ + agent_name = "unknown" + + # Try to get agent card from our stored attribute first, then fallback to SDK attribute + agent_card = getattr(a2a_client, "_litellm_agent_card", None) + if agent_card is None: + agent_card = getattr(a2a_client, "agent_card", None) + + if agent_card is not None: + agent_name = getattr(agent_card, "name", "unknown") or "unknown" + + # Build model string + model = f"a2a_agent/{agent_name}" + custom_llm_provider = "a2a_agent" + + # Set on litellm_logging_obj if available (for standard logging payload) + litellm_logging_obj = kwargs.get("litellm_logging_obj") + if litellm_logging_obj is not None: + litellm_logging_obj.model = model + litellm_logging_obj.custom_llm_provider = custom_llm_provider + litellm_logging_obj.model_call_details["model"] = model + litellm_logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + + return agent_name + + +@client +async def asend_message( + a2a_client: "A2AClientType", + request: "SendMessageRequest", + **kwargs: Any, +) -> LiteLLMSendMessageResponse: + """ + Async: Send a message to an A2A agent. + + Uses the @client decorator for LiteLLM logging and tracking. + + Args: + a2a_client: An initialized a2a.client.A2AClient instance + request: SendMessageRequest from a2a.types + **kwargs: Additional arguments passed to the client decorator + + Returns: + LiteLLMSendMessageResponse (wraps a2a SendMessageResponse with _hidden_params) + + Example: + ```python + from litellm.a2a_protocol import asend_message, create_a2a_client + from a2a.types import SendMessageRequest, MessageSendParams + from uuid import uuid4 + + # Create client once + a2a_client = await create_a2a_client(base_url="http://localhost:10001") + + # Use it for multiple requests + request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams( + message={ + "role": "user", + "parts": [{"kind": "text", "text": "Hello!"}], + "messageId": uuid4().hex, + } + ) + ) + response = await asend_message(a2a_client=a2a_client, request=request) + ``` + """ + agent_name = _get_a2a_model_info(a2a_client, kwargs) + + verbose_logger.info(f"A2A send_message request_id={request.id}, agent={agent_name}") + + a2a_response = await a2a_client.send_message(request) + + verbose_logger.info(f"A2A send_message completed, request_id={request.id}") + + # Wrap in LiteLLM response type for _hidden_params support + response = LiteLLMSendMessageResponse.from_a2a_response(a2a_response) + + return response + + +@client +def send_message( + a2a_client: "A2AClientType", + request: "SendMessageRequest", + **kwargs: Any, +) -> Union[LiteLLMSendMessageResponse, Coroutine[Any, Any, LiteLLMSendMessageResponse]]: + """ + Sync: Send a message to an A2A agent. + + Uses the @client decorator for LiteLLM logging and tracking. + + Args: + a2a_client: An initialized a2a.client.A2AClient instance + request: SendMessageRequest from a2a.types + **kwargs: Additional arguments passed to the client decorator + + Returns: + LiteLLMSendMessageResponse (wraps a2a SendMessageResponse with _hidden_params) + """ + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + + if loop is not None: + return asend_message(a2a_client=a2a_client, request=request, **kwargs) + else: + return asyncio.run(asend_message(a2a_client=a2a_client, request=request, **kwargs)) + + +async def asend_message_streaming( + a2a_client: "A2AClientType", + request: "SendStreamingMessageRequest", +) -> AsyncIterator["SendStreamingMessageResponse"]: + """ + Async: Send a streaming message to an A2A agent. + + Args: + a2a_client: An initialized a2a.client.A2AClient instance + request: SendStreamingMessageRequest from a2a.types + + Yields: + SendStreamingMessageResponse chunks from the agent + """ + verbose_logger.info(f"A2A send_message_streaming request_id={request.id}") + + stream = a2a_client.send_message_streaming(request) + + chunk_count = 0 + async for chunk in stream: + chunk_count += 1 + yield chunk + + verbose_logger.info( + f"A2A send_message_streaming completed, request_id={request.id}, chunks={chunk_count}" + ) + + +async def create_a2a_client( + base_url: str, + timeout: float = 60.0, + extra_headers: Optional[Dict[str, str]] = None, +) -> "A2AClientType": + """ + Create an A2A client for the given agent URL. + + This resolves the agent card and returns a ready-to-use A2A client. + The client can be reused for multiple requests. + + Args: + base_url: The base URL of the A2A agent (e.g., "http://localhost:10001") + timeout: Request timeout in seconds (default: 60.0) + extra_headers: Optional additional headers to include in requests + + Returns: + An initialized a2a.client.A2AClient instance + + Example: + ```python + from litellm.a2a_protocol import create_a2a_client, asend_message + + # Create client once + client = await create_a2a_client(base_url="http://localhost:10001") + + # Reuse for multiple requests + response1 = await asend_message(a2a_client=client, request=request1) + response2 = await asend_message(a2a_client=client, request=request2) + ``` + """ + if not A2A_SDK_AVAILABLE: + raise ImportError( + "The 'a2a' package is required for A2A agent invocation. " + "Install it with: pip install a2a" + ) + + verbose_logger.info(f"Creating A2A client for {base_url}") + + # Use LiteLLM's cached httpx client + http_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.A2A, + params={"timeout": timeout}, + ) + httpx_client = http_handler.client + + # Resolve agent card + resolver = A2ACardResolver( + httpx_client=httpx_client, + base_url=base_url, + ) + agent_card = await resolver.get_agent_card() + + verbose_logger.debug( + f"Resolved agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}" + ) + + # Create A2A client + a2a_client = _A2AClient( + httpx_client=httpx_client, + agent_card=agent_card, + ) + + # Store agent_card on client for later retrieval (SDK doesn't expose it) + a2a_client._litellm_agent_card = agent_card # type: ignore[attr-defined] + + verbose_logger.info(f"A2A client created for {base_url}") + + return a2a_client + + +async def aget_agent_card( + base_url: str, + timeout: float = 60.0, + extra_headers: Optional[Dict[str, str]] = None, +) -> "AgentCard": + """ + Fetch the agent card from an A2A agent. + + Args: + base_url: The base URL of the A2A agent (e.g., "http://localhost:10001") + timeout: Request timeout in seconds (default: 60.0) + extra_headers: Optional additional headers to include in requests + + Returns: + AgentCard from the A2A agent + """ + if not A2A_SDK_AVAILABLE: + raise ImportError( + "The 'a2a' package is required for A2A agent invocation. " + "Install it with: pip install a2a" + ) + + verbose_logger.info(f"Fetching agent card from {base_url}") + + # Use LiteLLM's cached httpx client + http_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.A2A, + params={"timeout": timeout}, + ) + httpx_client = http_handler.client + + resolver = A2ACardResolver( + httpx_client=httpx_client, + base_url=base_url, + ) + agent_card = await resolver.get_agent_card() + + verbose_logger.info( + f"Fetched agent card: {agent_card.name if hasattr(agent_card, 'name') else 'unknown'}" + ) + return agent_card diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 8289801ee30..50b48321db5 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -14,8 +14,7 @@ from litellm.utils import token_counter async def calculate_batch_cost_and_usage( file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai"], - model_name: Optional[str] = None, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm"], ) -> Tuple[float, Usage, List[str]]: """ Calculate the cost and usage of a batch @@ -37,8 +36,7 @@ async def calculate_batch_cost_and_usage( async def _handle_completed_batch( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"], - model_name: Optional[str] = None, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm"], ) -> Tuple[float, Usage, List[str]]: """Helper function to process a completed batch and handle logging""" # Get batch results @@ -84,8 +82,7 @@ def _get_batch_models_from_file_content( def _batch_cost_calculator( file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - model_name: Optional[str] = None, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm"] = "openai", ) -> float: """ Calculate the cost of a batch based on the output file id @@ -186,7 +183,7 @@ def calculate_vertex_ai_batch_cost_and_usage( async def _get_batch_output_file_content_as_dictionary( batch: Batch, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm"] = "openai", ) -> List[dict]: """ Get the batch output file content as a list of dictionaries @@ -225,7 +222,7 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]: def _get_batch_job_cost_from_file_content( file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm"] = "openai", ) -> float: """ Get the cost of a batch job from the file content @@ -253,8 +250,7 @@ def _get_batch_job_cost_from_file_content( def _get_batch_job_total_usage_from_file_content( file_content_dictionary: List[dict], - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", - model_name: Optional[str] = None, + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm"] = "openai", ) -> Usage: """ Get the tokens of a batch job from the file content diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 57a9857dd6a..b99f4a628dc 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -23,6 +23,7 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.batches.handler import AzureBatchesAPI +from litellm.llms.bedrock.batches.handler import BedrockBatchesHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.openai import OpenAIBatchesAPI @@ -35,7 +36,11 @@ from litellm.types.llms.openai import ( RetrieveBatchRequest, ) from litellm.types.router import GenericLiteLLMParams -from litellm.types.utils import LiteLLMBatch, LlmProviders +from litellm.types.utils import ( + OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, + LiteLLMBatch, + LlmProviders, +) from litellm.utils import ( ProviderConfigManager, client, @@ -100,7 +105,7 @@ async def acreate_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -148,7 +153,7 @@ def create_batch( completion_window: Literal["24h"], endpoint: Literal["/v1/chat/completions", "/v1/embeddings", "/v1/completions"], input_file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -235,7 +240,7 @@ def create_batch( ) return response api_base: Optional[str] = None - if custom_llm_provider == "openai": + if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( optional_params.api_base @@ -350,7 +355,7 @@ def create_batch( @client async def aretrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -396,10 +401,10 @@ def _handle_retrieve_batch_providers_without_provider_config( litellm_params: dict, _retrieve_batch_request: RetrieveBatchRequest, _is_async: bool, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", ): api_base: Optional[str] = None - if custom_llm_provider == "openai": + if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( optional_params.api_base @@ -512,7 +517,7 @@ def _handle_retrieve_batch_providers_without_provider_config( @client def retrieve_batch( batch_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -576,7 +581,7 @@ def retrieve_batch( async_kwargs = kwargs.copy() async_kwargs.pop("aws_region_name", None) - return _handle_async_invoke_status( + return BedrockBatchesHandler._handle_async_invoke_status( batch_id=batch_id, aws_region_name=kwargs.get("aws_region_name", "us-east-1"), logging_obj=litellm_logging_obj, @@ -644,7 +649,7 @@ def retrieve_batch( async def alist_batches( after: Optional[str] = None, limit: Optional[int] = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai", metadata: Optional[Dict[str, str]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -687,7 +692,7 @@ async def alist_batches( def list_batches( after: Optional[str] = None, limit: Optional[int] = None, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "hosted_vllm", "vertex_ai"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -727,7 +732,7 @@ def list_batches( timeout = 600.0 _is_async = kwargs.pop("alist_batches", False) is True - if custom_llm_provider == "openai": + if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( optional_params.api_base @@ -928,7 +933,7 @@ def cancel_batch( _is_async = kwargs.pop("acancel_batch", False) is True api_base: Optional[str] = None - if custom_llm_provider == "openai": + if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: api_base = ( optional_params.api_base or litellm.api_base @@ -1043,19 +1048,20 @@ def _handle_async_invoke_status( ) # Transform response to a LiteLLMBatch object + from litellm.types.llms.openai import BatchJobStatus from litellm.types.utils import LiteLLMBatch # Normalize status to lowercase (AWS returns 'Completed', 'Failed', etc.) aws_status_raw = status_response.get("status", "") aws_status_lower = aws_status_raw.lower() # Map AWS status values to LiteLLM expected values - status_mapping = { + status_mapping: dict[str, BatchJobStatus] = { "completed": "completed", "failed": "failed", "inprogress": "in_progress", "in_progress": "in_progress", } - normalized_status = status_mapping.get(aws_status_lower, aws_status_lower) + normalized_status: BatchJobStatus = status_mapping.get(aws_status_lower, "failed") # Default to "failed" if unknown status # Get output S3 URI safely output_s3_uri = "" @@ -1065,13 +1071,15 @@ def _handle_async_invoke_status( pass # Use BedrockBatchesConfig's timestamp parsing method (expects raw AWS status string) + import time + from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig created_at, in_progress_at, completed_at, failed_at, _, _ = BedrockBatchesConfig()._parse_timestamps_and_status(status_response, aws_status_raw) result = LiteLLMBatch( id=status_response["invocationArn"], object="batch", status=normalized_status, - created_at=created_at, + created_at=created_at or int(time.time()), # Provide default timestamp if None in_progress_at=in_progress_at, completed_at=completed_at, failed_at=failed_at, diff --git a/litellm/constants.py b/litellm/constants.py index e3de7368c8a..fa9f1d527af 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1,4 +1,5 @@ import os +import sys from typing import List, Literal DEFAULT_HEALTH_CHECK_PROMPT = str( @@ -99,10 +100,18 @@ RUNWAYML_POLLING_TIMEOUT = int( ########## Networking constants ############################################################## _DEFAULT_TTL_FOR_HTTPX_CLIENTS = 3600 # 1 hour, re-use the same httpx client for 1 hour -# Aiohttp connection pooling constants -AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 0)) +# Aiohttp connection pooling - prevents memory leaks from unbounded connection growth +# Set to 0 for unlimited (not recommended for production) +AIOHTTP_CONNECTOR_LIMIT = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT", 300)) +AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(os.getenv("AIOHTTP_CONNECTOR_LIMIT_PER_HOST", 50)) AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120)) AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300)) +# enable_cleanup_closed is only needed for Python versions with the SSL leak bug +# Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960) +# Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78 +AIOHTTP_NEEDS_CLEANUP_CLOSED = ( + (3, 13, 0) <= sys.version_info < (3, 13, 1) or sys.version_info < (3, 12, 7) +) # WebSocket constants # Default to None (unlimited) to match OpenAI's official agents SDK behavior @@ -255,7 +264,9 @@ TOGETHER_AI_EMBEDDING_350_M = int(os.getenv("TOGETHER_AI_EMBEDDING_350_M", 350)) QDRANT_SCALAR_QUANTILE = float(os.getenv("QDRANT_SCALAR_QUANTILE", 0.99)) QDRANT_VECTOR_SIZE = int(os.getenv("QDRANT_VECTOR_SIZE", 1536)) CACHED_STREAMING_CHUNK_DELAY = float(os.getenv("CACHED_STREAMING_CHUNK_DELAY", 0.02)) -AUDIO_SPEECH_CHUNK_SIZE = 8192 # chunk_size for audio speech streaming. Balance between latency and memory usage +AUDIO_SPEECH_CHUNK_SIZE = int( + os.getenv("AUDIO_SPEECH_CHUNK_SIZE", 8192) +) # chunk_size for audio speech streaming. Balance between latency and memory usage MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB = int( os.getenv("MAX_SIZE_PER_ITEM_IN_MEMORY_CACHE_IN_KB", 512) ) @@ -278,10 +289,16 @@ REDACTED_BY_LITELM_STRING = "REDACTED_BY_LITELM" MAX_LANGFUSE_INITIALIZED_CLIENTS = int( os.getenv("MAX_LANGFUSE_INITIALIZED_CLIENTS", 50) ) -LOGGING_WORKER_CONCURRENCY = int(os.getenv("LOGGING_WORKER_CONCURRENCY", 100)) # Must be above 0 +LOGGING_WORKER_CONCURRENCY = int( + os.getenv("LOGGING_WORKER_CONCURRENCY", 100) +) # Must be above 0 LOGGING_WORKER_MAX_QUEUE_SIZE = int(os.getenv("LOGGING_WORKER_MAX_QUEUE_SIZE", 50_000)) -LOGGING_WORKER_MAX_TIME_PER_COROUTINE = float(os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0)) -LOGGING_WORKER_CLEAR_PERCENTAGE = int(os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50)) # Percentage of queue to clear (default: 50%) +LOGGING_WORKER_MAX_TIME_PER_COROUTINE = float( + os.getenv("LOGGING_WORKER_MAX_TIME_PER_COROUTINE", 20.0) +) +LOGGING_WORKER_CLEAR_PERCENTAGE = int( + os.getenv("LOGGING_WORKER_CLEAR_PERCENTAGE", 50) +) # Percentage of queue to clear (default: 50%) MAX_ITERATIONS_TO_CLEAR_QUEUE = int(os.getenv("MAX_ITERATIONS_TO_CLEAR_QUEUE", 200)) MAX_TIME_TO_CLEAR_QUEUE = float(os.getenv("MAX_TIME_TO_CLEAR_QUEUE", 5.0)) LOGGING_WORKER_AGGRESSIVE_CLEAR_COOLDOWN_SECONDS = float( @@ -586,6 +603,7 @@ openai_compatible_providers: List = [ "cometapi", "clarifai", "docker_model_runner", + "ragflow", ] openai_text_completion_compatible_providers: List = ( [ # providers that support `/v1/completions` @@ -858,7 +876,9 @@ BEDROCK_INVOKE_PROVIDERS_LITERAL = Literal[ "nova", "deepseek_r1", "qwen3", + "qwen2", "twelvelabs", + "openai", ] BEDROCK_EMBEDDING_PROVIDERS_LITERAL = Literal[ @@ -908,6 +928,9 @@ BEDROCK_CONVERSE_MODELS = [ "meta.llama3-2-3b-instruct-v1:0", "meta.llama3-2-11b-instruct-v1:0", "meta.llama3-2-90b-instruct-v1:0", + "amazon.nova-lite-v1:0", + "amazon.nova-2-lite-v1:0", + "amazon.nova-pro-v1:0", ] diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9ef26d23ce2..57eab6d29a3 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -95,6 +95,7 @@ from litellm.utils import ( EmbeddingResponse, ImageResponse, ModelResponse, + ModelResponseStream, ProviderConfigManager, TextCompletionResponse, TranscriptionResponse, @@ -654,7 +655,9 @@ def _infer_call_type( if completion_response is None: return None - if isinstance(completion_response, ModelResponse): + if isinstance(completion_response, ModelResponse) or isinstance( + completion_response, ModelResponseStream + ): return "completion" elif isinstance(completion_response, EmbeddingResponse): return "embedding" @@ -934,6 +937,17 @@ def completion_cost( # noqa: PLR0915 prompt_tokens = token_counter(model=model, text=prompt) completion_tokens = token_counter(model=model, text=completion) + # Handle A2A calls before model check - A2A doesn't require a model + if call_type in ( + CallTypes.asend_message.value, + CallTypes.send_message.value, + ): + from litellm.a2a_protocol.cost_calculator import A2ACostCalculator + + return A2ACostCalculator.calculate_a2a_cost( + litellm_logging_obj=litellm_logging_obj + ) + if model is None: raise ValueError( f"Model is None and does not exist in passed completion_response. Passed completion_response={completion_response}, model={model}" @@ -1046,29 +1060,33 @@ def completion_cost( # noqa: PLR0915 number_of_queries = len(query) elif query is not None: number_of_queries = 1 - + search_model = model or "" if custom_llm_provider and "/" not in search_model: # If model is like "tavily-search", construct "tavily/search" for cost lookup search_model = f"{custom_llm_provider}/search" - - prompt_cost, completion_cost_result = search_provider_cost_per_query( - model=search_model, - custom_llm_provider=custom_llm_provider, - number_of_queries=number_of_queries, - optional_params=optional_params, + + prompt_cost, completion_cost_result = ( + search_provider_cost_per_query( + model=search_model, + custom_llm_provider=custom_llm_provider, + number_of_queries=number_of_queries, + optional_params=optional_params, + ) ) - + # Return the total cost (prompt_cost + completion_cost, but for search it's just prompt_cost) _final_cost = prompt_cost + completion_cost_result - + # Apply discount original_cost = _final_cost - _final_cost, discount_percent, discount_amount = _apply_cost_discount( - base_cost=_final_cost, - custom_llm_provider=custom_llm_provider, + _final_cost, discount_percent, discount_amount = ( + _apply_cost_discount( + base_cost=_final_cost, + custom_llm_provider=custom_llm_provider, + ) ) - + # Store cost breakdown in logging object if available _store_cost_breakdown_in_logging_obj( litellm_logging_obj=litellm_logging_obj, @@ -1080,7 +1098,7 @@ def completion_cost( # noqa: PLR0915 discount_percent=discount_percent, discount_amount=discount_amount, ) - + return _final_cost elif call_type == CallTypes.arealtime.value and isinstance( completion_response, LiteLLMRealtimeStreamLoggingObject diff --git a/litellm/files/main.py b/litellm/files/main.py index 535772fa42c..9378715a472 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -18,6 +18,7 @@ from litellm import get_secret_str from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI +from litellm.llms.bedrock.files.handler import BedrockFilesHandler from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai.openai import FileDeleted, FileObject, OpenAIFilesAPI @@ -30,7 +31,10 @@ from litellm.types.llms.openai import ( OpenAIFileObject, ) from litellm.types.router import * -from litellm.types.utils import LlmProviders +from litellm.types.utils import ( + OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, + LlmProviders, +) from litellm.utils import ( ProviderConfigManager, client, @@ -44,6 +48,7 @@ base_llm_http_handler = BaseLLMHTTPHandler() openai_files_instance = OpenAIFilesAPI() azure_files_instance = AzureOpenAIFilesAPI() vertex_ai_files_instance = VertexAIFilesHandler() +bedrock_files_instance = BedrockFilesHandler() ################################################# @@ -51,7 +56,7 @@ vertex_ai_files_instance = VertexAIFilesHandler() async def acreate_file( file: FileTypes, purpose: Literal["assistants", "batch", "fine-tune"], - custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -95,9 +100,7 @@ async def acreate_file( def create_file( file: FileTypes, purpose: Literal["assistants", "batch", "fine-tune"], - custom_llm_provider: Optional[ - Literal["openai", "azure", "vertex_ai", "bedrock"] - ] = None, + custom_llm_provider: Optional[Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"]] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -165,7 +168,7 @@ def create_file( ), timeout=timeout, ) - elif custom_llm_provider == "openai": + elif custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( optional_params.api_base @@ -276,7 +279,7 @@ def create_file( @client async def afile_retrieve( file_id: str, - custom_llm_provider: Literal["openai", "azure"] = "openai", + custom_llm_provider: Literal["openai", "azure", "hosted_vllm"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -317,7 +320,7 @@ async def afile_retrieve( @client def file_retrieve( file_id: str, - custom_llm_provider: Literal["openai", "azure"] = "openai", + custom_llm_provider: Literal["openai", "azure", "hosted_vllm"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -347,7 +350,7 @@ def file_retrieve( _is_async = kwargs.pop("is_async", False) is True - if custom_llm_provider == "openai": + if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( optional_params.api_base @@ -514,7 +517,7 @@ def file_delete( elif timeout is None: timeout = 600.0 _is_async = kwargs.pop("is_async", False) is True - if custom_llm_provider == "openai": + if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( optional_params.api_base @@ -670,7 +673,7 @@ def file_list( timeout = 600.0 _is_async = kwargs.pop("is_async", False) is True - if custom_llm_provider == "openai": + if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( optional_params.api_base @@ -754,7 +757,7 @@ def file_list( @client async def afile_content( file_id: str, - custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -799,7 +802,7 @@ def file_content( file_id: str, model: Optional[str] = None, custom_llm_provider: Optional[ - Union[Literal["openai", "azure", "vertex_ai"], str] + Union[Literal["openai", "azure", "vertex_ai", "bedrock", "hosted_vllm"], str] ] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, @@ -846,7 +849,7 @@ def file_content( _is_async = kwargs.pop("afile_content", False) is True - if custom_llm_provider == "openai": + if custom_llm_provider in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: # for deepinfra/perplexity/anyscale/groq we check in get_llm_provider and pass in the api base from there api_base = ( optional_params.api_base @@ -937,9 +940,18 @@ def file_content( timeout=timeout, max_retries=optional_params.max_retries, ) + elif custom_llm_provider == "bedrock": + response = bedrock_files_instance.file_content( + _is_async=_is_async, + file_content_request=_file_content_request, + api_base=optional_params.api_base, + optional_params=litellm_params_dict, + timeout=timeout, + max_retries=optional_params.max_retries, + ) else: raise litellm.exceptions.BadRequestError( - message="LiteLLM doesn't support {} for 'custom_llm_provider'. Supported providers are 'openai', 'azure', 'vertex_ai'.".format( + message="LiteLLM doesn't support {} for 'custom_llm_provider'. Supported providers are 'openai', 'azure', 'vertex_ai', 'bedrock'.".format( custom_llm_provider ), model="n/a", diff --git a/litellm/images/main.py b/litellm/images/main.py index eacd4778299..770b16c1ed2 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -6,7 +6,9 @@ from typing import Any, Coroutine, Dict, List, Literal, Optional, Union, cast, o import httpx import litellm -from litellm import client, exception_type, get_litellm_params +from litellm.utils import exception_type, get_litellm_params +# client is imported from litellm as it's a decorator +from litellm import client from litellm.constants import DEFAULT_IMAGE_ENDPOINT_MODEL from litellm.constants import request_timeout as DEFAULT_REQUEST_TIMEOUT from litellm.exceptions import LiteLLMUnknownProvider diff --git a/litellm/integrations/SlackAlerting/budget_alert_types.py b/litellm/integrations/SlackAlerting/budget_alert_types.py index 1e9ad286e37..dadfef3fc40 100644 --- a/litellm/integrations/SlackAlerting/budget_alert_types.py +++ b/litellm/integrations/SlackAlerting/budget_alert_types.py @@ -50,6 +50,14 @@ class TeamBudgetAlert(BaseBudgetAlertType): return user_info.team_id or "default_id" +class OrganizationBudgetAlert(BaseBudgetAlertType): + def get_event_message(self) -> str: + return "Organization Budget: " + + def get_id(self, user_info: CallInfo) -> str: + return user_info.organization_id or "default_id" + + class TokenBudgetAlert(BaseBudgetAlertType): def get_event_message(self) -> str: return "Key Budget: " @@ -72,6 +80,7 @@ def get_budget_alert_type( "soft_budget", "user_budget", "team_budget", + "organization_budget", "proxy_budget", "projected_limit_exceeded", ], @@ -83,6 +92,7 @@ def get_budget_alert_type( "soft_budget": SoftBudgetAlert(), "user_budget": UserBudgetAlert(), "team_budget": TeamBudgetAlert(), + "organization_budget": OrganizationBudgetAlert(), "token_budget": TokenBudgetAlert(), "projected_limit_exceeded": ProjectedLimitExceededAlert(), } diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 3efe5873786..0e691e2c43f 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -134,19 +134,25 @@ class SlackAlerting(CustomBatchLogger): if llm_router is not None: self.llm_router = llm_router - def _prepare_outage_value_for_cache(self, outage_value: Union[dict, ProviderRegionOutageModel, OutageModel]) -> dict: + def _prepare_outage_value_for_cache( + self, outage_value: Union[dict, ProviderRegionOutageModel, OutageModel] + ) -> dict: """ Helper method to prepare outage value for Redis caching. Converts set objects to lists for JSON serialization. """ # Convert to dict for processing cache_value = dict(outage_value) - - if "deployment_ids" in cache_value and isinstance(cache_value["deployment_ids"], set): + + if "deployment_ids" in cache_value and isinstance( + cache_value["deployment_ids"], set + ): cache_value["deployment_ids"] = list(cache_value["deployment_ids"]) return cache_value - def _restore_outage_value_from_cache(self, outage_value: Optional[dict]) -> Optional[dict]: + def _restore_outage_value_from_cache( + self, outage_value: Optional[dict] + ) -> Optional[dict]: """ Helper method to restore outage value after retrieving from cache. Converts list objects back to sets for proper handling. @@ -528,6 +534,7 @@ class SlackAlerting(CustomBatchLogger): "soft_budget", "user_budget", "team_budget", + "organization_budget", "proxy_budget", "projected_limit_exceeded", ], @@ -1338,7 +1345,7 @@ Model Info: subject=email_event["subject"], html=email_event["html"], ) - if webhook_event.event_group == "team": + if webhook_event.event_group == Litellm_EntityType.TEAM: from litellm.integrations.email_alerting import send_team_budget_alert await send_team_budget_alert(webhook_event=webhook_event) @@ -1399,7 +1406,7 @@ Model Info: current_time = datetime.now().strftime("%H:%M:%S") _proxy_base_url = os.getenv("PROXY_BASE_URL", None) # Use .name if it's an enum, otherwise use as is - alert_type_name = getattr(alert_type, 'name', alert_type) + alert_type_name = getattr(alert_type, "name", alert_type) alert_type_formatted = f"Alert type: `{alert_type_name}`" if alert_type == "daily_reports" or alert_type == "new_model_added": formatted_message = alert_type_formatted + message diff --git a/litellm/integrations/arize/arize_phoenix.py b/litellm/integrations/arize/arize_phoenix.py index ab70dd9d0e2..4a6e0cec8ca 100644 --- a/litellm/integrations/arize/arize_phoenix.py +++ b/litellm/integrations/arize/arize_phoenix.py @@ -1,18 +1,20 @@ import os -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, Any, Optional, Union +from datetime import datetime from litellm._logging import verbose_logger from litellm.integrations.arize import _utils from litellm.integrations.arize._utils import ArizeOTELAttributes from litellm.types.integrations.arize_phoenix import ArizePhoenixConfig +from litellm.types.services import ServiceLoggerPayload +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 - from .opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig - Protocol = _Protocol OpenTelemetryConfig = _OpenTelemetryConfig Span = Union[_Span, Any] @@ -25,7 +27,11 @@ else: ARIZE_HOSTED_PHOENIX_ENDPOINT = "https://otlp.arize.com/v1/traces" -class ArizePhoenixLogger: +class ArizePhoenixLogger(OpenTelemetry): + def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]): + ArizePhoenixLogger.set_arize_phoenix_attributes(span, kwargs, response_obj) + return + @staticmethod def set_arize_phoenix_attributes(span: Span, kwargs, response_obj): _utils.set_attributes(span, kwargs, response_obj, ArizeOTELAttributes) @@ -97,3 +103,46 @@ class ArizePhoenixLogger: endpoint=endpoint, project_name=project_name, ) + + async def async_service_success_hook( + self, + payload: ServiceLoggerPayload, + parent_otel_span: Optional[Span] = None, + start_time: Optional[Union[datetime, float]] = None, + end_time: Optional[Union[datetime, float]] = None, + event_metadata: Optional[dict] = None, + ): + pass # suppress additional spans + + async def async_service_failure_hook( + self, + payload: ServiceLoggerPayload, + error: Optional[str] = "", + parent_otel_span: Optional[Span] = None, + start_time: Optional[Union[datetime, float]] = None, + end_time: Optional[Union[float, datetime]] = None, + event_metadata: Optional[dict] = None, + ): + pass # suppress additional spans + + def create_litellm_proxy_request_started_span( + self, + start_time: datetime, + headers: dict, + ): + pass # suppress additional spans + + async def async_health_check(self): + + config = self.get_arize_phoenix_config() + + if not config.otlp_auth_headers: + return { + "status": "unhealthy", + "error_message": "PHOENIX_API_KEY environment variable not set", + } + + return { + "status": "healthy", + "message": "Arize-Phoenix credentials are configured properly", + } \ No newline at end of file diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 7f74f5d2157..782f9460044 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -17,10 +17,10 @@ from litellm.caching import DualCache from litellm.integrations.custom_logger import CustomLogger from litellm.types.guardrails import ( DynamicGuardrailParams, + GenericGuardrailAPIInputs, GuardrailEventHooks, LitellmParams, Mode, - PiiEntityType, ) from litellm.types.llms.openai import AllMessageValues from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -450,20 +450,22 @@ class CustomGuardrail(CustomLogger): async def apply_guardrail( self, - texts: List[str], + inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, - images: Optional[List[str]] = None, - ) -> Tuple[List[str], Optional[List[str]]]: + ) -> GenericGuardrailAPIInputs: """ - Apply your guardrail logic to the given text + Apply your guardrail logic to the given inputs Args: - texts: The texts to apply the guardrail to - images: The images to apply the guardrail to + inputs: Dictionary containing: + - texts: List of texts to apply the guardrail to + - images: Optional list of images to apply the guardrail to + - tool_calls: Optional list of tool calls to apply the guardrail to request_data: The request data dictionary - containing user api key metadata (e.g. user_id, team_id, etc.) input_type: The type of input to apply the guardrail to - "request" or "response" + logging_obj: Optional logging object for tracking the guardrail execution Any of the custom guardrails can override this method to provide custom guardrail logic @@ -474,7 +476,7 @@ class CustomGuardrail(CustomLogger): - If the guardrail raises an exception """ - return texts, images + return inputs def _process_response( self, diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 46e1a2c201f..21e1d562224 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -65,11 +65,11 @@ class DataDogLogger( `DD_SITE` - your datadog site, example = `"us5.datadoghq.com"` Optional environment variables (DataDog Agent): - `DD_AGENT_HOST` - hostname or IP of DataDog agent, example = `"localhost"` - `DD_AGENT_PORT` - port of DataDog agent (default: 10518 for logs) + `LITELLM_DD_AGENT_HOST` - hostname or IP of DataDog agent, example = `"localhost"` + `LITELLM_DD_AGENT_PORT` - port of DataDog agent (default: 10518 for logs) - Note: If DD_AGENT_HOST is set, logs will be sent to the agent instead of directly to DataDog API. - In this case, DD_API_KEY and DD_SITE are not required (agent handles authentication). + Note: We use LITELLM_DD_AGENT_HOST instead of DD_AGENT_HOST to avoid conflicts + with ddtrace which automatically sets DD_AGENT_HOST for APM tracing. """ try: verbose_logger.debug("Datadog: in init datadog logger") @@ -85,7 +85,8 @@ class DataDogLogger( ) # Configure DataDog endpoint (Agent or Direct API) - dd_agent_host = os.getenv("DD_AGENT_HOST") + # Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST + dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST") if dd_agent_host: self._configure_dd_agent(dd_agent_host=dd_agent_host) else: @@ -127,7 +128,7 @@ class DataDogLogger( Args: dd_agent_host: Hostname or IP of DataDog agent """ - dd_agent_port = os.getenv("DD_AGENT_PORT", "10518") # default port for logs + dd_agent_port = os.getenv("LITELLM_DD_AGENT_PORT", "10518") # default port for logs self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs" self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}") diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 90bf19b21fe..9f9d45d0e7d 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1065,14 +1065,7 @@ class OpenTelemetry(CustomLogger): self, span: Span, kwargs, response_obj: Optional[Any] ): try: - if self.callback_name == "arize_phoenix": - from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger - - ArizePhoenixLogger.set_arize_phoenix_attributes( - span, kwargs, response_obj - ) - return - elif self.callback_name == "langtrace": + if self.callback_name == "langtrace": from litellm.integrations.langtrace import LangtraceAttributes LangtraceAttributes().set_langtrace_attributes( @@ -1088,6 +1081,11 @@ class OpenTelemetry(CustomLogger): span, kwargs, response_obj ) return + elif self.callback_name == "weave_otel": + from litellm.integrations.weave.weave_otel import set_weave_otel_attributes + + set_weave_otel_attributes(span, kwargs, response_obj) + return from litellm.proxy._types import SpanAttributes optional_params = kwargs.get("optional_params", {}) diff --git a/litellm/integrations/weave/__init__.py b/litellm/integrations/weave/__init__.py new file mode 100644 index 00000000000..49af77b55e8 --- /dev/null +++ b/litellm/integrations/weave/__init__.py @@ -0,0 +1,7 @@ +""" +Weave (W&B) integration for LiteLLM via OpenTelemetry. +""" + +from litellm.integrations.weave.weave_otel import WeaveOtelLogger + +__all__ = ["WeaveOtelLogger"] diff --git a/litellm/integrations/weave/weave_otel.py b/litellm/integrations/weave/weave_otel.py new file mode 100644 index 00000000000..167deaf2cdc --- /dev/null +++ b/litellm/integrations/weave/weave_otel.py @@ -0,0 +1,329 @@ +from __future__ import annotations + +import base64 +import json +import os +from typing import TYPE_CHECKING, Any, Optional + +from opentelemetry.trace import Status, StatusCode +from typing_extensions import override + +from litellm._logging import verbose_logger +from litellm.integrations._types.open_inference import SpanAttributes as OpenInferenceSpanAttributes +from litellm.integrations.arize import _utils +from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig +from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes import ( + BaseLLMObsOTELAttributes, + safe_set_attribute, +) +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps +from litellm.types.integrations.weave_otel import WeaveOtelConfig, WeaveSpanAttributes +from litellm.types.utils import StandardCallbackDynamicParams + +if TYPE_CHECKING: + from opentelemetry.trace import Span + + +# Weave OTEL endpoint +# Multi-tenant cloud: https://trace.wandb.ai/otel/v1/traces +# Dedicated cloud: https://.wandb.io/traces/otel/v1/traces +WEAVE_BASE_URL = "https://trace.wandb.ai" +WEAVE_OTEL_ENDPOINT = "/otel/v1/traces" + + +class WeaveLLMObsOTELAttributes(BaseLLMObsOTELAttributes): + """ + Weave-specific LLM observability OTEL attributes. + + Weave automatically maps attributes from multiple frameworks including + GenAI, OpenInference, Langfuse, and others. + """ + + @staticmethod + @override + def set_messages(span: "Span", kwargs: dict[str, Any]): + """Set input messages as span attributes using OpenInference conventions.""" + + messages = kwargs.get("messages") or [] + optional_params = kwargs.get("optional_params") or {} + + prompt = {"messages": messages} + functions = optional_params.get("functions") + tools = optional_params.get("tools") + if functions is not None: + prompt["functions"] = functions + if tools is not None: + prompt["tools"] = tools + safe_set_attribute(span, OpenInferenceSpanAttributes.INPUT_VALUE, json.dumps(prompt)) + + +def _set_weave_specific_attributes(span: Span, kwargs: dict[str, Any], response_obj: Any): + """ + Sets Weave-specific metadata attributes onto the OTEL span. + + Based on Weave's OTEL attribute mappings from: + https://github.com/wandb/weave/blob/master/weave/trace_server/opentelemetry/constants.py + """ + + # Extract all needed data upfront + litellm_params = kwargs.get("litellm_params") or {} + # optional_params = kwargs.get("optional_params") or {} + metadata = kwargs.get("metadata") or {} + model = kwargs.get("model") or "" + custom_llm_provider = litellm_params.get("custom_llm_provider") or "" + + # Weave supports a custom display name and will default to the model name if not provided. + display_name = metadata.get("display_name") + if not display_name and model: + if custom_llm_provider: + display_name = f"{custom_llm_provider}/{model}" + else: + display_name = model + if display_name: + display_name = display_name.replace("/", "__") + safe_set_attribute(span, WeaveSpanAttributes.DISPLAY_NAME.value, display_name) + + # Weave threads are OpenInference sessions. + if (session_id := metadata.get("session_id")) is not None: + if isinstance(session_id, (list, dict)): + session_id = safe_dumps(session_id) + safe_set_attribute(span, WeaveSpanAttributes.THREAD_ID.value, session_id) + safe_set_attribute(span, WeaveSpanAttributes.IS_TURN.value, True) + + # Response attributes are already set by _utils.set_attributes, + # but we override them here to better match Weave's expectations + if response_obj: + output_dict = None + if hasattr(response_obj, "model_dump"): + output_dict = response_obj.model_dump() + elif hasattr(response_obj, "get"): + output_dict = response_obj + + if output_dict: + safe_set_attribute(span, OpenInferenceSpanAttributes.OUTPUT_VALUE, safe_dumps(output_dict)) + + +def _get_weave_authorization_header(api_key: str) -> str: + """ + Get the authorization header for Weave OpenTelemetry. + + Weave uses Basic auth with format: api: + """ + auth_string = f"api:{api_key}" + auth_header = base64.b64encode(auth_string.encode()).decode() + return f"Basic {auth_header}" + + +def get_weave_otel_config() -> WeaveOtelConfig: + """ + Retrieves the Weave OpenTelemetry configuration based on environment variables. + + Environment Variables: + WANDB_API_KEY: Required. W&B API key for authentication. + WANDB_PROJECT_ID: Required. Project ID in format /. + WANDB_HOST: Optional. Custom Weave host URL. Defaults to cloud endpoint. + + Returns: + WeaveOtelConfig: A Pydantic model containing Weave OTEL configuration. + + Raises: + ValueError: If required environment variables are missing. + """ + api_key = os.getenv("WANDB_API_KEY") + project_id = os.getenv("WANDB_PROJECT_ID") + host = os.getenv("WANDB_HOST") + + if not api_key: + raise ValueError("WANDB_API_KEY must be set for Weave OpenTelemetry integration.") + + if not project_id: + raise ValueError( + "WANDB_PROJECT_ID must be set for Weave OpenTelemetry integration. Format: /" + ) + + if host: + if not host.startswith("http"): + host = "https://" + host + # Self-managed instances use a different path + endpoint = host.rstrip("/") + WEAVE_OTEL_ENDPOINT + verbose_logger.debug(f"Using Weave OTEL endpoint from host: {endpoint}") + else: + endpoint = WEAVE_BASE_URL + WEAVE_OTEL_ENDPOINT + verbose_logger.debug(f"Using Weave cloud endpoint: {endpoint}") + + # Weave uses Basic auth with format: api: + auth_header = _get_weave_authorization_header(api_key=api_key) + otlp_auth_headers = f"Authorization={auth_header},project_id={project_id}" + + # Set standard OTEL environment variables + os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] = endpoint + os.environ["OTEL_EXPORTER_OTLP_HEADERS"] = otlp_auth_headers + + return WeaveOtelConfig( + otlp_auth_headers=otlp_auth_headers, + endpoint=endpoint, + project_id=project_id, + protocol="otlp_http", + ) + + +def set_weave_otel_attributes(span: Span, kwargs: dict[str, Any], response_obj: Any): + """ + Sets OpenTelemetry span attributes for Weave observability. + Uses the same attribute setting logic as other OTEL integrations for consistency. + """ + _utils.set_attributes(span, kwargs, response_obj, WeaveLLMObsOTELAttributes) + _set_weave_specific_attributes(span=span, kwargs=kwargs, response_obj=response_obj) + + +class WeaveOtelLogger(OpenTelemetry): + """ + Weave (W&B) OpenTelemetry Logger for LiteLLM. + + Sends LLM traces to Weave via the OpenTelemetry Protocol (OTLP). + + Environment Variables: + WANDB_API_KEY: Required. Weights & Biases API key for authentication. + WANDB_PROJECT_ID: Required. Project ID in format /. + WANDB_HOST: Optional. Custom Weave host URL. Defaults to cloud endpoint. + + Usage: + litellm.callbacks = ["weave_otel"] + + Or manually: + from litellm.integrations.weave.weave_otel import WeaveOtelLogger + weave_logger = WeaveOtelLogger(callback_name="weave_otel") + litellm.callbacks = [weave_logger] + + Reference: + https://docs.wandb.ai/weave/guides/tracking/otel + """ + + def __init__( + self, + config: Optional[OpenTelemetryConfig] = None, + callback_name: Optional[str] = "weave_otel", + **kwargs, + ): + """ + Initialize WeaveOtelLogger. + + If config is not provided, automatically configures from environment variables + (WANDB_API_KEY, WANDB_PROJECT_ID, WANDB_HOST) via get_weave_otel_config(). + """ + if config is None: + # Auto-configure from Weave environment variables + weave_config = get_weave_otel_config() + + config = OpenTelemetryConfig( + exporter=weave_config.protocol, + endpoint=weave_config.endpoint, + headers=weave_config.otlp_auth_headers, + ) + + super().__init__(config=config, callback_name=callback_name, **kwargs) + + def _maybe_log_raw_request(self, kwargs, response_obj, start_time, end_time, parent_span): + """ + Override to skip creating the raw_gen_ai_request child span. + + For Weave, we only want a single span per LLM call. The parent span + already contains all the necessary attributes, so the child span + is redundant. + """ + pass + + def _start_primary_span( + self, + kwargs, + response_obj, + start_time, + end_time, + context, + parent_span=None, + ): + """ + Override to always create a child span instead of reusing the parent span. + + This ensures that wrapper spans (like "B", "C", "D", "E") remain separate + from the LiteLLM LLM call spans, creating proper nesting in Weave. + """ + + otel_tracer = self.get_tracer_to_use_for_request(kwargs) + # Always create a new child span, even if parent_span is provided + # This ensures wrapper spans remain separate from LLM call spans + 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)) + return span + + def _handle_success(self, kwargs, response_obj, start_time, end_time): + """ + Override to prevent ending externally created parent spans. + + When wrapper spans (like "B", "C", "D", "E") are provided as parent spans, + they should be managed by the user code, not ended by LiteLLM. + """ + + verbose_logger.debug( + "Weave OpenTelemetry Logger: Logging kwargs: %s, OTEL config settings=%s", + kwargs, + self.config, + ) + ctx, parent_span = self._get_span_context(kwargs) + + # Always create a child span (handled by _start_primary_span override) + primary_span_parent = None + + # 1. Primary span + span = self._start_primary_span(kwargs, response_obj, start_time, end_time, ctx, primary_span_parent) + + # 2. Raw-request sub-span (skipped for Weave via _maybe_log_raw_request override) + self._maybe_log_raw_request(kwargs, response_obj, start_time, end_time, span) + + # 3. Guardrail span + self._create_guardrail_span(kwargs=kwargs, context=ctx) + + # 4. Metrics & cost recording + self._record_metrics(kwargs, response_obj, start_time, end_time) + + # 5. Semantic logs. + if self.config.enable_events: + self._emit_semantic_logs(kwargs, response_obj, span) + + # 6. Don't end parent span - it's managed by user code + # Since we always create a child span (never reuse parent), the parent span + # lifecycle is owned by the user. This prevents double-ending of wrapper spans + # like "B", "C", "D", "E" that users create and manage themselves. + + def construct_dynamic_otel_headers( + self, standard_callback_dynamic_params: StandardCallbackDynamicParams + ) -> dict | None: + """ + Construct dynamic Weave headers from standard callback dynamic params. + + This is used for team/key based logging. + + Returns: + dict: A dictionary of dynamic Weave headers + """ + dynamic_headers = {} + + dynamic_wandb_api_key = standard_callback_dynamic_params.get("wandb_api_key") + dynamic_weave_project_id = standard_callback_dynamic_params.get("weave_project_id") + + if dynamic_wandb_api_key: + auth_header = _get_weave_authorization_header( + api_key=dynamic_wandb_api_key, + ) + dynamic_headers["Authorization"] = auth_header + + if dynamic_weave_project_id: + dynamic_headers["project_id"] = dynamic_weave_project_id + + return dynamic_headers if dynamic_headers else None diff --git a/litellm/litellm_core_utils/README.md b/litellm/litellm_core_utils/README.md index 6494041291b..b61c8982762 100644 --- a/litellm/litellm_core_utils/README.md +++ b/litellm/litellm_core_utils/README.md @@ -9,4 +9,5 @@ Core files: - `default_encoding.py`: code for loading the default encoding (tiktoken) - `get_llm_provider_logic.py`: code for inferring the LLM provider from a given model name. - `duration_parser.py`: code for parsing durations - e.g. "1d", "1mo", "10s" +- `api_route_to_call_types.py`: mapping of API routes to their corresponding CallTypes (e.g., `/chat/completions` -> [acompletion, completion]) diff --git a/litellm/litellm_core_utils/api_route_to_call_types.py b/litellm/litellm_core_utils/api_route_to_call_types.py new file mode 100644 index 00000000000..35f83de1dd7 --- /dev/null +++ b/litellm/litellm_core_utils/api_route_to_call_types.py @@ -0,0 +1,38 @@ +""" +Dictionary mapping API routes to their corresponding CallTypes in LiteLLM. + +This dictionary maps each API endpoint to the CallTypes that can be used for that route. +Each route can have both async (prefixed with 'a') and sync call types. +""" + +from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes + + +def get_call_types_for_route(route: str) -> list: + """ + Get the list of CallTypes for a given API route. + + Args: + route: API route path (e.g., "/chat/completions") + + Returns: + List of CallTypes for that route, or empty list if route not found + """ + return API_ROUTE_TO_CALL_TYPES.get(route, []) + + +def get_routes_for_call_type(call_type: CallTypes) -> list: + """ + Get all routes that use a specific CallType. + + Args: + call_type: The CallType to search for + + Returns: + List of routes that use this CallType + """ + routes = [] + for route, types in API_ROUTE_TO_CALL_TYPES.items(): + if call_type in types: + routes.append(route) + return routes diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py index 538ef6be283..fdc9f374553 100644 --- a/litellm/litellm_core_utils/custom_logger_registry.py +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -75,6 +75,7 @@ class CustomLoggerRegistry: "langfuse_otel": OpenTelemetry, "arize_phoenix": OpenTelemetry, "langtrace": OpenTelemetry, + "weave_otel": OpenTelemetry, "mlflow": MlflowLogger, "langfuse": LangfusePromptManagement, "otel": OpenTelemetry, diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index b10011befcd..a90d16dba49 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -840,6 +840,16 @@ def _get_openai_compatible_provider_info( # noqa: PLR0915 ) = litellm.ClarifaiConfig()._get_openai_compatible_provider_info( api_base, api_key ) + elif custom_llm_provider == "ragflow": + full_model = f"ragflow/{model}" + ( + api_base, + dynamic_api_key, + _, + ) = litellm.RAGFlowConfig()._get_openai_compatible_provider_info( + full_model, api_base, api_key, "ragflow" + ) + model = full_model if api_base is not None and not isinstance(api_base, str): raise Exception("api base needs to be a string. api_base={}".format(api_base)) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0b0f483ff74..2b52d58e29d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -71,6 +71,7 @@ from litellm.litellm_core_utils.redact_messages import ( from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.llms.base_llm.search.transformation import SearchResponse from litellm.responses.utils import ResponseAPILoggingUtils +from litellm.types.agents import LiteLLMSendMessageResponse from litellm.types.containers.main import ContainerObject from litellm.types.llms.openai import ( AllMessageValues, @@ -1738,6 +1739,7 @@ class Logging(LiteLLMLoggingBaseClass): and logging_result.get("object") == "search" # Search API (dict format) or isinstance(logging_result, VideoObject) or isinstance(logging_result, ContainerObject) + or isinstance(logging_result, LiteLLMSendMessageResponse) # A2A or (self.call_type == CallTypes.call_mcp_tool.value) ): return True @@ -3617,15 +3619,15 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 for callback in _in_memory_loggers: if ( - isinstance(callback, OpenTelemetry) + isinstance(callback, ArizePhoenixLogger) and callback.callback_name == "arize_phoenix" ): return callback # type: ignore - _otel_logger = OpenTelemetry( + _arize_phoenix_otel_logger = ArizePhoenixLogger( config=otel_config, callback_name="arize_phoenix" ) - _in_memory_loggers.append(_otel_logger) - return _otel_logger # type: ignore + _in_memory_loggers.append(_arize_phoenix_otel_logger) + return _arize_phoenix_otel_logger # type: ignore elif logging_integration == "otel": from litellm.integrations.opentelemetry import OpenTelemetry @@ -3800,6 +3802,31 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 ) _in_memory_loggers.append(_otel_logger) return _otel_logger # type: ignore + elif logging_integration == "weave_otel": + from litellm.integrations.opentelemetry import ( + OpenTelemetryConfig, + ) + from litellm.integrations.weave.weave_otel import WeaveOtelLogger, get_weave_otel_config + + weave_otel_config = get_weave_otel_config() + + otel_config = OpenTelemetryConfig( + exporter=weave_otel_config.protocol, + endpoint=weave_otel_config.endpoint, + headers=weave_otel_config.otlp_auth_headers, + ) + + for callback in _in_memory_loggers: + if ( + isinstance(callback, WeaveOtelLogger) + and callback.callback_name == "weave_otel" + ): + return callback # type: ignore + _otel_logger = WeaveOtelLogger( + config=otel_config, callback_name="weave_otel" + ) + _in_memory_loggers.append(_otel_logger) + return _otel_logger # type: ignore elif logging_integration == "pagerduty": for callback in _in_memory_loggers: if isinstance(callback, PagerDutyAlerting): diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 262692d6d1a..04c8c235557 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3446,8 +3446,25 @@ class BedrockConverseMessagesProcessor: @staticmethod def _initial_message_setup( messages: List, + model: str, + llm_provider: str, user_continue_message: Optional[ChatCompletionUserMessage] = None, ) -> List: + # gracefully handle base case of no messages at all + if len(messages) == 0: + if user_continue_message is not None: + messages.append(user_continue_message) + elif litellm.modify_params: + messages.append(DEFAULT_USER_CONTINUE_MESSAGE) + else: + raise litellm.BadRequestError( + message=BAD_MESSAGE_ERROR_STR + + "bedrock requires at least one non-system message", + model=model, + llm_provider=llm_provider, + ) + + # if initial message is assistant message if messages[0].get("role") is not None and messages[0]["role"] == "assistant": if user_continue_message is not None: messages.insert(0, user_continue_message) @@ -3475,18 +3492,8 @@ class BedrockConverseMessagesProcessor: contents: List[BedrockMessageBlock] = [] msg_i = 0 - ## BASE CASE ## - if len(messages) == 0: - raise litellm.BadRequestError( - message=BAD_MESSAGE_ERROR_STR - + "bedrock requires at least one non-system message", - model=model, - llm_provider=llm_provider, - ) - - # if initial message is assistant message messages = BedrockConverseMessagesProcessor._initial_message_setup( - messages, user_continue_message + messages, model, llm_provider, user_continue_message ) while msg_i < len(messages): @@ -3847,28 +3854,9 @@ def _bedrock_converse_messages_pt( # noqa: PLR0915 contents: List[BedrockMessageBlock] = [] msg_i = 0 - ## BASE CASE ## - if len(messages) == 0: - raise litellm.BadRequestError( - message=BAD_MESSAGE_ERROR_STR - + "bedrock requires at least one non-system message", - model=model, - llm_provider=llm_provider, - ) - - # if initial message is assistant message - if messages[0].get("role") is not None and messages[0]["role"] == "assistant": - if user_continue_message is not None: - messages.insert(0, user_continue_message) - elif litellm.modify_params: - messages.insert(0, DEFAULT_USER_CONTINUE_MESSAGE) - - # if final message is assistant message - if messages[-1].get("role") is not None and messages[-1]["role"] == "assistant": - if user_continue_message is not None: - messages.append(user_continue_message) - elif litellm.modify_params: - messages.append(DEFAULT_USER_CONTINUE_MESSAGE) + messages = BedrockConverseMessagesProcessor._initial_message_setup( + messages, model, llm_provider, user_continue_message + ) while msg_i < len(messages): user_content: List[BedrockContentBlock] = [] diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 4d8e109d882..a7f460fab59 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -96,9 +96,9 @@ class CustomStreamWrapper: self.system_fingerprint: Optional[str] = None self.received_finish_reason: Optional[str] = None - self.intermittent_finish_reason: Optional[str] = ( - None # finish reasons that show up mid-stream - ) + self.intermittent_finish_reason: Optional[ + str + ] = None # finish reasons that show up mid-stream self.special_tokens = [ "<|assistant|>", "<|system|>", @@ -735,7 +735,7 @@ class CustomStreamWrapper: and completion_obj["function_call"] is not None ) or ( - "tool_calls" in model_response.choices[0].delta + "tool_calls" in model_response.choices[0].delta and model_response.choices[0].delta["tool_calls"] is not None ) or ( @@ -889,7 +889,6 @@ class CustomStreamWrapper: ## check if openai/azure chunk original_chunk = response_obj.get("original_chunk", None) if original_chunk: - if len(original_chunk.choices) > 0: choices = [] for choice in original_chunk.choices: @@ -906,7 +905,6 @@ class CustomStreamWrapper: print_verbose(f"choices in streaming: {choices}") setattr(model_response, "choices", choices) else: - return model_response.system_fingerprint = ( original_chunk.system_fingerprint @@ -1435,9 +1433,9 @@ class CustomStreamWrapper: _json_delta = delta.model_dump() print_verbose(f"_json_delta: {_json_delta}") if "role" not in _json_delta or _json_delta["role"] is None: - _json_delta["role"] = ( - "assistant" # mistral's api returns role as None - ) + _json_delta[ + "role" + ] = "assistant" # mistral's api returns role as None if "tool_calls" in _json_delta and isinstance( _json_delta["tool_calls"], list ): @@ -1533,7 +1531,7 @@ class CustomStreamWrapper: async def _call_post_streaming_deployment_hook(self, chunk): """ Call the post-call streaming deployment hook for callbacks. - + This allows callbacks to modify streaming chunks before they're returned. """ try: @@ -1544,15 +1542,17 @@ class CustomStreamWrapper: # Get request kwargs from logging object request_data = self.logging_obj.model_call_details call_type_str = self.logging_obj.call_type - + try: typed_call_type = CallTypes(call_type_str) except ValueError: typed_call_type = None - + # Call hooks for all callbacks for callback in litellm.callbacks: - if isinstance(callback, CustomLogger) and hasattr(callback, "async_post_call_streaming_deployment_hook"): + if isinstance(callback, CustomLogger) and hasattr( + callback, "async_post_call_streaming_deployment_hook" + ): result = await callback.async_post_call_streaming_deployment_hook( request_data=request_data, response_chunk=chunk, @@ -1560,11 +1560,14 @@ class CustomStreamWrapper: ) if result is not None: chunk = result - + return chunk except Exception as e: from litellm._logging import verbose_logger - verbose_logger.exception(f"Error in post-call streaming deployment hook: {str(e)}") + + verbose_logger.exception( + f"Error in post-call streaming deployment hook: {str(e)}" + ) return chunk def cache_streaming_response(self, processed_chunk, cache_hit: bool): @@ -1687,7 +1690,7 @@ class CustomStreamWrapper: response, "usage" ): # remove usage from chunk, only send on final chunk # Convert the object to a dictionary - obj_dict = response.dict() + obj_dict = response.model_dump() # Remove an attribute (e.g., 'attr2') if "usage" in obj_dict: @@ -1852,7 +1855,7 @@ class CustomStreamWrapper: processed_chunk, "usage" ): # remove usage from chunk, only send on final chunk # Convert the object to a dictionary - obj_dict = processed_chunk.dict() + obj_dict = processed_chunk.model_dump() # Remove an attribute (e.g., 'attr2') if "usage" in obj_dict: @@ -1872,11 +1875,15 @@ class CustomStreamWrapper: if self.sent_last_chunk is True and self.stream_options is None: usage = calculate_total_usage(chunks=self.chunks) processed_chunk._hidden_params["usage"] = usage - + # Call post-call streaming deployment hook for final chunk if self.sent_last_chunk is True: - processed_chunk = await self._call_post_streaming_deployment_hook(processed_chunk) - + processed_chunk = ( + await self._call_post_streaming_deployment_hook( + processed_chunk + ) + ) + return processed_chunk raise StopAsyncIteration else: # temporary patch for non-aiohttp async calls @@ -1890,9 +1897,9 @@ class CustomStreamWrapper: chunk = next(self.completion_stream) if chunk is not None and chunk != b"": print_verbose(f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}") - processed_chunk: Optional[ModelResponseStream] = ( - self.chunk_creator(chunk=chunk) - ) + processed_chunk: Optional[ + ModelResponseStream + ] = self.chunk_creator(chunk=chunk) print_verbose( f"PROCESSED CHUNK POST CHUNK CREATOR: {processed_chunk}" ) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 383a4941238..9b6511f151c 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -12,10 +12,17 @@ Pattern Overview: 4. Apply guardrail responses back to the original structure """ +import json from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast from litellm._logging import verbose_proxy_logger +from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, +) from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.guardrails import GenericGuardrailAPIInputs +from litellm.types.llms.anthropic import AllAnthropicToolsValues +from litellm.types.llms.openai import ChatCompletionToolParam if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -36,6 +43,10 @@ class AnthropicMessagesHandler(BaseTranslation): Methods can be overridden to customize behavior for different message formats. """ + def __init__(self): + super().__init__() + self.adapter = LiteLLMAnthropicMessagesAdapter() + async def process_input_messages( self, data: dict, @@ -46,11 +57,13 @@ class AnthropicMessagesHandler(BaseTranslation): Process input messages by applying guardrails to text content. """ messages = data.get("messages") + tools = data.get("tools", None) if messages is None: return data texts_to_check: List[str] = [] images_to_check: List[str] = [] + tools_to_check: List[ChatCompletionToolParam] = [] task_mappings: List[Tuple[int, Optional[int]]] = [] # Track (message_index, content_index) for each text # content_index is None for string content, int for list content @@ -65,18 +78,28 @@ class AnthropicMessagesHandler(BaseTranslation): task_mappings=task_mappings, ) + if tools is not None: + self._extract_input_tools( + tools=tools, + tools_to_check=tools_to_check, + ) + # Step 2: Apply guardrail to all texts in batch if texts_to_check: - guardrailed_texts, guardrailed_images = ( - await guardrail_to_apply.apply_guardrail( - texts=texts_to_check, - request_data=data, - input_type="request", - images=images_to_check if images_to_check else None, - logging_obj=litellm_logging_obj, - ) + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + if tools_to_check: + inputs["tools"] = tools_to_check + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) + # Step 3: Map guardrail responses back to original message structure await self._apply_guardrail_responses_to_input( messages=messages, @@ -104,15 +127,17 @@ class AnthropicMessagesHandler(BaseTranslation): Override this method to customize text/image extraction logic. """ content = message.get("content", None) - if content is None: + tools = message.get("tools", None) + if content is None and tools is None: return - if isinstance(content, str): + ## CHECK FOR TEXT + IMAGES + if content is not None and isinstance(content, str): # Simple string content texts_to_check.append(content) task_mappings.append((msg_idx, None)) - elif isinstance(content, list): + elif content is not None and isinstance(content, list): # List content (e.g., multimodal with text and images) for content_idx, content_item in enumerate(content): # Extract text @@ -130,6 +155,22 @@ class AnthropicMessagesHandler(BaseTranslation): if data: images_to_check.append(data) + def _extract_input_tools( + self, + tools: List[Dict[str, Any]], + tools_to_check: List[ChatCompletionToolParam], + ) -> None: + """ + Extract tools from a message. + """ + ## CHECK FOR TOOLS + if tools is not None and isinstance(tools, list): + # TRANSFORM ANTHROPIC TOOLS TO OPENAI TOOLS + openai_tools = self.adapter.translate_anthropic_tools_to_openai( + tools=cast(List[AllAnthropicToolsValues], tools) + ) + tools_to_check.extend(openai_tools) + async def _apply_guardrail_responses_to_input( self, messages: List[Dict[str, Any]], @@ -223,16 +264,18 @@ class AnthropicMessagesHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata - guardrailed_texts, guardrailed_images = ( - await guardrail_to_apply.apply_guardrail( - texts=texts_to_check, - request_data=request_data, - input_type="response", - images=images_to_check if images_to_check else None, - logging_obj=litellm_logging_obj, - ) + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) + # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( response=response, @@ -246,6 +289,112 @@ class AnthropicMessagesHandler(BaseTranslation): return response + async def process_output_streaming_response( + self, + responses_so_far: List[Any], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + user_api_key_dict: Optional[Any] = None, + ) -> List[Any]: + """ + Process output streaming response by applying guardrails to text content. + + Get the string so far, check the apply guardrail to the string so far, and return the list of responses so far. + """ + string_so_far = self.get_streaming_string_so_far(responses_so_far) + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid + inputs={"texts": [string_so_far]}, + request_data={}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + def get_streaming_string_so_far(self, responses_so_far: List[Any]) -> str: + """ + Parse streaming responses and extract accumulated text content. + + Handles two formats: + 1. Raw bytes in SSE (Server-Sent Events) format from Anthropic API + 2. Parsed dict objects (for backwards compatibility) + + SSE format example: + b'event: content_block_delta\\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" curious"}}\\n\\n' + + Dict format example: + { + "type": "content_block_delta", + "index": 0, + "delta": { + "type": "text_delta", + "text": " curious" + } + } + """ + text_so_far = "" + for response in responses_so_far: + # Handle raw bytes in SSE format + if isinstance(response, bytes): + text_so_far += self._extract_text_from_sse(response) + # Handle already-parsed dict format + elif isinstance(response, dict): + delta = response.get("delta") if response.get("delta") else None + if delta and delta.get("type") == "text_delta": + text = delta.get("text", "") + if text: + text_so_far += text + return text_so_far + + def _extract_text_from_sse(self, sse_bytes: bytes) -> str: + """ + Extract text content from Server-Sent Events (SSE) format. + + Args: + sse_bytes: Raw bytes in SSE format + + Returns: + Accumulated text from all content_block_delta events + """ + text = "" + try: + # Decode bytes to string + sse_string = sse_bytes.decode("utf-8") + + # Split by double newline to get individual events + events = sse_string.split("\n\n") + + for event in events: + if not event.strip(): + continue + + # Parse event lines + lines = event.strip().split("\n") + event_type = None + data_line = None + + for line in lines: + if line.startswith("event:"): + event_type = line[6:].strip() + elif line.startswith("data:"): + data_line = line[5:].strip() + + # Only process content_block_delta events + if event_type == "content_block_delta" and data_line: + try: + data = json.loads(data_line) + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + text += delta.get("text", "") + except json.JSONDecodeError: + verbose_proxy_logger.warning( + f"Failed to parse JSON from SSE data: {data_line}" + ) + + except Exception as e: + verbose_proxy_logger.error(f"Error extracting text from SSE: {e}") + + return text + def _has_text_content(self, response: "AnthropicMessagesResponse") -> bool: """ Check if response has any text content to process. diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index b363b747de5..36156d56a59 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -436,9 +436,7 @@ class AnthropicChatCompletion(BaseLLM): else: if client is None or not isinstance(client, HTTPHandler): - client = _get_httpx_client( - params={"timeout": timeout} - ) + client = _get_httpx_client(params={"timeout": timeout}) else: client = client @@ -528,9 +526,7 @@ class ModelResponseIterator: usage_object=cast(dict, anthropic_usage_chunk), reasoning_content=None ) - def _content_block_delta_helper( - self, chunk: dict - ) -> Tuple[ + def _content_block_delta_helper(self, chunk: dict) -> Tuple[ str, Optional[ChatCompletionToolCallChunk], List[Union[ChatCompletionThinkingBlock, ChatCompletionRedactedThinkingBlock]], diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 4221dfacf34..bdc986ae27f 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -786,6 +786,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): valid_content: bool = False system_message_block = ChatCompletionSystemMessage(**message) if isinstance(system_message_block["content"], str): + # Skip empty text blocks - Anthropic API raises errors for empty text + if not system_message_block["content"]: + continue anthropic_system_message_content = AnthropicSystemMessageContent( type="text", text=system_message_block["content"], @@ -800,10 +803,14 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): valid_content = True elif isinstance(message["content"], list): for _content in message["content"]: + # Skip empty text blocks - Anthropic API raises errors for empty text + text_value = _content.get("text") + if _content.get("type") == "text" and not text_value: + continue anthropic_system_message_content = ( AnthropicSystemMessageContent( type=_content.get("type"), - text=_content.get("text"), + text=text_value, ) ) if "cache_control" in _content: diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index e7aa93ac882..994afa26e9c 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -1020,7 +1020,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): headers: dict, client=None, timeout=None, - ) -> litellm.ImageResponse: + ) -> ImageResponse: response: Optional[dict] = None try: diff --git a/litellm/llms/azure/chat/gpt_5_transformation.py b/litellm/llms/azure/chat/gpt_5_transformation.py index 209475730f8..87f81d117f0 100644 --- a/litellm/llms/azure/chat/gpt_5_transformation.py +++ b/litellm/llms/azure/chat/gpt_5_transformation.py @@ -40,7 +40,11 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): or optional_params.get("reasoning_effort") ) - if reasoning_effort_value == "none": + # gpt-5.1 supports reasoning_effort='none', but other gpt-5 models don't + # See: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning + is_gpt_5_1 = self.is_model_gpt_5_1_model(model) + + if reasoning_effort_value == "none" and not is_gpt_5_1: if litellm.drop_params is True or ( drop_params is not None and drop_params is True ): @@ -54,7 +58,7 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): raise UnsupportedParamsError( status_code=400, message=( - "Azure OpenAI does not support reasoning_effort='none'. " + "Azure OpenAI does not support reasoning_effort='none' for this model. " "Supported values are: 'low', 'medium', and 'high'. " "To drop this parameter, set `litellm.drop_params=True` or for proxy:\n\n" "`litellm_settings:\n drop_params: true`\n" @@ -70,7 +74,8 @@ class AzureOpenAIGPT5Config(AzureOpenAIConfig, OpenAIGPT5Config): drop_params=drop_params, ) - if result.get("reasoning_effort") == "none": + # Only drop reasoning_effort='none' for non-gpt-5.1 models + if result.get("reasoning_effort") == "none" and not is_gpt_5_1: result.pop("reasoning_effort") return result diff --git a/litellm/llms/azure_ai/embed/handler.py b/litellm/llms/azure_ai/embed/handler.py index 13b8cc4cf29..67733d1ccb5 100644 --- a/litellm/llms/azure_ai/embed/handler.py +++ b/litellm/llms/azure_ai/embed/handler.py @@ -58,7 +58,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): data: ImageEmbeddingRequest, timeout: float, logging_obj, - model_response: litellm.EmbeddingResponse, + model_response: EmbeddingResponse, optional_params: dict, api_key: Optional[str], api_base: Optional[str], @@ -138,7 +138,7 @@ class AzureAIEmbedding(OpenAIChatCompletion): input: List, timeout: float, logging_obj, - model_response: litellm.EmbeddingResponse, + model_response: EmbeddingResponse, optional_params: dict, api_key: Optional[str] = None, api_base: Optional[str] = None, diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 5acbf4e9f4f..7106c207bd6 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,5 +1,5 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, Dict, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -84,3 +84,17 @@ class BaseTranslation(ABC): user_api_key_dict: User API key metadata (passed separately since response doesn't contain it) """ pass + + async def process_output_streaming_response( + self, + responses_so_far: List[Any], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, + user_api_key_dict: Optional["UserAPIKeyAuth"] = None, + ) -> Any: + """ + Process output streaming response with guardrails. + + Optional to override in subclasses. + """ + return responses_so_far diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index ed658c793af..816b93edd20 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -353,6 +353,10 @@ class BaseAWSLLM: model_id = BaseAWSLLM._get_model_id_from_model_with_spec( model_id, spec="deepseek_r1" ) + elif provider == "openai" and "openai/" in model_id: + model_id = BaseAWSLLM._get_model_id_from_model_with_spec( + model_id, spec="openai" + ) return model_id @staticmethod diff --git a/litellm/llms/bedrock/batches/handler.py b/litellm/llms/bedrock/batches/handler.py new file mode 100644 index 00000000000..4a26bd43348 --- /dev/null +++ b/litellm/llms/bedrock/batches/handler.py @@ -0,0 +1,96 @@ +from openai.types.batch import BatchRequestCounts +from openai.types.batch import Metadata as OpenAIBatchMetadata + +from litellm.types.utils import LiteLLMBatch + + +class BedrockBatchesHandler: + """ + Handler for Bedrock Batches. + + Specific providers/models needed some special handling. + + E.g. Twelve Labs Embedding Async Invoke + """ + @staticmethod + def _handle_async_invoke_status( + batch_id: str, aws_region_name: str, logging_obj=None, **kwargs + ) -> "LiteLLMBatch": + """ + Handle async invoke status check for AWS Bedrock. + + This is for Twelve Labs Embedding Async Invoke. + + Args: + batch_id: The async invoke ARN + aws_region_name: AWS region name + **kwargs: Additional parameters + + Returns: + dict: Status information including status, output_file_id (S3 URL), etc. + """ + import asyncio + + from litellm.llms.bedrock.embed.embedding import BedrockEmbedding + + async def _async_get_status(): + # Create embedding handler instance + embedding_handler = BedrockEmbedding() + + # Get the status of the async invoke job + status_response = await embedding_handler._get_async_invoke_status( + invocation_arn=batch_id, + aws_region_name=aws_region_name, + logging_obj=logging_obj, + **kwargs, + ) + + # Transform response to a LiteLLMBatch object + from litellm.types.utils import LiteLLMBatch + + openai_batch_metadata: OpenAIBatchMetadata = { + "output_file_id": status_response["outputDataConfig"][ + "s3OutputDataConfig" + ]["s3Uri"], + "failure_message": status_response.get("failureMessage") or "", + "model_arn": status_response["modelArn"], + } + + result = LiteLLMBatch( + id=status_response["invocationArn"], + object="batch", + status=status_response["status"], + created_at=status_response["submitTime"], + in_progress_at=status_response["lastModifiedTime"], + completed_at=status_response.get("endTime"), + failed_at=status_response.get("endTime") + if status_response["status"] == "failed" + else None, + request_counts=BatchRequestCounts( + total=1, + completed=1 if status_response["status"] == "completed" else 0, + failed=1 if status_response["status"] == "failed" else 0, + ), + metadata=openai_batch_metadata, + completion_window="24h", + endpoint="/v1/embeddings", + input_file_id="", + ) + + return result + + # Since this function is called from within an async context via run_in_executor, + # we need to create a new event loop in a thread to avoid conflicts + import concurrent.futures + + def run_in_thread(): + new_loop = asyncio.new_event_loop() + asyncio.set_event_loop(new_loop) + try: + return new_loop.run_until_complete(_async_get_status()) + finally: + new_loop.close() + + with concurrent.futures.ThreadPoolExecutor() as executor: + future = executor.submit(run_in_thread) + return future.result() diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index fd1f6f0c893..d5bd054118d 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -29,6 +29,7 @@ def make_sync_call( logging_obj: LiteLLMLoggingObject, json_mode: Optional[bool] = False, fake_stream: bool = False, + stream_chunk_size: int = 1024, ): if client is None: client = _get_httpx_client() # Create a new client if none provided @@ -66,7 +67,7 @@ def make_sync_call( ) else: decoder = AWSEventStreamDecoder(model=model) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) # LOGGING logging_obj.post_call( @@ -102,6 +103,7 @@ class BedrockConverseLLM(BaseAWSLLM): fake_stream: bool = False, json_mode: Optional[bool] = False, api_key: Optional[str] = None, + stream_chunk_size: int = 1024, ) -> CustomStreamWrapper: request_data = await litellm.AmazonConverseConfig()._async_transform_request( model=model, @@ -143,6 +145,7 @@ class BedrockConverseLLM(BaseAWSLLM): logging_obj=logging_obj, fake_stream=fake_stream, json_mode=json_mode, + stream_chunk_size=stream_chunk_size, ) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, @@ -260,6 +263,7 @@ class BedrockConverseLLM(BaseAWSLLM): ): ## SETUP ## stream = optional_params.pop("stream", None) + stream_chunk_size = optional_params.pop("stream_chunk_size", 1024) unencoded_model_id = optional_params.pop("model_id", None) fake_stream = optional_params.pop("fake_stream", False) json_mode = optional_params.get("json_mode", False) @@ -356,7 +360,8 @@ class BedrockConverseLLM(BaseAWSLLM): json_mode=json_mode, fake_stream=fake_stream, credentials=credentials, - api_key=api_key + api_key=api_key, + stream_chunk_size=stream_chunk_size, ) # type: ignore ### ASYNC COMPLETION return self.async_completion( @@ -433,6 +438,7 @@ class BedrockConverseLLM(BaseAWSLLM): logging_obj=logging_obj, json_mode=json_mode, fake_stream=fake_stream, + stream_chunk_size=stream_chunk_size, ) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 3b3a138ec67..705f3c9e630 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -246,6 +246,93 @@ class AmazonConverseConfig(BaseConfig): llm_provider="bedrock", ) + def _is_nova_lite_2_model(self, model: str) -> bool: + """ + Check if the model is a Nova Lite 2 model that supports reasoningConfig. + + Nova Lite 2 models use a different reasoning configuration structure compared to + Anthropic's thinking parameter and GPT-OSS's reasoning_effort parameter. + + Supported models: + - amazon.nova-2-lite-v1:0 + - us.amazon.nova-2-lite-v1:0 + - eu.amazon.nova-2-lite-v1:0 + - apac.amazon.nova-2-lite-v1:0 + + Args: + model: The model identifier + + Returns: + True if the model is a Nova Lite 2 model, False otherwise + + Examples: + >>> config = AmazonConverseConfig() + >>> config._is_nova_lite_2_model("amazon.nova-2-lite-v1:0") + True + >>> config._is_nova_lite_2_model("us.amazon.nova-2-lite-v1:0") + True + >>> config._is_nova_lite_2_model("amazon.nova-pro-1-5-v1:0") + False + >>> config._is_nova_lite_2_model("amazon.nova-pro-v1:0") + False + """ + # Remove regional prefix if present (us., eu., apac.) + model_without_region = model + for prefix in ["us.", "eu.", "apac."]: + if model.startswith(prefix): + model_without_region = model[len(prefix) :] + break + + # Check if the model is specifically Nova Lite 2 + return "nova-2-lite" in model_without_region + + def _transform_reasoning_effort_to_reasoning_config( + self, reasoning_effort: str + ) -> dict: + """ + Transform reasoning_effort parameter to Nova 2 reasoningConfig structure. + + Nova 2 models use a reasoningConfig structure in additionalModelRequestFields + that differs from both Anthropic's thinking parameter and GPT-OSS's reasoning_effort. + + Args: + reasoning_effort: The reasoning effort level, must be "low" or "high" + + Returns: + dict: A dictionary containing the reasoningConfig structure: + { + "reasoningConfig": { + "type": "enabled", + "maxReasoningEffort": "low" | "medium" |"high" + } + } + + Raises: + BadRequestError: If reasoning_effort is not "low", "medium" or "high" + + Examples: + >>> config = AmazonConverseConfig() + >>> config._transform_reasoning_effort_to_reasoning_config("high") + {'reasoningConfig': {'type': 'enabled', 'maxReasoningEffort': 'high'}} + >>> config._transform_reasoning_effort_to_reasoning_config("low") + {'reasoningConfig': {'type': 'enabled', 'maxReasoningEffort': 'low'}} + """ + valid_values = ["low", "medium", "high"] + if reasoning_effort not in valid_values: + raise litellm.exceptions.BadRequestError( + message=f"Invalid reasoning_effort value '{reasoning_effort}' for Nova 2 models. " + f"Supported values: {valid_values}", + model="amazon.nova-2-lite-v1:0", + llm_provider="bedrock_converse", + ) + + return { + "reasoningConfig": { + "type": "enabled", + "maxReasoningEffort": reasoning_effort, + } + } + def get_supported_openai_params(self, model: str) -> List[str]: from litellm.utils import supports_function_calling @@ -299,6 +386,10 @@ class AmazonConverseConfig(BaseConfig): if "gpt-oss" in model: supported_params.append("reasoning_effort") + elif self._is_nova_lite_2_model(model): + # Nova Lite 2 models support reasoning_effort (transformed to reasoningConfig) + # These models use a different reasoning structure than Anthropic's thinking parameter + supported_params.append("reasoning_effort") elif ( "claude-3-7" in model or "claude-sonnet-4" in model @@ -564,6 +655,12 @@ class AmazonConverseConfig(BaseConfig): # GPT-OSS models: keep reasoning_effort as-is # It will be passed through to additionalModelRequestFields optional_params["reasoning_effort"] = value + elif self._is_nova_lite_2_model(model): + # Nova Lite 2 models: transform to reasoningConfig + reasoning_config = ( + self._transform_reasoning_effort_to_reasoning_config(value) + ) + optional_params.update(reasoning_config) else: # Anthropic and other models: convert to thinking parameter optional_params["thinking"] = AnthropicConfig._map_reasoning_effort( @@ -574,8 +671,9 @@ class AmazonConverseConfig(BaseConfig): self._validate_request_metadata(value) # type: ignore optional_params["requestMetadata"] = value - # Only update thinking tokens for non-GPT-OSS models - if "gpt-oss" not in model: + # Only update thinking tokens for non-GPT-OSS models and non-Nova-Lite-2 models + # Nova Lite 2 handles token budgeting differently through reasoningConfig + if "gpt-oss" not in model and not self._is_nova_lite_2_model(model): self.update_optional_params_with_thinking_tokens( non_default_params=non_default_params, optional_params=optional_params ) diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index b35e86cabd2..5e33a266449 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -73,6 +73,9 @@ bedrock_tool_name_mappings: InMemoryCache = InMemoryCache( max_size_in_memory=50, default_ttl=600 ) from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig +from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import ( + AmazonBedrockOpenAIConfig, +) converse_config = AmazonConverseConfig() @@ -189,6 +192,7 @@ async def make_call( fake_stream: bool = False, json_mode: Optional[bool] = False, bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None, + stream_chunk_size: int = 1024, ): try: if client is None: @@ -232,7 +236,7 @@ async def make_call( json_mode=json_mode, ) completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) + response.aiter_bytes(chunk_size=stream_chunk_size) ) elif bedrock_invoke_provider == "deepseek_r1": decoder = AmazonDeepSeekR1StreamDecoder( @@ -240,12 +244,12 @@ async def make_call( sync_stream=False, ) completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) + response.aiter_bytes(chunk_size=stream_chunk_size) ) else: decoder = AWSEventStreamDecoder(model=model) completion_stream = decoder.aiter_bytes( - response.aiter_bytes(chunk_size=1024) + response.aiter_bytes(chunk_size=stream_chunk_size) ) # LOGGING @@ -278,6 +282,7 @@ def make_sync_call( fake_stream: bool = False, json_mode: Optional[bool] = False, bedrock_invoke_provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL] = None, + stream_chunk_size: int = 1024, ): try: if client is None: @@ -318,16 +323,16 @@ def make_sync_call( sync_stream=True, json_mode=json_mode, ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) elif bedrock_invoke_provider == "deepseek_r1": decoder = AmazonDeepSeekR1StreamDecoder( model=model, sync_stream=True, ) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) else: decoder = AWSEventStreamDecoder(model=model) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) # LOGGING logging_obj.post_call( @@ -401,6 +406,10 @@ class BedrockLLM(BaseAWSLLM): prompt = prompt_factory( model=model, messages=messages, custom_llm_provider="bedrock" ) + elif provider == "openai": + # OpenAI uses messages directly, no prompt conversion needed + # Return empty prompt as it won't be used + prompt = "" elif provider == "cohere": prompt, chat_history = cohere_message_pt(messages=messages) else: @@ -578,6 +587,30 @@ class BedrockLLM(BaseAWSLLM): ) elif provider == "meta" or provider == "llama": outputText = completion_response["generation"] + elif provider == "openai": + # OpenAI imported models use OpenAI Chat Completions format + if "choices" in completion_response and len(completion_response["choices"]) > 0: + choice = completion_response["choices"][0] + if "message" in choice: + outputText = choice["message"].get("content") + elif "text" in choice: # fallback for completion format + outputText = choice["text"] + + # Set finish reason + if "finish_reason" in choice: + model_response.choices[0].finish_reason = map_finish_reason( + choice["finish_reason"] + ) + + # Set usage if available + if "usage" in completion_response: + usage = completion_response["usage"] + _usage = litellm.Usage( + prompt_tokens=usage.get("prompt_tokens", 0), + completion_tokens=usage.get("completion_tokens", 0), + total_tokens=usage.get("total_tokens", 0), + ) + setattr(model_response, "usage", _usage) elif provider == "mistral": outputText = completion_response["outputs"][0]["text"] model_response.choices[0].finish_reason = completion_response[ @@ -698,6 +731,7 @@ class BedrockLLM(BaseAWSLLM): ## SETUP ## stream = optional_params.pop("stream", None) + stream_chunk_size = optional_params.pop("stream_chunk_size", 1024) provider = self.get_bedrock_invoke_provider(model) modelId = self.get_bedrock_model_id( @@ -895,6 +929,20 @@ class BedrockLLM(BaseAWSLLM): ): # completion(top_k=3) > anthropic_config(top_k=3) <- allows for dynamic variables to be passed in inference_params[k] = v data = json.dumps({"prompt": prompt, **inference_params}) + elif provider == "openai": + ## OpenAI imported models use OpenAI Chat Completions format (messages-based) + # Use AmazonBedrockOpenAIConfig for proper OpenAI transformation + openai_config = AmazonBedrockOpenAIConfig() + supported_params = openai_config.get_supported_openai_params(model=model) + + # Filter to only supported OpenAI params + filtered_params = { + k: v for k, v in inference_params.items() + if k in supported_params + } + + # OpenAI uses messages format, not prompt + data = json.dumps({"messages": messages, **filtered_params}) else: ## LOGGING logging_obj.pre_call( @@ -958,6 +1006,7 @@ class BedrockLLM(BaseAWSLLM): headers=prepped.headers, timeout=timeout, client=client, + stream_chunk_size=stream_chunk_size, ) # type: ignore ### ASYNC COMPLETION return self.async_completion( @@ -1003,7 +1052,7 @@ class BedrockLLM(BaseAWSLLM): decoder = AWSEventStreamDecoder(model=model) - completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=1024)) + completion_stream = decoder.iter_bytes(response.iter_bytes(chunk_size=stream_chunk_size)) streaming_response = CustomStreamWrapper( completion_stream=completion_stream, model=model, @@ -1123,6 +1172,7 @@ class BedrockLLM(BaseAWSLLM): logger_fn=None, headers={}, client: Optional[AsyncHTTPHandler] = None, + stream_chunk_size: int = 1024, ) -> CustomStreamWrapper: # The call is not made here; instead, we prepare the necessary objects for the stream. @@ -1138,6 +1188,7 @@ class BedrockLLM(BaseAWSLLM): messages=messages, logging_obj=logging_obj, fake_stream=True if "ai21" in api_base else False, + stream_chunk_size=stream_chunk_size, ), model=model, custom_llm_provider="bedrock", diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py index a81d55f0ad2..3506c8f1cc0 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_nova_transformation.py @@ -10,7 +10,6 @@ from typing import Any, List, Optional import httpx -import litellm from litellm.litellm_core_utils.litellm_logging import Logging from litellm.types.llms.bedrock import BedrockInvokeNovaRequest from litellm.types.llms.openai import AllMessageValues @@ -80,7 +79,7 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig): encoding: Any, api_key: Optional[str] = None, json_mode: Optional[bool] = None, - ) -> litellm.ModelResponse: + ) -> ModelResponse: return AmazonConverseConfig.transform_response( self, model, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py new file mode 100644 index 00000000000..c532d8ea27c --- /dev/null +++ b/litellm/llms/bedrock/chat/invoke_transformations/amazon_qwen2_transformation.py @@ -0,0 +1,98 @@ +""" +Handles transforming requests for `bedrock/invoke/{qwen2} models` + +Inherits from `AmazonQwen3Config` since Qwen2 and Qwen3 architectures are mostly similar. +The main difference is in the response format: Qwen2 uses "text" field while Qwen3 uses "generation" field. + +Qwen2 + Invoke API Tutorial: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html +""" + +from typing import Any, List, Optional + +import httpx + +from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen3_transformation import ( + AmazonQwen3Config, +) +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + LiteLLMLoggingObj, +) +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import ModelResponse + + +class AmazonQwen2Config(AmazonQwen3Config): + """ + Config for sending `qwen2` requests to `/bedrock/invoke/` + + Inherits from AmazonQwen3Config since Qwen2 and Qwen3 architectures are mostly similar. + The main difference is in the response format: Qwen2 uses "text" field while Qwen3 uses "generation" field. + + Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/invoke-imported-model.html + """ + + def transform_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ModelResponse, + logging_obj: LiteLLMLoggingObj, + request_data: dict, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + encoding: Any, + api_key: Optional[str] = None, + json_mode: Optional[bool] = None, + ) -> ModelResponse: + """ + Transform Qwen2 Bedrock response to OpenAI format + + Qwen2 uses "text" field, but we also support "generation" field for compatibility. + """ + try: + if hasattr(raw_response, 'json'): + response_data = raw_response.json() + else: + response_data = raw_response + + # Extract the generated text - Qwen2 uses "text" field, but also support "generation" for compatibility + generated_text = response_data.get("generation", "") or response_data.get("text", "") + + # Clean up the response (remove assistant start token if present) + if generated_text.startswith("<|im_start|>assistant\n"): + generated_text = generated_text[len("<|im_start|>assistant\n"):] + if generated_text.endswith("<|im_end|>"): + generated_text = generated_text[:-len("<|im_end|>")] + + # Set the content in the existing model_response structure + if hasattr(model_response, 'choices') and len(model_response.choices) > 0: + choice = model_response.choices[0] + if hasattr(choice, 'message'): + choice.message.content = generated_text + choice.finish_reason = "stop" + else: + # Handle streaming choices + choice.delta.content = generated_text + choice.finish_reason = "stop" + + # Set usage information if available in response + if "usage" in response_data: + usage_data = response_data["usage"] + if hasattr(model_response, 'usage'): + model_response.usage.prompt_tokens = usage_data.get("prompt_tokens", 0) + model_response.usage.completion_tokens = usage_data.get("completion_tokens", 0) + model_response.usage.total_tokens = usage_data.get("total_tokens", 0) + + return model_response + + except Exception as e: + if logging_obj: + logging_obj.post_call( + input=messages, + api_key=api_key, + original_response=raw_response, + additional_args={"error": str(e)}, + ) + raise e + diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index 6c389ff3b7d..bcb4cae1c8b 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -258,6 +258,15 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): litellm_params=litellm_params, headers=headers, ) + elif provider == "openai": + # OpenAI imported models use OpenAI Chat Completions format + return litellm.AmazonBedrockOpenAIConfig().transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) else: raise BedrockError( status_code=404, diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index ad115bc7e92..21a78c30343 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -27,6 +27,25 @@ class BedrockError(BaseLLMException): pass +# Lazy import cache to avoid circular imports and performance impact +_get_model_info = None + + +def get_cached_model_info(): + """ + Lazy import and cache get_model_info to avoid circular imports. + + This function is used by bedrock transformation classes that need get_model_info + but cannot import it at module level due to circular import issues. + The function is cached after first use to avoid performance impact. + """ + global _get_model_info + if _get_model_info is None: + from litellm import get_model_info + _get_model_info = get_model_info + return _get_model_info + + class AmazonBedrockGlobalConfig: def __init__(self): pass @@ -616,6 +635,8 @@ def get_bedrock_chat_config(model: str): return litellm.AmazonInvokeNovaConfig() elif bedrock_invoke_provider == "qwen3": return litellm.AmazonQwen3Config() + elif bedrock_invoke_provider == "qwen2": + return litellm.AmazonQwen2Config() elif bedrock_invoke_provider == "twelvelabs": return litellm.AmazonTwelveLabsPegasusConfig() else: diff --git a/litellm/llms/bedrock/files/handler.py b/litellm/llms/bedrock/files/handler.py new file mode 100644 index 00000000000..d6177e090d5 --- /dev/null +++ b/litellm/llms/bedrock/files/handler.py @@ -0,0 +1,206 @@ +import asyncio +import base64 +from typing import Any, Coroutine, Optional, Tuple, Union + +import httpx + +from litellm import LlmProviders +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.types.llms.openai import ( + FileContentRequest, + HttpxBinaryResponseContent, +) +from litellm.types.utils import SpecialEnums + +from ..base_aws_llm import BaseAWSLLM + + +class BedrockFilesHandler(BaseAWSLLM): + """ + Handles downloading files from S3 for Bedrock batch processing. + + This implementation downloads files from S3 buckets where Bedrock + stores batch output files. + """ + + def __init__(self): + super().__init__() + self.async_httpx_client = get_async_httpx_client( + llm_provider=LlmProviders.BEDROCK, + ) + + def _extract_s3_uri_from_file_id(self, file_id: str) -> str: + """ + Extract S3 URI from encoded file ID. + + The file ID can be in two formats: + 1. Base64-encoded unified file ID containing: llm_output_file_id,s3://bucket/path + 2. Direct S3 URI: s3://bucket/path + + Args: + file_id: Encoded file ID or direct S3 URI + + Returns: + S3 URI (e.g., "s3://bucket-name/path/to/file") + """ + # First, try to decode if it's a base64-encoded unified file ID + try: + # Add padding if needed + padded = file_id + "=" * (-len(file_id) % 4) + decoded = base64.urlsafe_b64decode(padded).decode() + + # Check if it's a unified file ID format + if decoded.startswith(SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value): + # Extract llm_output_file_id from the decoded string + if "llm_output_file_id," in decoded: + s3_uri = decoded.split("llm_output_file_id,")[1].split(";")[0] + return s3_uri + except Exception: + pass + + # If not base64 encoded or doesn't contain llm_output_file_id, assume it's already an S3 URI + if file_id.startswith("s3://"): + return file_id + + # If it doesn't start with s3://, assume it's a direct S3 URI and add the prefix + return f"s3://{file_id}" + + def _parse_s3_uri(self, s3_uri: str) -> Tuple[str, str]: + """ + Parse S3 URI to extract bucket name and object key. + + Args: + s3_uri: S3 URI (e.g., "s3://bucket-name/path/to/file") + + Returns: + Tuple of (bucket_name, object_key) + """ + if not s3_uri.startswith("s3://"): + raise ValueError(f"Invalid S3 URI format: {s3_uri}. Expected format: s3://bucket-name/path/to/file") + + # Remove 's3://' prefix + path = s3_uri[5:] + + if "/" in path: + bucket_name, object_key = path.split("/", 1) + else: + bucket_name = path + object_key = "" + + return bucket_name, object_key + + async def afile_content( + self, + file_content_request: FileContentRequest, + optional_params: dict, + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + ) -> HttpxBinaryResponseContent: + """ + Download file content from S3 bucket for Bedrock files. + + Args: + file_content_request: Contains file_id (encoded or S3 URI) + optional_params: Optional parameters containing AWS credentials + timeout: Request timeout + max_retries: Max retry attempts + + Returns: + HttpxBinaryResponseContent: Binary content wrapped in compatible response format + """ + import boto3 + from botocore.credentials import Credentials + + file_id = file_content_request.get("file_id") + if not file_id: + raise ValueError("file_id is required in file_content_request") + + # Extract S3 URI from file ID + s3_uri = self._extract_s3_uri_from_file_id(file_id) + bucket_name, object_key = self._parse_s3_uri(s3_uri) + + # Get AWS credentials + aws_region_name = self._get_aws_region_name( + optional_params=optional_params, model="" + ) + credentials: Credentials = self.get_credentials( + aws_access_key_id=optional_params.get("aws_access_key_id"), + aws_secret_access_key=optional_params.get("aws_secret_access_key"), + aws_session_token=optional_params.get("aws_session_token"), + aws_region_name=aws_region_name, + aws_session_name=optional_params.get("aws_session_name"), + aws_profile_name=optional_params.get("aws_profile_name"), + aws_role_name=optional_params.get("aws_role_name"), + aws_web_identity_token=optional_params.get("aws_web_identity_token"), + aws_sts_endpoint=optional_params.get("aws_sts_endpoint"), + ) + + # Create S3 client + s3_client = boto3.client( + "s3", + aws_access_key_id=credentials.access_key, + aws_secret_access_key=credentials.secret_key, + aws_session_token=credentials.token, + region_name=aws_region_name, + ) + + # Download file from S3 + try: + response = s3_client.get_object(Bucket=bucket_name, Key=object_key) + file_content = response["Body"].read() + except Exception as e: + raise ValueError(f"Failed to download file from S3: {s3_uri}. Error: {str(e)}") + + # Create mock HTTP response + mock_response = httpx.Response( + status_code=200, + content=file_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request(method="GET", url=s3_uri), + ) + + return HttpxBinaryResponseContent(response=mock_response) + + def file_content( + self, + _is_async: bool, + file_content_request: FileContentRequest, + api_base: Optional[str], + optional_params: dict, + timeout: Union[float, httpx.Timeout], + max_retries: Optional[int], + ) -> Union[ + HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent] + ]: + """ + Download file content from S3 bucket for Bedrock files. + Supports both sync and async operations. + + Args: + _is_async: Whether to run asynchronously + file_content_request: Contains file_id (encoded or S3 URI) + api_base: API base (unused for S3 operations) + optional_params: Optional parameters containing AWS credentials + timeout: Request timeout + max_retries: Max retry attempts + + Returns: + HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format + """ + if _is_async: + return self.afile_content( + file_content_request=file_content_request, + optional_params=optional_params, + timeout=timeout, + max_retries=max_retries, + ) + else: + return asyncio.run( + self.afile_content( + file_content_request=file_content_request, + optional_params=optional_params, + timeout=timeout, + max_retries=max_retries, + ) + ) + diff --git a/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py b/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py index f2b94b617c0..18366999583 100644 --- a/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py +++ b/litellm/llms/bedrock/image/amazon_nova_canvas_transformation.py @@ -3,7 +3,6 @@ from typing import Any, Dict, List, Optional from openai.types.image import Image -from litellm import get_model_info from litellm.types.llms.bedrock import ( AmazonNovaCanvasColorGuidedGenerationParams, AmazonNovaCanvasColorGuidedRequest, @@ -15,6 +14,7 @@ from litellm.types.llms.bedrock import ( AmazonNovaCanvasTextToImageRequest, AmazonNovaCanvasTextToImageResponse, ) +from litellm.llms.bedrock.common_utils import get_cached_model_info from litellm.types.utils import ImageResponse @@ -207,6 +207,7 @@ class AmazonNovaCanvasConfig: size: Optional[str] = None, optional_params: Optional[dict] = None, ) -> float: + get_model_info = get_cached_model_info() model_info = get_model_info( model=model, custom_llm_provider="bedrock", diff --git a/litellm/llms/bedrock/image/amazon_stability1_transformation.py b/litellm/llms/bedrock/image/amazon_stability1_transformation.py index 63af32f3f56..07f82cec232 100644 --- a/litellm/llms/bedrock/image/amazon_stability1_transformation.py +++ b/litellm/llms/bedrock/image/amazon_stability1_transformation.py @@ -5,7 +5,7 @@ from typing import List, Optional from openai.types.image import Image -from litellm import get_model_info +from litellm.llms.bedrock.common_utils import get_cached_model_info from litellm.types.utils import ImageResponse @@ -151,6 +151,7 @@ class AmazonStabilityConfig: size = size or "1024-x-1024" model = f"{size}/{steps}/{model}" + get_model_info = get_cached_model_info() model_info = get_model_info( model=model, custom_llm_provider="bedrock", diff --git a/litellm/llms/bedrock/image/amazon_stability3_transformation.py b/litellm/llms/bedrock/image/amazon_stability3_transformation.py index 445a2fe1100..160d0af8e80 100644 --- a/litellm/llms/bedrock/image/amazon_stability3_transformation.py +++ b/litellm/llms/bedrock/image/amazon_stability3_transformation.py @@ -3,12 +3,12 @@ from typing import List, Optional from openai.types.image import Image -from litellm import get_model_info from litellm.llms.bedrock.common_utils import BedrockError from litellm.types.llms.bedrock import ( AmazonStability3TextToImageRequest, AmazonStability3TextToImageResponse, ) +from litellm.llms.bedrock.common_utils import get_cached_model_info from litellm.types.utils import ImageResponse @@ -115,6 +115,7 @@ class AmazonStability3Config: size: Optional[str] = None, optional_params: Optional[dict] = None, ) -> float: + get_model_info = get_cached_model_info() model_info = get_model_info( model=model, custom_llm_provider="bedrock", diff --git a/litellm/llms/bedrock/image/amazon_titan_transformation.py b/litellm/llms/bedrock/image/amazon_titan_transformation.py index bed9ad0c300..65411cabdcf 100644 --- a/litellm/llms/bedrock/image/amazon_titan_transformation.py +++ b/litellm/llms/bedrock/image/amazon_titan_transformation.py @@ -7,7 +7,7 @@ from typing import List, Optional from openai.types.image import Image -from litellm import get_model_info +from litellm.utils import get_model_info from litellm.types.llms.bedrock import ( AmazonNovaCanvasImageGenerationConfig, AmazonTitanImageGenerationRequestBody, diff --git a/litellm/llms/cohere/rerank/guardrail_translation/handler.py b/litellm/llms/cohere/rerank/guardrail_translation/handler.py index e042295ab0b..6893a5991c3 100644 --- a/litellm/llms/cohere/rerank/guardrail_translation/handler.py +++ b/litellm/llms/cohere/rerank/guardrail_translation/handler.py @@ -49,12 +49,13 @@ class CohereRerankHandler(BaseTranslation): # Process query only query = data.get("query") if query is not None and isinstance(query, str): - guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail( - texts=[query], + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [query]}, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) data["query"] = guardrailed_texts[0] if guardrailed_texts else query verbose_proxy_logger.debug( diff --git a/litellm/llms/custom_httpx/aiohttp_transport.py b/litellm/llms/custom_httpx/aiohttp_transport.py index 6997afafd8d..f845bf7cb90 100644 --- a/litellm/llms/custom_httpx/aiohttp_transport.py +++ b/litellm/llms/custom_httpx/aiohttp_transport.py @@ -82,9 +82,7 @@ class AiohttpResponseStream(httpx.AsyncByteStream): async def __aiter__(self) -> typing.AsyncIterator[bytes]: try: - async for chunk in self._aiohttp_response.content.iter_chunked( - self.CHUNK_SIZE - ): + async for chunk in self._aiohttp_response.content.iter_chunked(self.CHUNK_SIZE): yield chunk except ( aiohttp.ClientPayloadError, @@ -120,16 +118,13 @@ class AiohttpResponseStream(httpx.AsyncByteStream): class AiohttpTransport(httpx.AsyncBaseTransport): - def __init__( - self, client: Union[ClientSession, Callable[[], ClientSession]] - ) -> None: + def __init__(self, client: Union[ClientSession, Callable[[], ClientSession]]) -> None: self.client = client ######################################################### # Class variables for proxy settings ######################################################### - self.proxy: Optional[str] = None - self.checked_proxy_env_settings: bool = False + self.proxy_cache: Dict[str, Optional[str]] = {} async def aclose(self) -> None: if isinstance(self.client, ClientSession): @@ -184,11 +179,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): current_loop = asyncio.get_running_loop() # If session is from a different or closed loop, recreate it - if ( - session_loop is None - or session_loop != current_loop - or session_loop.is_closed() - ): + if session_loop is None or session_loop != current_loop or session_loop.is_closed(): # Close old session to prevent leaks old_session = self.client try: @@ -215,7 +206,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): self.client = ClientSession() return self.client - + async def _make_aiohttp_request( self, client_session: ClientSession, @@ -226,20 +217,20 @@ class LiteLLMAiohttpTransport(AiohttpTransport): ) -> ClientResponse: """ Helper function to make an aiohttp request with the given parameters. - + Args: client_session: The aiohttp ClientSession to use request: The httpx Request to send timeout: Timeout settings dict with 'connect', 'read', 'pool' keys proxy: Optional proxy URL sni_hostname: Optional SNI hostname for SSL - + Returns: ClientResponse from aiohttp """ from aiohttp import ClientTimeout from yarl import URL as YarlURL - + try: data = request.content except httpx.RequestNotRead: @@ -262,9 +253,9 @@ class LiteLLMAiohttpTransport(AiohttpTransport): proxy=proxy, server_hostname=sni_hostname, ).__aenter__() - + return response - + async def handle_async_request( self, request: httpx.Request, @@ -297,7 +288,7 @@ class LiteLLMAiohttpTransport(AiohttpTransport): else: self.client = ClientSession() client_session = self.client - + # Retry the request with the new session with map_aiohttp_exceptions(): response = await self._make_aiohttp_request( @@ -317,45 +308,41 @@ class LiteLLMAiohttpTransport(AiohttpTransport): content=AiohttpResponseStream(response), request=request, ) - async def _get_proxy_settings(self, request: httpx.Request): proxy = None - if not ( - litellm.disable_aiohttp_trust_env - or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False")) - ): + if not (litellm.disable_aiohttp_trust_env or str_to_bool(os.getenv("DISABLE_AIOHTTP_TRUST_ENV", "False"))): try: proxy = self._proxy_from_env(request.url) except Exception as e: # pragma: no cover - best effort verbose_logger.debug(f"Error reading proxy env: {e}") return proxy - def _proxy_from_env(self, url: httpx.URL) -> typing.Optional[str]: """ Return proxy URL from env for the given request URL Only check the proxy env settings once, this is a costly operation for CPU % usage - + .""" ######################################################### # Check if we've already checked the proxy env settings ######################################################### - if self.checked_proxy_env_settings is True: - return self.proxy - - ######################################################### - # set self.checked_proxy_env_settings to True - ######################################################### - self.checked_proxy_env_settings = True + proxy_cache_key = url.host + + if proxy_cache_key in self.proxy_cache: + return self.proxy_cache[proxy_cache_key] + proxies = urllib.request.getproxies() if urllib.request.proxy_bypass(url.host): - return None + proxy_url = None + else: + proxy = proxies.get(url.scheme) or proxies.get("all") + if proxy and "://" not in proxy: + proxy = f"http://{proxy}" + proxy_url = proxy - proxy = proxies.get(url.scheme) or proxies.get("all") - if proxy and "://" not in proxy: - proxy = f"http://{proxy}" - self.proxy = proxy - return self.proxy + self.proxy_cache[proxy_cache_key] = proxy_url + + return proxy_url diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index c35e910ab08..d9a9d4f9dc1 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -16,7 +16,9 @@ from litellm._logging import verbose_logger from litellm.constants import ( _DEFAULT_TTL_FOR_HTTPX_CLIENTS, AIOHTTP_CONNECTOR_LIMIT, + AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, + AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, DEFAULT_SSL_CIPHERS, ) @@ -792,15 +794,20 @@ class AsyncHTTPHandler: verbose_logger.debug( "NEW SESSION: Creating new ClientSession (no shared session provided)" ) + transport_connector_kwargs = { + "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, + "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, + "enable_cleanup_closed": True, + **connector_kwargs, + } + if AIOHTTP_CONNECTOR_LIMIT > 0: + transport_connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT + if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: + transport_connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST + return LiteLLMAiohttpTransport( client=lambda: ClientSession( - connector=TCPConnector( - limit=AIOHTTP_CONNECTOR_LIMIT, - keepalive_timeout=AIOHTTP_KEEPALIVE_TIMEOUT, - ttl_dns_cache=AIOHTTP_TTL_DNS_CACHE, - enable_cleanup_closed=True, - **connector_kwargs, - ), + connector=TCPConnector(**transport_connector_kwargs), trust_env=trust_env, ), ) diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index 20e0d412edc..a75ecd8cc7b 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -218,22 +218,47 @@ class GroqChatConfig(OpenAILikeChatConfig): When using tools in this way: - https://docs.anthropic.com/en/docs/build-with-claude/tool-use#json-mode - You usually want to provide a single tool - You should set tool_choice (see Forcing tool use) to instruct the model to explicitly use that tool - - Remember that the model will pass the input to the tool, so the name of the tool and description should be from the model’s perspective. + - Remember that the model will pass the input to the tool, so the name of the tool and description should be from the model's perspective. + + Note: This workaround is only for models that don't support native json_schema. + Models like gpt-oss-120b, llama-4, kimi-k2 support native json_schema and should + pass response_format directly to Groq. + See: https://console.groq.com/docs/structured-outputs#supported-models """ if json_schema is not None: - _tool_choice = { - "type": "function", - "function": {"name": "json_tool_call"}, - } - _tool = self._create_json_tool_call_for_response_format( - json_schema=json_schema, - ) - optional_params["tools"] = [_tool] - optional_params["tool_choice"] = _tool_choice - optional_params["json_mode"] = True - non_default_params.pop( - "response_format", None - ) # only remove if it's a json_schema - handled via using groq's tool calling params. + # Check if model supports native response_schema + if not litellm.supports_response_schema( + model=model, custom_llm_provider="groq" + ): + # Check if user is also passing tools - this combination won't work + # See: https://console.groq.com/docs/structured-outputs + # "Streaming and tool use are not currently supported with Structured Outputs" + if "tools" in non_default_params: + raise litellm.BadRequestError( + message=f"Groq model '{model}' does not support native structured outputs. " + "LiteLLM uses a tool-calling workaround for structured outputs on this model, " + "which is incompatible with user-provided tools. " + "Either use a model that supports native structured outputs " + "(e.g., gpt-oss-120b, llama-4, kimi-k2), or remove the tools parameter. " + "See: https://console.groq.com/docs/structured-outputs#supported-models", + model=model, + llm_provider="groq", + ) + # Use workaround only for models without native support + _tool_choice = { + "type": "function", + "function": {"name": "json_tool_call"}, + } + _tool = self._create_json_tool_call_for_response_format( + json_schema=json_schema, + ) + optional_params["tools"] = [_tool] + optional_params["tool_choice"] = _tool_choice + optional_params["json_mode"] = True + non_default_params.pop( + "response_format", None + ) # only remove if it's a json_schema - handled via using groq's tool calling params. + # else: model supports native json_schema, let response_format pass through optional_params = super().map_openai_params( non_default_params, optional_params, model, drop_params ) diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 0abc94012ee..76aa6f730c2 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,16 +14,18 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation -from litellm.types.utils import Choices +from litellm.types.guardrails import GenericGuardrailAPIInputs +from litellm.types.llms.openai import ChatCompletionToolParam +from litellm.types.utils import Choices, StreamingChoices if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail - from litellm.types.utils import ModelResponse + from litellm.types.utils import ModelResponse, ModelResponseStream class OpenAIChatCompletionsHandler(BaseTranslation): @@ -52,38 +54,60 @@ class OpenAIChatCompletionsHandler(BaseTranslation): texts_to_check: List[str] = [] images_to_check: List[str] = [] - task_mappings: List[Tuple[int, Optional[int]]] = [] - # Track (message_index, content_index) for each text + tool_calls_to_check: List[ChatCompletionToolParam] = [] + text_task_mappings: List[Tuple[int, Optional[int]]] = [] + tool_call_task_mappings: List[Tuple[int, int]] = [] + # text_task_mappings: Track (message_index, content_index) for each text # content_index is None for string content, int for list content + # tool_call_task_mappings: Track (message_index, tool_call_index) for each tool call - # Step 1: Extract all text content and images + # Step 1: Extract all text content, images, and tool calls for msg_idx, message in enumerate(messages): - self._extract_input_text_and_images( + self._extract_inputs( message=message, msg_idx=msg_idx, texts_to_check=texts_to_check, images_to_check=images_to_check, - task_mappings=task_mappings, + tool_calls_to_check=tool_calls_to_check, + text_task_mappings=text_task_mappings, + tool_call_task_mappings=tool_call_task_mappings, ) - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: - guardrailed_texts, guardrailed_images = ( - await guardrail_to_apply.apply_guardrail( - texts=texts_to_check, - request_data=data, - input_type="request", - images=images_to_check if images_to_check else None, - logging_obj=litellm_logging_obj, - ) + # Step 2: Apply guardrail to all texts and tool calls in batch + if texts_to_check or tool_calls_to_check: + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + if tool_calls_to_check: + inputs["tool_calls"] = tool_calls_to_check # type: ignore + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) + guardrailed_tool_calls = guardrailed_inputs.get("tools", []) + # Step 3: Map guardrail responses back to original message structure - await self._apply_guardrail_responses_to_input( - messages=messages, - responses=guardrailed_texts, - task_mappings=task_mappings, - ) + if guardrailed_texts and texts_to_check: + await self._apply_guardrail_responses_to_input_texts( + messages=messages, + responses=guardrailed_texts, + task_mappings=text_task_mappings, + ) + + # Step 4: Apply guardrailed tool calls back to messages + if guardrailed_tool_calls: + # Note: The guardrail may modify tool_calls_to_check in place + # or we may need to handle returned tool calls differently + await self._apply_guardrail_responses_to_input_tool_calls( + messages=messages, + tool_calls=guardrailed_tool_calls, # type: ignore + task_mappings=tool_call_task_mappings, + ) verbose_proxy_logger.debug( "OpenAI Chat Completions: Processed input messages: %s", messages @@ -91,61 +115,71 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return data - def _extract_input_text_and_images( + def _extract_inputs( self, message: Dict[str, Any], msg_idx: int, texts_to_check: List[str], images_to_check: List[str], - task_mappings: List[Tuple[int, Optional[int]]], + tool_calls_to_check: List[ChatCompletionToolParam], + text_task_mappings: List[Tuple[int, Optional[int]]], + tool_call_task_mappings: List[Tuple[int, int]], ) -> None: """ - Extract text content and images from a message. + Extract text content, images, and tool calls from a message. - Override this method to customize text/image extraction logic. + Override this method to customize text/image/tool call extraction logic. """ content = message.get("content", None) - if content is None: - return + if content is not None: + if isinstance(content, str): + # Simple string content + texts_to_check.append(content) + text_task_mappings.append((msg_idx, None)) - if isinstance(content, str): - # Simple string content - texts_to_check.append(content) - task_mappings.append((msg_idx, None)) + elif isinstance(content, list): + # List content (e.g., multimodal with text and images) + for content_idx, content_item in enumerate(content): + # Extract text + text_str = content_item.get("text", None) + if text_str is not None: + texts_to_check.append(text_str) + text_task_mappings.append((msg_idx, int(content_idx))) - elif isinstance(content, list): - # List content (e.g., multimodal with text and images) - for content_idx, content_item in enumerate(content): - # Extract text - text_str = content_item.get("text", None) - if text_str is not None: - texts_to_check.append(text_str) - task_mappings.append((msg_idx, int(content_idx))) + # Extract images (image_url) + if content_item.get("type") == "image_url": + image_url = content_item.get("image_url", {}) + if isinstance(image_url, dict): + url = image_url.get("url") + if url: + images_to_check.append(url) - # Extract images (image_url) - if content_item.get("type") == "image_url": - image_url = content_item.get("image_url", {}) - if isinstance(image_url, dict): - url = image_url.get("url") - if url: - images_to_check.append(url) + # Extract tool calls (typically in assistant messages) + tool_calls = message.get("tools", None) + if tool_calls is not None and isinstance(tool_calls, list): + for tool_call_idx, tool_call in enumerate(tool_calls): + if isinstance(tool_call, dict): + # Add the full tool call object to the list + tool_calls_to_check.append(ChatCompletionToolParam(**tool_call)) + tool_call_task_mappings.append((msg_idx, int(tool_call_idx))) - async def _apply_guardrail_responses_to_input( + async def _apply_guardrail_responses_to_input_texts( self, messages: List[Dict[str, Any]], responses: List[str], task_mappings: List[Tuple[int, Optional[int]]], ) -> None: """ - Apply guardrail responses back to input messages. + Apply guardrail responses back to input message text content. - Override this method to customize how responses are applied. + Override this method to customize how text responses are applied. """ for task_idx, guardrail_response in enumerate(responses): mapping = task_mappings[task_idx] msg_idx = cast(int, mapping[0]) content_idx_optional = cast(Optional[int], mapping[1]) + # Handle content content = messages[msg_idx].get("content", None) if content is None: continue @@ -160,6 +194,31 @@ class OpenAIChatCompletionsHandler(BaseTranslation): "text" ] = guardrail_response + async def _apply_guardrail_responses_to_input_tool_calls( + self, + messages: List[Dict[str, Any]], + tool_calls: List[Dict[str, Any]], + task_mappings: List[Tuple[int, int]], + ) -> None: + """ + Apply guardrailed tool calls back to input messages. + + The guardrail may have modified the tool_calls list in place, + so we apply the modified tool calls back to the original messages. + + Override this method to customize how tool call responses are applied. + """ + for task_idx, (msg_idx, tool_call_idx) in enumerate(task_mappings): + if task_idx < len(tool_calls): + guardrailed_tool_call = tool_calls[task_idx] + message_tool_calls = messages[msg_idx].get("tool_calls", None) + if message_tool_calls is not None and isinstance( + message_tool_calls, list + ): + if tool_call_idx < len(message_tool_calls): + # Replace the tool call with the guardrailed version + message_tool_calls[tool_call_idx] = guardrailed_tool_call + async def process_output_response( self, response: "ModelResponse", @@ -193,21 +252,27 @@ class OpenAIChatCompletionsHandler(BaseTranslation): texts_to_check: List[str] = [] images_to_check: List[str] = [] - task_mappings: List[Tuple[int, Optional[int]]] = [] - # Track (choice_index, content_index) for each text + tool_calls_to_check: List[Dict[str, Any]] = [] + text_task_mappings: List[Tuple[int, Optional[int]]] = [] + tool_call_task_mappings: List[Tuple[int, int]] = [] + # text_task_mappings: Track (choice_index, content_index) for each text + # content_index is None for string content, int for list content + # tool_call_task_mappings: Track (choice_index, tool_call_index) for each tool call - # Step 1: Extract all text content and images from response choices + # Step 1: Extract all text content, images, and tool calls from response choices for choice_idx, choice in enumerate(response.choices): self._extract_output_text_and_images( choice=choice, choice_idx=choice_idx, texts_to_check=texts_to_check, images_to_check=images_to_check, - task_mappings=task_mappings, + tool_calls_to_check=tool_calls_to_check, + text_task_mappings=text_task_mappings, + tool_call_task_mappings=tool_call_task_mappings, ) - # Step 2: Apply guardrail to all texts in batch - if texts_to_check: + # Step 2: Apply guardrail to all texts and tool calls in batch + if texts_to_check or tool_calls_to_check: # Create a request_data dict with response info and user API key metadata request_data: dict = {"response": response} @@ -218,22 +283,36 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata - guardrailed_texts, guardrailed_images = ( - await guardrail_to_apply.apply_guardrail( - texts=texts_to_check, - request_data=request_data, - input_type="response", - images=images_to_check if images_to_check else None, - logging_obj=litellm_logging_obj, - ) + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + if tool_calls_to_check: + inputs["tool_calls"] = tool_calls_to_check # type: ignore + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) + # Step 3: Map guardrail responses back to original response structure - await self._apply_guardrail_responses_to_output( - response=response, - responses=guardrailed_texts, - task_mappings=task_mappings, - ) + if guardrailed_texts and texts_to_check: + await self._apply_guardrail_responses_to_output_texts( + response=response, + responses=guardrailed_texts, + task_mappings=text_task_mappings, + ) + + # Step 4: Apply guardrailed tool calls back to response + if tool_calls_to_check: + await self._apply_guardrail_responses_to_output_tool_calls( + response=response, + tool_calls=tool_calls_to_check, + task_mappings=tool_call_task_mappings, + ) verbose_proxy_logger.debug( "OpenAI Chat Completions: Processed output response: %s", response @@ -241,51 +320,210 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return response - def _has_text_content(self, response: "ModelResponse") -> bool: + async def process_output_streaming_response( + self, + responses_so_far: List["ModelResponseStream"], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + user_api_key_dict: Optional[Any] = None, + ) -> List["ModelResponseStream"]: """ - Check if response has any text content to process. + Process output streaming responses by applying guardrails to text content. + + Args: + responses_so_far: List of LiteLLM ModelResponseStream objects + guardrail_to_apply: The guardrail instance to apply + litellm_logging_obj: Optional logging object + user_api_key_dict: User API key metadata to pass to guardrails + + Returns: + Modified list of responses with guardrail applied to content + + Response Format Support: + - String content: choice.message.content = "text here" + - List content: choice.message.content = [{"type": "text", "text": "text here"}, ...] + """ + + # Step 0: Check if any response has text content to process + has_any_text_content = False + for response in responses_so_far: + if self._has_text_content(response): + has_any_text_content = True + break + + if not has_any_text_content: + verbose_proxy_logger.warning( + "OpenAI Chat Completions: No text content in streaming responses, skipping guardrail" + ) + return responses_so_far + + # Step 1: Combine all streaming chunks into complete text per choice + # For streaming, we need to concatenate all delta.content across all chunks + # Key: (choice_idx, content_idx), Value: combined text + combined_texts: Dict[Tuple[int, Optional[int]], str] = {} + + for response_idx, response in enumerate(responses_so_far): + for choice_idx, choice in enumerate(response.choices): + if isinstance(choice, litellm.StreamingChoices): + content = choice.delta.content + elif isinstance(choice, litellm.Choices): + content = choice.message.content + else: + continue + + if content is None: + continue + + if isinstance(content, str): + # String content - accumulate for this choice + key = (choice_idx, None) + if key not in combined_texts: + combined_texts[key] = "" + combined_texts[key] += content + + elif isinstance(content, list): + # List content - accumulate for each content item + for content_idx, content_item in enumerate(content): + text_str = content_item.get("text") + if text_str: + key = (choice_idx, content_idx) + if key not in combined_texts: + combined_texts[key] = "" + combined_texts[key] += text_str + + # Step 2: Create lists for guardrail processing + texts_to_check: List[str] = [] + images_to_check: List[str] = [] + task_mappings: List[Tuple[int, Optional[int]]] = [] + # Track (choice_index, content_index) for each combined text + + for (choice_idx, content_idx), combined_text in combined_texts.items(): + texts_to_check.append(combined_text) + task_mappings.append((choice_idx, content_idx)) + + # Step 3: Apply guardrail to all combined texts in batch + if texts_to_check: + # Create a request_data dict with response info and user API key metadata + request_data: dict = {"responses": responses_so_far} + + # Add user API key metadata with prefixed keys + user_metadata = self.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + request_data["litellm_metadata"] = user_metadata + + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, + ) + + guardrailed_texts = guardrailed_inputs.get("texts", []) + + # Step 4: Apply guardrailed text back to all streaming chunks + # For each choice, replace the combined text across all chunks + await self._apply_guardrail_responses_to_output_streaming( + responses=responses_so_far, + guardrailed_texts=guardrailed_texts, + task_mappings=task_mappings, + ) + + verbose_proxy_logger.debug( + "OpenAI Chat Completions: Processed output streaming responses: %s", + responses_so_far, + ) + + return responses_so_far + + def _has_text_content( + self, response: Union["ModelResponse", "ModelResponseStream"] + ) -> bool: + """ + Check if response has any text content or tool calls to process. Override this method to customize text content detection. """ - for choice in response.choices: - if isinstance(choice, litellm.Choices): - if choice.message.content and isinstance(choice.message.content, str): - return True + from litellm.types.utils import ModelResponse, ModelResponseStream + + if isinstance(response, ModelResponse): + for choice in response.choices: + if isinstance(choice, litellm.Choices): + # Check for text content + if choice.message.content and isinstance( + choice.message.content, str + ): + return True + # Check for tool calls + if choice.message.tool_calls and isinstance( + choice.message.tool_calls, list + ): + if len(choice.message.tool_calls) > 0: + return True + elif isinstance(response, ModelResponseStream): + for choice in response.choices: + if isinstance(choice, litellm.StreamingChoices): + # Check for text content + if choice.delta.content and isinstance(choice.delta.content, str): + return True + # Check for tool calls + if choice.delta.tool_calls and isinstance( + choice.delta.tool_calls, list + ): + if len(choice.delta.tool_calls) > 0: + return True return False def _extract_output_text_and_images( self, - choice: Any, + choice: Union[Choices, StreamingChoices], choice_idx: int, texts_to_check: List[str], images_to_check: List[str], - task_mappings: List[Tuple[int, Optional[int]]], + tool_calls_to_check: List[Dict[str, Any]], + text_task_mappings: List[Tuple[int, Optional[int]]], + tool_call_task_mappings: List[Tuple[int, int]], ) -> None: """ - Extract text content and images from a response choice. + Extract text content, images, and tool calls from a response choice. - Override this method to customize text/image extraction logic. + Override this method to customize text/image/tool call extraction logic. """ - if not isinstance(choice, litellm.Choices): - return - verbose_proxy_logger.debug( "OpenAI Chat Completions: Processing choice: %s", choice ) - if choice.message.content and isinstance(choice.message.content, str): - # Simple string content - texts_to_check.append(choice.message.content) - task_mappings.append((choice_idx, None)) + # Determine content source and tool calls based on choice type + content = None + tool_calls = None + if isinstance(choice, litellm.Choices): + content = choice.message.content + tool_calls = choice.message.tool_calls + elif isinstance(choice, litellm.StreamingChoices): + content = choice.delta.content + tool_calls = choice.delta.tool_calls + else: + # Unknown choice type, skip processing + return - elif choice.message.content and isinstance(choice.message.content, list): + # Process content if it exists + if content and isinstance(content, str): + # Simple string content + texts_to_check.append(content) + text_task_mappings.append((choice_idx, None)) + + elif content and isinstance(content, list): # List content (e.g., multimodal response) - for content_idx, content_item in enumerate(choice.message.content): + for content_idx, content_item in enumerate(content): # Extract text content_text = content_item.get("text") if content_text: texts_to_check.append(content_text) - task_mappings.append((choice_idx, int(content_idx))) + text_task_mappings.append((choice_idx, int(content_idx))) # Extract images if content_item.get("type") == "image_url": @@ -295,36 +533,181 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if url: images_to_check.append(url) - async def _apply_guardrail_responses_to_output( + # Process tool calls if they exist + if tool_calls is not None and isinstance(tool_calls, list): + for tool_call_idx, tool_call in enumerate(tool_calls): + # Convert tool call to dict format for guardrail processing + tool_call_dict = self._convert_tool_call_to_dict(tool_call) + if tool_call_dict: + tool_calls_to_check.append(tool_call_dict) + tool_call_task_mappings.append((choice_idx, int(tool_call_idx))) + + def _convert_tool_call_to_dict( + self, tool_call: Union[Dict[str, Any], Any] + ) -> Optional[Dict[str, Any]]: + """ + Convert a tool call object to dictionary format. + + Tool calls can be either dict or object depending on the type. + """ + if isinstance(tool_call, dict): + return tool_call + elif hasattr(tool_call, "id") and hasattr(tool_call, "function"): + # Convert object to dict + function = tool_call.function + function_dict = {} + if hasattr(function, "name"): + function_dict["name"] = function.name + if hasattr(function, "arguments"): + function_dict["arguments"] = function.arguments + + tool_call_dict = { + "id": tool_call.id if hasattr(tool_call, "id") else None, + "type": tool_call.type if hasattr(tool_call, "type") else "function", + "function": function_dict, + } + return tool_call_dict + return None + + async def _apply_guardrail_responses_to_output_texts( self, response: "ModelResponse", responses: List[str], task_mappings: List[Tuple[int, Optional[int]]], ) -> None: """ - Apply guardrail responses back to output response. + Apply guardrail text responses back to output response. - Override this method to customize how responses are applied. + Override this method to customize how text responses are applied. """ for task_idx, guardrail_response in enumerate(responses): mapping = task_mappings[task_idx] choice_idx = cast(int, mapping[0]) content_idx_optional = cast(Optional[int], mapping[1]) - content = cast(Choices, response.choices[choice_idx]).message.content + choice = cast(Choices, response.choices[choice_idx]) + + # Handle content + content = choice.message.content if content is None: continue if isinstance(content, str) and content_idx_optional is None: # Replace string content with guardrail response - cast(Choices, response.choices[choice_idx]).message.content = ( - guardrail_response - ) + choice.message.content = guardrail_response elif isinstance(content, list) and content_idx_optional is not None: # Replace specific text item in list content - cast(Choices, response.choices[choice_idx]).message.content[ # type: ignore - content_idx_optional - ][ - "text" - ] = guardrail_response + choice.message.content[content_idx_optional]["text"] = guardrail_response # type: ignore + + async def _apply_guardrail_responses_to_output_tool_calls( + self, + response: "ModelResponse", + tool_calls: List[Dict[str, Any]], + task_mappings: List[Tuple[int, int]], + ) -> None: + """ + Apply guardrailed tool calls back to output response. + + The guardrail may have modified the tool_calls list in place, + so we apply the modified tool calls back to the original response. + + Override this method to customize how tool call responses are applied. + """ + for task_idx, (choice_idx, tool_call_idx) in enumerate(task_mappings): + if task_idx < len(tool_calls): + guardrailed_tool_call = tool_calls[task_idx] + choice = cast(Choices, response.choices[choice_idx]) + choice_tool_calls = choice.message.tool_calls + + if choice_tool_calls is not None and isinstance( + choice_tool_calls, list + ): + if tool_call_idx < len(choice_tool_calls): + # Update the tool call with guardrailed version + existing_tool_call = choice_tool_calls[tool_call_idx] + # Update object attributes (output responses always have typed objects) + if "function" in guardrailed_tool_call: + func_dict = guardrailed_tool_call["function"] + if "arguments" in func_dict: + existing_tool_call.function.arguments = func_dict[ + "arguments" + ] + if "name" in func_dict: + existing_tool_call.function.name = func_dict["name"] + + async def _apply_guardrail_responses_to_output_streaming( + self, + responses: List["ModelResponseStream"], + guardrailed_texts: List[str], + task_mappings: List[Tuple[int, Optional[int]]], + ) -> None: + """ + Apply guardrail responses back to output streaming responses. + + For streaming responses, the guardrailed text (which is the combined text from all chunks) + is placed in the first chunk, and subsequent chunks are cleared. + + Args: + responses: List of ModelResponseStream objects to modify + guardrailed_texts: List of guardrailed text responses (combined from all chunks) + task_mappings: List of tuples (choice_idx, content_idx) + + Override this method to customize how responses are applied to streaming responses. + """ + # Build a mapping of what guardrailed text to use for each (choice_idx, content_idx) + guardrail_map: Dict[Tuple[int, Optional[int]], str] = {} + for task_idx, guardrail_response in enumerate(guardrailed_texts): + mapping = task_mappings[task_idx] + choice_idx = cast(int, mapping[0]) + content_idx_optional = cast(Optional[int], mapping[1]) + guardrail_map[(choice_idx, content_idx_optional)] = guardrail_response + + # Track which choices we've already set the guardrailed text for + # Key: (choice_idx, content_idx), Value: boolean (True if already set) + already_set: Dict[Tuple[int, Optional[int]], bool] = {} + + # Iterate through all responses and update content + for response_idx, response in enumerate(responses): + for choice_idx_in_response, choice in enumerate(response.choices): + if isinstance(choice, litellm.StreamingChoices): + content = choice.delta.content + elif isinstance(choice, litellm.Choices): + content = choice.message.content + else: + continue + + if content is None: + continue + + if isinstance(content, str): + # String content + key = (choice_idx_in_response, None) + if key in guardrail_map: + if key not in already_set: + # First chunk - set the complete guardrailed text + if isinstance(choice, litellm.StreamingChoices): + choice.delta.content = guardrail_map[key] + elif isinstance(choice, litellm.Choices): + choice.message.content = guardrail_map[key] + already_set[key] = True + else: + # Subsequent chunks - clear the content + if isinstance(choice, litellm.StreamingChoices): + choice.delta.content = "" + elif isinstance(choice, litellm.Choices): + choice.message.content = "" + + elif isinstance(content, list): + # List content - handle each content item + for content_idx, content_item in enumerate(content): + if "text" in content_item: + key = (choice_idx_in_response, content_idx) + if key in guardrail_map: + if key not in already_set: + # First chunk - set the complete guardrailed text + content_item["text"] = guardrail_map[key] + already_set[key] = True + else: + # Subsequent chunks - clear the text + content_item["text"] = "" diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index f8b733567d3..73d08cfead4 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -53,12 +53,13 @@ class OpenAITextCompletionHandler(BaseTranslation): if isinstance(prompt, str): # Single string prompt - guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail( - texts=[prompt], + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [prompt]}, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) data["prompt"] = guardrailed_texts[0] if guardrailed_texts else prompt verbose_proxy_logger.debug( @@ -79,12 +80,13 @@ class OpenAITextCompletionHandler(BaseTranslation): text_indices.append(idx) if texts_to_check: - guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail( - texts=texts_to_check, + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": texts_to_check}, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) # Replace guardrailed texts back for guardrail_idx, prompt_idx in enumerate(text_indices): @@ -152,12 +154,13 @@ class OpenAITextCompletionHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata - guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail( - texts=texts_to_check, + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": texts_to_check}, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) # Apply guardrailed texts back to choices for guardrail_idx, choice_idx in enumerate(choice_indices): diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index 53ee994c48f..842a64b1878 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -52,12 +52,13 @@ class OpenAIImageGenerationHandler(BaseTranslation): # Apply guardrail to the prompt if isinstance(prompt, str): - guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail( - texts=[prompt], + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [prompt]}, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) data["prompt"] = guardrailed_texts[0] if guardrailed_texts else prompt verbose_proxy_logger.debug( diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 3282b7665c0..20842525e59 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -444,6 +444,11 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): else: headers = {} response = raw_response.parse() + if not hasattr(response, "model_dump"): + raise OpenAIError( + status_code=500, + message=f"Empty or invalid response from LLM endpoint. Received: {response!r}. Check the reverse proxy or model server configuration.", + ) return headers, response except openai.APITimeoutError as e: end_time = time.time() @@ -477,7 +482,14 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): else: headers = {} response = raw_response.parse() + if not hasattr(response, "model_dump"): + raise OpenAIError( + status_code=500, + message=f"Empty or invalid response from LLM endpoint. Received: {response!r}. Check the reverse proxy or model server configuration.", + ) return headers, response + except OpenAIError: + raise except Exception as e: if raw_response is not None: raise Exception( diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 667a72a426a..9377eb4e193 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -28,10 +28,19 @@ Output: response.output is List[GenericResponseOutputItem] where each has: - text: str """ -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast + +from openai import BaseModel + +from openai import BaseModel from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, +) +from litellm.types.guardrails import GenericGuardrailAPIInputs +from litellm.types.llms.openai import ChatCompletionToolParam from litellm.types.responses.main import GenericResponseOutputItem, OutputText if TYPE_CHECKING: @@ -63,17 +72,28 @@ class OpenAIResponsesHandler(BaseTranslation): Handles both string input and list of message objects. """ input_data: Optional[Union[str, "ResponseInputParam"]] = data.get("input") + tools_to_check: List[ChatCompletionToolParam] = [] if input_data is None: return data # Handle simple string input if isinstance(input_data, str): - guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail( - texts=[input_data], + inputs = GenericGuardrailAPIInputs(texts=[input_data]) + + # Extract and transform tools if present + + if "tools" in data and data["tools"]: + self._extract_and_transform_tools(data["tools"], tools_to_check) + if tools_to_check: + inputs["tools"] = tools_to_check + + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) data["input"] = guardrailed_texts[0] if guardrailed_texts else input_data verbose_proxy_logger.debug("OpenAI Responses API: Processed string input") return data @@ -88,7 +108,7 @@ class OpenAIResponsesHandler(BaseTranslation): # Track (message_index, content_index) for each text # content_index is None for string content, int for list content - # Step 1: Extract all text content and images + # Step 1: Extract all text content, images, and tools for msg_idx, message in enumerate(input_data): self._extract_input_text_and_images( message=message, @@ -98,18 +118,26 @@ class OpenAIResponsesHandler(BaseTranslation): task_mappings=task_mappings, ) + # Extract and transform tools if present + if "tools" in data and data["tools"]: + self._extract_and_transform_tools(data["tools"], tools_to_check) + # Step 2: Apply guardrail to all texts in batch if texts_to_check: - guardrailed_texts, guardrailed_images = ( - await guardrail_to_apply.apply_guardrail( - texts=texts_to_check, - request_data=data, - input_type="request", - images=images_to_check if images_to_check else None, - logging_obj=litellm_logging_obj, - ) + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + if tools_to_check: + inputs["tools"] = tools_to_check + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) + # Step 3: Map guardrail responses back to original input structure await self._apply_guardrail_responses_to_input( messages=input_data, @@ -123,6 +151,29 @@ class OpenAIResponsesHandler(BaseTranslation): return data + def _extract_and_transform_tools( + self, + tools: List[Dict[str, Any]], + tools_to_check: List[ChatCompletionToolParam], + ) -> None: + """ + Extract and transform tools from Responses API format to Chat Completion format. + + Uses the LiteLLM transformation function to convert Responses API tools + to Chat Completion tools that can be passed to guardrails. + """ + if tools is not None and isinstance(tools, list): + # Transform Responses API tools to Chat Completion tools + ( + transformed_tools, + _, + ) = LiteLLMCompletionResponsesConfig.transform_responses_api_tools_to_chat_completion_tools( + tools # type: ignore + ) + tools_to_check.extend( + cast(List[ChatCompletionToolParam], transformed_tools) + ) + def _extract_input_text_and_images( self, message: Any, # Can be Dict[str, Any] or ResponseInputParam @@ -252,16 +303,18 @@ class OpenAIResponsesHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata - guardrailed_texts, guardrailed_images = ( - await guardrail_to_apply.apply_guardrail( - texts=texts_to_check, - request_data=request_data, - input_type="response", - images=images_to_check if images_to_check else None, - logging_obj=litellm_logging_obj, - ) + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + if images_to_check: + inputs["images"] = images_to_check + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="response", + logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) + # Step 3: Map guardrail responses back to original response structure await self._apply_guardrail_responses_to_output( response=response, @@ -275,6 +328,31 @@ class OpenAIResponsesHandler(BaseTranslation): return response + async def process_output_streaming_response( + self, + responses_so_far: List[Any], + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: Optional[Any] = None, + user_api_key_dict: Optional[Any] = None, + ) -> List[Any]: + """ + Process output streaming response by applying guardrails to text content. + """ + string_so_far = self.get_streaming_string_so_far(responses_so_far) + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [string_so_far]}, + request_data={}, + input_type="response", + logging_obj=litellm_logging_obj, + ) + return responses_so_far + + def get_streaming_string_so_far(self, responses_so_far: List[Any]) -> str: + """ + Get the string so far from the responses so far. + """ + return "".join([response.get("text", "") for response in responses_so_far]) + def _has_text_content(self, response: "ResponsesAPIResponse") -> bool: """ Check if response has any text content to process. @@ -285,6 +363,17 @@ class OpenAIResponsesHandler(BaseTranslation): return False for output_item in response.output: + if isinstance(output_item, BaseModel): + try: + generic_response_output_item = ( + GenericResponseOutputItem.model_validate( + output_item.model_dump() + ) + ) + if generic_response_output_item.content: + output_item = generic_response_output_item + except Exception: + continue if isinstance(output_item, (GenericResponseOutputItem, dict)): content = ( output_item.content @@ -296,9 +385,11 @@ class OpenAIResponsesHandler(BaseTranslation): # Check if it's an OutputText with text if isinstance(content_item, OutputText): if content_item.text: + return True elif isinstance(content_item, dict): if content_item.get("text"): + return True return False @@ -316,8 +407,16 @@ class OpenAIResponsesHandler(BaseTranslation): Override this method to customize text/image extraction logic. """ # Handle both GenericResponseOutputItem and dict - if isinstance(output_item, GenericResponseOutputItem): - content = output_item.content + content: Optional[Union[List[OutputText], List[dict]]] = None + if isinstance(output_item, BaseModel): + try: + generic_response_output_item = GenericResponseOutputItem.model_validate( + output_item.model_dump() + ) + if generic_response_output_item.content: + content = generic_response_output_item.content + except Exception: + return elif isinstance(output_item, dict): content = output_item.get("content", []) else: diff --git a/litellm/llms/openai/speech/guardrail_translation/handler.py b/litellm/llms/openai/speech/guardrail_translation/handler.py index 71edd4f2801..4c2f71477be 100644 --- a/litellm/llms/openai/speech/guardrail_translation/handler.py +++ b/litellm/llms/openai/speech/guardrail_translation/handler.py @@ -50,12 +50,13 @@ class OpenAITextToSpeechHandler(BaseTranslation): return data if isinstance(input_text, str): - guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail( - texts=[input_text], + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [input_text]}, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) data["input"] = guardrailed_texts[0] if guardrailed_texts else input_text verbose_proxy_logger.debug( diff --git a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py index 18678a9878b..ac416f42c81 100644 --- a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py +++ b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py @@ -88,12 +88,13 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata - guardrailed_texts, _ = await guardrail_to_apply.apply_guardrail( - texts=[original_text], + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [original_text]}, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, ) + guardrailed_texts = guardrailed_inputs.get("texts", []) response.text = guardrailed_texts[0] if guardrailed_texts else original_text verbose_proxy_logger.debug( diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index f1eafe4e294..b5610852fd2 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -10,6 +10,7 @@ from enum import Enum from typing import Any, AsyncIterator, Iterator, List, Optional, Tuple, Union, cast import httpx +import litellm from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException @@ -28,6 +29,20 @@ class CacheControlSupportedModels(str, Enum): class OpenrouterConfig(OpenAIGPTConfig): + def get_supported_openai_params(self, model: str) -> list: + """ + Allow reasoning parameters for models flagged as reasoning-capable. + """ + supported_params = super().get_supported_openai_params(model=model) + try: + if litellm.supports_reasoning( + model=model, custom_llm_provider="openrouter" + ) or litellm.supports_reasoning(model=model): + supported_params.append("reasoning_effort") + except Exception: + pass + return list(dict.fromkeys(supported_params)) + def map_openai_params( self, non_default_params: dict, diff --git a/litellm/llms/ovhcloud/chat/transformation.py b/litellm/llms/ovhcloud/chat/transformation.py index 6bdc28620ff..e9dc5be3eed 100644 --- a/litellm/llms/ovhcloud/chat/transformation.py +++ b/litellm/llms/ovhcloud/chat/transformation.py @@ -7,7 +7,9 @@ More information on our website: https://endpoints.ai.cloud.ovh.net from typing import Optional, Union, List import httpx -from litellm import ModelResponseStream, OpenAIGPTConfig, get_model_info, verbose_logger +from litellm.utils import ModelResponseStream, get_model_info +from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig +from litellm._logging import verbose_logger from litellm.llms.ovhcloud.utils import OVHCloudException from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index ae96335cec3..0702086a578 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -118,8 +118,8 @@ class PassThroughEndpointHandler(BaseTranslation): return data # Apply guardrail (pass-through doesn't modify the text, just checks it) - await guardrail_to_apply.apply_guardrail( - texts=[text_to_check], + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [text_to_check]}, request_data=data, input_type="request", logging_obj=litellm_logging_obj, @@ -178,8 +178,8 @@ class PassThroughEndpointHandler(BaseTranslation): request_data["litellm_metadata"] = user_metadata # Apply guardrail (pass-through doesn't modify the text, just checks it) - await guardrail_to_apply.apply_guardrail( - texts=[text_to_check], + guardrailed_inputs = await guardrail_to_apply.apply_guardrail( + inputs={"texts": [text_to_check]}, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/ragflow/__init__.py b/litellm/llms/ragflow/__init__.py new file mode 100644 index 00000000000..17d12bed31c --- /dev/null +++ b/litellm/llms/ragflow/__init__.py @@ -0,0 +1,8 @@ +""" +RAGFlow provider for LiteLLM. + +RAGFlow provides OpenAI-compatible APIs with unique path structures: +- Chat endpoint: /api/v1/chats_openai/{chat_id}/chat/completions +- Agent endpoint: /api/v1/agents_openai/{agent_id}/chat/completions +""" + diff --git a/litellm/llms/ragflow/chat/__init__.py b/litellm/llms/ragflow/chat/__init__.py new file mode 100644 index 00000000000..0e0f47d07b6 --- /dev/null +++ b/litellm/llms/ragflow/chat/__init__.py @@ -0,0 +1,4 @@ +""" +RAGFlow chat completion configuration. +""" + diff --git a/litellm/llms/ragflow/chat/transformation.py b/litellm/llms/ragflow/chat/transformation.py new file mode 100644 index 00000000000..58fbfa83c98 --- /dev/null +++ b/litellm/llms/ragflow/chat/transformation.py @@ -0,0 +1,264 @@ +""" +RAGFlow provider configuration for OpenAI-compatible API. + +RAGFlow provides OpenAI-compatible APIs with unique path structures: +- Chat endpoint: /api/v1/chats_openai/{chat_id}/chat/completions +- Agent endpoint: /api/v1/agents_openai/{agent_id}/chat/completions + +Model name format: +- Chat: ragflow/chat/{chat_id}/{model_name} +- Agent: ragflow/agent/{agent_id}/{model_name} +""" + +from typing import List, Optional, Tuple + +import litellm +from litellm.llms.openai.openai import OpenAIConfig +from litellm.secret_managers.main import get_secret, get_secret_str +from litellm.types.llms.openai import AllMessageValues + + +class RAGFlowConfig(OpenAIConfig): + """ + Configuration for RAGFlow OpenAI-compatible API. + + Handles both chat and agent endpoints by parsing the model name format: + - ragflow/chat/{chat_id}/{model_name} for chat endpoints + - ragflow/agent/{agent_id}/{model_name} for agent endpoints + """ + + def _parse_ragflow_model(self, model: str) -> Tuple[str, str, str]: + """ + Parse RAGFlow model name format: ragflow/{endpoint_type}/{id}/{model_name} + + Args: + model: Model name in format ragflow/chat/{chat_id}/{model} or ragflow/agent/{agent_id}/{model} + + Returns: + Tuple of (endpoint_type, id, model_name) + + Raises: + ValueError: If model format is invalid + """ + parts = model.split("/") + if len(parts) < 4: + raise ValueError( + f"Invalid RAGFlow model format: {model}. " + f"Expected format: ragflow/chat/{{chat_id}}/{{model}} or ragflow/agent/{{agent_id}}/{{model}}" + ) + + if parts[0] != "ragflow": + raise ValueError( + f"Invalid RAGFlow model format: {model}. Must start with 'ragflow/'" + ) + + endpoint_type = parts[1] + if endpoint_type not in ["chat", "agent"]: + raise ValueError( + f"Invalid RAGFlow endpoint type: {endpoint_type}. Must be 'chat' or 'agent'" + ) + + entity_id = parts[2] + model_name = "/".join(parts[3:]) # Handle model names that might contain slashes + + return endpoint_type, entity_id, model_name + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete URL for the RAGFlow API call. + + Constructs URL based on endpoint type: + - Chat: /api/v1/chats_openai/{chat_id}/chat/completions + - Agent: /api/v1/agents_openai/{agent_id}/chat/completions + + Args: + api_base: Base API URL (e.g., http://ragflow-server:port or http://ragflow-server:port/v1) + api_key: API key (not used in URL construction) + model: Model name in format ragflow/{endpoint_type}/{id}/{model} + optional_params: Optional parameters + litellm_params: LiteLLM parameters (may contain api_base) + stream: Whether streaming is enabled + + Returns: + Complete URL for the API call + """ + # Get api_base from multiple sources: input param, litellm_params, environment, or global litellm setting + if litellm_params and hasattr(litellm_params, 'api_base') and litellm_params.api_base: + api_base = api_base or litellm_params.api_base + + api_base = ( + api_base + or litellm.api_base + or get_secret("RAGFLOW_API_BASE") + or get_secret_str("RAGFLOW_API_BASE") + ) + + if api_base is None: + raise ValueError("api_base is required for RAGFlow provider. Set it via api_base parameter, RAGFLOW_API_BASE environment variable, or litellm.api_base") + + # Parse model name to extract endpoint type and ID + endpoint_type, entity_id, _ = self._parse_ragflow_model(model) + + # Remove trailing slash from api_base if present + api_base = api_base.rstrip("/") + + # Strip /v1 or /api/v1 from api_base if present, since we'll add the full path + # Check /api/v1 first because /api/v1 ends with /v1 + if api_base.endswith("/api/v1"): + api_base = api_base[:-7] # Remove /api/v1 + elif api_base.endswith("/v1"): + api_base = api_base[:-3] # Remove /v1 + + # Construct the RAGFlow-specific path + if endpoint_type == "chat": + path = f"/api/v1/chats_openai/{entity_id}/chat/completions" + else: # agent + path = f"/api/v1/agents_openai/{entity_id}/chat/completions" + + # Ensure path starts with / + if not path.startswith("/"): + path = "/" + path + + return f"{api_base}{path}" + + def _get_openai_compatible_provider_info( + self, + model: str, + api_base: Optional[str], + api_key: Optional[str], + custom_llm_provider: str, + ) -> Tuple[Optional[str], Optional[str], str]: + """ + Get OpenAI-compatible provider information for RAGFlow. + + Args: + model: Model name (will be parsed to extract actual model name) + api_base: Base API URL (from input params) + api_key: API key (from input params) + custom_llm_provider: Custom LLM provider name + + Returns: + Tuple of (api_base, api_key, custom_llm_provider) + """ + # Parse model to extract the actual model name + # The model name will be stored in litellm_params for use in requests + _, _, actual_model = self._parse_ragflow_model(model) + + # Get api_base from multiple sources: input param, environment, or global litellm setting + dynamic_api_base = ( + api_base + or litellm.api_base + or get_secret("RAGFLOW_API_BASE") + or get_secret_str("RAGFLOW_API_BASE") + ) + + # Get api_key from multiple sources: input param, environment, or global litellm setting + dynamic_api_key = ( + api_key + or litellm.api_key + or get_secret_str("RAGFLOW_API_KEY") + ) + + return dynamic_api_base, dynamic_api_key, custom_llm_provider + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate environment and set up headers for RAGFlow API. + + Args: + headers: Request headers + model: Model name + messages: Chat messages + optional_params: Optional parameters + litellm_params: LiteLLM parameters (may contain api_key) + api_key: API key (from input params) + api_base: Base API URL + + Returns: + Updated headers dictionary + """ + # Use api_key from litellm_params if available, otherwise fall back to other sources + if litellm_params and hasattr(litellm_params, 'api_key') and litellm_params.api_key: + api_key = api_key or litellm_params.api_key + + # Get api_key from multiple sources: input param, litellm_params, environment, or global litellm setting + api_key = ( + api_key + or litellm.api_key + or get_secret_str("RAGFLOW_API_KEY") + ) + + if api_key is not None: + headers["Authorization"] = f"Bearer {api_key}" + + # Ensure Content-Type is set to application/json + if "content-type" not in headers and "Content-Type" not in headers: + headers["Content-Type"] = "application/json" + + # Parse model to extract actual model name and store it + # The actual model name should be used in the request body + try: + _, _, actual_model = self._parse_ragflow_model(model) + # Store the actual model name in litellm_params for use in transform_request + litellm_params["_ragflow_actual_model"] = actual_model + except ValueError: + # If parsing fails, use the original model name + pass + + return headers + + def transform_request( + self, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Transform request for RAGFlow API. + + Uses the actual model name extracted from the RAGFlow model format. + + Args: + model: Model name in RAGFlow format + messages: Chat messages + optional_params: Optional parameters + litellm_params: LiteLLM parameters (may contain _ragflow_actual_model) + headers: Request headers + + Returns: + Transformed request dictionary + """ + # Get the actual model name from litellm_params if available + actual_model = litellm_params.get("_ragflow_actual_model") + if actual_model is None: + # Fallback: try to parse the model name + try: + _, _, actual_model = self._parse_ragflow_model(model) + except ValueError: + # If parsing fails, use the original model name + actual_model = model + + # Use parent's transform_request with the actual model name + return super().transform_request( + actual_model, messages, optional_params, litellm_params, headers + ) + diff --git a/litellm/llms/ragflow/vector_stores/__init__.py b/litellm/llms/ragflow/vector_stores/__init__.py new file mode 100644 index 00000000000..3be29310b39 --- /dev/null +++ b/litellm/llms/ragflow/vector_stores/__init__.py @@ -0,0 +1,2 @@ +# RAGFlow vector stores module + diff --git a/litellm/llms/ragflow/vector_stores/transformation.py b/litellm/llms/ragflow/vector_stores/transformation.py new file mode 100644 index 00000000000..b6401a4b8d7 --- /dev/null +++ b/litellm/llms/ragflow/vector_stores/transformation.py @@ -0,0 +1,249 @@ +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + +import httpx + +from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.router import GenericLiteLLMParams +from litellm.types.vector_stores import ( + BaseVectorStoreAuthCredentials, + VectorStoreCreateOptionalRequestParams, + VectorStoreCreateResponse, + VectorStoreFileCounts, + VectorStoreIndexEndpoints, + VectorStoreSearchOptionalRequestParams, + VectorStoreSearchResponse, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class RAGFlowVectorStoreConfig(BaseVectorStoreConfig): + """Vector store configuration for RAGFlow datasets.""" + + def get_auth_credentials( + self, litellm_params: dict + ) -> BaseVectorStoreAuthCredentials: + api_key = litellm_params.get("api_key") + if api_key is None: + # Try to get from environment variable + api_key = get_secret_str("RAGFLOW_API_KEY") + if api_key is None: + raise ValueError("api_key is required (set RAGFLOW_API_KEY env var or pass in litellm_params)") + return { + "headers": { + "Authorization": f"Bearer {api_key}", + }, + } + + def get_vector_store_endpoints_by_type(self) -> VectorStoreIndexEndpoints: + """RAGFlow vector stores are management-only, no search support.""" + return { + "read": [], + "write": [], + } + + def validate_environment( + self, headers: dict, litellm_params: Optional[GenericLiteLLMParams] + ) -> dict: + """Validate environment and set headers for RAGFlow API.""" + litellm_params = litellm_params or GenericLiteLLMParams() + api_key = ( + litellm_params.api_key + or get_secret_str("RAGFLOW_API_KEY") + ) + + if api_key is None: + raise ValueError("RAGFLOW_API_KEY is required (set env var or pass in litellm_params)") + + headers.update( + { + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + } + ) + return headers + + def get_complete_url( + self, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + """ + Get the complete URL for RAGFlow datasets API. + + Supports: + - RAGFLOW_API_BASE env var + - api_base in litellm_params + - Default: http://localhost:9380 + """ + api_base = ( + api_base + or litellm_params.get("api_base") + or get_secret_str("RAGFLOW_API_BASE") + or "http://localhost:9380" + ) + + # Remove trailing slashes + api_base = api_base.rstrip("/") + + # RAGFlow datasets API endpoint + return f"{api_base}/api/v1/datasets" + + def transform_search_vector_store_request( + self, + vector_store_id: str, + query: Union[str, List[str]], + vector_store_search_optional_params: VectorStoreSearchOptionalRequestParams, + api_base: str, + litellm_logging_obj: LiteLLMLoggingObj, + litellm_params: dict, + ) -> Tuple[str, Dict]: + """RAGFlow vector stores are management-only, search is not supported.""" + raise NotImplementedError( + "RAGFlow vector stores support dataset management only, not search/retrieval" + ) + + def transform_search_vector_store_response( + self, response: httpx.Response, litellm_logging_obj: LiteLLMLoggingObj + ) -> VectorStoreSearchResponse: + """RAGFlow vector stores are management-only, search is not supported.""" + raise NotImplementedError( + "RAGFlow vector stores support dataset management only, not search/retrieval" + ) + + def transform_create_vector_store_request( + self, + vector_store_create_optional_params: VectorStoreCreateOptionalRequestParams, + api_base: str, + ) -> Tuple[str, Dict]: + """ + Transform create request to RAGFlow POST /api/v1/datasets format. + + Maps LiteLLM params to RAGFlow dataset creation parameters. + RAGFlow-specific fields can be passed via metadata. + """ + url = api_base # Already includes /api/v1/datasets from get_complete_url + + # Extract name (required by RAGFlow) + name = vector_store_create_optional_params.get("name") + if not name: + raise ValueError("name is required for RAGFlow dataset creation") + + # Build request body + request_body: Dict[str, Any] = { + "name": name, + } + + # Extract RAGFlow-specific fields from metadata + metadata = vector_store_create_optional_params.get("metadata") + if metadata: + # RAGFlow-specific fields that can be in metadata + ragflow_fields = [ + "avatar", + "description", + "embedding_model", + "permission", + "chunk_method", + "parser_config", + "parse_type", + "pipeline_id", + ] + + for field in ragflow_fields: + if field in metadata: + request_body[field] = metadata[field] + + # Validate: chunk_method and pipeline_id are mutually exclusive + if "chunk_method" in request_body and "pipeline_id" in request_body: + raise ValueError( + "chunk_method and pipeline_id are mutually exclusive. " + "Specify either chunk_method or pipeline_id, not both." + ) + + # If neither chunk_method nor pipeline_id is specified, default to naive + if "chunk_method" not in request_body and "pipeline_id" not in request_body: + request_body["chunk_method"] = "naive" + + return url, request_body + + def transform_create_vector_store_response( + self, response: httpx.Response + ) -> VectorStoreCreateResponse: + """ + Transform RAGFlow response to VectorStoreCreateResponse format. + + RAGFlow response format: + { + "code": 0, + "data": { + "id": "...", + "name": "...", + "create_time": 1745836841611, # milliseconds + ... + } + } + """ + try: + response_json = response.json() + + # Check for RAGFlow error response + if response_json.get("code") != 0: + error_message = response_json.get("message", "Unknown error") + raise self.get_error_class( + error_message=error_message, + status_code=response.status_code, + headers=response.headers, + ) + + data = response_json.get("data", {}) + + # Extract dataset ID + dataset_id = data.get("id") + if not dataset_id: + raise ValueError("RAGFlow response missing dataset id") + + # Extract name + name = data.get("name") + + # Convert create_time from milliseconds to seconds (Unix timestamp) + create_time_ms = data.get("create_time", 0) + created_at = int(create_time_ms / 1000) if create_time_ms else None + + # Build VectorStoreCreateResponse + return VectorStoreCreateResponse( + id=dataset_id, + object="vector_store", + created_at=created_at or 0, + name=name, + bytes=0, # RAGFlow doesn't provide bytes in response + file_counts=VectorStoreFileCounts( + in_progress=0, + completed=0, + failed=0, + cancelled=0, + total=0, + ), + status="completed", + expires_after=None, + expires_at=None, + last_active_at=None, + metadata=None, + ) + except Exception as e: + # If it's already a ValueError we raised, re-raise it + if isinstance(e, ValueError) and "RAGFlow response" in str(e): + raise + # If it's already our error class (has status_code), re-raise + if hasattr(e, "status_code"): + raise + # Otherwise, wrap in our error class + raise self.get_error_class( + error_message=str(e), + status_code=response.status_code, + headers=response.headers, + ) + diff --git a/litellm/llms/together_ai/chat.py b/litellm/llms/together_ai/chat.py index 06d33f69750..e8a784d2779 100644 --- a/litellm/llms/together_ai/chat.py +++ b/litellm/llms/together_ai/chat.py @@ -8,7 +8,8 @@ Docs: https://docs.together.ai/reference/completions-1 from typing import Optional -from litellm import get_model_info, verbose_logger +from litellm.utils import get_model_info +from litellm._logging import verbose_logger from ..openai.chat.gpt_transformation import OpenAIGPTConfig diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index dc6a3170afe..a3ea8afe40d 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -5,7 +5,8 @@ from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, get_ty import httpx import litellm -from litellm import supports_response_schema, supports_system_messages, verbose_logger +from litellm.utils import supports_response_schema, supports_system_messages +from litellm._logging import verbose_logger from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_defs from litellm.llms.base_llm.base_utils import BaseLLMModelInfo, BaseTokenCounter 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 5fef8c1ec49..665661b9d22 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 @@ -1091,6 +1091,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if "thoughtSignature" in part: part_copy = part.copy() part_copy.pop("thoughtSignature") + + text_content = part_copy.get("text") + if isinstance(text_content, str) and text_content.strip() == "": + continue + thinking_blocks.append( ChatCompletionThinkingBlock( type="thinking", @@ -1205,14 +1210,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): } # Embed thought signature in ID for OpenAI client compatibility if thought_signature: - _tool_response_chunk[ - "id" - ] = _encode_tool_call_id_with_signature( - _tool_response_chunk["id"] or "", thought_signature - ) _tool_response_chunk["provider_specific_fields"] = { # type: ignore "thought_signature": thought_signature } + # Only embed in ID if preview features are enabled + if litellm.enable_preview_features: + _tool_response_chunk[ + "id" + ] = _encode_tool_call_id_with_signature( + _tool_response_chunk["id"] or "", thought_signature + ) _tools.append(_tool_response_chunk) cumulative_tool_call_idx += 1 if len(_tools) == 0: diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py index af9af71fef4..859bb0a6984 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_handler.py @@ -8,7 +8,7 @@ from typing import Any, Literal, Optional, Union import httpx import litellm -from litellm import EmbeddingResponse +from litellm.types.utils import EmbeddingResponse from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 2c0f5dad228..455ec1d18f5 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -6,7 +6,7 @@ Why separate file? Make it easy to see how transformation works from typing import List -from litellm import EmbeddingResponse +from litellm.types.utils import EmbeddingResponse from litellm.types.llms.openai import EmbeddingInput from litellm.types.llms.vertex_ai import ( ContentType, diff --git a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py index 04be4de8e32..e14cfe3be0b 100644 --- a/litellm/llms/vertex_ai/image_generation/image_generation_handler.py +++ b/litellm/llms/vertex_ai/image_generation/image_generation_handler.py @@ -176,7 +176,7 @@ class VertexImageGeneration(VertexLLM): vertex_project: Optional[str], vertex_location: Optional[str], vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES], - model_response: litellm.ImageResponse, + model_response: ImageResponse, logging_obj: Any, model: str = "imagegeneration", # vertex ai uses imagegeneration as the default model client: Optional[AsyncHTTPHandler] = None, diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py index 582d7a4c569..d0ffc7be0a6 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/embedding_handler.py @@ -147,13 +147,13 @@ class VertexMultimodalEmbedding(VertexLLM): optional_params: dict, litellm_params: dict, data: dict, - model_response: litellm.EmbeddingResponse, + model_response: EmbeddingResponse, timeout: Optional[Union[float, httpx.Timeout]], logging_obj: LiteLLMLoggingObj, headers={}, client: Optional[AsyncHTTPHandler] = None, api_key: Optional[str] = None, - ) -> litellm.EmbeddingResponse: + ) -> EmbeddingResponse: if client is None: _params = {} if timeout is not None: diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py new file mode 100644 index 00000000000..aff14b1004f --- /dev/null +++ b/litellm/llms/vertex_ai/text_to_speech/transformation.py @@ -0,0 +1,472 @@ +""" +Vertex AI Text-to-Speech transformation + +Maps OpenAI TTS spec to Google Cloud Text-to-Speech API +Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/synthesize +""" + +import base64 +from typing import TYPE_CHECKING, Any, Coroutine, Dict, Optional, Tuple, Union + +import httpx + +from litellm.llms.base_llm.text_to_speech.transformation import ( + BaseTextToSpeechConfig, + TextToSpeechRequestData, +) +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES +from litellm.types.llms.vertex_ai_text_to_speech import ( + VertexTextToSpeechAudioConfig, + VertexTextToSpeechInput, + VertexTextToSpeechVoice, +) + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.llms.openai import HttpxBinaryResponseContent +else: + LiteLLMLoggingObj = Any + HttpxBinaryResponseContent = Any + + +class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase): + """ + Configuration for Google Cloud/Vertex AI Text-to-Speech + + Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/synthesize + """ + + # Default values + DEFAULT_LANGUAGE_CODE = "en-US" + DEFAULT_VOICE_NAME = "en-US-Studio-O" + DEFAULT_AUDIO_ENCODING = "LINEAR16" + DEFAULT_SPEAKING_RATE = "1" + + # API endpoint + TTS_API_URL = "https://texttospeech.googleapis.com/v1/text:synthesize" + + # Voice name mappings from OpenAI voices to Google Cloud voices + # Users can pass either: + # 1. OpenAI voice names (alloy, echo, fable, onyx, nova, shimmer) - will be mapped + # 2. Google Cloud/Vertex AI voice names (en-US-Studio-O, en-US-Wavenet-D, etc.) - used directly + VOICE_MAPPINGS = { + "alloy": "en-US-Studio-O", + "echo": "en-US-Studio-M", + "fable": "en-GB-Studio-B", + "onyx": "en-US-Wavenet-D", + "nova": "en-US-Studio-O", + "shimmer": "en-US-Wavenet-F", + } + + # Response format mappings from OpenAI to Google Cloud audio encoding + FORMAT_MAPPINGS = { + "mp3": "MP3", + "opus": "OGG_OPUS", + "aac": "MP3", # Google doesn't have AAC, use MP3 + "flac": "FLAC", + "wav": "LINEAR16", + "pcm": "LINEAR16", + } + + def __init__(self) -> None: + BaseTextToSpeechConfig.__init__(self) + VertexBase.__init__(self) + + def _map_voice_to_vertex_format( + self, + voice: Optional[Union[str, Dict]], + ) -> Tuple[Optional[str], Optional[Dict]]: + """ + Map voice to Vertex AI format. + + Supports both: + 1. OpenAI voice names (alloy, echo, fable, onyx, nova, shimmer) - will be mapped + 2. Vertex AI voice names (en-US-Studio-O, en-US-Wavenet-D, etc.) - used directly + 3. Dict with languageCode and name - used as-is + + Returns: + Tuple of (voice_str, voice_dict) where: + - voice_str: Original string voice (for interface compatibility) + - voice_dict: Vertex AI format dict with languageCode and name + """ + if voice is None: + return None, None + + if isinstance(voice, dict): + # Already in Vertex AI format + return None, voice + + # voice is a string + voice_str = voice + + # Map OpenAI voice if it's a known OpenAI voice, otherwise use directly + if voice in self.VOICE_MAPPINGS: + mapped_voice_name = self.VOICE_MAPPINGS[voice] + else: + # Assume it's already a Vertex AI voice name + mapped_voice_name = voice + + # Extract language code from voice name (e.g., "en-US-Studio-O" -> "en-US") + parts = mapped_voice_name.split("-") + if len(parts) >= 2: + language_code = f"{parts[0]}-{parts[1]}" + else: + language_code = self.DEFAULT_LANGUAGE_CODE + + voice_dict = { + "languageCode": language_code, + "name": mapped_voice_name, + } + + return voice_str, voice_dict + + def dispatch_text_to_speech( + self, + model: str, + input: str, + voice: Optional[Union[str, Dict]], + optional_params: Dict, + litellm_params_dict: Dict, + logging_obj: "LiteLLMLoggingObj", + timeout: Union[float, httpx.Timeout], + extra_headers: Optional[Dict[str, Any]], + base_llm_http_handler: Any, + aspeech: bool, + api_base: Optional[str], + api_key: Optional[str], + **kwargs: Any, + ) -> Union[ + "HttpxBinaryResponseContent", + Coroutine[Any, Any, "HttpxBinaryResponseContent"], + ]: + """ + Dispatch method to handle Vertex AI TTS requests + + This method encapsulates Vertex AI-specific credential resolution and parameter handling. + Voice mapping is handled in map_openai_params (similar to Azure AVA pattern). + + Args: + base_llm_http_handler: The BaseLLMHTTPHandler instance from main.py + """ + # Resolve Vertex AI credentials using VertexBase helpers + vertex_credentials = self.safe_get_vertex_ai_credentials(litellm_params_dict) + vertex_project = self.safe_get_vertex_ai_project(litellm_params_dict) + vertex_location = self.safe_get_vertex_ai_location(litellm_params_dict) + + # Convert voice to string if it's a dict (extract name) + # Actual voice mapping happens in map_openai_params + voice_str: Optional[str] = None + if isinstance(voice, str): + voice_str = voice + elif isinstance(voice, dict): + # Extract voice name from dict if needed + voice_str = voice.get("name") if voice else None + + # Store credentials in litellm_params for use in transform methods + litellm_params_dict.update({ + "vertex_credentials": vertex_credentials, + "vertex_project": vertex_project, + "vertex_location": vertex_location, + "api_base": api_base, + }) + + # Call the text_to_speech_handler + response = base_llm_http_handler.text_to_speech_handler( + model=model, + input=input, + voice=voice_str, + text_to_speech_provider_config=self, + text_to_speech_optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params_dict, + logging_obj=logging_obj, + timeout=timeout, + extra_headers=extra_headers, + client=None, + _is_async=aspeech, + ) + + return response + + def get_supported_openai_params(self, model: str) -> list: + """ + Vertex AI TTS supports these OpenAI parameters + + Note: Vertex AI also supports additional parameters like audioConfig + which can be passed but are not part of the OpenAI spec + """ + return ["voice", "response_format", "speed"] + + def map_openai_params( + self, + model: str, + optional_params: Dict, + voice: Optional[Union[str, Dict]] = None, + drop_params: bool = False, + kwargs: Dict = {}, + ) -> Tuple[Optional[str], Dict]: + """ + Map OpenAI parameters to Vertex AI TTS parameters + + Voice handling (similar to Azure AVA): + - If voice is an OpenAI voice name (alloy, echo, etc.), it maps to a Vertex AI voice + - If voice is already a Vertex AI voice name (en-US-Studio-O, etc.), it's used directly + - If voice is a dict with languageCode and name, it's used as-is + + Note: For Vertex AI, voice dict is stored in mapped_params["vertex_voice_dict"] + because the base class interface expects voice to be a string. + + Returns: + Tuple of (mapped_voice_str, mapped_params) + """ + mapped_params = {} + + ########################################################## + # Map voice using helper + ########################################################## + mapped_voice_str, voice_dict = self._map_voice_to_vertex_format(voice) + if voice_dict is not None: + mapped_params["vertex_voice_dict"] = voice_dict + + # Map response format + if "response_format" in optional_params: + format_name = optional_params["response_format"] + if format_name in self.FORMAT_MAPPINGS: + mapped_params["audioEncoding"] = self.FORMAT_MAPPINGS[format_name] + else: + # Try to use it directly as Google Cloud format + mapped_params["audioEncoding"] = format_name + else: + # Default to LINEAR16 + mapped_params["audioEncoding"] = self.DEFAULT_AUDIO_ENCODING + + # Map speed (OpenAI: 0.25-4.0, Vertex AI: speakingRate 0.25-4.0) + if "speed" in optional_params: + speed = optional_params["speed"] + if speed is not None: + mapped_params["speakingRate"] = str(speed) + + # Pass through Vertex AI-specific parameters from kwargs + if "audioConfig" in kwargs: + mapped_params["audioConfig"] = kwargs["audioConfig"] + + if "use_ssml" in kwargs: + mapped_params["use_ssml"] = kwargs["use_ssml"] + + return mapped_voice_str, mapped_params + + def validate_environment( + self, + headers: dict, + model: str, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate Vertex AI environment and set up authentication headers + + Note: Actual authentication is handled in transform_text_to_speech_request + because Vertex AI requires OAuth2 token refresh + """ + validated_headers = headers.copy() + + # Content-Type for JSON + validated_headers["Content-Type"] = "application/json" + validated_headers["charset"] = "UTF-8" + + return validated_headers + + def get_complete_url( + self, + model: str, + api_base: Optional[str], + litellm_params: dict, + ) -> str: + """ + Get the complete URL for Vertex AI TTS request + + Google Cloud TTS endpoint: https://texttospeech.googleapis.com/v1/text:synthesize + """ + if api_base: + return api_base + + return self.TTS_API_URL + + def _validate_vertex_input( + self, + input_data: VertexTextToSpeechInput, + optional_params: Dict, + ) -> VertexTextToSpeechInput: + """ + Validate and transform input for Vertex AI TTS + + Handles text vs SSML input detection and validation + """ + # Remove None values + if input_data.get("text") is None: + input_data.pop("text", None) + if input_data.get("ssml") is None: + input_data.pop("ssml", None) + + # Check if use_ssml is set + use_ssml = optional_params.get("use_ssml", False) + + if use_ssml: + if "text" in input_data: + input_data["ssml"] = input_data.pop("text") + elif "ssml" not in input_data: + raise ValueError("SSML input is required when use_ssml is True.") + else: + # LiteLLM will auto-detect if text is in ssml format + # check if "text" is an ssml - in this case we should pass it as ssml instead of text + if input_data: + _text = input_data.get("text", None) or "" + if "" in _text: + input_data["ssml"] = input_data.pop("text") + + if not input_data: + raise ValueError("Either 'text' or 'ssml' must be provided.") + if "text" in input_data and "ssml" in input_data: + raise ValueError("Only one of 'text' or 'ssml' should be provided, not both.") + + return input_data + + def transform_text_to_speech_request( + self, + model: str, + input: str, + voice: Optional[str], + optional_params: Dict, + litellm_params: Dict, + headers: dict, + ) -> TextToSpeechRequestData: + """ + Transform OpenAI TTS request to Vertex AI TTS format + + This method handles: + 1. Authentication with Vertex AI + 2. Building the request body + 3. Setting up headers + + Returns: + TextToSpeechRequestData: Contains dict_body and headers + """ + # Get Vertex AI credentials from litellm_params + vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES] = litellm_params.get( + "vertex_credentials" + ) + vertex_project: Optional[str] = litellm_params.get("vertex_project") + + ####### Authenticate with Vertex AI ######## + _auth_header, vertex_project = self._ensure_access_token( + credentials=vertex_credentials, + project_id=vertex_project, + custom_llm_provider="vertex_ai_beta", + ) + + auth_header, _ = self._get_token_and_url( + model="", + auth_header=_auth_header, + gemini_api_key=None, + vertex_credentials=vertex_credentials, + vertex_project=vertex_project, + vertex_location=litellm_params.get("vertex_location"), + stream=False, + custom_llm_provider="vertex_ai_beta", + api_base=litellm_params.get("api_base"), + ) + + # Set authentication headers + headers["Authorization"] = f"Bearer {auth_header}" + headers["x-goog-user-project"] = vertex_project + + ####### Build the request ################ + vertex_input = VertexTextToSpeechInput(text=input) + vertex_input = self._validate_vertex_input(vertex_input, optional_params) + + # Build voice configuration + # Check for voice dict stored in: + # 1. litellm_params by dispatch method + # 2. optional_params by map_openai_params + voice_dict = ( + litellm_params.get("vertex_voice_dict") + or optional_params.get("vertex_voice_dict") + ) + if voice_dict is not None and isinstance(voice_dict, dict): + vertex_voice = VertexTextToSpeechVoice(**voice_dict) + elif voice is not None and isinstance(voice, str): + # Handle string voice (shouldn't normally happen if dispatch was called) + parts = voice.split("-") + if len(parts) >= 2: + language_code = f"{parts[0]}-{parts[1]}" + else: + language_code = self.DEFAULT_LANGUAGE_CODE + vertex_voice = VertexTextToSpeechVoice( + languageCode=language_code, + name=voice, + ) + else: + # Use defaults + vertex_voice = VertexTextToSpeechVoice( + languageCode=self.DEFAULT_LANGUAGE_CODE, + name=self.DEFAULT_VOICE_NAME, + ) + + # Build audio configuration + audio_encoding = optional_params.get("audioEncoding", self.DEFAULT_AUDIO_ENCODING) + speaking_rate = optional_params.get("speakingRate", self.DEFAULT_SPEAKING_RATE) + + # Check for full audioConfig in optional_params + if "audioConfig" in optional_params: + vertex_audio_config = VertexTextToSpeechAudioConfig(**optional_params["audioConfig"]) + else: + vertex_audio_config = VertexTextToSpeechAudioConfig( + audioEncoding=audio_encoding, + speakingRate=speaking_rate, + ) + + request_body: Dict[str, Any] = { + "input": dict(vertex_input), + "voice": dict(vertex_voice), + "audioConfig": dict(vertex_audio_config), + } + + return TextToSpeechRequestData( + dict_body=request_body, + headers=headers, + ) + + def transform_text_to_speech_response( + self, + model: str, + raw_response: httpx.Response, + logging_obj: "LiteLLMLoggingObj", + ) -> "HttpxBinaryResponseContent": + """ + Transform Vertex AI TTS response to standard format + + Vertex AI returns JSON with base64-encoded audio content. + We decode it and return as HttpxBinaryResponseContent. + """ + from litellm.types.llms.openai import HttpxBinaryResponseContent + + # Parse JSON response + _json_response = raw_response.json() + + # Get base64-encoded audio content + response_content = _json_response.get("audioContent") + if not response_content: + raise ValueError("No audioContent in Vertex AI TTS response") + + # Decode base64 to get binary content + binary_data = base64.b64decode(response_content) + + # Create an httpx.Response object with the binary data + response = httpx.Response( + status_code=200, + content=binary_data, + ) + + # Initialize the HttpxBinaryResponseContent instance + return HttpxBinaryResponseContent(response) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py index a170e6cc7f2..aaa6a0bb95f 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/embedding_handler.py @@ -137,7 +137,7 @@ class VertexEmbedding(VertexBase): self, model: str, input: Union[list, str], - model_response: litellm.EmbeddingResponse, + model_response: EmbeddingResponse, logging_obj: LiteLLMLoggingObject, optional_params: dict, custom_llm_provider: Literal[ @@ -152,7 +152,7 @@ class VertexEmbedding(VertexBase): gemini_api_key: Optional[str] = None, extra_headers: Optional[dict] = None, encoding=None, - ) -> litellm.EmbeddingResponse: + ) -> EmbeddingResponse: """ Async embedding implementation """ diff --git a/litellm/main.py b/litellm/main.py index a09a9453017..0e19699309d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -52,13 +52,10 @@ from pydantic import BaseModel from typing_extensions import overload import litellm -from litellm import ( # type: ignore - client, - exception_type, - get_litellm_params, - get_optional_params, -) - +# client must be imported from litellm as it's a decorator used at function definition time +from litellm import client +# Other utils are imported directly to avoid circular imports +from litellm.utils import exception_type, get_litellm_params, get_optional_params # Logging is imported lazily when needed to avoid loading litellm_logging at import time if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging @@ -206,7 +203,6 @@ from .llms.vertex_ai.image_generation.image_generation_handler import ( from .llms.vertex_ai.multimodal_embeddings.embedding_handler import ( VertexMultimodalEmbedding, ) -from .llms.vertex_ai.text_to_speech.text_to_speech_handler import VertexTextToSpeechAPI from .llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels from .llms.vertex_ai.vertex_embeddings.embedding_handler import VertexEmbedding from .llms.vertex_ai.vertex_gemma_models.main import VertexAIGemmaModels @@ -277,7 +273,7 @@ google_batch_embeddings = GoogleBatchEmbeddings() vertex_partner_models_chat_completion = VertexAIPartnerModels() vertex_gemma_chat_completion = VertexAIGemmaModels() vertex_model_garden_chat_completion = VertexAIModelGardenModels() -vertex_text_to_speech = VertexTextToSpeechAPI() +# vertex_text_to_speech is now replaced by VertexAITextToSpeechConfig sagemaker_llm = SagemakerLLM() watsonx_chat_completion = WatsonXChatHandler() openai_like_embedding = OpenAILikeEmbeddingHandler() @@ -1989,6 +1985,36 @@ def completion( # type: ignore # noqa: PLR0915 ) raise e + elif custom_llm_provider == "ragflow": + ## COMPLETION CALL - RAGFlow uses HTTP handler to support custom URL paths + try: + response = base_llm_http_handler.completion( + model=model, + messages=messages, + headers=headers, + model_response=model_response, + api_key=api_key, + api_base=api_base, + acompletion=acompletion, + logging_obj=logging, + optional_params=optional_params, + litellm_params=litellm_params, + shared_session=shared_session, + timeout=timeout, + client=client, + custom_llm_provider=custom_llm_provider, + encoding=encoding, + stream=stream, + provider_config=provider_config, + ) + except Exception as e: + logging.post_call( + input=messages, + api_key=api_key, + original_response=str(e), + additional_args={"headers": headers}, + ) + raise e elif custom_llm_provider == "xai": ## COMPLETION CALL try: @@ -6145,30 +6171,13 @@ def speech( # noqa: PLR0915 _is_async=aspeech or False, ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta": + from litellm.llms.vertex_ai.text_to_speech.transformation import ( + VertexAITextToSpeechConfig, + ) + generic_optional_params = GenericLiteLLMParams(**kwargs) - api_base = generic_optional_params.api_base or "" - vertex_ai_project = ( - generic_optional_params.vertex_project - or litellm.vertex_project - or get_secret_str("VERTEXAI_PROJECT") - ) - vertex_ai_location = ( - generic_optional_params.vertex_location - or litellm.vertex_location - or get_secret_str("VERTEXAI_LOCATION") - ) - vertex_credentials = ( - generic_optional_params.vertex_credentials - or get_secret_str("VERTEXAI_CREDENTIALS") - ) - - if voice is not None and not isinstance(voice, dict): - raise litellm.BadRequestError( - message=f"'voice' is required to be passed as a dict for Vertex AI TTS, passed in voice={voice}", - model=model, - llm_provider=custom_llm_provider, - ) + # Handle Gemini models separately (they use speech_to_completion_bridge) if "gemini" in model: from .endpoints.speech.speech_to_completion_bridge.handler import ( speech_to_completion_bridge_handler, @@ -6184,19 +6193,37 @@ def speech( # noqa: PLR0915 logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, ) - response = vertex_text_to_speech.audio_speech( - _is_async=aspeech, - vertex_credentials=vertex_credentials, - vertex_project=vertex_ai_project, - vertex_location=vertex_ai_location, - timeout=timeout, - api_base=api_base, + + # Vertex AI Text-to-Speech (Google Cloud TTS) + if text_to_speech_provider_config is None: + text_to_speech_provider_config = VertexAITextToSpeechConfig() + + # Cast to specific Vertex AI config type to access dispatch method + vertex_config = cast( + VertexAITextToSpeechConfig, text_to_speech_provider_config + ) + + # Store Vertex AI specific params in litellm_params_dict + litellm_params_dict.update({ + "vertex_project": generic_optional_params.vertex_project, + "vertex_location": generic_optional_params.vertex_location, + "vertex_credentials": generic_optional_params.vertex_credentials, + }) + + response = vertex_config.dispatch_text_to_speech( model=model, input=input, voice=voice, optional_params=optional_params, - kwargs=kwargs, + litellm_params_dict=litellm_params_dict, logging_obj=logging_obj, + timeout=timeout, + extra_headers=headers, + base_llm_http_handler=base_llm_http_handler, + aspeech=aspeech or False, + api_base=generic_optional_params.api_base, + api_key=None, # Vertex AI uses OAuth, not API key + **kwargs, ) elif custom_llm_provider == "gemini": from .endpoints.speech.speech_to_completion_bridge.handler import ( diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9fdc1704f41..d02a01e3a67 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -269,6 +269,71 @@ "supports_response_schema": true, "supports_vision": true }, + "amazon.nova-2-lite-v1:0": { + "input_cost_per_token": 3e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_video_input": true, + "supports_vision": true + }, + "apac.amazon.nova-2-lite-v1:0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.75e-06, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_video_input": true, + "supports_vision": true + }, + "eu.amazon.nova-2-lite-v1:0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.75e-06, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_video_input": true, + "supports_vision": true + }, + "us.amazon.nova-2-lite-v1:0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.75e-06, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_video_input": true, + "supports_vision": true + }, + "amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, "litellm_provider": "bedrock_converse", @@ -360,6 +425,15 @@ "litellm_provider": "bedrock", "mode": "image_generation" }, + "amazon.titan-image-generator-v2:0": { + "input_cost_per_image": 0.0, + "output_cost_per_image": 0.008, + "output_cost_per_image_premium_image": 0.01, + "output_cost_per_image_above_1024_and_1024_pixels": 0.01, + "output_cost_per_image_above_1024_and_1024_pixels_and_premium_image": 0.012, + "litellm_provider": "bedrock", + "mode": "image_generation" + }, "twelvelabs.marengo-embed-2-7-v1:0": { "input_cost_per_token": 7e-05, "litellm_provider": "bedrock", @@ -6717,6 +6791,33 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "claude-opus-4-5": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_creation_input_token_cost_above_1hr": 1e-05, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "anthropic", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 159 + }, "claude-sonnet-4-20250514": { "deprecation_date": "2026-05-14", "cache_creation_input_token_cost": 3.75e-06, @@ -7824,26 +7925,298 @@ "source": "https://www.databricks.com/product/pricing/foundation-model-serving" }, "databricks/databricks-claude-3-7-sonnet": { - "input_cost_per_token": 2.5e-06, - "input_dbu_cost_per_token": 3.571e-05, + "input_cost_per_token": 2.9999900000000002e-06, + "input_dbu_cost_per_token": 4.2857e-05, "litellm_provider": "databricks", "max_input_tokens": 200000, "max_output_tokens": 128000, "max_tokens": 200000, "metadata": { - "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Claude 3.7 conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.7857e-05, - "output_db_cost_per_token": 0.000214286, - "source": "https://www.databricks.com/product/pricing/foundation-model-serving", + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", "supports_assistant_prefill": true, "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": true }, + "databricks/databricks-claude-haiku-4-5": { + "input_cost_per_token": 1.00002e-06, + "input_dbu_cost_per_token": 1.4286e-05, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 5.00003e-06, + "output_dbu_cost_per_token": 7.1429e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "databricks/databricks-claude-opus-4": { + "input_cost_per_token": 1.5000020000000002e-05, + "input_dbu_cost_per_token": 0.000214286, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 32000, + "max_tokens": 200000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 7.500003000000001e-05, + "output_dbu_cost_per_token": 0.001071429, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "databricks/databricks-claude-opus-4-1": { + "input_cost_per_token": 1.5000020000000002e-05, + "input_dbu_cost_per_token": 0.000214286, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 32000, + "max_tokens": 200000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 7.500003000000001e-05, + "output_dbu_cost_per_token": 0.001071429, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "databricks/databricks-claude-opus-4-5": { + "input_cost_per_token": 5.00003e-06, + "input_dbu_cost_per_token": 7.1429e-05, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.5000010000000002e-05, + "output_dbu_cost_per_token": 0.000357143, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "databricks/databricks-claude-sonnet-4": { + "input_cost_per_token": 2.9999900000000002e-06, + "input_dbu_cost_per_token": 4.2857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "databricks/databricks-claude-sonnet-4-1": { + "input_cost_per_token": 2.9999900000000002e-06, + "input_dbu_cost_per_token": 4.2857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "databricks/databricks-claude-sonnet-4-5": { + "input_cost_per_token": 2.9999900000000002e-06, + "input_dbu_cost_per_token": 4.2857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 200000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, + "databricks/databricks-gemini-2-5-flash": { + "input_cost_per_token": 3.0001999999999996e-07, + "input_dbu_cost_per_token": 4.285999999999999e-06, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65535, + "max_tokens": 1048576, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 2.49998e-06, + "output_dbu_cost_per_token": 3.5714e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "databricks/databricks-gemini-2-5-pro": { + "input_cost_per_token": 1.24999e-06, + "input_dbu_cost_per_token": 1.7857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 1048576, + "max_output_tokens": 65536, + "max_tokens": 1048576, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 9.999990000000002e-06, + "output_dbu_cost_per_token": 0.000142857, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving", + "supports_function_calling": true, + "supports_tool_choice": true + }, + "databricks/databricks-gemma-3-12b": { + "input_cost_per_token": 1.5000999999999998e-07, + "input_dbu_cost_per_token": 2.1429999999999996e-06, + "litellm_provider": "databricks", + "max_input_tokens": 128000, + "max_output_tokens": 32000, + "max_tokens": 128000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 5.0001e-07, + "output_dbu_cost_per_token": 7.143e-06, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving" + }, + "databricks/databricks-gpt-5": { + "input_cost_per_token": 1.24999e-06, + "input_dbu_cost_per_token": 1.7857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 400000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 9.999990000000002e-06, + "output_dbu_cost_per_token": 0.000142857, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-1": { + "input_cost_per_token": 1.24999e-06, + "input_dbu_cost_per_token": 1.7857e-05, + "litellm_provider": "databricks", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 400000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 9.999990000000002e-06, + "output_dbu_cost_per_token": 0.000142857, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-mini": { + "input_cost_per_token": 2.4997000000000006e-07, + "input_dbu_cost_per_token": 3.571e-06, + "litellm_provider": "databricks", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 400000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 1.9999700000000004e-06, + "output_dbu_cost_per_token": 2.8571e-05, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-5-nano": { + "input_cost_per_token": 4.998e-08, + "input_dbu_cost_per_token": 7.14e-07, + "litellm_provider": "databricks", + "max_input_tokens": 400000, + "max_output_tokens": 128000, + "max_tokens": 400000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 3.9998000000000007e-07, + "output_dbu_cost_per_token": 5.714000000000001e-06, + "source": "https://www.databricks.com/product/pricing/proprietary-foundation-model-serving" + }, + "databricks/databricks-gpt-oss-120b": { + "input_cost_per_token": 1.5000999999999998e-07, + "input_dbu_cost_per_token": 2.1429999999999996e-06, + "litellm_provider": "databricks", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 5.9997e-07, + "output_dbu_cost_per_token": 8.571e-06, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving" + }, + "databricks/databricks-gpt-oss-20b": { + "input_cost_per_token": 7e-08, + "input_dbu_cost_per_token": 1e-06, + "litellm_provider": "databricks", + "max_input_tokens": 131072, + "max_output_tokens": 131072, + "max_tokens": 131072, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 3.0001999999999996e-07, + "output_dbu_cost_per_token": 4.285999999999999e-06, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving" + }, "databricks/databricks-gte-large-en": { - "input_cost_per_token": 1.2999e-07, + "input_cost_per_token": 1.2999000000000001e-07, "input_dbu_cost_per_token": 1.857e-06, "litellm_provider": "databricks", "max_input_tokens": 8192, @@ -7868,14 +8241,14 @@ "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.5e-06, + "output_cost_per_token": 1.5000300000000002e-06, "output_dbu_cost_per_token": 2.1429e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", "supports_tool_choice": true }, "databricks/databricks-llama-4-maverick": { - "input_cost_per_token": 5e-06, - "input_dbu_cost_per_token": 7.143e-05, + "input_cost_per_token": 5.0001e-07, + "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", "max_input_tokens": 128000, "max_output_tokens": 128000, @@ -7884,13 +8257,13 @@ "notes": "Databricks documentation now provides both DBU costs (_dbu_cost_per_token) and dollar costs(_cost_per_token)." }, "mode": "chat", - "output_cost_per_token": 1.5e-05, - "output_dbu_cost_per_token": 0.00021429, + "output_cost_per_token": 1.5000300000000002e-06, + "output_dbu_cost_per_token": 2.1429e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", "supports_tool_choice": true }, "databricks/databricks-meta-llama-3-1-405b-instruct": { - "input_cost_per_token": 5e-06, + "input_cost_per_token": 5.00003e-06, "input_dbu_cost_per_token": 7.1429e-05, "litellm_provider": "databricks", "max_input_tokens": 128000, @@ -7900,14 +8273,29 @@ "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 1.500002e-05, - "output_db_cost_per_token": 0.000214286, + "output_cost_per_token": 1.5000020000000002e-05, + "output_dbu_cost_per_token": 0.000214286, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", "supports_tool_choice": true }, + "databricks/databricks-meta-llama-3-1-8b-instruct": { + "input_cost_per_token": 1.5000999999999998e-07, + "input_dbu_cost_per_token": 2.1429999999999996e-06, + "litellm_provider": "databricks", + "max_input_tokens": 200000, + "max_output_tokens": 128000, + "max_tokens": 200000, + "metadata": { + "notes": "Input/output cost per token is dbu cost * $0.070. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." + }, + "mode": "chat", + "output_cost_per_token": 4.5003000000000007e-07, + "output_dbu_cost_per_token": 6.429000000000001e-06, + "source": "https://www.databricks.com/product/pricing/foundation-model-serving" + }, "databricks/databricks-meta-llama-3-3-70b-instruct": { - "input_cost_per_token": 1.00002e-06, - "input_dbu_cost_per_token": 1.4286e-05, + "input_cost_per_token": 5.0001e-07, + "input_dbu_cost_per_token": 7.143e-06, "litellm_provider": "databricks", "max_input_tokens": 128000, "max_output_tokens": 128000, @@ -7916,8 +8304,8 @@ "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.99999e-06, - "output_dbu_cost_per_token": 4.2857e-05, + "output_cost_per_token": 1.5000300000000002e-06, + "output_dbu_cost_per_token": 2.1429e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", "supports_tool_choice": true }, @@ -7932,7 +8320,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 2.99999e-06, + "output_cost_per_token": 2.9999900000000002e-06, "output_dbu_cost_per_token": 4.2857e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", "supports_tool_choice": true @@ -7948,13 +8336,13 @@ "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 9.9902e-07, + "output_cost_per_token": 1.00002e-06, "output_dbu_cost_per_token": 1.4286e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", "supports_tool_choice": true }, "databricks/databricks-mpt-30b-instruct": { - "input_cost_per_token": 9.9902e-07, + "input_cost_per_token": 1.00002e-06, "input_dbu_cost_per_token": 1.4286e-05, "litellm_provider": "databricks", "max_input_tokens": 8192, @@ -7964,7 +8352,7 @@ "notes": "Input/output cost per token is dbu cost * $0.070, based on databricks Llama 3.1 70B conversion. Number provided for reference, '*_dbu_cost_per_token' used in actual calculation." }, "mode": "chat", - "output_cost_per_token": 9.9902e-07, + "output_cost_per_token": 1.00002e-06, "output_dbu_cost_per_token": 1.4286e-05, "source": "https://www.databricks.com/product/pricing/foundation-model-serving", "supports_tool_choice": true @@ -9250,6 +9638,21 @@ "supports_prompt_caching": true, "supports_tool_choice": true }, + "deepseek/deepseek-v3.2": { + "input_cost_per_token": 2.8e-07, + "input_cost_per_token_cache_hit": 2.8e-08, + "litellm_provider": "deepseek", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "deepseek.v3-v1:0": { "input_cost_per_token": 5.8e-07, "litellm_provider": "bedrock_converse", @@ -10107,6 +10510,19 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/accounts/fireworks/models/deepseek-v3p2": { + "input_cost_per_token": 1.2e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://fireworks.ai/models/fireworks/deepseek-v3p2", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "fireworks_ai/accounts/fireworks/models/firefunction-v2": { "input_cost_per_token": 9e-07, "litellm_provider": "fireworks_ai", @@ -10445,25 +10861,25 @@ "supports_tool_choice": true }, "ft:babbage-002": { - "input_cost_per_token": 4e-07, + "input_cost_per_token": 1.6e-06, "input_cost_per_token_batches": 2e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 16384, "mode": "completion", - "output_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 2e-07 }, "ft:davinci-002": { - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.2e-05, "input_cost_per_token_batches": 1e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 16384, "mode": "completion", - "output_cost_per_token": 2e-06, + "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 1e-06 }, "ft:gpt-3.5-turbo": { @@ -10526,6 +10942,7 @@ "supports_tool_choice": true }, "ft:gpt-4o-2024-08-06": { + "cache_read_input_token_cost": 1.875e-06, "input_cost_per_token": 3.75e-06, "input_cost_per_token_batches": 1.875e-06, "litellm_provider": "openai", @@ -10538,6 +10955,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -10558,8 +10976,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true + "supports_tool_choice": true }, "ft:gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 1.5e-07, @@ -10578,8 +10995,79 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true + "supports_tool_choice": true + }, + "ft:gpt-4.1-2025-04-14": { + "cache_read_input_token_cost": 7.5e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "ft:gpt-4.1-mini-2025-04-14": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-07, + "input_cost_per_token_batches": 4e-07, + "litellm_provider": "openai", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "output_cost_per_token_batches": 1.6e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "ft:gpt-4.1-nano-2025-04-14": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, + "litellm_provider": "openai", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 8e-07, + "output_cost_per_token_batches": 4e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "ft:o4-mini-2025-04-16": { + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 1.6e-05, + "output_cost_per_token_batches": 8e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true }, "gemini-1.0-pro": { "input_cost_per_character": 1.25e-07, @@ -16402,7 +16890,7 @@ "output_cost_per_token": 9.9e-07, "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/distil-whisper-large-v3-en": { @@ -16421,7 +16909,7 @@ "mode": "chat", "output_cost_per_token": 7e-08, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/gemma2-9b-it": { @@ -16433,7 +16921,7 @@ "mode": "chat", "output_cost_per_token": 2e-07, "supports_function_calling": false, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": false }, "groq/llama-3.1-405b-reasoning": { @@ -16445,7 +16933,7 @@ "mode": "chat", "output_cost_per_token": 7.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.1-70b-versatile": { @@ -16458,7 +16946,7 @@ "mode": "chat", "output_cost_per_token": 7.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.1-8b-instant": { @@ -16470,7 +16958,7 @@ "mode": "chat", "output_cost_per_token": 8e-08, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-11b-text-preview": { @@ -16483,7 +16971,7 @@ "mode": "chat", "output_cost_per_token": 1.8e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-11b-vision-preview": { @@ -16496,7 +16984,7 @@ "mode": "chat", "output_cost_per_token": 1.8e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true, "supports_vision": true }, @@ -16510,7 +16998,7 @@ "mode": "chat", "output_cost_per_token": 4e-08, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-3b-preview": { @@ -16523,7 +17011,7 @@ "mode": "chat", "output_cost_per_token": 6e-08, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-90b-text-preview": { @@ -16536,7 +17024,7 @@ "mode": "chat", "output_cost_per_token": 9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-90b-vision-preview": { @@ -16549,7 +17037,7 @@ "mode": "chat", "output_cost_per_token": 9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true, "supports_vision": true }, @@ -16573,7 +17061,7 @@ "mode": "chat", "output_cost_per_token": 7.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-guard-3-8b": { @@ -16594,7 +17082,7 @@ "mode": "chat", "output_cost_per_token": 8e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama3-groq-70b-8192-tool-use-preview": { @@ -16607,7 +17095,7 @@ "mode": "chat", "output_cost_per_token": 8.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama3-groq-8b-8192-tool-use-preview": { @@ -16620,7 +17108,7 @@ "mode": "chat", "output_cost_per_token": 1.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/meta-llama/llama-4-maverick-17b-128e-instruct": { @@ -16666,7 +17154,7 @@ "mode": "chat", "output_cost_per_token": 2.4e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/moonshotai/kimi-k2-instruct": { @@ -16742,7 +17230,7 @@ "output_cost_per_token": 5.9e-07, "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/whisper-large-v3": { @@ -18417,6 +18905,34 @@ "supports_tool_choice": true, "supports_web_search": true }, + "moonshot/kimi-k2-0905-preview": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_web_search": true + }, + "moonshot/kimi-k2-turbo-preview": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.15e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_web_search": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 2e-06, @@ -18474,14 +18990,15 @@ "supports_vision": true }, "moonshot/kimi-thinking-preview": { - "input_cost_per_token": 3e-05, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 6e-07, "litellm_provider": "moonshot", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-05, - "source": "https://platform.moonshot.ai/docs/pricing", + "output_cost_per_token": 2.5e-06, + "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", "supports_vision": true }, "moonshot/kimi-k2-thinking": { @@ -18498,6 +19015,20 @@ "supports_tool_choice": true, "supports_web_search": true }, + "moonshot/kimi-k2-thinking-turbo": { + "cache_read_input_token_cost": 1.5e-7, + "input_cost_per_token": 1.15e-6, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8e-6, + "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_web_search": true + }, "moonshot/moonshot-v1-128k": { "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", @@ -20173,6 +20704,21 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/deepseek/deepseek-v3.2": { + "input_cost_per_token": 2.8e-07, + "input_cost_per_token_cache_hit": 2.8e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/deepseek/deepseek-v3.2-exp": { "input_cost_per_token": 2e-07, "input_cost_per_token_cache_hit": 2e-08, @@ -23428,6 +23974,32 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "global.anthropic.claude-opus-4-5-20251101-v1:0": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 159 + }, "us.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -24511,6 +25083,15 @@ "supports_parallel_function_calling": true, "supports_tool_choice": true }, + "vertex_ai/chirp": { + "input_cost_per_character": 30e-06, + "litellm_provider": "vertex_ai", + "mode": "audio_speech", + "source": "https://cloud.google.com/text-to-speech/pricing", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, "vertex_ai/claude-3-5-haiku": { "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", @@ -25941,8 +26522,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.135, - "output_cost_per_token": 0.4, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.5e-06, "litellm_provider": "wandb", "mode": "chat" }, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 2c03cbdae31..d42fd80cbd7 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -17,16 +17,15 @@ from urllib.parse import urlparse from fastapi import HTTPException from httpx import HTTPStatusError from mcp import ReadResourceResult, Resource +from mcp.types import CallToolRequestParams as MCPCallToolRequestParams from mcp.types import ( - CallToolRequestParams as MCPCallToolRequestParams, + CallToolResult, GetPromptRequestParams, GetPromptResult, Prompt, ResourceTemplate, ) -from mcp.types import CallToolResult from mcp.types import Tool as MCPTool - from pydantic import AnyUrl import litellm @@ -1949,7 +1948,12 @@ class MCPServerManager: ) = split_server_prefix_from_name(tool_name) if original_tool_name in self.tool_name_to_mcp_server_name_mapping: for server in self.get_registry().values(): - if normalize_server_name(server.name) == normalize_server_name( + if server.server_name is None: + if normalize_server_name(server.name) == normalize_server_name( + server_name_from_prefix + ): + return server + elif normalize_server_name(server.server_name) == normalize_server_name( server_name_from_prefix ): return server diff --git a/litellm/proxy/_experimental/out/assets/logos/a2a_agent.png b/litellm/proxy/_experimental/out/assets/logos/a2a_agent.png new file mode 100644 index 00000000000..305ae1acaf4 Binary files /dev/null and b/litellm/proxy/_experimental/out/assets/logos/a2a_agent.png differ diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index d2bbe8ee6e3..f1916e99ff9 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -3,23 +3,19 @@ model_list: litellm_params: model: openai/gpt-3.5-turbo api_key: os.environ/OPENAI_API_KEY + - model_name: claude-sonnet-4-5-20250929 + litellm_params: + model: anthropic/claude-sonnet-4-5-20250929 + - model_name: gpt-4.1-mini + litellm_params: + model: openai/gpt-4.1-mini + guardrails: - # - guardrail_name: model-armor-shield - # litellm_params: - # guardrail: model_armor - # mode: "post_call" # Run on both input and output - # template_id: "test-prompt-template" # Required: Your Model Armor template ID - # project_id: "test-vector-store-db" # Your GCP project ID - # location: "us" # GCP location (default: us-central1) - # mask_request_content: true # Enable request content masking - # mask_response_content: true # Enable response content masking - # fail_on_error: true # Fail request if Model Armor errors (default: true) - # default_on: true # Run by default for all requests - guardrail_name: generic-guardrail litellm_params: guardrail: generic_guardrail_api - mode: [pre_call, post_call] + mode: ["pre_call", "post_call", "during_call"] headers: Authorization: Bearer mock-bedrock-token-12345 api_base: http://localhost:8080 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b5b0bd80602..53a8627bc8f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2444,6 +2444,7 @@ class CallInfo(LiteLLMPydanticObjectBase): user_id: Optional[str] = None team_id: Optional[str] = None team_alias: Optional[str] = None + organization_id: Optional[str] = None user_email: Optional[str] = None key_alias: Optional[str] = None projected_exceeded_date: Optional[str] = None @@ -3440,9 +3441,7 @@ class LiteLLM_JWTAuth(LiteLLMPydanticObjectBase): team_id_upsert: bool = False team_ids_jwt_field: Optional[str] = None upsert_sso_user_to_team: bool = False - team_allowed_routes: List[ - Literal["openai_routes", "info_routes", "management_routes"] - ] = ["openai_routes", "info_routes"] + team_allowed_routes: List[str] = ["openai_routes", "info_routes"] team_id_default: Optional[str] = Field( default=None, description="If no team_id given, default permissions/spend-tracking to this team.s", diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py new file mode 100644 index 00000000000..32af1232f6b --- /dev/null +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -0,0 +1,220 @@ +""" +A2A Protocol endpoints for LiteLLM Proxy. + +Allows clients to invoke agents through LiteLLM using the A2A protocol. +The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM. +""" + +import json +from typing import Any, Optional + +from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi.responses import JSONResponse, StreamingResponse + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + +router = APIRouter() + + +def _jsonrpc_error( + request_id: Optional[str], + code: int, + message: str, + status_code: int = 400, +) -> JSONResponse: + """Create a JSON-RPC 2.0 error response.""" + return JSONResponse( + content={ + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": code, "message": message}, + }, + status_code=status_code, + ) + + +def _get_agent(agent_id: str): + """Look up an agent by ID or name. Returns None if not found.""" + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) + if agent is None: + agent = global_agent_registry.get_agent_by_name(agent_name=agent_id) + return agent + + +async def _handle_stream_message( + a2a_client: Any, + request_id: str, + params: dict, +) -> StreamingResponse: + """Handle message/stream method.""" + from a2a.types import MessageSendParams, SendStreamingMessageRequest + + a2a_request = SendStreamingMessageRequest( + id=request_id, + params=MessageSendParams(**params), + ) + + async def stream_response(): + try: + async for chunk in a2a_client.send_message_streaming(a2a_request): + yield json.dumps(chunk.model_dump(mode="json", exclude_none=True)) + "\n" + except Exception as e: + verbose_proxy_logger.exception(f"Error streaming A2A response: {e}") + yield json.dumps({ + "jsonrpc": "2.0", + "id": request_id, + "error": {"code": -32603, "message": f"Streaming error: {str(e)}"}, + }) + "\n" + + return StreamingResponse(stream_response(), media_type="application/x-ndjson") + + +@router.get( + "/a2a/{agent_id}/.well-known/agent-card.json", + tags=["[beta] A2A Agents"], + dependencies=[Depends(user_api_key_auth)], +) +async def get_agent_card( + agent_id: str, + request: Request, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get the agent card for an agent (A2A discovery endpoint). + + The URL in the agent card is rewritten to point to the LiteLLM proxy, + so all subsequent A2A calls go through LiteLLM for logging and cost tracking. + """ + try: + agent = _get_agent(agent_id) + if agent is None: + raise HTTPException(status_code=404, detail=f"Agent '{agent_id}' not found") + + # Copy and rewrite URL to point to LiteLLM proxy + agent_card = dict(agent.agent_card_params) + agent_card["url"] = f"{str(request.base_url).rstrip('/')}/a2a/{agent_id}" + + verbose_proxy_logger.debug( + f"Returning agent card for '{agent_id}' with proxy URL: {agent_card['url']}" + ) + return JSONResponse(content=agent_card) + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception(f"Error getting agent card: {e}") + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/a2a/{agent_id}", + tags=["[beta] A2A Agents"], + dependencies=[Depends(user_api_key_auth)], +) +@router.post( + "/a2a/{agent_id}/message/send", + tags=["[beta] A2A Agents"], + dependencies=[Depends(user_api_key_auth)], +) +@router.post( + "/v1/a2a/{agent_id}/message/send", + tags=["[beta] A2A Agents"], + dependencies=[Depends(user_api_key_auth)], +) +async def invoke_agent_a2a( + agent_id: str, + request: Request, + fastapi_response: Response, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Invoke an agent using the A2A protocol (JSON-RPC 2.0). + + Supported methods: + - message/send: Send a message and get a response + - message/stream: Send a message and stream the response + """ + from litellm.a2a_protocol import asend_message, create_a2a_client + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + from litellm.proxy.proxy_server import ( + general_settings, + proxy_config, + version, + ) + + body = {} + try: + body = await request.json() + verbose_proxy_logger.debug(f"A2A request for agent '{agent_id}': {body}") + + # Validate JSON-RPC format + if body.get("jsonrpc") != "2.0": + return _jsonrpc_error(body.get("id"), -32600, "Invalid Request: jsonrpc must be '2.0'") + + request_id = body.get("id") + method = body.get("method") + params = body.get("params", {}) + + # Find the agent + agent = _get_agent(agent_id) + if agent is None: + return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' not found", 404) + + # Get backend URL and agent name + agent_url = agent.agent_card_params.get("url") + agent_name = agent.agent_card_params.get("name", agent_id) + if not agent_url: + return _jsonrpc_error(request_id, -32000, f"Agent '{agent_id}' has no URL configured", 500) + + verbose_proxy_logger.info(f"Proxying A2A request to agent '{agent_id}' at {agent_url}") + + # Set up data dict for litellm processing + body.update({ + "model": f"a2a_agent/{agent_name}", + "custom_llm_provider": "a2a_agent", + }) + + # Add litellm data (user_api_key, user_id, team_id, etc.) + data = await add_litellm_data_to_request( + data=body, + request=request, + user_api_key_dict=user_api_key_dict, + proxy_config=proxy_config, + general_settings=general_settings, + version=version, + ) + + # Create A2A client + a2a_client = await create_a2a_client(base_url=agent_url) + + if method == "message/send": + from a2a.types import MessageSendParams, SendMessageRequest + + a2a_request = SendMessageRequest( + id=request_id, + params=MessageSendParams(**params), + ) + + # Pass litellm data through kwargs for proper logging + response = await asend_message( + a2a_client=a2a_client, + request=a2a_request, + metadata=data.get("metadata", {}), + proxy_server_request=data.get("proxy_server_request"), + ) + return JSONResponse(content=response.model_dump(mode="json", exclude_none=True)) + + elif method == "message/stream": + return await _handle_stream_message(a2a_client, request_id, params) + else: + return _jsonrpc_error(request_id, -32601, f"Method '{method}' not found") + + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception(f"Error invoking agent: {e}") + return _jsonrpc_error(body.get("id"), -32603, f"Internal error: {str(e)}", 500) diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 489c8e82302..90688036d84 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -29,7 +29,7 @@ router = APIRouter() @router.get( "/v1/agents", - tags=["[beta] Agents"], + tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], response_model=List[AgentResponse], ) @@ -101,7 +101,7 @@ from litellm.proxy.agent_endpoints.agent_registry import ( @router.post( "/v1/agents", - tags=["[beta] Agents"], + tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], response_model=AgentResponse, ) @@ -196,7 +196,7 @@ async def create_agent( @router.get( "/v1/agents/{agent_id}", - tags=["[beta] Agents"], + tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], response_model=AgentResponse, ) @@ -239,7 +239,7 @@ async def get_agent_by_id(agent_id: str): @router.put( "/v1/agents/{agent_id}", - tags=["[beta] Agents"], + tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], response_model=AgentResponse, ) @@ -328,7 +328,7 @@ async def update_agent( @router.patch( "/v1/agents/{agent_id}", - tags=["[beta] Agents"], + tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], response_model=AgentResponse, ) @@ -471,7 +471,7 @@ async def delete_agent(agent_id: str): @router.post( "/v1/agents/{agent_id}/make_public", - tags=["[beta] Agents"], + tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], response_model=AgentMakePublicResponse, ) @@ -585,7 +585,7 @@ async def make_agent_public( @router.post( "/v1/agents/make_public", - tags=["[beta] Agents"], + tags=["[beta] A2A Agents"], dependencies=[Depends(user_api_key_auth)], response_model=AgentMakePublicResponse, ) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index abea9e6fee1..3317560904d 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -2,20 +2,14 @@ Unified /v1/messages endpoint - (Anthropic Spec) """ -import asyncio -from fastapi import APIRouter, Depends, HTTPException, Request, Response, status +from fastapi import APIRouter, Depends, HTTPException, Request, Response -import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth -from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - create_streaming_response, -) +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.common_utils.http_parsing_utils import _read_request_body -from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request from litellm.types.utils import TokenCountResponse router = APIRouter() @@ -49,169 +43,28 @@ async def anthropic_response( # noqa: PLR0915 version, ) - request_data = await _read_request_body(request=request) - data: dict = {**request_data} + data = await _read_request_body(request=request) + base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) try: - data["model"] = ( - general_settings.get("completion_model", None) # server default - or user_model # model name passed via cli args - or data.get("model", None) # default passed in http request - ) - if user_model: - data["model"] = user_model - - data = await add_litellm_data_to_request( - data=data, # type: ignore + result = await base_llm_response_processor.base_process_llm_request( request=request, - general_settings=general_settings, + fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, - version=version, + route_type="anthropic_messages", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, proxy_config=proxy_config, + select_data_generator=None, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, ) - - # override with user settings, these are params passed via cli - if user_temperature: - data["temperature"] = user_temperature - if user_request_timeout: - data["request_timeout"] = user_request_timeout - if user_max_tokens: - data["max_tokens"] = user_max_tokens - if user_api_base: - data["api_base"] = user_api_base - - ### MODEL ALIAS MAPPING ### - # check if model name in model alias map - # get the actual model name - if data["model"] in litellm.model_alias_map: - data["model"] = litellm.model_alias_map[data["model"]] - - ### CALL HOOKS ### - modify incoming data before calling the model - data = await proxy_logging_obj.pre_call_hook( # type: ignore - user_api_key_dict=user_api_key_dict, data=data, call_type=CallTypes.anthropic_messages.value - ) - - tasks = [] - tasks.append( - proxy_logging_obj.during_call_hook( - data=data, - user_api_key_dict=user_api_key_dict, - call_type=ProxyBaseLLMRequestProcessing._get_pre_call_type( - route_type="anthropic_messages" # type: ignore - ), - ) - ) - - ### ROUTE THE REQUESTs ### - router_model_names = llm_router.model_names if llm_router is not None else [] - - # skip router if user passed their key - if ( - llm_router is not None and data["model"] in router_model_names - ): # model in router model list - llm_coro = llm_router.aanthropic_messages(**data) - elif ( - llm_router is not None - and llm_router.model_group_alias is not None - and data["model"] in llm_router.model_group_alias - ): # model set in model_group_alias - llm_coro = llm_router.aanthropic_messages(**data) - elif ( - llm_router is not None and data["model"] in llm_router.deployment_names - ): # model in router deployments, calling a specific deployment on the router - llm_coro = llm_router.aanthropic_messages(**data, specific_deployment=True) - elif ( - llm_router is not None and llm_router.has_model_id(data["model"]) - ): # model in router model list - llm_coro = llm_router.aanthropic_messages(**data) - elif ( - llm_router is not None - and data["model"] not in router_model_names - and ( - llm_router.default_deployment is not None - or len(llm_router.pattern_router.patterns) > 0 - ) - ): # model in router deployments, calling a specific deployment on the router - llm_coro = llm_router.aanthropic_messages(**data) - elif user_model is not None: # `litellm --model ` - llm_coro = litellm.anthropic_messages(**data) - else: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "error": "completion: Invalid model name passed in model=" - + data.get("model", "") - }, - ) - - tasks.append(llm_coro) - - # wait for call to end - llm_responses = asyncio.gather( - *tasks - ) # run the moderation check in parallel to the actual llm api call - - responses = await llm_responses - - response = responses[1] - - # Extract model_id from request metadata (set by router during routing) - litellm_metadata = data.get("litellm_metadata", {}) or {} - model_info = litellm_metadata.get("model_info", {}) or {} - model_id = model_info.get("id", "") or "" - - # Get other metadata from hidden_params - hidden_params = getattr(response, "_hidden_params", {}) or {} - cache_key = hidden_params.get("cache_key", None) or "" - api_base = hidden_params.get("api_base", None) or "" - response_cost = hidden_params.get("response_cost", None) or "" - - ### ALERTING ### - asyncio.create_task( - proxy_logging_obj.update_request_status( - litellm_call_id=data.get("litellm_call_id", ""), status="success" - ) - ) - - verbose_proxy_logger.debug("final response: %s", response) - - fastapi_response.headers.update( - ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - request_data=data, - hidden_params=hidden_params, - ) - ) - - if ( - "stream" in data and data["stream"] is True - ): # use generate_responses to stream responses - selected_data_generator = ( - ProxyBaseLLMRequestProcessing.async_sse_data_generator( - response=response, - user_api_key_dict=user_api_key_dict, - request_data=data, - proxy_logging_obj=proxy_logging_obj, - ) - ) - - return await create_streaming_response( - generator=selected_data_generator, - media_type="text/event-stream", - headers=dict(fastapi_response.headers), - ) - - ### CALL HOOKS ### - modify outgoing data - response = await proxy_logging_obj.post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response # type: ignore - ) - - verbose_proxy_logger.debug("\nResponse from Litellm:\n{}".format(response)) - return response + return result except Exception as e: await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data @@ -280,35 +133,30 @@ async def count_tokens( Returns: {"input_tokens": } """ from litellm.proxy.proxy_server import token_counter as internal_token_counter - + try: request_data = await _read_request_body(request=request) data: dict = {**request_data} - + # Extract required fields model_name = data.get("model") messages = data.get("messages", []) - + if not model_name: raise HTTPException( - status_code=400, - detail={"error": "model parameter is required"} + status_code=400, detail={"error": "model parameter is required"} ) - + if not messages: raise HTTPException( - status_code=400, - detail={"error": "messages parameter is required"} + status_code=400, detail={"error": "messages parameter is required"} ) - + # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest - - token_request = TokenCountRequest( - model=model_name, - messages=messages - ) - + + token_request = TokenCountRequest(model=model_name, messages=messages) + # Call the internal token counter function with direct request flag set to False token_response = await internal_token_counter( request=token_request, @@ -319,17 +167,18 @@ async def count_tokens( _token_response_dict = token_response.model_dump() elif isinstance(token_response, dict): _token_response_dict = token_response - + # Convert the internal response to Anthropic API format return {"input_tokens": _token_response_dict.get("total_tokens", 0)} - + except HTTPException: raise except Exception as e: verbose_proxy_logger.exception( - "litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - {}".format(str(e)) + "litellm.proxy.anthropic_endpoints.count_tokens(): Exception occurred - {}".format( + str(e) + ) ) raise HTTPException( - status_code=500, - detail={"error": f"Internal server error: {str(e)}"} + status_code=500, detail={"error": f"Internal server error: {str(e)}"} ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c9774b18b88..fc79a4d3591 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -143,6 +143,14 @@ async def common_checks( valid_token=valid_token, ) + # 3.1. If organization is in budget + await _organization_max_budget_check( + valid_token=valid_token, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + await _tag_max_budget_check( request_body=request_body, prisma_client=prisma_client, @@ -182,7 +190,18 @@ async def common_checks( general_settings.get("enforce_user_param", None) is not None and general_settings["enforce_user_param"] is True ): - if RouteChecks.is_llm_api_route(route=route) and "user" not in request_body: + # Get HTTP method from request + http_method = request.method if hasattr(request, 'method') else None + + # Check if it's a POST request and if it's an OpenAI route but not MCP + is_post_method = http_method and http_method.upper() == "POST" + is_openai_route = RouteChecks.is_llm_api_route(route=route) + is_mcp_route = route in LiteLLMRoutes.mcp_routes.value or RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value + ) + + # Enforce user param only for POST requests on OpenAI routes (excluding MCP routes) + if is_post_method and is_openai_route and not is_mcp_route and "user" not in request_body: raise Exception( f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}" ) @@ -1893,6 +1912,7 @@ async def _virtual_key_max_budget_check( max_budget=valid_token.max_budget, user_id=valid_token.user_id, team_id=valid_token.team_id, + organization_id=valid_token.org_id, user_email=user_email, key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, @@ -1939,6 +1959,7 @@ async def _virtual_key_soft_budget_check( user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, user_email=None, key_alias=valid_token.key_alias, event_group=Litellm_EntityType.KEY, @@ -1977,6 +1998,7 @@ async def _team_max_budget_check( user_id=valid_token.user_id, team_id=valid_token.team_id, team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, event_group=Litellm_EntityType.TEAM, ) asyncio.create_task( @@ -1993,6 +2015,65 @@ async def _team_max_budget_check( ) +async def _organization_max_budget_check( + valid_token: Optional[UserAPIKeyAuth], + prisma_client: Optional[PrismaClient], + user_api_key_cache: DualCache, + proxy_logging_obj: ProxyLogging, +): + """ + Check if the organization is over its max budget. + + Raises: + BudgetExceededError if the organization is over its max budget. + Triggers a budget alert if the organization is over its max budget. + """ + # Only check if token has organization info and organization_max_budget is set + if ( + valid_token is None + or valid_token.org_id is None + or valid_token.organization_max_budget is None + or valid_token.organization_max_budget <= 0 + ): + return + + # Get organization object to check current spend + if prisma_client is not None: + org_table = await get_org_object( + org_id=valid_token.org_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + if ( + org_table is not None + and org_table.spend >= valid_token.organization_max_budget + ): + # Trigger budget alert + call_info = CallInfo( + token=valid_token.token, + spend=org_table.spend, + max_budget=valid_token.organization_max_budget, + user_id=valid_token.user_id, + team_id=valid_token.team_id, + team_alias=valid_token.team_alias, + organization_id=valid_token.org_id, + event_group=Litellm_EntityType.ORGANIZATION, + ) + asyncio.create_task( + proxy_logging_obj.budget_alerts( + type="organization_budget", + user_info=call_info, + ) + ) + + raise litellm.BudgetExceededError( + current_cost=org_table.spend, + max_budget=valid_token.organization_max_budget, + message=f"Budget has been exceeded! Organization={valid_token.org_id} Current cost: {org_table.spend}, Max budget: {valid_token.organization_max_budget}", + ) + + async def _tag_max_budget_check( request_body: dict, prisma_client: Optional[PrismaClient], diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py new file mode 100644 index 00000000000..6ef983b221c --- /dev/null +++ b/litellm/proxy/auth/login_utils.py @@ -0,0 +1,337 @@ +""" +Login utilities for handling user authentication in the proxy server. + +This module contains the core login logic that can be reused across different +login endpoints (e.g., /login and /v2/login). +""" + +import os +import secrets +from typing import Literal, Optional, cast + +import litellm +from fastapi import HTTPException + +from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.proxy._types import ( + LiteLLM_UserTable, + LitellmUserRoles, + ProxyErrorTypes, + ProxyException, + UpdateUserRequest, + UserAPIKeyAuth, + hash_token, +) +from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.key_management_endpoints import ( + generate_key_helper_fn, +) +from litellm.proxy.management_endpoints.ui_sso import ( + get_disabled_non_admin_personal_key_creation, +) +from litellm.proxy.utils import PrismaClient, get_server_root_path +from litellm.secret_managers.main import get_secret_bool +from litellm.types.proxy.ui_sso import ReturnedUITokenObject + + +def get_ui_credentials(master_key: Optional[str]) -> tuple[str, str]: + """ + Get UI username and password from environment variables or master key. + + Args: + master_key: Master key for the proxy (used as fallback for password) + + Returns: + tuple[str, str]: A tuple containing (ui_username, ui_password) + + Raises: + ProxyException: If neither UI_PASSWORD nor master_key is available + """ + ui_username = os.getenv("UI_USERNAME", "admin") + ui_password = os.getenv("UI_PASSWORD", None) + if ui_password is None: + ui_password = str(master_key) if master_key is not None else None + if ui_password is None: + raise ProxyException( + message="set Proxy master key to use UI. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.", + type=ProxyErrorTypes.auth_error, + param="UI_PASSWORD", + code=500, + ) + return ui_username, ui_password + + +class LoginResult: + """Result object containing authentication data from login.""" + + def __init__( + self, + user_id: str, + key: str, + user_email: Optional[str], + user_role: str, + login_method: str = "username_password", + ): + self.user_id = user_id + self.key = key + self.user_email = user_email + self.user_role = user_role + self.login_method = login_method + + +async def authenticate_user( + username: str, + password: str, + master_key: Optional[str], + prisma_client: Optional[PrismaClient], +) -> LoginResult: + """ + Authenticate a user and generate an API key for UI access. + + This function handles two login scenarios: + 1. Admin login using UI_USERNAME and UI_PASSWORD + 2. User login using email and password from database + + Args: + username: Username or email from the login form + password: Password from the login form + master_key: Master key for the proxy (required) + prisma_client: Prisma database client (optional) + + Returns: + LoginResult: Object containing authentication data + + Raises: + ProxyException: If authentication fails or required configuration is missing + """ + if master_key is None: + raise ProxyException( + message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.", + type=ProxyErrorTypes.auth_error, + param="master_key", + code=500, + ) + + ui_username, ui_password = get_ui_credentials(master_key) + + # Check if we can find the `username` in the db. On the UI, users can enter username=their email + _user_row: Optional[LiteLLM_UserTable] = None + user_role: Optional[ + Literal[ + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.INTERNAL_USER, + LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, + ] + ] = None + + if prisma_client is not None: + _user_row = cast( + Optional[LiteLLM_UserTable], + await prisma_client.db.litellm_usertable.find_first( + where={"user_email": {"equals": username}} + ), + ) + + """ + To login to Admin UI, we support the following + - Login with UI_USERNAME and UI_PASSWORD + - Login with Invite Link `user_email` and `password` combination + """ + if secrets.compare_digest(username, ui_username) and secrets.compare_digest( + password, ui_password + ): + # Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin + user_role = LitellmUserRoles.PROXY_ADMIN + user_id = LITELLM_PROXY_ADMIN_NAME + + # we want the key created to have PROXY_ADMIN_PERMISSIONS + key_user_id = LITELLM_PROXY_ADMIN_NAME + if ( + os.getenv("PROXY_ADMIN_ID", None) is not None + and os.environ["PROXY_ADMIN_ID"] == user_id + ) or user_id == LITELLM_PROXY_ADMIN_NAME: + # checks if user is admin + key_user_id = os.getenv("PROXY_ADMIN_ID", LITELLM_PROXY_ADMIN_NAME) + + # Admin is Authe'd in - generate key for the UI to access Proxy + + # ensure this user is set as the proxy admin, in this route there is no sso, we can assume this user is only the admin + await user_update( + data=UpdateUserRequest( + user_id=key_user_id, + user_role=user_role, + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + ), + ) + + if os.getenv("DATABASE_URL") is not None: + response = await generate_key_helper_fn( + request_type="key", + **{ + "user_role": LitellmUserRoles.PROXY_ADMIN, + "duration": "24hr", + "key_max_budget": litellm.max_ui_session_budget, + "models": [], + "aliases": {}, + "config": {}, + "spend": 0, + "user_id": key_user_id, + "team_id": "litellm-dashboard", + }, # type: ignore + ) + else: + raise ProxyException( + message="No Database connected. Set DATABASE_URL in .env. If set, use `--detailed_debug` to debug issue.", + type=ProxyErrorTypes.auth_error, + param="DATABASE_URL", + code=500, + ) + + key = response["token"] # type: ignore + + if get_secret_bool("EXPERIMENTAL_UI_LOGIN"): + user_info: Optional[LiteLLM_UserTable] = None + if _user_row is not None: + user_info = _user_row + elif ( + user_id is not None + ): # if user_id is not None, we are using the UI_USERNAME and UI_PASSWORD + from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken + + user_info = LiteLLM_UserTable( + user_id=user_id, + user_role=user_role, + models=[], + max_budget=litellm.max_ui_session_budget, + ) + if user_info is None: + raise HTTPException( + status_code=401, + detail={ + "error": "User Information is required for experimental UI login" + }, + ) + + key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( + user_info + ) + + return LoginResult( + user_id=user_id, + key=key, + user_email=None, + user_role=user_role, + login_method="username_password", + ) + + elif _user_row is not None: + """ + When sharing invite links + + -> if the user has no role in the DB assume they are only a viewer + """ + user_id = getattr(_user_row, "user_id", "unknown") + user_role = getattr( + _user_row, "user_role", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY + ) + user_email = getattr(_user_row, "user_email", "unknown") + _password = getattr(_user_row, "password", "unknown") + + if _password is None: + raise ProxyException( + message="User has no password set. Please set a password for the user via `/user/update`.", + type=ProxyErrorTypes.auth_error, + param="password", + code=401, + ) + + # check if password == _user_row.password + hash_password = hash_token(token=password) + if secrets.compare_digest(password, _password) or secrets.compare_digest( + hash_password, _password + ): + if os.getenv("DATABASE_URL") is not None: + response = await generate_key_helper_fn( + request_type="key", + **{ # type: ignore + "user_role": user_role, + "duration": "24hr", + "key_max_budget": litellm.max_ui_session_budget, + "models": [], + "aliases": {}, + "config": {}, + "spend": 0, + "user_id": user_id, + "team_id": "litellm-dashboard", + }, + ) + else: + raise ProxyException( + message="No Database connected. Set DATABASE_URL in .env. If set, use `--detailed_debug` to debug issue.", + type=ProxyErrorTypes.auth_error, + param="DATABASE_URL", + code=500, + ) + + key = response["token"] # type: ignore + + return LoginResult( + user_id=user_id, + key=key, + user_email=user_email, + user_role=cast(str, user_role), + login_method="username_password", + ) + else: + raise ProxyException( + message=f"Invalid credentials used to access UI.\nNot valid credentials for {username}", + type=ProxyErrorTypes.auth_error, + param="invalid_credentials", + code=401, + ) + else: + raise ProxyException( + message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file", + type=ProxyErrorTypes.auth_error, + param="invalid_credentials", + code=401, + ) + + +def create_ui_token_object( + login_result: LoginResult, + general_settings: dict, + premium_user: bool, +) -> ReturnedUITokenObject: + """ + Create a ReturnedUITokenObject from a LoginResult. + + Args: + login_result: The result from authenticate_user + general_settings: General proxy settings dictionary + premium_user: Whether premium features are enabled + + Returns: + ReturnedUITokenObject: Token object ready for JWT encoding + """ + disabled_non_admin_personal_key_creation = ( + get_disabled_non_admin_personal_key_creation() + ) + + return ReturnedUITokenObject( + user_id=login_result.user_id, + key=login_result.key, + user_email=login_result.user_email, + user_role=login_result.user_role, + login_method=login_result.login_method, + premium_user=premium_user, + auth_header_name=general_settings.get( + "litellm_key_header_name", "Authorization" + ), + disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, + server_root_path=get_server_root_path(), + ) + diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index d2b04410026..0d2ffc70f29 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -337,6 +337,7 @@ class ProxyBaseLLMRequestProcessing: "alist_skills", "aget_skill", "adelete_skill", + "anthropic_messages", ], version: Optional[str] = None, user_model: Optional[str] = None, @@ -461,11 +462,12 @@ class ProxyBaseLLMRequestProcessing: "alist_skills", "aget_skill", "adelete_skill", + "anthropic_messages", ], proxy_logging_obj: ProxyLogging, general_settings: dict, proxy_config: ProxyConfig, - select_data_generator: Callable, + select_data_generator: Optional[Callable] = None, llm_router: Optional[Router] = None, model: Optional[str] = None, user_model: Optional[str] = None, @@ -605,7 +607,21 @@ class ProxyBaseLLMRequestProcessing: status_code=response.status_code, headers=custom_headers, ) - else: + elif route_type == "anthropic_messages": + selected_data_generator = ( + ProxyBaseLLMRequestProcessing.async_sse_data_generator( + response=response, + user_api_key_dict=user_api_key_dict, + request_data=self.data, + proxy_logging_obj=proxy_logging_obj, + ) + ) + return await create_streaming_response( + generator=selected_data_generator, + media_type="text/event-stream", + headers=custom_headers, + ) + elif select_data_generator: selected_data_generator = select_data_generator( response=response, user_api_key_dict=user_api_key_dict, @@ -804,12 +820,38 @@ class ProxyBaseLLMRequestProcessing: # This matches the original behavior before the refactor in commit 511d435f6f error_body = await e.response.aread() error_text = error_body.decode("utf-8") - + raise HTTPException( status_code=e.response.status_code, detail={"error": error_text}, ) error_msg = f"{str(e)}" + # Check for AttributeError in various places: + # 1. Direct AttributeError (already handled above) + # 2. In underlying exception (__cause__, __context__, original_exception) + has_attribute_error = ( + ( + isinstance(e, Exception) + and isinstance(getattr(e, "__cause__", None), AttributeError) + ) + or ( + isinstance(e, Exception) + and isinstance(getattr(e, "__context__", None), AttributeError) + ) + or ( + isinstance(e, Exception) + and isinstance(getattr(e, "original_exception", None), AttributeError) + ) + ) + + if has_attribute_error: + raise ProxyException( + message=f"Invalid request format: {error_msg}", + type="invalid_request_error", + param=None, + code=status.HTTP_400_BAD_REQUEST, + headers=headers, + ) raise ProxyException( message=getattr(e, "message", error_msg), type=getattr(e, "type", "None"), @@ -1093,7 +1135,9 @@ class ProxyBaseLLMRequestProcessing: return obj return None - def maybe_get_model_id(self, _logging_obj: Optional[LiteLLMLoggingObj]) -> Optional[str]: + def maybe_get_model_id( + self, _logging_obj: Optional[LiteLLMLoggingObj] + ) -> Optional[str]: """ Get model_id from logging object or request metadata. @@ -1103,10 +1147,7 @@ class ProxyBaseLLMRequestProcessing: model_id = None if _logging_obj: # 1. Try getting from litellm_params (updated during call) - if ( - hasattr(_logging_obj, "litellm_params") - and _logging_obj.litellm_params - ): + if hasattr(_logging_obj, "litellm_params") and _logging_obj.litellm_params: # First check direct model_info path (set by router.py with selected deployment) model_info = _logging_obj.litellm_params.get("model_info") or {} model_id = model_info.get("id", None) diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index c448742f6dc..69472c2cda4 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -72,7 +72,7 @@ class CustomOpenAPISpec: openapi_schema["components"]["schemas"] = {} # Add the schema - openapi_schema["components"]["schemas"][schema_name] = schema_def + CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def}) @staticmethod def add_request_body_to_paths(openapi_schema: Dict[str, Any], paths: List[str], schema_ref: str) -> None: diff --git a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py index 9a6032fd084..9aaa2fb8381 100644 --- a/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -1,4 +1,5 @@ #### Analytics Endpoints ##### +import os from fastapi import APIRouter from litellm.types.proxy.discovery_endpoints.ui_discovery_endpoints import ( @@ -14,8 +15,12 @@ router = APIRouter() ) # if mounted at root path async def get_ui_config(): from litellm.proxy.utils import get_proxy_base_url, get_server_root_path + from litellm.proxy.auth.auth_utils import _has_user_setup_sso + + auto_redirect_ui_login_to_sso = os.getenv("AUTO_REDIRECT_UI_LOGIN_TO_SSO", "true").lower() == "true" return UiDiscoveryEndpoints( server_root_path=get_server_root_path(), proxy_base_url=get_proxy_base_url(), + auto_redirect_to_sso=_has_user_setup_sso() and auto_redirect_ui_login_to_sso, ) diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index 395fdb249d3..bb78383ce44 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1214,13 +1214,15 @@ async def apply_guardrail( detail=f"Guardrail '{request.guardrail_name}' not found. Please ensure the guardrail is configured in your LiteLLM proxy.", ) - response_text = await active_guardrail.apply_guardrail( - texts=[request.text], + guardrailed_inputs = await active_guardrail.apply_guardrail( + inputs={"texts": [request.text]}, request_data={}, input_type="request", - images=None, ) + response_text = guardrailed_inputs.get("texts", []) - return ApplyGuardrailResponse(response_text=response_text[0][0]) + return ApplyGuardrailResponse( + response_text=response_text[0] if response_text else request.text + ) except Exception as e: raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index f9c3caf4944..e0fb3192401 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -20,6 +20,7 @@ from typing import ( AsyncGenerator, List, Literal, + NamedTuple, Optional, Tuple, Union, @@ -41,7 +42,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.secret_managers.main import get_secret_str -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GenericGuardrailAPIInputs, GuardrailEventHooks from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockContentItem, @@ -53,6 +54,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.utils import ( CallTypes, CallTypesLiteral, @@ -728,9 +730,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 @@ -800,9 +802,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 @@ -1244,12 +1246,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): async def apply_guardrail( self, - texts: List[str], + inputs: "GenericGuardrailAPIInputs", request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, - images: Optional[List[str]] = None, - ) -> Tuple[List[str], Optional[List[str]]]: + ) -> "GenericGuardrailAPIInputs": """ Apply Bedrock guardrail to a batch of texts for testing purposes. @@ -1257,17 +1258,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): It creates mock messages to test the guardrail functionality. Args: - texts: List of texts to analyze + inputs: Dictionary containing texts and optional images request_data: Request data dictionary for logging metadata input_type: Whether this is a "request" or "response" - images: Optional list of images (not processed separately) + logging_obj: Optional logging object Returns: - Tuple of (processed_texts, images) - texts may be masked, images unchanged + GenericGuardrailAPIInputs - processed_texts may be masked, images unchanged Raises: Exception: If content is blocked by Bedrock guardrail """ + texts = inputs.get("texts", []) try: verbose_proxy_logger.debug( f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)" @@ -1275,10 +1277,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): masked_texts = [] - for text in texts: - mock_messages: List[AllMessageValues] = [ - ChatCompletionUserMessage(role="user", content=text) - ] + mock_messages: List[AllMessageValues] = [ + ChatCompletionUserMessage(role="user", content=text) for text in texts + ] + request_messages = mock_messages filter_result = self._prepare_guardrail_messages_for_role( messages=request_messages @@ -1291,44 +1293,37 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data=request_data, ) - bedrock_response = await self.make_bedrock_api_request( - source="INPUT", - messages=mock_messages, - request_data=request_data, + if bedrock_response.get("action") == "BLOCKED": + raise Exception( + f"Content blocked by Bedrock guardrail: {bedrock_response.get('reason', 'Unknown reason')}" ) - 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 - # Apply any masking that was applied by the guardrail - masked_text = text - output_list = bedrock_response.get("output") - if output_list: - # If the guardrail returned modified content, use that - for output_item in output_list: + output_list = bedrock_response.get("output") + if output_list: + # If the guardrail returned modified content, use that + for output_item in output_list: + text_content = output_item.get("text") + if text_content: + masked_text = str(text_content) + masked_texts.append(masked_text) + else: + outputs_list = bedrock_response.get("outputs") + if outputs_list: + # Fallback to outputs field if output is not available + for output_item in outputs_list: text_content = output_item.get("text") if text_content: masked_text = str(text_content) - break - else: - outputs_list = bedrock_response.get("outputs") - if outputs_list: - # Fallback to outputs field if output is not available - for output_item in outputs_list: - text_content = output_item.get("text") - if text_content: - masked_text = str(text_content) - break - - masked_texts.append(masked_text) + masked_texts.append(masked_text) verbose_proxy_logger.debug( "Bedrock Guardrail: Successfully applied guardrail" ) - return masked_texts, images + inputs["texts"] = masked_texts + return inputs except Exception as e: verbose_proxy_logger.error( diff --git a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py index 49401ca7921..493d432eebb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/enkryptai/enkryptai.py @@ -15,7 +15,6 @@ from typing import ( List, Literal, Optional, - Tuple, Union, ) @@ -30,7 +29,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.guardrails import GenericGuardrailAPIInputs, GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import ( EnkryptAIProcessedResult, EnkryptAIResponse, @@ -481,27 +480,28 @@ class EnkryptAIGuardrails(CustomGuardrail): async def apply_guardrail( self, - texts: List[str], + inputs: "GenericGuardrailAPIInputs", request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, - images: Optional[List[str]] = None, - ) -> Tuple[List[str], Optional[List[str]]]: + ) -> "GenericGuardrailAPIInputs": """ Apply EnkryptAI guardrail to a batch of texts. Args: - texts: List of texts to check for attacks + inputs: Dictionary containing texts and optional images request_data: Request data dictionary containing metadata input_type: Whether this is a "request" or "response" - images: Optional list of images (not used by EnkryptAI) + logging_obj: Optional logging object Returns: - Tuple of (texts, images) - texts unchanged if passed, images unchanged + GenericGuardrailAPIInputs - texts unchanged if passed, images unchanged Raises: ValueError: If any attacks are detected """ + texts = inputs.get("texts", []) + # Check each text for attacks for text in texts: result = await self._call_enkryptai_guardrails( @@ -517,7 +517,7 @@ class EnkryptAIGuardrails(CustomGuardrail): error_message = self._create_error_message(processed_result) raise ValueError(error_message) - return texts, images + return inputs async def async_post_call_streaming_iterator_hook( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index f806e8b4802..33912b8fcd9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -14,7 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) -from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.guardrails import GenericGuardrailAPIInputs, GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import ( GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest, @@ -144,22 +144,24 @@ class GenericGuardrailAPI(CustomGuardrail): async def apply_guardrail( self, - texts: List[str], + inputs: GenericGuardrailAPIInputs, request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, - images: Optional[List[str]] = None, - ) -> Tuple[List[str], Optional[List[str]]]: + ) -> GenericGuardrailAPIInputs: """ - Apply the Generic Guardrail API to the given text. + Apply the Generic Guardrail API to the given inputs. This is the main method that gets called by the framework. Args: - texts: List of texts to check + inputs: Dictionary containing: + - texts: List of texts to check + - images: Optional list of images to check + - tool_calls: Optional list of tool calls to check request_data: Request data dictionary containing user_api_key_dict and other metadata input_type: Whether this is a "request" or "response" guardrail - images: Optional list of images to check + logging_obj: Optional logging object for tracking the guardrail execution Returns: Tuple of (processed texts, processed images) @@ -169,6 +171,11 @@ class GenericGuardrailAPI(CustomGuardrail): """ verbose_proxy_logger.debug("Generic Guardrail API: Applying guardrail to text") + # Extract texts and images from inputs + texts = inputs.get("texts", []) + images = inputs.get("images") + tools = inputs.get("tools") + # Use provided request_data or create an empty dict if request_data is None: request_data = {} @@ -193,6 +200,7 @@ class GenericGuardrailAPI(CustomGuardrail): texts=texts, request_data=user_metadata, images=images, + tools=tools, additional_provider_specific_params=additional_params, input_type=input_type, ) @@ -230,17 +238,19 @@ class GenericGuardrailAPI(CustomGuardrail): ) raise Exception(f"Content blocked by guardrail: {error_message}") - elif guardrail_response.action == "GUARDRAIL_INTERVENED": - # Content was modified by the guardrail - if guardrail_response.texts: - verbose_proxy_logger.debug("Generic Guardrail API modified text") - return guardrail_response.texts, guardrail_response.images - # Action is NONE or no modifications needed - return ( - guardrail_response.texts or texts, - guardrail_response.images or images, - ) + return_inputs = GenericGuardrailAPIInputs(texts=texts) + if guardrail_response.texts: + return_inputs["texts"] = guardrail_response.texts + if guardrail_response.images: + return_inputs["images"] = guardrail_response.images + elif images: + return_inputs["images"] = images + if guardrail_response.tools: + return_inputs["tools"] = guardrail_response.tools + elif tools: + return_inputs["tools"] = tools + return return_inputs except Exception as e: # Check if it's already an exception we raised diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index 776b9d36d5e..6756b188847 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -27,6 +27,7 @@ from litellm.integrations.custom_guardrail import CustomGuardrail if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.guardrails import GenericGuardrailAPIInputs from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import ( BlockedWord, @@ -304,12 +305,11 @@ class ContentFilterGuardrail(CustomGuardrail): async def apply_guardrail( self, - texts: List[str], + inputs: "GenericGuardrailAPIInputs", request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, - images: Optional[List[str]] = None, - ) -> Tuple[List[str], Optional[List[str]]]: + ) -> "GenericGuardrailAPIInputs": """ Apply content filtering guardrail to a batch of texts. @@ -317,17 +317,19 @@ class ContentFilterGuardrail(CustomGuardrail): either blocking the request or masking the sensitive content. Args: - texts: List of texts to apply the guardrail to + inputs: Dictionary containing texts and optional images request_data: Request data dictionary for logging metadata input_type: Whether this is a "request" or "response" - images: Optional list of images (not processed) + logging_obj: Optional logging object Returns: - Tuple of (processed_texts, images) - texts may be masked, images unchanged + GenericGuardrailAPIInputs - processed_texts may be masked, images unchanged Raises: HTTPException: If sensitive content is detected and action is BLOCK """ + texts = inputs.get("texts", []) + verbose_proxy_logger.debug( f"ContentFilterGuardrail: Applying guardrail to {len(texts)} text(s)" ) @@ -386,7 +388,8 @@ class ContentFilterGuardrail(CustomGuardrail): verbose_proxy_logger.debug( "ContentFilterGuardrail: Guardrail applied successfully" ) - return processed_texts, images + inputs["texts"] = processed_texts + return inputs async def async_post_call_streaming_iterator_hook( self, @@ -423,12 +426,17 @@ class ContentFilterGuardrail(CustomGuardrail): if isinstance(choice.delta.content, str): # Check the chunk content using apply_guardrail try: - processed_content = await self.apply_guardrail( - texts=[choice.delta.content], + guardrailed_inputs = await self.apply_guardrail( + inputs={"texts": [choice.delta.content]}, input_type="response", - images=None, request_data=request_data, ) + processed_texts = guardrailed_inputs.get("texts", []) + processed_content = ( + processed_texts[0] + if processed_texts + else choice.delta.content + ) if processed_content != choice.delta.content: choice.delta.content = processed_content verbose_proxy_logger.debug( diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 6b19685f07e..d183b688edd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -29,9 +29,11 @@ import aiohttp import litellm # noqa: E401 from litellm import get_secret from litellm._logging import verbose_proxy_logger +from litellm.types.guardrails import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm._uuid import uuid from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError @@ -713,17 +715,18 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): async def apply_guardrail( self, - texts: List[str], + inputs: "GenericGuardrailAPIInputs", request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, - images: Optional[List[str]] = None, - ) -> Tuple[List[str], Optional[List[str]]]: + ) -> "GenericGuardrailAPIInputs": """ UI will call this function to check: 1. If the connection to the guardrail is working 2. When Testing the guardrail with some text, this function will be called with the input text and returns a text after applying the guardrail """ + texts = inputs.get("texts", []) + new_texts = [] for text in texts: modified_text = await self.check_pii( @@ -733,7 +736,8 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): request_data=request_data or {}, ) new_texts.append(modified_text) - return new_texts, images + inputs["texts"] = new_texts + return inputs def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: """ 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 aeae19a8270..c16cb89b785 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -6,13 +6,14 @@ Unified Guardrail, leveraging LiteLLM's /applyGuardrail endpoint 3. Implements a way to call /applyGuardrail endpoint for `/chat/completions` + `/v1/messages` requests on async_post_call_streaming_iterator_hook """ -from typing import Any, AsyncGenerator, Union +from typing import Any, AsyncGenerator, List, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.cost_calculator import _infer_call_type from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks @@ -125,16 +126,21 @@ class UnifiedLLMGuardrails(CustomLogger): ) is not True ): - return verbose_proxy_logger.debug( "async_post_call_success_hook response: %s", response ) - call_type = _infer_call_type(call_type=None, completion_response=response) + 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: + call_type = call_types[0] if call_type is None: + call_type = _infer_call_type(call_type=None, completion_response=response) + if call_type is None: return response if endpoint_guardrail_translation_mappings is None: @@ -176,6 +182,111 @@ class UnifiedLLMGuardrails(CustomLogger): See Aim guardrail implementation for an example - https://github.com/BerriAI/litellm/blob/d0e022cfacb8e9ebc5409bb652059b6fd97b45c0/litellm/proxy/guardrails/guardrail_hooks/aim.py#L168 Triggered by mode: 'post_call' + + Supports sampling_rate parameter to control how often chunks are processed. + sampling_rate=1 means every chunk, sampling_rate=5 means every 5th chunk, etc. """ + + global endpoint_guardrail_translation_mappings + + guardrail_to_apply: CustomGuardrail = request_data.pop( + "guardrail_to_apply", None + ) + + + # Get sampling rate from guardrail config or optional_params, default to 5 + sampling_rate = 5 + if guardrail_to_apply is not None: + # Check guardrail config first + guardrail_config = getattr(guardrail_to_apply, "guardrail_config", {}) + sampling_rate = guardrail_config.get( + "streaming_sampling_rate", sampling_rate + ) + + # Also check optional_params as fallback + sampling_rate = self.optional_params.get( + "streaming_sampling_rate", sampling_rate + ) + + if guardrail_to_apply is None: + async for item in response: + yield item + return + + event_type: GuardrailEventHooks = GuardrailEventHooks.post_call + if ( + guardrail_to_apply.should_run_guardrail( + data=request_data, event_type=event_type + ) + is not True + ): + verbose_proxy_logger.debug( + "UnifiedLLMGuardrails: Post-call streaming scanning disabled for %s", + guardrail_to_apply.guardrail_name, + ) + async for item in response: + yield item + return + + # Initialize translation mappings if needed + if endpoint_guardrail_translation_mappings is None: + endpoint_guardrail_translation_mappings = ( + load_guardrail_translation_mappings() + ) + + # Infer call type from first chunk + call_type = None + chunk_counter = 0 + responses_so_far: List[Any] = [] + async for item in response: - yield item + chunk_counter += 1 + responses_so_far.append(item) + + # Infer call type from first chunk if not already done + if call_type is None and 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: + call_type = call_types[0] + + if call_type is None: + call_type = _infer_call_type(call_type=None, completion_response=item) + + # If call type not supported, just pass through all chunks + if ( + call_type is None + or CallTypes(call_type) not in endpoint_guardrail_translation_mappings + ): + yield item + async for remaining_item in response: + yield remaining_item + return + + # 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, + sampling_rate, + guardrail_to_apply.guardrail_name, + ) + + endpoint_translation = endpoint_guardrail_translation_mappings[ + CallTypes(call_type) + ]() + + processed_items = ( + await endpoint_translation.process_output_streaming_response( + responses_so_far=responses_so_far, + guardrail_to_apply=guardrail_to_apply, + litellm_logging_obj=request_data.get("litellm_logging_obj"), + user_api_key_dict=user_api_key_dict, + ) + ) + + last_item = processed_items[-1] + + yield last_item + else: + yield item diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index 562f0d40272..bebb87f8d21 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -4,7 +4,7 @@ # # +-------------------------------------------------------------+ import os -from typing import TYPE_CHECKING, List, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Literal, Optional from fastapi import HTTPException @@ -14,6 +14,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.types.guardrails import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -71,27 +72,27 @@ class ZscalerAIGuard(CustomGuardrail): async def apply_guardrail( self, - texts: List[str], + inputs: "GenericGuardrailAPIInputs", request_data: dict, input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, - images: Optional[List[str]] = None, - ) -> Tuple[List[str], Optional[List[str]]]: + ) -> "GenericGuardrailAPIInputs": """ Apply Zscaler AI Guard guardrail to batch of texts. Args: - texts: List of texts to check + inputs: Dictionary containing texts and optional images request_data: Request data dictionary containing metadata input_type: Whether this is a "request" or "response" - images: Optional list of images (not used by Zscaler) + logging_obj: Optional logging object Returns: - Tuple of (processed_texts, images) - texts unchanged if passed, images unchanged + GenericGuardrailAPIInputs - texts unchanged if passed, images unchanged Raises: Exception: If content is blocked by Zscaler AI Guard """ + texts = inputs.get("texts", []) try: verbose_proxy_logger.debug(f"ZscalerAIGuard: Checking {len(texts)} text(s)") @@ -143,7 +144,7 @@ class ZscalerAIGuard(CustomGuardrail): raise e verbose_proxy_logger.debug("ZscalerAIGuard: Successfully applied guardrail.") - return texts, images + return inputs def extract_blocking_info(self, response): """ diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 9436df585a3..2abba1d4976 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1343,19 +1343,19 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): "INSIDE parallel request limiter ASYNC SUCCESS LOGGING" ) - # Get metadata from kwargs - litellm_metadata = kwargs["litellm_params"].get( - get_metadata_variable_name_from_kwargs(kwargs), {} + # Get metadata from standard_logging_object - this correctly handles both + # 'metadata' and 'litellm_metadata' fields from litellm_params + standard_logging_object = kwargs.get("standard_logging_object") or {} + standard_logging_metadata = standard_logging_object.get("metadata") or {} + + # user_api_key_hash is the same as user_api_key (it's the hash) + user_api_key = standard_logging_metadata.get("user_api_key_hash") + user_api_key_user_id = standard_logging_metadata.get("user_api_key_user_id") + user_api_key_team_id = standard_logging_metadata.get("user_api_key_team_id") + user_api_key_organization_id = standard_logging_metadata.get( + "user_api_key_org_id" ) - if litellm_metadata is None: - return - user_api_key = litellm_metadata.get("user_api_key") - user_api_key_user_id = litellm_metadata.get("user_api_key_user_id") - user_api_key_team_id = litellm_metadata.get("user_api_key_team_id") - user_api_key_organization_id = litellm_metadata.get( - "user_api_key_organization_id" - ) - user_api_key_end_user_id = kwargs.get("user") or litellm_metadata.get( + user_api_key_end_user_id = kwargs.get("user") or standard_logging_metadata.get( "user_api_key_end_user_id" ) model_group = get_model_group_from_litellm_kwargs(kwargs) @@ -1501,10 +1501,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): litellm_parent_otel_span: Union[ Span, None ] = _get_parent_otel_span_from_kwargs(kwargs) - litellm_metadata = kwargs["litellm_params"]["metadata"] - user_api_key = ( - litellm_metadata.get("user_api_key") if litellm_metadata else None - ) + # Get metadata from standard_logging_object - this correctly handles both + # 'metadata' and 'litellm_metadata' fields from litellm_params + standard_logging_object = kwargs.get("standard_logging_object") or {} + standard_logging_metadata = standard_logging_object.get("metadata") or {} + user_api_key = standard_logging_metadata.get("user_api_key_hash") + pipeline_operations: List[RedisPipelineIncrementOperation] = [] if user_api_key: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0a7fc62a42a..9dc255bd79a 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -611,6 +611,8 @@ class LiteLLMProxyRequestSetup: data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr( user_api_key_dict, "end_user_max_budget", None ) + # Add the full UserAPIKeyAuth object for MCP server access control + data[_metadata_variable_name]["user_api_key_auth"] = user_api_key_dict return data @staticmethod diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 7d047ca2c29..7c93c8424ab 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -563,10 +563,15 @@ async def user_info( user_id = user_api_key_dict.user_id ## GET USER ROW ## + user_info = None if user_id is not None: user_info = await prisma_client.get_data(user_id=user_id) - else: - user_info = None + + if user_info is None: + raise HTTPException( + status_code=404, + detail=f"User {user_id} not found", + ) ## GET ALL TEAMS ## team_list = [] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 99789162318..e1d5a90dc79 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -31,7 +31,9 @@ from typing import ( from litellm._uuid import uuid from litellm.constants import ( AIOHTTP_CONNECTOR_LIMIT, + AIOHTTP_CONNECTOR_LIMIT_PER_HOST, AIOHTTP_KEEPALIVE_TIMEOUT, + AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, BASE_MCP_ROUTE, @@ -186,6 +188,7 @@ from litellm.proxy._experimental.mcp_server.tool_registry import ( global_mcp_tool_registry, ) from litellm.proxy._types import * +from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router from litellm.proxy.analytics_endpoints.analytics_endpoints import ( @@ -553,10 +556,11 @@ else: global_max_parallel_request_retry_timeout_env ) -ui_link = f"{server_root_path}/ui/" +ui_link = f"{server_root_path}/ui" +fallback_login_link = f"{server_root_path}/fallback/login" model_hub_link = f"{server_root_path}/ui/model_hub_table" ui_message = ( - f"👉 [```LiteLLM Admin Panel on /ui```]({ui_link}). Create, Edit Keys with SSO" + f"👉 [```LiteLLM Admin Panel on /ui```]({ui_link}). Create, Edit Keys with SSO. Having issues? Try [```Fallback Login```]({fallback_login_link})" ) ui_message += "\n\n💸 [```LiteLLM Model Cost Map```](https://models.litellm.ai/)." @@ -626,21 +630,26 @@ async def proxy_shutdown_event(): async def _initialize_shared_aiohttp_session(): - """Initialize shared aiohttp session for connection reuse.""" + """Initialize shared aiohttp session for connection reuse with connection limits.""" try: from aiohttp import ClientSession, TCPConnector - # Create connector with connection pooling settings optimized for long-lived connections - connector = TCPConnector( - limit=AIOHTTP_CONNECTOR_LIMIT, - keepalive_timeout=AIOHTTP_KEEPALIVE_TIMEOUT, - ttl_dns_cache=AIOHTTP_TTL_DNS_CACHE, - enable_cleanup_closed=True, - ) - + connector_kwargs = { + "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, + "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, + "enable_cleanup_closed": True, + } + if AIOHTTP_CONNECTOR_LIMIT > 0: + connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT + if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0: + connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST + + connector = TCPConnector(**connector_kwargs) session = ClientSession(connector=connector) + verbose_proxy_logger.info( - f"SESSION REUSE: Created shared aiohttp session for connection pooling (ID: {id(session)})" + f"SESSION REUSE: Created shared aiohttp session for connection pooling (ID: {id(session)}, " + f"limit={AIOHTTP_CONNECTOR_LIMIT}, limit_per_host={AIOHTTP_CONNECTOR_LIMIT_PER_HOST})" ) return session except Exception as e: @@ -8266,257 +8275,97 @@ async def fallback_login(request: Request): ) # hidden since this is a helper for UI sso login async def login(request: Request): # noqa: PLR0915 global premium_user, general_settings, master_key - from litellm.types.proxy.ui_sso import ReturnedUITokenObject + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.utils import get_custom_url - if master_key is None: - raise ProxyException( - message="Master Key not set for Proxy. Please set Master Key to use Admin UI. Set `LITELLM_MASTER_KEY` in .env or set general_settings:master_key in config.yaml. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.", - type=ProxyErrorTypes.auth_error, - param="master_key", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) form = await request.form() username = str(form.get("username")) password = str(form.get("password")) - ui_username = os.getenv("UI_USERNAME", "admin") - ui_password = os.getenv("UI_PASSWORD", None) - if ui_password is None: - ui_password = str(master_key) if master_key is not None else None - if ui_password is None: - raise ProxyException( - message="set Proxy master key to use UI. https://docs.litellm.ai/docs/proxy/virtual_keys. If set, use `--detailed_debug` to debug issue.", - type=ProxyErrorTypes.auth_error, - param="UI_PASSWORD", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - # check if we can find the `username` in the db. on the ui, users can enter username=their email - _user_row: Optional[LiteLLM_UserTable] = None - user_role: Optional[ - Literal[ - LitellmUserRoles.PROXY_ADMIN, - LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, - LitellmUserRoles.INTERNAL_USER, - LitellmUserRoles.INTERNAL_USER_VIEW_ONLY, - ] - ] = None - if prisma_client is not None: - _user_row = cast( - Optional[LiteLLM_UserTable], - await prisma_client.db.litellm_usertable.find_first( - where={"user_email": {"equals": username}} - ), - ) - disabled_non_admin_personal_key_creation = ( - get_disabled_non_admin_personal_key_creation() + # Authenticate user and get login result + login_result = await authenticate_user( + username=username, + password=password, + master_key=master_key, + prisma_client=prisma_client, ) - """ - To login to Admin UI, we support the following - - Login with UI_USERNAME and UI_PASSWORD - - Login with Invite Link `user_email` and `password` combination - """ - if secrets.compare_digest(username, ui_username) and secrets.compare_digest( - password, ui_password - ): - # Non SSO -> If user is using UI_USERNAME and UI_PASSWORD they are Proxy admin - user_role = LitellmUserRoles.PROXY_ADMIN - user_id = litellm_proxy_admin_name - # we want the key created to have PROXY_ADMIN_PERMISSIONS - key_user_id = litellm_proxy_admin_name - if ( - os.getenv("PROXY_ADMIN_ID", None) is not None - and os.environ["PROXY_ADMIN_ID"] == user_id - ) or user_id == litellm_proxy_admin_name: - # checks if user is admin - key_user_id = os.getenv("PROXY_ADMIN_ID", litellm_proxy_admin_name) + # Create UI token object + returned_ui_token_object = create_ui_token_object( + login_result=login_result, + general_settings=general_settings, + premium_user=premium_user, + ) - # Admin is Authe'd in - generate key for the UI to access Proxy + # Generate JWT token + import jwt - # ensure this user is set as the proxy admin, in this route there is no sso, we can assume this user is only the admin - await user_update( - data=UpdateUserRequest( - user_id=key_user_id, - user_role=user_role, - ), - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, - ), - ) - if os.getenv("DATABASE_URL") is not None: - response = await generate_key_helper_fn( - request_type="key", - **{ - "user_role": LitellmUserRoles.PROXY_ADMIN, - "duration": "24hr", - "key_max_budget": litellm.max_ui_session_budget, - "models": [], - "aliases": {}, - "config": {}, - "spend": 0, - "user_id": key_user_id, - "team_id": "litellm-dashboard", - }, # type: ignore - ) - else: - raise ProxyException( - message="No Database connected. Set DATABASE_URL in .env. If set, use `--detailed_debug` to debug issue.", - type=ProxyErrorTypes.auth_error, - param="DATABASE_URL", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - key = response["token"] # type: ignore - litellm_dashboard_ui = get_custom_url(str(request.base_url)) - if litellm_dashboard_ui.endswith("/"): - litellm_dashboard_ui += "ui/" - else: - litellm_dashboard_ui += "/ui/" - import jwt + jwt_token = jwt.encode( # type: ignore + cast(dict, returned_ui_token_object), + master_key, + algorithm="HS256", + ) - if get_secret_bool("EXPERIMENTAL_UI_LOGIN"): - user_info: Optional[LiteLLM_UserTable] = None - if _user_row is not None: - user_info = _user_row - elif ( - user_id is not None - ): # if user_id is not None, we are using the UI_USERNAME and UI_PASSWORD - user_info = LiteLLM_UserTable( - user_id=user_id, - user_role=user_role, - models=[], - max_budget=litellm.max_ui_session_budget, - ) - if user_info is None: - raise HTTPException( - status_code=401, - detail={ - "error": "User Information is required for experimental UI login" - }, - ) - - key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( - user_info - ) - - returned_ui_token_object = ReturnedUITokenObject( - user_id=user_id, - key=key, - user_email=None, - user_role=user_role, - login_method="username_password", - premium_user=premium_user, - auth_header_name=general_settings.get( - "litellm_key_header_name", "Authorization" - ), - disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, - server_root_path=get_server_root_path(), - ) - - jwt_token = jwt.encode( # type: ignore - cast(dict, returned_ui_token_object), - master_key, - algorithm="HS256", - ) - litellm_dashboard_ui += "?login=success" - redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token) - return redirect_response - elif _user_row is not None: - """ - When sharing invite links - - -> if the user has no role in the DB assume they are only a viewer - """ - user_id = getattr(_user_row, "user_id", "unknown") - user_role = getattr( - _user_row, "user_role", LitellmUserRoles.INTERNAL_USER_VIEW_ONLY - ) - user_email = getattr(_user_row, "user_email", "unknown") - _password = getattr(_user_row, "password", "unknown") - - if _password is None: - raise ProxyException( - message="User has no password set. Please set a password for the user via `/user/update`.", - type=ProxyErrorTypes.auth_error, - param="password", - code=status.HTTP_401_UNAUTHORIZED, - ) - - # check if password == _user_row.password - hash_password = hash_token(token=password) - if secrets.compare_digest(password, _password) or secrets.compare_digest( - hash_password, _password - ): - if os.getenv("DATABASE_URL") is not None: - response = await generate_key_helper_fn( - request_type="key", - **{ # type: ignore - "user_role": user_role, - "duration": "24hr", - "key_max_budget": litellm.max_ui_session_budget, - "models": [], - "aliases": {}, - "config": {}, - "spend": 0, - "user_id": user_id, - "team_id": "litellm-dashboard", - }, - ) - else: - raise ProxyException( - message="No Database connected. Set DATABASE_URL in .env. If set, use `--detailed_debug` to debug issue.", - type=ProxyErrorTypes.auth_error, - param="DATABASE_URL", - code=status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - key = response["token"] # type: ignore - litellm_dashboard_ui = get_custom_url(str(request.base_url)) - if litellm_dashboard_ui.endswith("/"): - litellm_dashboard_ui += "ui/" - else: - litellm_dashboard_ui += "/ui/" - import jwt - - returned_ui_token_object = ReturnedUITokenObject( - user_id=user_id, - key=key, - user_email=user_email, - user_role=cast(str, user_role), - login_method="username_password", - premium_user=premium_user, - auth_header_name=general_settings.get( - "litellm_key_header_name", "Authorization" - ), - disabled_non_admin_personal_key_creation=disabled_non_admin_personal_key_creation, - server_root_path=get_server_root_path(), - ) - - jwt_token = jwt.encode( # type: ignore - cast(dict, returned_ui_token_object), - master_key, - algorithm="HS256", - ) - litellm_dashboard_ui += "?login=success" - redirect_response = RedirectResponse( - url=litellm_dashboard_ui, status_code=303 - ) - redirect_response.set_cookie(key="token", value=jwt_token) - return redirect_response - else: - raise ProxyException( - message=f"Invalid credentials used to access UI.\nNot valid credentials for {username}", - type=ProxyErrorTypes.auth_error, - param="invalid_credentials", - code=status.HTTP_401_UNAUTHORIZED, - ) + # Build redirect URL + litellm_dashboard_ui = get_custom_url(str(request.base_url)) + if litellm_dashboard_ui.endswith("/"): + litellm_dashboard_ui += "ui/" else: - raise ProxyException( - message="Invalid credentials used to access UI.\nCheck 'UI_USERNAME', 'UI_PASSWORD' in .env file", - type=ProxyErrorTypes.auth_error, - param="invalid_credentials", - code=status.HTTP_401_UNAUTHORIZED, - ) + litellm_dashboard_ui += "/ui/" + litellm_dashboard_ui += "?login=success" + # Create redirect response with cookie + redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) + redirect_response.set_cookie(key="token", value=jwt_token) + return redirect_response + + +@router.post( + "/v2/login", include_in_schema=False +) # hidden helper for UI logins via API +async def login_v2(request: Request): # noqa: PLR0915 + global premium_user, general_settings, master_key + from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.utils import get_custom_url + + body = await request.json() + username = str(body.get("username")) + password = str(body.get("password")) + + login_result = await authenticate_user( + username=username, + password=password, + master_key=master_key, + prisma_client=prisma_client, + ) + + returned_ui_token_object = create_ui_token_object( + login_result=login_result, + general_settings=general_settings, + premium_user=premium_user, + ) + + import jwt + + jwt_token = jwt.encode( # type: ignore + cast(dict, returned_ui_token_object), + master_key, + algorithm="HS256", + ) + + litellm_dashboard_ui = get_custom_url(str(request.base_url)) + if litellm_dashboard_ui.endswith("/"): + litellm_dashboard_ui += "ui/" + else: + litellm_dashboard_ui += "/ui/" + litellm_dashboard_ui += "?login=success" + + json_response = JSONResponse( + content={"redirect_url": litellm_dashboard_ui}, + status_code=status.HTTP_200_OK, + ) + json_response.set_cookie(key="token", value=jwt_token) + return json_response @app.get("/onboarding/get_token", include_in_schema=False) async def onboarding(invite_link: str, request: Request): @@ -10221,6 +10070,7 @@ app.include_router(user_agent_analytics_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) app.include_router(agent_endpoints_router) +app.include_router(a2a_router) ######################################################## # MCP Server ######################################################## diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index 8e08bd8d991..a23775f122d 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2607,6 +2607,16 @@ "options": null, "default_value": null }, + { + "key": "api_base", + "label": "API Base (Optional)", + "placeholder": null, + "tooltip": "Supports custom API base for Google Private API endpoints.", + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, { "key": "vertex_credentials", "label": "Vertex Credentials", diff --git a/litellm/proxy/route_llm_request.py b/litellm/proxy/route_llm_request.py index 79e316d4702..221aa16f912 100644 --- a/litellm/proxy/route_llm_request.py +++ b/litellm/proxy/route_llm_request.py @@ -75,12 +75,13 @@ def add_shared_session_to_data(data: dict) -> None: """ Add shared aiohttp session for connection reuse (prevents cold starts). Silently continues without session reuse if import fails or session is unavailable. - + Args: data: Dictionary to add the shared session to """ try: from litellm.proxy.proxy_server import shared_aiohttp_session + if shared_aiohttp_session is not None and not shared_aiohttp_session.closed: data["shared_session"] = shared_aiohttp_session except Exception: @@ -136,13 +137,14 @@ async def route_request( "aget_skill", "adelete_skill", "aingest", + "anthropic_messages", ], ): """ Common helper to route the request """ add_shared_session_to_data(data) - + team_id = get_team_id_from_data(data) router_model_names = llm_router.model_names if llm_router is not None else [] @@ -177,7 +179,12 @@ async def route_request( return llm_router.abatch_completion(models=models, **data) elif llm_router is not None: # Skip model-based routing for container operations - if route_type in ["acreate_container", "alist_containers", "aretrieve_container", "adelete_container"]: + if route_type in [ + "acreate_container", + "alist_containers", + "aretrieve_container", + "adelete_container", + ]: return getattr(llm_router, f"{route_type}")(**data) if route_type in [ "avideo_list", @@ -196,7 +203,7 @@ async def route_request( ] and (data.get("model") is None or data.get("model") == ""): # These endpoints don't need a model, use custom_llm_provider directly return getattr(litellm, f"{route_type}")(**data) - + team_model_name = ( llm_router.map_team_model(data["model"], team_id) if team_id is not None @@ -206,9 +213,8 @@ async def route_request( data["model"] = team_model_name return getattr(llm_router, f"{route_type}")(**data) - elif ( - data["model"] in router_model_names - or llm_router.has_model_id(data["model"]) + elif data["model"] in router_model_names or llm_router.has_model_id( + data["model"] ): return getattr(llm_router, f"{route_type}")(**data) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 9594a55962c..28a7ef001d5 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1072,6 +1072,7 @@ class ProxyLogging: "user_budget", "soft_budget", "team_budget", + "organization_budget", "proxy_budget", "projected_limit_exceeded", ], @@ -1416,6 +1417,7 @@ class ProxyLogging: 3. /image/generation 4. /files """ + from litellm.types.guardrails import GuardrailEventHooks guardrail_callbacks: List[CustomGuardrail] = [] @@ -1450,6 +1452,7 @@ class ProxyLogging: continue guardrail_response: Optional[Any] = None + if "apply_guardrail" in type(callback).__dict__: data["guardrail_to_apply"] = callback guardrail_response = ( @@ -1559,7 +1562,9 @@ class ProxyLogging: Covers: 1. /chat/completions """ + 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( @@ -1573,11 +1578,22 @@ class ProxyLogging: ) or _callback.should_run_guardrail( data=request_data, event_type=GuardrailEventHooks.post_call ): - response = _callback.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=response, - request_data=request_data, - ) + + if "apply_guardrail" in type(callback).__dict__: + request_data["guardrail_to_apply"] = callback + response = ( + unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + request_data=request_data, + response=response, + ) + ) + else: + response = _callback.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=response, + request_data=request_data, + ) return response def _init_response_taking_too_long_task(self, data: Optional[dict] = None): diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 950ea7063f5..e837346df23 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -22,6 +22,9 @@ import litellm from litellm._logging import verbose_logger from litellm.constants import request_timeout from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + update_responses_input_with_model_file_ids, +) from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( @@ -38,9 +41,6 @@ from litellm.types.llms.openai import ( ToolChoice, ToolParam, ) -from litellm.litellm_core_utils.prompt_templates.common_utils import ( - update_responses_input_with_model_file_ids, -) # Handle ResponseText import with fallback if TYPE_CHECKING: @@ -168,7 +168,8 @@ async def aresponses_api_with_mcp( ) = LiteLLM_Proxy_MCP_Handler._parse_mcp_tools(tools) # Process MCP tools through the complete pipeline (fetch + filter + deduplicate + transform) - user_api_key_auth = kwargs.get("user_api_key_auth") + # Extract user_api_key_auth from litellm_metadata (where it's added by add_user_api_key_auth_to_request_metadata) + user_api_key_auth = kwargs.get("user_api_key_auth") or kwargs.get("litellm_metadata", {}).get("user_api_key_auth") # Get original MCP tools (for events) and OpenAI tools (for LLM) by reusing existing methods ( diff --git a/litellm/responses/utils.py b/litellm/responses/utils.py index 198182cf118..ad99609e905 100644 --- a/litellm/responses/utils.py +++ b/litellm/responses/utils.py @@ -24,7 +24,12 @@ from litellm.types.llms.openai import ( ResponseText, ) from litellm.types.responses.main import DecodedResponseId -from litellm.types.utils import PromptTokensDetails, SpecialEnums, Usage +from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + PromptTokensDetails, + SpecialEnums, + Usage, +) class ResponsesAPIRequestUtils: @@ -366,10 +371,11 @@ class ResponsesAPIRequestUtils: Headers from tools.headers in request body should be passed to MCP server. """ from starlette.datastructures import Headers + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, ) - + # Extract headers from secret_fields which contains the original request headers raw_headers_from_request: Optional[Dict[str, str]] = None if secret_fields and isinstance(secret_fields, dict): @@ -445,11 +451,20 @@ class ResponseAPILoggingUtils: cached_tokens=response_api_usage.input_tokens_details.cached_tokens, audio_tokens=response_api_usage.input_tokens_details.audio_tokens, ) + completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None + if response_api_usage.output_tokens_details: + completion_tokens_details = CompletionTokensDetailsWrapper( + reasoning_tokens=getattr( + response_api_usage.output_tokens_details, "reasoning_tokens", None + ) + ) + chat_usage = Usage( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, prompt_tokens_details=prompt_tokens_details, + completion_tokens_details=completion_tokens_details, ) # Preserve cost attribute if it exists on ResponseAPIUsage diff --git a/litellm/router.py b/litellm/router.py index a52e0260bad..9488f341cbf 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -153,11 +153,7 @@ from litellm.types.utils import ( ) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import ( - ModelResponseStream, - StandardLoggingPayload, - Usage, -) +from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -779,6 +775,9 @@ class Router: self.aanthropic_messages = self.factory_function( litellm.anthropic_messages, call_type="anthropic_messages" ) + self.anthropic_messages = self.factory_function( + litellm.anthropic_messages, call_type="anthropic_messages" + ) self.agenerate_content = self.factory_function( litellm.agenerate_content, call_type="agenerate_content" ) @@ -886,6 +885,7 @@ class Router: from litellm.vector_store_files.main import ( update as vector_store_file_update_fn, ) + self.avector_store_file_create = self.factory_function( avector_store_file_create_fn, call_type="avector_store_file_create" ) @@ -3865,6 +3865,7 @@ class Router: "retrieve_container", "delete_container", ): + def sync_wrapper( custom_llm_provider: Optional[str] = None, client: Optional[Any] = None, @@ -5944,36 +5945,38 @@ class Router: """ Get API credentials and provider info from a model name in model_list. Useful for passthrough endpoints (files, batches, etc.) that need credentials. - + This method tries to find a deployment by model_id first, and if not found, it tries to find by model_group_name (model_name). - + Args: model_id: Model ID or model name from model_list (e.g., "gpt-4o-litellm") - + Returns: Dictionary containing api_key, api_base, custom_llm_provider, etc. Returns None if model not found. - + Example: credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm") # Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", ...} """ # Try to get deployment by model_id first deployment = self.get_deployment(model_id=model_id) - + # If not found, try by model_group_name if deployment is None: - deployment = self.get_deployment_by_model_group_name(model_group_name=model_id) - + deployment = self.get_deployment_by_model_group_name( + model_group_name=model_id + ) + if deployment is None: return None - + # Get basic credentials credentials = CredentialLiteLLMParams( **deployment.litellm_params.model_dump(exclude_none=True) ).model_dump(exclude_none=True) - + # Add custom_llm_provider if deployment.litellm_params.custom_llm_provider: credentials["custom_llm_provider"] = ( @@ -5986,7 +5989,7 @@ class Router: )[0] else: credentials["custom_llm_provider"] = "openai" # default - + return credentials @overload diff --git a/litellm/secret_managers/cyberark_secret_manager.py b/litellm/secret_managers/cyberark_secret_manager.py index 5f1e8b5cd07..2745df00778 100644 --- a/litellm/secret_managers/cyberark_secret_manager.py +++ b/litellm/secret_managers/cyberark_secret_manager.py @@ -16,6 +16,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.proxy._types import KeyManagementSystem from .base_secret_manager import BaseSecretManager +from .main import str_to_bool class CyberArkSecretManager(BaseSecretManager): @@ -32,6 +33,11 @@ class CyberArkSecretManager(BaseSecretManager): self.tls_cert_path = os.getenv("CYBERARK_CLIENT_CERT", "") self.tls_key_path = os.getenv("CYBERARK_CLIENT_KEY", "") + # SSL verification - can be disabled for self-signed certificates + # Set CYBERARK_SSL_VERIFY=false to disable SSL verification + ssl_verify_env = str_to_bool(os.getenv("CYBERARK_SSL_VERIFY")) + self.ssl_verify: bool = ssl_verify_env if ssl_verify_env is not None else True + # Validate environment if not self.conjur_api_key and not ( self.tls_cert_path and self.tls_key_path @@ -52,6 +58,11 @@ class CyberArkSecretManager(BaseSecretManager): f"CyberArk secret manager is only available for premium users. {CommonProxyErrors.not_premium_user.value}" ) + if not self.ssl_verify: + verbose_logger.warning( + "CyberArk SSL verification is disabled. This is insecure and should only be used for testing with self-signed certificates." + ) + def _authenticate(self) -> str: """ Authenticate with CyberArk Conjur and get a session token. @@ -71,13 +82,16 @@ class CyberArkSecretManager(BaseSecretManager): try: if self.tls_cert_path and self.tls_key_path: - # Certificate-based authentication - http_client = httpx.Client(cert=(self.tls_cert_path, self.tls_key_path)) + # Certificate-based authentication - need custom client for cert + http_client = httpx.Client( + cert=(self.tls_cert_path, self.tls_key_path), + verify=self.ssl_verify, + ) resp = http_client.post(auth_url, content=self.conjur_api_key) else: # API key authentication - http_handler = _get_httpx_client() - resp = http_handler.post(auth_url, content=self.conjur_api_key) + http_handler = _get_httpx_client(params={"ssl_verify": self.ssl_verify}) + resp = http_handler.client.post(auth_url, content=self.conjur_api_key) resp.raise_for_status() @@ -117,8 +131,8 @@ class CyberArkSecretManager(BaseSecretManager): policy_yaml = f"- !variable {secret_name}\n" try: - client = _get_httpx_client() - resp = client.post( + client = _get_httpx_client(params={"ssl_verify": self.ssl_verify}) + resp = client.client.post( policy_url, headers={ **self._get_request_headers(), @@ -180,6 +194,7 @@ class CyberArkSecretManager(BaseSecretManager): async_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.SecretManager, + params={"ssl_verify": self.ssl_verify}, ) try: @@ -227,11 +242,11 @@ class CyberArkSecretManager(BaseSecretManager): if self.cache.get_cache(secret_name) is not None: return self.cache.get_cache(secret_name) - sync_client = _get_httpx_client() + sync_client = _get_httpx_client(params={"ssl_verify": self.ssl_verify}) try: url = self.get_url(secret_name) - response = sync_client.get(url, headers=self._get_request_headers()) + response = sync_client.client.get(url, headers=self._get_request_headers()) response.raise_for_status() # CyberArk Conjur returns the raw secret value as text @@ -278,7 +293,7 @@ class CyberArkSecretManager(BaseSecretManager): """ async_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.SecretManager, - params={"timeout": timeout}, + params={"ssl_verify": self.ssl_verify}, ) try: diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 850dab0ea7f..2eb26dc6227 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,9 +1,14 @@ from datetime import datetime -from typing import Any, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union -from pydantic import BaseModel +from pydantic import BaseModel, PrivateAttr from typing_extensions import Required, TypedDict +from litellm.types.llms.base import LiteLLMPydanticObjectBase + +if TYPE_CHECKING: + from a2a.types import SendMessageResponse + # AgentProvider class AgentProvider(TypedDict, total=False): @@ -200,3 +205,41 @@ class AgentMakePublicResponse(BaseModel): class MakeAgentsPublicRequest(BaseModel): agent_ids: List[str] + + +class LiteLLMSendMessageResponse(LiteLLMPydanticObjectBase): + """ + LiteLLM wrapper for A2A SendMessageResponse. + + Wraps the a2a SDK's SendMessageResponse with LiteLLM's _hidden_params + for cost tracking and logging integration. + """ + + # A2A response fields + id: str + jsonrpc: str = "2.0" + result: Optional[Dict[str, Any]] = None + error: Optional[Dict[str, Any]] = None + + model_config = {"extra": "allow"} + + # LiteLLM private attributes for logging/cost tracking + _hidden_params: dict = PrivateAttr(default_factory=dict) + + @classmethod + def from_a2a_response( + cls, response: "SendMessageResponse" + ) -> "LiteLLMSendMessageResponse": + """ + Create a LiteLLMSendMessageResponse from an a2a SDK SendMessageResponse. + + Args: + response: The a2a SDK SendMessageResponse + + Returns: + LiteLLMSendMessageResponse with _hidden_params support + """ + # Convert the a2a response to a dict + response_dict = response.model_dump(mode="json", exclude_none=True) + + return cls(**response_dict) diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 2c0b10ec67d..ea3a6cc4ede 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -5,6 +5,7 @@ from typing import Any, Dict, List, Literal, Optional, Union from pydantic import BaseModel, ConfigDict, Field from typing_extensions import Required, TypedDict +from litellm.types.llms.openai import ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import ( EnkryptAIGuardrailConfigs, ) @@ -738,3 +739,9 @@ class PatchGuardrailRequest(BaseModel): guardrail_name: Optional[str] = None litellm_params: Optional[BaseLitellmParams] = None guardrail_info: Optional[Dict[str, Any]] = None + + +class GenericGuardrailAPIInputs(TypedDict, total=False): + texts: List[str] + images: List[str] + tools: List[ChatCompletionToolParam] diff --git a/litellm/types/integrations/weave_otel.py b/litellm/types/integrations/weave_otel.py new file mode 100644 index 00000000000..5b40ff85340 --- /dev/null +++ b/litellm/types/integrations/weave_otel.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +from enum import Enum +from typing import Literal + +from pydantic import BaseModel + + +class WeaveOtelConfig(BaseModel): + """Configuration for Weave OpenTelemetry integration.""" + + otlp_auth_headers: str | None = None + endpoint: str | None = None + project_id: str | None = None + protocol: Literal["otlp_grpc", "otlp_http"] = "otlp_http" + + +class WeaveSpanAttributes(str, Enum): + """ + Weave-specific span attributes for OpenTelemetry traces. + + Based on Weave's OTEL attribute mappings from: + https://github.com/wandb/weave/blob/master/weave/trace_server/opentelemetry/constants.py + """ + + DISPLAY_NAME = "wandb.display_name" + + # Thread organization, similar to OpenInference session_id. + THREAD_ID = "wandb.thread_id" + IS_TURN = "wandb.is_turn" + diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 4d5ca02f032..9ed25005c05 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -24,6 +24,7 @@ class httpxSpecialProvider(str, Enum): Search = "search" MCP = "mcp" RAG = "rag" + A2A = "a2a" VerifyTypes = Union[str, bool, ssl.SSLContext] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 61d58e4c86d..9fb97d47d17 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -303,7 +303,7 @@ class OpenAIFileObject(BaseModel): `fine-tune`, `fine-tune-results`, `vision`, and `user_data`. """ - status: Literal["uploaded", "processed", "error"] + status: Optional[Literal["uploaded", "processed", "error"]] = None """Deprecated. The current status of the file, which can be either `uploaded`, `processed`, or diff --git a/litellm/types/llms/vertex_ai_text_to_speech.py b/litellm/types/llms/vertex_ai_text_to_speech.py new file mode 100644 index 00000000000..e65b75356bf --- /dev/null +++ b/litellm/types/llms/vertex_ai_text_to_speech.py @@ -0,0 +1,54 @@ +""" +Type definitions for Vertex AI Text-to-Speech API + +Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/synthesize +""" + +from typing import Optional + +from typing_extensions import TypedDict + + +class VertexTextToSpeechInput(TypedDict, total=False): + """ + Input for Vertex AI Text-to-Speech synthesis. + + Exactly one of text or ssml must be provided. + """ + text: Optional[str] + ssml: Optional[str] + + +class VertexTextToSpeechVoice(TypedDict, total=False): + """ + Voice configuration for Vertex AI Text-to-Speech. + + Attributes: + languageCode: The language code (e.g., "en-US", "de-DE") + name: The voice name (e.g., "en-US-Studio-O", "en-US-Wavenet-D") + """ + languageCode: str + name: str + + +class VertexTextToSpeechAudioConfig(TypedDict, total=False): + """ + Audio configuration for Vertex AI Text-to-Speech. + + Attributes: + audioEncoding: The audio encoding format (e.g., "LINEAR16", "MP3", "OGG_OPUS") + speakingRate: The speaking rate (0.25 to 4.0, default "1") + """ + audioEncoding: str + speakingRate: str + + +class VertexTextToSpeechRequest(TypedDict, total=False): + """ + Request body for Vertex AI Text-to-Speech API. + + Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/synthesize + """ + input: VertexTextToSpeechInput + voice: VertexTextToSpeechVoice + audioConfig: Optional[VertexTextToSpeechAudioConfig] diff --git a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py index 7132077b5fa..f100dd35fa6 100644 --- a/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py +++ b/litellm/types/proxy/discovery_endpoints/ui_discovery_endpoints.py @@ -6,3 +6,4 @@ from pydantic import BaseModel class UiDiscoveryEndpoints(BaseModel): server_root_path: str proxy_base_url: Optional[str] + auto_redirect_to_sso: bool diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index d2acc3e3403..66823270e87 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -3,6 +3,7 @@ from typing import Any, Dict, List, Literal, Optional from pydantic import BaseModel, Field from typing_extensions import TypedDict +from litellm.types.llms.openai import ChatCompletionToolParam from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -59,6 +60,7 @@ class GenericGuardrailAPIRequest: litellm_trace_id: Optional[str], additional_provider_specific_params: Optional[Dict[str, Any]] = None, images: Optional[List[str]] = None, + tools: Optional[List[ChatCompletionToolParam]] = None, ): self.texts = texts self.request_data = request_data @@ -69,12 +71,14 @@ class GenericGuardrailAPIRequest: self.input_type = input_type self.litellm_call_id = litellm_call_id self.litellm_trace_id = litellm_trace_id + self.tools = tools def to_dict(self) -> dict: return { "texts": self.texts, "request_data": self.request_data, "images": self.images, + "tools": self.tools, "additional_provider_specific_params": self.additional_provider_specific_params, "input_type": self.input_type, "litellm_call_id": self.litellm_call_id, @@ -87,6 +91,7 @@ class GenericGuardrailAPIResponse: texts: Optional[List[str]] images: Optional[List[str]] + tools: Optional[List[ChatCompletionToolParam]] action: str blocked_reason: Optional[str] @@ -96,11 +101,13 @@ class GenericGuardrailAPIResponse: texts: Optional[List[str]] = None, blocked_reason: Optional[str] = None, images: Optional[List[str]] = None, + tools: Optional[List[ChatCompletionToolParam]] = None, ): self.action = action self.blocked_reason = blocked_reason self.texts = texts self.images = images + self.tools = tools @classmethod def from_dict(cls, data: dict) -> "GenericGuardrailAPIResponse": diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 58267fdfea9..4861510da24 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -340,11 +340,23 @@ class CallTypes(str, Enum): generate_content_stream = "generate_content_stream" agenerate_content_stream = "agenerate_content_stream" + ######################################################### + # OCR Call Types + ######################################################### + ocr = "ocr" + aocr = "aocr" + ######################################################### # MCP Call Types ######################################################### call_mcp_tool = "call_mcp_tool" + ######################################################### + # A2A Call Types + ######################################################### + asend_message = "asend_message" + send_message = "send_message" + CallTypesLiteral = Literal[ "embedding", @@ -397,10 +409,337 @@ CallTypesLiteral = Literal[ "vector_store_file_delete", "avector_store_file_delete", "call_mcp_tool", + "asend_message", + "send_message", "aresponses", "responses", ] +# Mapping of API routes to their corresponding call types +API_ROUTE_TO_CALL_TYPES = { + # Chat Completions + "/chat/completions": [CallTypes.acompletion, CallTypes.completion], + "/v1/chat/completions": [CallTypes.acompletion, CallTypes.completion], + "/engines/{model}/chat/completions": [CallTypes.acompletion, CallTypes.completion], + "/openai/deployments/{model}/chat/completions": [ + CallTypes.acompletion, + CallTypes.completion, + ], + # Text Completions + "/completions": [CallTypes.atext_completion, CallTypes.text_completion], + "/v1/completions": [CallTypes.atext_completion, CallTypes.text_completion], + "/engines/{model}/completions": [ + CallTypes.atext_completion, + CallTypes.text_completion, + ], + "/openai/deployments/{model}/completions": [ + CallTypes.atext_completion, + CallTypes.text_completion, + ], + # Embeddings + "/embeddings": [CallTypes.aembedding, CallTypes.embedding], + "/v1/embeddings": [CallTypes.aembedding, CallTypes.embedding], + "/engines/{model}/embeddings": [CallTypes.aembedding, CallTypes.embedding], + "/openai/deployments/{model}/embeddings": [ + CallTypes.aembedding, + CallTypes.embedding, + ], + # Image Generation + "/images/generations": [CallTypes.aimage_generation, CallTypes.image_generation], + "/v1/images/generations": [CallTypes.aimage_generation, CallTypes.image_generation], + "/engines/{model}/images/generations": [ + CallTypes.aimage_generation, + CallTypes.image_generation, + ], + "/openai/deployments/{model}/images/generations": [ + CallTypes.aimage_generation, + CallTypes.image_generation, + ], + # Image Edits + "/images/edits": [CallTypes.aimage_edit, CallTypes.image_edit], + "/v1/images/edits": [CallTypes.aimage_edit, CallTypes.image_edit], + # Audio Transcriptions + "/audio/transcriptions": [CallTypes.atranscription, CallTypes.transcription], + "/v1/audio/transcriptions": [CallTypes.atranscription, CallTypes.transcription], + # Audio Speech + "/audio/speech": [CallTypes.aspeech, CallTypes.speech], + "/v1/audio/speech": [CallTypes.aspeech, CallTypes.speech], + # Moderations + "/moderations": [CallTypes.amoderation, CallTypes.moderation], + "/v1/moderations": [CallTypes.amoderation, CallTypes.moderation], + # Rerank + "/rerank": [CallTypes.arerank, CallTypes.rerank], + "/v1/rerank": [CallTypes.arerank, CallTypes.rerank], + "/v2/rerank": [CallTypes.arerank, CallTypes.rerank], + # Search + "/search": [CallTypes.asearch, CallTypes.search], + "/v1/search": [CallTypes.asearch, CallTypes.search], + # Batches + "/batches": [CallTypes.acreate_batch, CallTypes.create_batch], + "/v1/batches": [CallTypes.acreate_batch, CallTypes.create_batch], + "/batches/{batch_id}": [CallTypes.aretrieve_batch, CallTypes.retrieve_batch], + "/v1/batches/{batch_id}": [CallTypes.aretrieve_batch, CallTypes.retrieve_batch], + # Files + "/files": [ + CallTypes.acreate_file, + CallTypes.create_file, + CallTypes.afile_list, + CallTypes.file_list, + ], + "/v1/files": [ + CallTypes.acreate_file, + CallTypes.create_file, + CallTypes.afile_list, + CallTypes.file_list, + ], + "/files/{file_id}": [ + CallTypes.afile_retrieve, + CallTypes.file_retrieve, + CallTypes.afile_delete, + CallTypes.file_delete, + ], + "/v1/files/{file_id}": [ + CallTypes.afile_retrieve, + CallTypes.file_retrieve, + CallTypes.afile_delete, + CallTypes.file_delete, + ], + "/files/{file_id}/content": [CallTypes.afile_content, CallTypes.file_content], + "/v1/files/{file_id}/content": [CallTypes.afile_content, CallTypes.file_content], + # Assistants + "/assistants": [ + CallTypes.aget_assistants, + CallTypes.get_assistants, + CallTypes.acreate_assistants, + CallTypes.create_assistants, + ], + "/v1/assistants": [ + CallTypes.aget_assistants, + CallTypes.get_assistants, + CallTypes.acreate_assistants, + CallTypes.create_assistants, + ], + "/assistants/{assistant_id}": [ + CallTypes.adelete_assistant, + CallTypes.delete_assistant, + ], + "/v1/assistants/{assistant_id}": [ + CallTypes.adelete_assistant, + CallTypes.delete_assistant, + ], + # Threads + "/threads": [CallTypes.acreate_thread, CallTypes.create_thread], + "/v1/threads": [CallTypes.acreate_thread, CallTypes.create_thread], + "/threads/{thread_id}": [CallTypes.aget_thread, CallTypes.get_thread], + "/v1/threads/{thread_id}": [CallTypes.aget_thread, CallTypes.get_thread], + # Thread Messages + "/threads/{thread_id}/messages": [ + CallTypes.a_add_message, + CallTypes.add_message, + CallTypes.aget_messages, + CallTypes.get_messages, + ], + "/v1/threads/{thread_id}/messages": [ + CallTypes.a_add_message, + CallTypes.add_message, + CallTypes.aget_messages, + CallTypes.get_messages, + ], + # Thread Runs + "/threads/{thread_id}/runs": [ + CallTypes.arun_thread, + CallTypes.run_thread, + CallTypes.arun_thread_stream, + CallTypes.run_thread_stream, + ], + "/v1/threads/{thread_id}/runs": [ + CallTypes.arun_thread, + CallTypes.run_thread, + CallTypes.arun_thread_stream, + CallTypes.run_thread_stream, + ], + # Fine-tuning Jobs + "/fine_tuning/jobs": [ + CallTypes.acreate_fine_tuning_job, + CallTypes.create_fine_tuning_job, + CallTypes.alist_fine_tuning_jobs, + CallTypes.list_fine_tuning_jobs, + ], + "/v1/fine_tuning/jobs": [ + CallTypes.acreate_fine_tuning_job, + CallTypes.create_fine_tuning_job, + CallTypes.alist_fine_tuning_jobs, + CallTypes.list_fine_tuning_jobs, + ], + "/fine_tuning/jobs/{fine_tuning_job_id}": [ + CallTypes.aretrieve_fine_tuning_job, + CallTypes.retrieve_fine_tuning_job, + ], + "/v1/fine_tuning/jobs/{fine_tuning_job_id}": [ + CallTypes.aretrieve_fine_tuning_job, + CallTypes.retrieve_fine_tuning_job, + ], + "/fine_tuning/jobs/{fine_tuning_job_id}/cancel": [ + CallTypes.acancel_fine_tuning_job, + CallTypes.cancel_fine_tuning_job, + ], + "/v1/fine_tuning/jobs/{fine_tuning_job_id}/cancel": [ + CallTypes.acancel_fine_tuning_job, + CallTypes.cancel_fine_tuning_job, + ], + # Video Generation + "/videos": [ + CallTypes.acreate_video, + CallTypes.create_video, + CallTypes.avideo_list, + CallTypes.video_list, + ], + "/v1/videos": [ + CallTypes.acreate_video, + CallTypes.create_video, + CallTypes.avideo_list, + CallTypes.video_list, + ], + "/videos/{video_id}": [ + CallTypes.avideo_retrieve, + CallTypes.video_retrieve, + CallTypes.avideo_delete, + CallTypes.video_delete, + ], + "/v1/videos/{video_id}": [ + CallTypes.avideo_retrieve, + CallTypes.video_retrieve, + CallTypes.avideo_delete, + CallTypes.video_delete, + ], + "/videos/{video_id}/content": [CallTypes.avideo_content, CallTypes.video_content], + "/v1/videos/{video_id}/content": [ + CallTypes.avideo_content, + CallTypes.video_content, + ], + "/videos/{video_id}/remix": [CallTypes.avideo_remix, CallTypes.video_remix], + "/v1/videos/{video_id}/remix": [CallTypes.avideo_remix, CallTypes.video_remix], + # Vector Stores + "/vector_stores": [CallTypes.avector_store_create, CallTypes.vector_store_create], + "/v1/vector_stores": [ + CallTypes.avector_store_create, + CallTypes.vector_store_create, + ], + "/vector_stores/{vector_store_id}/search": [ + CallTypes.avector_store_search, + CallTypes.vector_store_search, + ], + "/v1/vector_stores/{vector_store_id}/search": [ + CallTypes.avector_store_search, + CallTypes.vector_store_search, + ], + "/vector_stores/{vector_store_id}/files": [ + CallTypes.avector_store_file_create, + CallTypes.vector_store_file_create, + CallTypes.avector_store_file_list, + CallTypes.vector_store_file_list, + ], + "/v1/vector_stores/{vector_store_id}/files": [ + CallTypes.avector_store_file_create, + CallTypes.vector_store_file_create, + CallTypes.avector_store_file_list, + CallTypes.vector_store_file_list, + ], + "/vector_stores/{vector_store_id}/files/{file_id}": [ + CallTypes.avector_store_file_retrieve, + CallTypes.vector_store_file_retrieve, + CallTypes.avector_store_file_delete, + CallTypes.vector_store_file_delete, + ], + "/v1/vector_stores/{vector_store_id}/files/{file_id}": [ + CallTypes.avector_store_file_retrieve, + CallTypes.vector_store_file_retrieve, + CallTypes.avector_store_file_delete, + CallTypes.vector_store_file_delete, + ], + "/vector_stores/{vector_store_id}/files/{file_id}/content": [ + CallTypes.avector_store_file_content, + CallTypes.vector_store_file_content, + ], + "/v1/vector_stores/{vector_store_id}/files/{file_id}/content": [ + CallTypes.avector_store_file_content, + CallTypes.vector_store_file_content, + ], + "/vector_stores/{vector_store_id}/files/{file_id}/update": [ + CallTypes.avector_store_file_update, + CallTypes.vector_store_file_update, + ], + "/v1/vector_stores/{vector_store_id}/files/{file_id}/update": [ + CallTypes.avector_store_file_update, + CallTypes.vector_store_file_update, + ], + # Containers + "/containers": [ + CallTypes.acreate_container, + CallTypes.create_container, + CallTypes.alist_containers, + CallTypes.list_containers, + ], + "/v1/containers": [ + CallTypes.acreate_container, + CallTypes.create_container, + CallTypes.alist_containers, + CallTypes.list_containers, + ], + "/containers/{container_id}": [ + CallTypes.aretrieve_container, + CallTypes.retrieve_container, + CallTypes.adelete_container, + CallTypes.delete_container, + ], + "/v1/containers/{container_id}": [ + CallTypes.aretrieve_container, + CallTypes.retrieve_container, + CallTypes.adelete_container, + CallTypes.delete_container, + ], + # Responses API + "/responses": [CallTypes.aresponses, CallTypes.responses], + "/v1/responses": [CallTypes.aresponses, CallTypes.responses], + "/responses/{response_id}": [CallTypes.aresponses, CallTypes.responses], + "/v1/responses/{response_id}": [CallTypes.aresponses, CallTypes.responses], + "/responses/{response_id}/input_items": [CallTypes.alist_input_items], + "/v1/responses/{response_id}/input_items": [CallTypes.alist_input_items], + # Realtime API + "/realtime": [CallTypes.arealtime], + "/v1/realtime": [CallTypes.arealtime], + # Provider-specific routes + "/anthropic/v1/messages": [CallTypes.anthropic_messages], + # Google GenAI routes + "/generate_content": [CallTypes.agenerate_content, CallTypes.generate_content], + "/models/{model}:generateContent": [ + CallTypes.agenerate_content, + CallTypes.generate_content, + ], + "/generate_content_stream": [ + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream, + ], + "/models/{model}:streamGenerateContent": [ + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream, + ], + # MCP (Model Context Protocol) + "/mcp/call_tool": [CallTypes.call_mcp_tool], + # A2A (Agent-to-Agent) + "/a2a/{agent_id}": [CallTypes.asend_message, CallTypes.send_message], + # Passthrough endpoints + "/llm_passthrough": [ + CallTypes.llm_passthrough_route, + CallTypes.allm_passthrough_route, + ], + "/v1/llm_passthrough": [ + CallTypes.llm_passthrough_route, + CallTypes.allm_passthrough_route, + ], + "/v1/messages": [CallTypes.anthropic_messages], +} + class PassthroughCallTypes(Enum): passthrough_image_generation = "passthrough-image-generation" @@ -1060,7 +1399,10 @@ class Usage(CompletionUsage): # Auto-calculate text_tokens only if provider didn't set it explicitly # Formula: text_tokens = completion_tokens - reasoning_tokens - image_tokens - audio_tokens - if _completion_tokens_details.text_tokens is None and completion_tokens is not None: + if ( + _completion_tokens_details.text_tokens is None + and completion_tokens is not None + ): calculated_text_tokens = completion_tokens - reasoning_tokens # Subtract other modality tokens if present @@ -2340,6 +2682,10 @@ class StandardCallbackDynamicParams(TypedDict, total=False): posthog_api_key: Optional[str] posthog_api_url: Optional[str] + # Weave (W&B) dynamic params + wandb_api_key: Optional[str] + weave_project_id: Optional[str] + # Logging settings turn_off_message_logging: Optional[bool] # when true will not log messages litellm_disabled_callbacks: Optional[List[str]] @@ -2618,6 +2964,7 @@ class LlmProviders(str, Enum): DATABRICKS = "databricks" EMPOWER = "empower" GITHUB = "github" + RAGFLOW = "ragflow" COMPACTIFAI = "compactifai" DOCKER_MODEL_RUNNER = "docker_model_runner" CUSTOM = "custom" @@ -2655,11 +3002,18 @@ class LlmProviders(str, Enum): WANDB = "wandb" OVHCLOUD = "ovhcloud" LEMONADE = "lemonade" + A2A_AGENT = "a2a_agent" # Create a set of all provider values for quick lookup LlmProvidersSet = {provider.value for provider in LlmProviders} +# File and Batch API providers that are OpenAI-compatible +OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS: set[str] = { + LlmProviders.OPENAI.value, + LlmProviders.HOSTED_VLLM.value, +} + class SearchProviders(str, Enum): """ diff --git a/litellm/types/vector_stores.py b/litellm/types/vector_stores.py index 6ae0b4bd2fd..a4ceb2c9ac7 100644 --- a/litellm/types/vector_stores.py +++ b/litellm/types/vector_stores.py @@ -3,17 +3,15 @@ from datetime import datetime from enum import Enum from typing import Any, Dict, List, Literal, Optional, Tuple, Union -from annotated_types import Ge from pydantic import BaseModel from typing_extensions import TypedDict -from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams - class SupportedVectorStoreIntegrations(str, Enum): """Supported vector store integrations.""" BEDROCK = "bedrock" + RAGFLOW = "ragflow" class LiteLLM_VectorStoreConfig(TypedDict, total=False): diff --git a/litellm/utils.py b/litellm/utils.py index 6c50afc5f49..b77c0e62e7d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7110,6 +7110,8 @@ class ProviderConfigManager: return litellm.CompactifAIChatConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: return litellm.GithubCopilotConfig() + elif litellm.LlmProviders.RAGFLOW == provider: + return litellm.RAGFlowConfig() elif ( litellm.LlmProviders.CUSTOM == provider or litellm.LlmProviders.CUSTOM_OPENAI == provider @@ -7631,6 +7633,12 @@ class ProviderConfigManager: ) return GeminiVectorStoreConfig() + elif litellm.LlmProviders.RAGFLOW == provider: + from litellm.llms.ragflow.vector_stores.transformation import ( + RAGFlowVectorStoreConfig, + ) + + return RAGFlowVectorStoreConfig() return None @staticmethod @@ -7910,6 +7918,12 @@ class ProviderConfigManager: ) return RunwayMLTextToSpeechConfig() + elif litellm.LlmProviders.VERTEX_AI == provider: + from litellm.llms.vertex_ai.text_to_speech.transformation import ( + VertexAITextToSpeechConfig, + ) + + return VertexAITextToSpeechConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f28e9b1290f..d02a01e3a67 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -269,6 +269,71 @@ "supports_response_schema": true, "supports_vision": true }, + "amazon.nova-2-lite-v1:0": { + "input_cost_per_token": 3e-07, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_video_input": true, + "supports_vision": true + }, + "apac.amazon.nova-2-lite-v1:0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.75e-06, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_video_input": true, + "supports_vision": true + }, + "eu.amazon.nova-2-lite-v1:0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.75e-06, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_video_input": true, + "supports_vision": true + }, + "us.amazon.nova-2-lite-v1:0": { + "input_cost_per_token": 6e-08, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.75e-06, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_video_input": true, + "supports_vision": true + }, + "amazon.nova-micro-v1:0": { "input_cost_per_token": 3.5e-08, "litellm_provider": "bedrock_converse", @@ -360,6 +425,15 @@ "litellm_provider": "bedrock", "mode": "image_generation" }, + "amazon.titan-image-generator-v2:0": { + "input_cost_per_image": 0.0, + "output_cost_per_image": 0.008, + "output_cost_per_image_premium_image": 0.01, + "output_cost_per_image_above_1024_and_1024_pixels": 0.01, + "output_cost_per_image_above_1024_and_1024_pixels_and_premium_image": 0.012, + "litellm_provider": "bedrock", + "mode": "image_generation" + }, "twelvelabs.marengo-embed-2-7-v1:0": { "input_cost_per_token": 7e-05, "litellm_provider": "bedrock", @@ -9564,6 +9638,21 @@ "supports_prompt_caching": true, "supports_tool_choice": true }, + "deepseek/deepseek-v3.2": { + "input_cost_per_token": 2.8e-07, + "input_cost_per_token_cache_hit": 2.8e-08, + "litellm_provider": "deepseek", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "deepseek.v3-v1:0": { "input_cost_per_token": 5.8e-07, "litellm_provider": "bedrock_converse", @@ -10421,6 +10510,19 @@ "supports_response_schema": true, "supports_tool_choice": true }, + "fireworks_ai/accounts/fireworks/models/deepseek-v3p2": { + "input_cost_per_token": 1.2e-06, + "litellm_provider": "fireworks_ai", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 163840, + "mode": "chat", + "output_cost_per_token": 1.2e-06, + "source": "https://fireworks.ai/models/fireworks/deepseek-v3p2", + "supports_function_calling": true, + "supports_response_schema": true, + "supports_tool_choice": true + }, "fireworks_ai/accounts/fireworks/models/firefunction-v2": { "input_cost_per_token": 9e-07, "litellm_provider": "fireworks_ai", @@ -10759,25 +10861,25 @@ "supports_tool_choice": true }, "ft:babbage-002": { - "input_cost_per_token": 4e-07, + "input_cost_per_token": 1.6e-06, "input_cost_per_token_batches": 2e-07, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 16384, "mode": "completion", - "output_cost_per_token": 4e-07, + "output_cost_per_token": 1.6e-06, "output_cost_per_token_batches": 2e-07 }, "ft:davinci-002": { - "input_cost_per_token": 2e-06, + "input_cost_per_token": 1.2e-05, "input_cost_per_token_batches": 1e-06, "litellm_provider": "text-completion-openai", "max_input_tokens": 16384, "max_output_tokens": 4096, "max_tokens": 16384, "mode": "completion", - "output_cost_per_token": 2e-06, + "output_cost_per_token": 1.2e-05, "output_cost_per_token_batches": 1e-06 }, "ft:gpt-3.5-turbo": { @@ -10840,6 +10942,7 @@ "supports_tool_choice": true }, "ft:gpt-4o-2024-08-06": { + "cache_read_input_token_cost": 1.875e-06, "input_cost_per_token": 3.75e-06, "input_cost_per_token_batches": 1.875e-06, "litellm_provider": "openai", @@ -10852,6 +10955,7 @@ "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, + "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, @@ -10872,8 +10976,7 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true + "supports_tool_choice": true }, "ft:gpt-4o-mini-2024-07-18": { "cache_read_input_token_cost": 1.5e-07, @@ -10892,8 +10995,79 @@ "supports_prompt_caching": true, "supports_response_schema": true, "supports_system_messages": true, - "supports_tool_choice": true, - "supports_vision": true + "supports_tool_choice": true + }, + "ft:gpt-4.1-2025-04-14": { + "cache_read_input_token_cost": 7.5e-07, + "input_cost_per_token": 3e-06, + "input_cost_per_token_batches": 1.5e-06, + "litellm_provider": "openai", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 1.2e-05, + "output_cost_per_token_batches": 6e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "ft:gpt-4.1-mini-2025-04-14": { + "cache_read_input_token_cost": 2e-07, + "input_cost_per_token": 8e-07, + "input_cost_per_token_batches": 4e-07, + "litellm_provider": "openai", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 3.2e-06, + "output_cost_per_token_batches": 1.6e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "ft:gpt-4.1-nano-2025-04-14": { + "cache_read_input_token_cost": 5e-08, + "input_cost_per_token": 2e-07, + "input_cost_per_token_batches": 1e-07, + "litellm_provider": "openai", + "max_input_tokens": 1047576, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "chat", + "output_cost_per_token": 8e-07, + "output_cost_per_token_batches": 4e-07, + "supports_function_calling": true, + "supports_parallel_function_calling": true, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true + }, + "ft:o4-mini-2025-04-16": { + "cache_read_input_token_cost": 1e-06, + "input_cost_per_token": 4e-06, + "input_cost_per_token_batches": 2e-06, + "litellm_provider": "openai", + "max_input_tokens": 200000, + "max_output_tokens": 100000, + "max_tokens": 100000, + "mode": "chat", + "output_cost_per_token": 1.6e-05, + "output_cost_per_token_batches": 8e-06, + "supports_function_calling": true, + "supports_parallel_function_calling": false, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true }, "gemini-1.0-pro": { "input_cost_per_character": 1.25e-07, @@ -16716,7 +16890,7 @@ "output_cost_per_token": 9.9e-07, "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/distil-whisper-large-v3-en": { @@ -16735,7 +16909,7 @@ "mode": "chat", "output_cost_per_token": 7e-08, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/gemma2-9b-it": { @@ -16747,7 +16921,7 @@ "mode": "chat", "output_cost_per_token": 2e-07, "supports_function_calling": false, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": false }, "groq/llama-3.1-405b-reasoning": { @@ -16759,7 +16933,7 @@ "mode": "chat", "output_cost_per_token": 7.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.1-70b-versatile": { @@ -16772,7 +16946,7 @@ "mode": "chat", "output_cost_per_token": 7.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.1-8b-instant": { @@ -16784,7 +16958,7 @@ "mode": "chat", "output_cost_per_token": 8e-08, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-11b-text-preview": { @@ -16797,7 +16971,7 @@ "mode": "chat", "output_cost_per_token": 1.8e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-11b-vision-preview": { @@ -16810,7 +16984,7 @@ "mode": "chat", "output_cost_per_token": 1.8e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true, "supports_vision": true }, @@ -16824,7 +16998,7 @@ "mode": "chat", "output_cost_per_token": 4e-08, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-3b-preview": { @@ -16837,7 +17011,7 @@ "mode": "chat", "output_cost_per_token": 6e-08, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-90b-text-preview": { @@ -16850,7 +17024,7 @@ "mode": "chat", "output_cost_per_token": 9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-3.2-90b-vision-preview": { @@ -16863,7 +17037,7 @@ "mode": "chat", "output_cost_per_token": 9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true, "supports_vision": true }, @@ -16887,7 +17061,7 @@ "mode": "chat", "output_cost_per_token": 7.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama-guard-3-8b": { @@ -16908,7 +17082,7 @@ "mode": "chat", "output_cost_per_token": 8e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama3-groq-70b-8192-tool-use-preview": { @@ -16921,7 +17095,7 @@ "mode": "chat", "output_cost_per_token": 8.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/llama3-groq-8b-8192-tool-use-preview": { @@ -16934,7 +17108,7 @@ "mode": "chat", "output_cost_per_token": 1.9e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/meta-llama/llama-4-maverick-17b-128e-instruct": { @@ -16980,7 +17154,7 @@ "mode": "chat", "output_cost_per_token": 2.4e-07, "supports_function_calling": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/moonshotai/kimi-k2-instruct": { @@ -17056,7 +17230,7 @@ "output_cost_per_token": 5.9e-07, "supports_function_calling": true, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_tool_choice": true }, "groq/whisper-large-v3": { @@ -18731,6 +18905,34 @@ "supports_tool_choice": true, "supports_web_search": true }, + "moonshot/kimi-k2-0905-preview": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 6e-07, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 2.5e-06, + "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_web_search": true + }, + "moonshot/kimi-k2-turbo-preview": { + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 1.15e-06, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8e-06, + "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_web_search": true + }, "moonshot/kimi-latest": { "cache_read_input_token_cost": 1.5e-07, "input_cost_per_token": 2e-06, @@ -18788,14 +18990,15 @@ "supports_vision": true }, "moonshot/kimi-thinking-preview": { - "input_cost_per_token": 3e-05, + "cache_read_input_token_cost": 1.5e-07, + "input_cost_per_token": 6e-07, "litellm_provider": "moonshot", "max_input_tokens": 131072, "max_output_tokens": 131072, "max_tokens": 131072, "mode": "chat", - "output_cost_per_token": 3e-05, - "source": "https://platform.moonshot.ai/docs/pricing", + "output_cost_per_token": 2.5e-06, + "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", "supports_vision": true }, "moonshot/kimi-k2-thinking": { @@ -18812,6 +19015,20 @@ "supports_tool_choice": true, "supports_web_search": true }, + "moonshot/kimi-k2-thinking-turbo": { + "cache_read_input_token_cost": 1.5e-7, + "input_cost_per_token": 1.15e-6, + "litellm_provider": "moonshot", + "max_input_tokens": 262144, + "max_output_tokens": 262144, + "max_tokens": 262144, + "mode": "chat", + "output_cost_per_token": 8e-6, + "source": "https://platform.moonshot.ai/docs/pricing/chat#generation-model-kimi-k2", + "supports_function_calling": true, + "supports_tool_choice": true, + "supports_web_search": true + }, "moonshot/moonshot-v1-128k": { "input_cost_per_token": 2e-06, "litellm_provider": "moonshot", @@ -20487,6 +20704,21 @@ "supports_reasoning": true, "supports_tool_choice": true }, + "openrouter/deepseek/deepseek-v3.2": { + "input_cost_per_token": 2.8e-07, + "input_cost_per_token_cache_hit": 2.8e-08, + "litellm_provider": "openrouter", + "max_input_tokens": 163840, + "max_output_tokens": 163840, + "max_tokens": 8192, + "mode": "chat", + "output_cost_per_token": 4e-07, + "supports_assistant_prefill": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_tool_choice": true + }, "openrouter/deepseek/deepseek-v3.2-exp": { "input_cost_per_token": 2e-07, "input_cost_per_token_cache_hit": 2e-08, @@ -23742,6 +23974,32 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159 }, + "global.anthropic.claude-opus-4-5-20251101-v1:0": { + "cache_creation_input_token_cost": 6.25e-06, + "cache_read_input_token_cost": 5e-07, + "input_cost_per_token": 5e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 200000, + "max_output_tokens": 64000, + "max_tokens": 64000, + "mode": "chat", + "output_cost_per_token": 2.5e-05, + "search_context_cost_per_query": { + "search_context_size_high": 0.01, + "search_context_size_low": 0.01, + "search_context_size_medium": 0.01 + }, + "supports_assistant_prefill": true, + "supports_computer_use": true, + "supports_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_tool_choice": true, + "supports_vision": true, + "tool_use_system_prompt_tokens": 159 + }, "us.anthropic.claude-sonnet-4-20250514-v1:0": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -24825,6 +25083,15 @@ "supports_parallel_function_calling": true, "supports_tool_choice": true }, + "vertex_ai/chirp": { + "input_cost_per_character": 30e-06, + "litellm_provider": "vertex_ai", + "mode": "audio_speech", + "source": "https://cloud.google.com/text-to-speech/pricing", + "supported_endpoints": [ + "/v1/audio/speech" + ] + }, "vertex_ai/claude-3-5-haiku": { "input_cost_per_token": 1e-06, "litellm_provider": "vertex_ai-anthropic_models", @@ -26255,8 +26522,8 @@ "max_tokens": 128000, "max_input_tokens": 128000, "max_output_tokens": 128000, - "input_cost_per_token": 0.135, - "output_cost_per_token": 0.4, + "input_cost_per_token": 6e-07, + "output_cost_per_token": 2.5e-06, "litellm_provider": "wandb", "mode": "chat" }, diff --git a/poetry.lock b/poetry.lock index 15140856781..711ab0bbeae 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.1.4 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.2.1 and should not be changed by hand. [[package]] name = "aiofiles" diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index b5bde3e5ce4..64c135ee85d 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -17,7 +17,8 @@ "rerank": "Supports /rerank endpoint", "ocr": "Supports /ocr endpoint", "search": "Supports /search endpoint", - "skills": "Supports /skills endpoint" + "skills": "Supports /skills endpoint", + "a2a_(Agent Gateway)": "Supports /a2a/{agent}/message/send endpoint (A2A Protocol)" } } }, @@ -35,7 +36,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "ai21": { @@ -51,7 +53,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "ai21_chat": { @@ -67,7 +70,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "anthropic": { @@ -84,7 +88,8 @@ "moderations": false, "batches": true, "rerank": false, - "skills": true + "skills": true, + "a2a": true } }, "anthropic_text": { @@ -101,7 +106,8 @@ "moderations": false, "batches": true, "rerank": false, - "skills": true + "skills": true, + "a2a": true } }, "assemblyai": { @@ -117,7 +123,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "auto_router": { @@ -133,7 +140,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "bedrock": { @@ -149,7 +157,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": true + "rerank": true, + "a2a": true } }, "sagemaker": { @@ -165,7 +174,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "azure": { @@ -181,7 +191,8 @@ "audio_speech": true, "moderations": true, "batches": true, - "rerank": false + "rerank": false, + "a2a": true } }, "azure_ai": { @@ -198,7 +209,8 @@ "moderations": true, "batches": true, "rerank": false, - "ocr": true + "ocr": true, + "a2a": true } }, "azure_ai/doc-intelligence": { @@ -231,7 +243,8 @@ "audio_speech": true, "moderations": true, "batches": true, - "rerank": false + "rerank": false, + "a2a": true } }, "baseten": { @@ -247,7 +260,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "bytez": { @@ -263,7 +277,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "cerebras": { @@ -279,7 +294,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "clarifai": { @@ -295,7 +311,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "cloudflare": { @@ -311,7 +328,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "codestral": { @@ -327,7 +345,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "cohere": { @@ -343,7 +362,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": true + "rerank": true, + "a2a": true } }, "cohere_chat": { @@ -359,7 +379,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "cometapi": { @@ -375,7 +396,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "compactifai": { @@ -391,7 +413,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "custom": { @@ -407,7 +430,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "custom_openai": { @@ -423,7 +447,8 @@ "audio_speech": true, "moderations": true, "batches": true, - "rerank": false + "rerank": false, + "a2a": true } }, "dashscope": { @@ -439,7 +464,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "databricks": { @@ -455,7 +481,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "dataforseo": { @@ -488,7 +515,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "deepgram": { @@ -504,7 +532,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "deepinfra": { @@ -520,7 +549,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "deepseek": { @@ -536,7 +566,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "elevenlabs": { @@ -552,7 +583,8 @@ "audio_speech": true, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "exa_ai": { @@ -585,7 +617,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "fal_ai": { @@ -601,7 +634,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "featherless_ai": { @@ -617,7 +651,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "fireworks_ai": { @@ -633,7 +668,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "firecrawl": { @@ -666,7 +702,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "galadriel": { @@ -682,7 +719,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "github_copilot": { @@ -698,7 +736,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "github": { @@ -714,7 +753,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "vertex_ai": { @@ -731,7 +771,24 @@ "moderations": false, "batches": false, "rerank": false, - "ocr": true + "ocr": true, + "a2a": true + } + }, + "vertex_ai/chirp": { + "display_name": "Google - Vertex AI Chirp3 HD (`vertex_ai/chirp`)", + "url": "https://docs.litellm.ai/docs/providers/vertex_speech", + "endpoints": { + "chat_completions": false, + "messages": false, + "responses": false, + "embeddings": false, + "image_generations": false, + "audio_transcriptions": false, + "audio_speech": true, + "moderations": false, + "batches": false, + "rerank": false } }, "gemini": { @@ -747,7 +804,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "gradient_ai": { @@ -763,7 +821,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "groq": { @@ -779,7 +838,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "heroku": { @@ -795,7 +855,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "hosted_vllm": { @@ -811,7 +872,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "huggingface": { @@ -827,7 +889,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": true + "rerank": true, + "a2a": true } }, "hyperbolic": { @@ -843,7 +906,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "watsonx": { @@ -859,7 +923,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "infinity": { @@ -907,7 +972,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "lemonade": { @@ -923,7 +989,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "litellm_proxy": { @@ -939,7 +1006,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "llamafile": { @@ -955,7 +1023,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "lm_studio": { @@ -971,7 +1040,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "maritalk": { @@ -987,7 +1057,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "meta_llama": { @@ -1003,7 +1074,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "mistral": { @@ -1020,7 +1092,8 @@ "moderations": false, "batches": false, "rerank": false, - "ocr": true + "ocr": true, + "a2a": true } }, "moonshot": { @@ -1036,7 +1109,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "docker_model_runner": { @@ -1052,7 +1126,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "morph": { @@ -1068,7 +1143,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "nebius": { @@ -1084,7 +1160,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "nlp_cloud": { @@ -1100,7 +1177,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "novita": { @@ -1116,7 +1194,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "nscale": { @@ -1132,7 +1211,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "nvidia_nim": { @@ -1148,7 +1228,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "oci": { @@ -1164,7 +1245,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "ollama": { @@ -1180,7 +1262,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "ollama_chat": { @@ -1196,7 +1279,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "oobabooga": { @@ -1212,7 +1296,8 @@ "audio_speech": true, "moderations": true, "batches": true, - "rerank": false + "rerank": false, + "a2a": true } }, "openai": { @@ -1228,7 +1313,8 @@ "audio_speech": true, "moderations": true, "batches": true, - "rerank": false + "rerank": false, + "a2a": true } }, "openai_like": { @@ -1260,7 +1346,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "ovhcloud": { @@ -1276,7 +1363,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "parallel_ai": { @@ -1310,7 +1398,8 @@ "moderations": false, "batches": false, "rerank": false, - "search": true + "search": true, + "a2a": true } }, "petals": { @@ -1326,7 +1415,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "publicai": { @@ -1342,7 +1432,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "predibase": { @@ -1358,7 +1449,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "recraft": { @@ -1390,7 +1482,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "runwayml": { @@ -1423,7 +1516,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "searxng": { @@ -1456,7 +1550,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "snowflake": { @@ -1472,7 +1567,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "text-completion-codestral": { @@ -1488,7 +1584,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "text-completion-openai": { @@ -1504,7 +1601,8 @@ "audio_speech": true, "moderations": true, "batches": true, - "rerank": false + "rerank": false, + "a2a": true } }, "together_ai": { @@ -1520,7 +1618,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "topaz": { @@ -1536,7 +1635,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "tavily": { @@ -1569,7 +1669,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "v0": { @@ -1585,7 +1686,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "vercel_ai_gateway": { @@ -1601,7 +1703,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "vllm": { @@ -1617,7 +1720,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "volcengine": { @@ -1633,7 +1737,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "voyage": { @@ -1665,7 +1770,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "watsonx_text": { @@ -1681,7 +1787,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "xai": { @@ -1697,7 +1804,8 @@ "audio_speech": false, "moderations": false, "batches": false, - "rerank": false + "rerank": false, + "a2a": true } }, "xinference": { diff --git a/tests/agent_tests/test_a2a.py b/tests/agent_tests/test_a2a.py new file mode 100644 index 00000000000..eeab2680564 --- /dev/null +++ b/tests/agent_tests/test_a2a.py @@ -0,0 +1,167 @@ +""" +Test for LiteLLM A2A module. + +Run with: + pytest tests/agent_tests/test_a2a.py -v -s +""" + +import asyncio +import os +import sys +import json +from typing import Optional +from uuid import uuid4 + +import pytest + +import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import StandardLoggingPayload + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +from a2a.types import MessageSendParams, SendMessageRequest + + +@pytest.mark.asyncio +async def test_asend_message_with_client_decorator(): + """ + Test asend_message standalone function with @client decorator. + This tests the LiteLLM logging integration. + """ + litellm._turn_on_debug() + from litellm.a2a_protocol import asend_message, create_a2a_client + + # Create the A2A client first + a2a_client = await create_a2a_client(base_url="http://localhost:10001") + + # Build the request matching A2A SDK spec + send_message_payload = { + "message": { + "role": "user", + "parts": [ + { + "kind": "text", + "text": "Hello from @client decorated function!", + } + ], + "messageId": uuid4().hex, + }, + } + + request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams(**send_message_payload), + ) + + # Send message using standalone function with @client decorator + response = await asend_message(a2a_client=a2a_client, request=request) + + # Print response for debugging + print("\n=== A2A Response (standalone with @client) ===") + print(response.model_dump(mode="json", exclude_none=True)) + + # Basic assertions + assert response is not None + + +class TestA2ALogger(CustomLogger): + """Custom logger to capture A2A logging payloads for testing.""" + + def __init__(self): + self.standard_logging_payload: Optional[StandardLoggingPayload] = None + self.logged_kwargs: Optional[dict] = None + self.log_success_called = False + super().__init__() + + async def async_log_success_event( + self, kwargs, response_obj, start_time, end_time + ): + print("TestA2ALogger: async_log_success_event called") + self.log_success_called = True + self.logged_kwargs = kwargs + self.standard_logging_payload = kwargs.get("standard_logging_object", None) + print(f"Captured standard_logging_payload: {self.standard_logging_payload}") + + +@pytest.mark.asyncio +async def test_a2a_logging_payload(): + """ + Test that A2A calls create a standard logging payload. + Validates the @client decorator integration with LiteLLM logging. + """ + # Reset callbacks and set up custom logger + litellm.logging_callback_manager._reset_all_callbacks() + test_logger = TestA2ALogger() + litellm.callbacks = [test_logger] + + from litellm.a2a_protocol import asend_message, create_a2a_client + + # Create the A2A client first + a2a_client = await create_a2a_client(base_url="http://localhost:10001") + + # Build the request + send_message_payload = { + "message": { + "role": "user", + "parts": [ + { + "kind": "text", + "text": "Hello! Testing logging payload.", + } + ], + "messageId": uuid4().hex, + }, + } + + request = SendMessageRequest( + id=str(uuid4()), + params=MessageSendParams(**send_message_payload), + ) + + # Send message + response = await asend_message(a2a_client=a2a_client, request=request) + + # Give async logging time to complete + await asyncio.sleep(1) + + # Print debug info + print("\n=== Logging Validation ===") + print(f"log_success_called: {test_logger.log_success_called}") + print(f"standard_logging_payload: {test_logger.standard_logging_payload}") + print(f"logged kwargs: {json.dumps(test_logger.logged_kwargs, indent=4, default=str)}") + + # Verify logging was called + assert test_logger.log_success_called is True + assert test_logger.standard_logging_payload is not None + + # Verify standard_logging_payload exists + slp = test_logger.standard_logging_payload + assert slp is not None + + # Get values from standard logging payload + logged_model = slp.get("model") if isinstance(slp, dict) else getattr(slp, "model", None) + logged_provider = slp.get("custom_llm_provider") if isinstance(slp, dict) else getattr(slp, "custom_llm_provider", None) + call_type = slp.get("call_type") if isinstance(slp, dict) else getattr(slp, "call_type", None) + response_cost = slp.get("response_cost") if isinstance(slp, dict) else getattr(slp, "response_cost", None) + + print(f"\n=== Standard Logging Payload Validation ===") + print(f"model: {logged_model}") + print(f"custom_llm_provider: {logged_provider}") + print(f"call_type: {call_type}") + print(f"response_cost: {response_cost}") + + # Verify model and custom_llm_provider are set correctly + assert logged_model is not None, "model should be set" + assert "a2a_agent/" in logged_model, f"model should contain 'a2a_agent/', got: {logged_model}" + assert logged_provider == "a2a_agent", f"custom_llm_provider should be 'a2a_agent', got: {logged_provider}" + + # Verify call_type is correct for A2A + assert call_type == "asend_message", f"call_type should be 'asend_message', got: {call_type}" + + # Verify response_cost is set to 0.0 (not None, not an error) + # This confirms the A2A cost calculator is working + assert response_cost is not None, "response_cost should not be None" + assert response_cost == 0.0, f"response_cost should be 0.0 for A2A, got: {response_cost}" diff --git a/tests/audio_tests/speech_vertex.mp3 b/tests/audio_tests/speech_vertex.mp3 new file mode 100644 index 00000000000..91efaea150c Binary files /dev/null and b/tests/audio_tests/speech_vertex.mp3 differ diff --git a/tests/batches_tests/test_hosted_vllm_batches_and_files.py b/tests/batches_tests/test_hosted_vllm_batches_and_files.py new file mode 100644 index 00000000000..aa4d847a45d --- /dev/null +++ b/tests/batches_tests/test_hosted_vllm_batches_and_files.py @@ -0,0 +1,106 @@ +""" +Unit Tests for hosted_vllm Batches and Files API + +Tests the integration of hosted_vllm provider with LiteLLM's batch and file operations. +Tests against a real OpenAI-compatible endpoint. +""" +import json +import os +import sys +import time +import uuid + +import httpx +import pytest +from dotenv import load_dotenv + +load_dotenv() +sys.path.insert( + 0, os.path.abspath("../..") +) + +import litellm + + +SERVER_URL = "https://exampleopenaiendpoint-production-0ee2.up.railway.app/v1" + + +@pytest.mark.asyncio() +@pytest.mark.skip(reason="Local only test") +async def test_hosted_vllm_full_workflow(): + """ + Test the complete workflow: create file -> create batch -> retrieve batch -> retrieve file. + Tests against real OpenAI-compatible endpoint. + """ + litellm._turn_on_debug() + file_name = "openai_batch_completions.jsonl" + _current_dir = os.path.dirname(os.path.abspath(__file__)) + file_path = os.path.join(_current_dir, file_name) + + # Step 1: Create file + print("\n=== Step 1: Creating file ===") + file_obj = await litellm.acreate_file( + file=open(file_path, "rb"), + purpose="batch", + custom_llm_provider="hosted_vllm", + api_base=SERVER_URL, + api_key="test-api-key", + ) + + print(f"✓ Created file: {file_obj.id}") + assert file_obj.id is not None + assert file_obj.object == "file" + assert file_obj.purpose == "batch" + + # Step 2: Create batch + print("\n=== Step 2: Creating batch ===") + batch_obj = await litellm.acreate_batch( + completion_window="24h", + endpoint="/v1/chat/completions", + input_file_id=file_obj.id, + custom_llm_provider="hosted_vllm", + metadata={"test": "hosted_vllm_integration"}, + api_base=SERVER_URL, + api_key="test-api-key", + ) + + print(f"✓ Created batch: {batch_obj.id}") + print(f" Status: {batch_obj.status}") + print(f" Input file: {batch_obj.input_file_id}") + assert batch_obj.id is not None + assert batch_obj.object == "batch" + assert batch_obj.input_file_id == file_obj.id + assert batch_obj.endpoint == "/v1/chat/completions" + + # Step 3: Retrieve batch + print("\n=== Step 3: Retrieving batch ===") + retrieved_batch = await litellm.aretrieve_batch( + batch_id=batch_obj.id, + custom_llm_provider="hosted_vllm", + api_base=SERVER_URL, + api_key="test-api-key", + ) + + print(f"✓ Retrieved batch: {retrieved_batch.id}") + print(f" Status: {retrieved_batch.status}") + print(f" Output file: {retrieved_batch.output_file_id}") + assert retrieved_batch.id == batch_obj.id + assert retrieved_batch.object == "batch" + assert retrieved_batch.input_file_id == file_obj.id + + # Step 4: Retrieve file (verify file still accessible) + print("\n=== Step 4: Retrieving original file ===") + retrieved_file = await litellm.afile_retrieve( + file_id=file_obj.id, + custom_llm_provider="hosted_vllm", + api_base=SERVER_URL, + api_key="test-api-key", + ) + + print(f"✓ Retrieved file: {retrieved_file.id}") + print(f" Filename: {retrieved_file.filename}") + print(f" Bytes: {retrieved_file.bytes}") + assert retrieved_file.id == file_obj.id + assert retrieved_file.object == "file" + + print("\n✅ Full workflow test completed successfully!") diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py index 186056dac90..7ce99abdd15 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_apply_guardrail_endpoint.py @@ -1,6 +1,7 @@ """ Test the /guardrails/apply_guardrail endpoint """ + import os import sys from unittest.mock import AsyncMock, Mock, patch @@ -22,37 +23,45 @@ async def test_apply_guardrail_endpoint_returns_correct_response(): from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail # Mock the guardrail registry - with patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY") as mock_registry: + with patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry: # Create a mock guardrail mock_guardrail = Mock(spec=CustomGuardrail) - mock_guardrail.apply_guardrail = AsyncMock(return_value="Redacted text: [REDACTED] and [REDACTED]") - + # Apply guardrail now returns a tuple (List[str], Optional[List[str]]) + mock_guardrail.apply_guardrail = AsyncMock( + return_value=(["Redacted text: [REDACTED] and [REDACTED]"], None) + ) + # Configure the registry to return our mock guardrail mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail - + # Create the request request = ApplyGuardrailRequest( guardrail_name="test-guardrail", text="Test text with PII", language="en", - entities=["EMAIL_ADDRESS", "PERSON"] + entities=["EMAIL_ADDRESS", "PERSON"], ) - + # Create a mock user API key user_api_key_dict = UserAPIKeyAuth(api_key="test-key") - + # Call the endpoint - response = await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict) - + response = await apply_guardrail( + request=request, user_api_key_dict=user_api_key_dict + ) + # Verify the response is of the correct type assert isinstance(response, ApplyGuardrailResponse) assert response.response_text == "Redacted text: [REDACTED] and [REDACTED]" - - # Verify the guardrail was called with correct parameters + + # Verify the guardrail was called with correct parameters (new signature) mock_guardrail.apply_guardrail.assert_called_once_with( - text="Test text with PII", - language="en", - entities=["EMAIL_ADDRESS", "PERSON"] + texts=["Test text with PII"], + request_data={}, + input_type="request", + images=None, ) @@ -63,23 +72,23 @@ async def test_apply_guardrail_endpoint_guardrail_not_found(): from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail # Mock the guardrail registry to return None - with patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY") as mock_registry: + with patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry: mock_registry.get_initialized_guardrail_callback.return_value = None - + # Create the request request = ApplyGuardrailRequest( - guardrail_name="non-existent-guardrail", - text="Test text", - language="en" + guardrail_name="non-existent-guardrail", text="Test text", language="en" ) - + # Create a mock user API key user_api_key_dict = UserAPIKeyAuth(api_key="test-key") - + # Verify exception is raised with pytest.raises(ProxyException) as exc_info: await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict) - + assert "non-existent-guardrail" in exc_info.value.message assert "not found" in exc_info.value.message @@ -90,34 +99,41 @@ async def test_apply_guardrail_endpoint_with_presidio_guardrail(): from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail # Mock the guardrail registry - with patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY") as mock_registry: + with patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry: # Create a mock guardrail that simulates Presidio behavior mock_guardrail = Mock(spec=CustomGuardrail) - # Simulate masking PII entities + # Simulate masking PII entities - returns tuple (List[str], Optional[List[str]]) mock_guardrail.apply_guardrail = AsyncMock( - return_value="My name is [PERSON] and my email is [EMAIL_ADDRESS]" + return_value=(["My name is [PERSON] and my email is [EMAIL_ADDRESS]"], None) ) - + # Configure the registry to return our mock guardrail mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail - + # Create the request request = ApplyGuardrailRequest( guardrail_name="pii-detection-guard", text="My name is John Doe and my email is john@example.com", language="en", - entities=["EMAIL_ADDRESS", "PERSON"] + entities=["EMAIL_ADDRESS", "PERSON"], ) - + # Create a mock user API key user_api_key_dict = UserAPIKeyAuth(api_key="test-key") - + # Call the endpoint - response = await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict) - + response = await apply_guardrail( + request=request, user_api_key_dict=user_api_key_dict + ) + # Verify the response is of the correct type assert isinstance(response, ApplyGuardrailResponse) - assert response.response_text == "My name is [PERSON] and my email is [EMAIL_ADDRESS]" + assert ( + response.response_text + == "My name is [PERSON] and my email is [EMAIL_ADDRESS]" + ) assert "john@example.com" not in response.response_text assert "John Doe" not in response.response_text @@ -128,33 +144,37 @@ async def test_apply_guardrail_endpoint_without_optional_params(): from litellm.proxy.guardrails.guardrail_endpoints import apply_guardrail # Mock the guardrail registry - with patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY") as mock_registry: + with patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry: # Create a mock guardrail mock_guardrail = Mock(spec=CustomGuardrail) - mock_guardrail.apply_guardrail = AsyncMock(return_value="Processed text") - + # Returns tuple (List[str], Optional[List[str]]) + mock_guardrail.apply_guardrail = AsyncMock( + return_value=(["Processed text"], None) + ) + # Configure the registry to return our mock guardrail mock_registry.get_initialized_guardrail_callback.return_value = mock_guardrail - + # Create the request without optional parameters request = ApplyGuardrailRequest( - guardrail_name="test-guardrail", - text="Test text" + guardrail_name="test-guardrail", text="Test text" ) - + # Create a mock user API key user_api_key_dict = UserAPIKeyAuth(api_key="test-key") - + # Call the endpoint - response = await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict) - + response = await apply_guardrail( + request=request, user_api_key_dict=user_api_key_dict + ) + # Verify the response is of the correct type assert isinstance(response, ApplyGuardrailResponse) assert response.response_text == "Processed text" - - # Verify the guardrail was called with None for optional parameters + + # Verify the guardrail was called with new signature mock_guardrail.apply_guardrail.assert_called_once_with( - text="Test text", - language=None, - entities=None + texts=["Test text"], request_data={}, input_type="request", images=None ) diff --git a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py index e65d01e41e4..8d98f56cd3a 100644 --- a/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py +++ b/tests/enterprise/litellm_enterprise/proxy/guardrails/test_bedrock_apply_guardrail.py @@ -1,6 +1,7 @@ """ Test the Bedrock guardrail apply_guardrail functionality """ + import os import sys from unittest.mock import AsyncMock, Mock, patch @@ -23,32 +24,30 @@ async def test_bedrock_apply_guardrail_success(): guardrail = BedrockGuardrail( guardrail_name="test-bedrock-guard", guardrailIdentifier="test-guard-id", - guardrailVersion="DRAFT" + guardrailVersion="DRAFT", ) - + # Mock the make_bedrock_api_request method - with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request: + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api_request: # Mock a successful response from Bedrock mock_response = { "action": "ALLOWED", - "content": [ - { - "text": { - "text": "This is a test message with some content" - } - } - ] + "content": [{"text": {"text": "This is a test message with some content"}}], } mock_api_request.return_value = mock_response - - # Test the apply_guardrail method - result = await guardrail.apply_guardrail( - text="This is a test message with some content", - language="en" + + # Test the apply_guardrail method with new signature + guardrailed_inputs = await guardrail.apply_guardrail( + inputs={"texts": ["This is a test message with some content"]}, + request_data={}, + input_type="request", ) - + result = guardrailed_inputs.get("texts", []) + # Verify the result - assert result == "This is a test message with some content" + assert result == ["This is a test message with some content"] mock_api_request.assert_called_once() @@ -59,25 +58,25 @@ async def test_bedrock_apply_guardrail_blocked(): guardrail = BedrockGuardrail( guardrail_name="test-bedrock-guard", guardrailIdentifier="test-guard-id", - guardrailVersion="DRAFT" + guardrailVersion="DRAFT", ) - + # Mock the make_bedrock_api_request method - with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request: + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api_request: # Mock a blocked response from Bedrock - mock_response = { - "action": "BLOCKED", - "reason": "Content violates policy" - } + mock_response = {"action": "BLOCKED", "reason": "Content violates policy"} mock_api_request.return_value = mock_response - + # Test the apply_guardrail method should raise an exception with pytest.raises(Exception) as exc_info: await guardrail.apply_guardrail( - text="This is blocked content", - language="en" + inputs={"texts": ["This is blocked content"]}, + request_data={}, + input_type="request", ) - + assert "Content blocked by Bedrock guardrail" in str(exc_info.value) assert "Content violates policy" in str(exc_info.value) @@ -89,30 +88,30 @@ async def test_bedrock_apply_guardrail_with_masking(): guardrail = BedrockGuardrail( guardrail_name="test-bedrock-guard", guardrailIdentifier="test-guard-id", - guardrailVersion="DRAFT" + guardrailVersion="DRAFT", ) - + # Mock the make_bedrock_api_request method - with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request: + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api_request: # Mock a response with masked content mock_response = { "action": "ALLOWED", - "outputs": [ - { - "text": "This is a test message with [REDACTED] content" - } - ] + "outputs": [{"text": "This is a test message with [REDACTED] content"}], } mock_api_request.return_value = mock_response - - # Test the apply_guardrail method - result = await guardrail.apply_guardrail( - text="This is a test message with sensitive content", - language="en" + + # Test the apply_guardrail method with new signature + guardrailed_inputs = await guardrail.apply_guardrail( + inputs={"texts": ["This is a test message with sensitive content"]}, + request_data={}, + input_type="request", ) - + result = guardrailed_inputs.get("texts", []) + # Verify the result contains the masked content - assert result == "This is a test message with [REDACTED] content" + assert result == ["This is a test message with [REDACTED] content"] mock_api_request.assert_called_once() @@ -123,21 +122,24 @@ async def test_bedrock_apply_guardrail_api_failure(): guardrail = BedrockGuardrail( guardrail_name="test-bedrock-guard", guardrailIdentifier="test-guard-id", - guardrailVersion="DRAFT" + guardrailVersion="DRAFT", ) - + # Mock the make_bedrock_api_request method to raise an exception - with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request: + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api_request: mock_api_request.side_effect = Exception("API connection failed") - + # Test the apply_guardrail method should raise an exception with pytest.raises(Exception) as exc_info: await guardrail.apply_guardrail( - text="This is a test message", - language="en" + inputs={"texts": ["This is a test message"]}, + request_data={}, + input_type="request", ) - - assert "Bedrock guardrail failed" in str(exc_info.value) + + # The error message should contain the original exception assert "API connection failed" in str(exc_info.value) @@ -150,44 +152,50 @@ async def test_bedrock_apply_guardrail_endpoint_integration(): guardrail = BedrockGuardrail( guardrail_name="test-bedrock-guard", guardrailIdentifier="test-guard-id", - guardrailVersion="DRAFT" + guardrailVersion="DRAFT", ) - + # Mock the guardrail registry - with patch("litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY") as mock_registry: + with patch( + "litellm.proxy.guardrails.guardrail_endpoints.GUARDRAIL_REGISTRY" + ) as mock_registry: # Mock the make_bedrock_api_request method - with patch.object(guardrail, 'make_bedrock_api_request', new_callable=AsyncMock) as mock_api_request: + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api_request: # Mock a successful response from Bedrock mock_response = { "action": "ALLOWED", - "outputs": [ - { - "text": "This is a test message with processed content" - } - ] + "outputs": [{"text": "This is a test message with processed content"}], } mock_api_request.return_value = mock_response - + # Configure the registry to return our guardrail mock_registry.get_initialized_guardrail_callback.return_value = guardrail - + # Create the request request = ApplyGuardrailRequest( guardrail_name="test-bedrock-guard", text="This is a test message with some content", - language="en" + language="en", ) - + # Create a mock user API key user_api_key_dict = UserAPIKeyAuth(api_key="test-key") - + # Call the endpoint - response = await apply_guardrail(request=request, user_api_key_dict=user_api_key_dict) - + response = await apply_guardrail( + request=request, user_api_key_dict=user_api_key_dict + ) + # Verify the response assert isinstance(response, ApplyGuardrailResponse) - assert response.response_text == "This is a test message with processed content" - mock_api_request.assert_called_once() + assert ( + response.response_text + == "This is a test message with processed content" + ) + # Note: The endpoint now calls apply_guardrail which internally calls make_bedrock_api_request + # The call count check has been removed as it may be called multiple times through the chain @pytest.mark.asyncio @@ -208,18 +216,22 @@ async def test_bedrock_apply_guardrail_filters_request_messages_when_flag_enable request_data = {"messages": request_messages} - with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api: mock_api.return_value = {"action": "ALLOWED"} - result = await guardrail.apply_guardrail( - text="latest question", + guardrailed_inputs = await guardrail.apply_guardrail( + inputs={"texts": ["latest question"]}, request_data=request_data, + input_type="request", ) + result = guardrailed_inputs.get("texts", []) assert mock_api.called _, kwargs = mock_api.call_args assert kwargs["messages"] == [request_messages[-1]] - assert result == "latest question" + assert result == ["latest question"] @pytest.mark.asyncio @@ -238,19 +250,23 @@ async def test_bedrock_apply_guardrail_filters_request_messages_when_flag_enable request_data = {"messages": request_messages} - with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api: mock_api.return_value = {"action": "BLOCKED", "reason": "policy"} with pytest.raises(Exception, match="policy") as exc_info: await guardrail.apply_guardrail( - text="blocked", + inputs={"texts": ["blocked"]}, request_data=request_data, + input_type="request", ) assert mock_api.called _, kwargs = mock_api.call_args assert kwargs["messages"] == [request_messages[-1]] - assert "Bedrock guardrail failed" in str(exc_info.value) + assert "Content blocked by Bedrock guardrail" in str(exc_info.value) + def test_bedrock_guardrail_filters_latest_user_message_when_enabled(): guardrail = BedrockGuardrail( diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index b0e192c00a7..1795ff23605 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -2,6 +2,7 @@ import sys import os import io, asyncio import pytest + sys.path.insert(0, os.path.abspath("../..")) import litellm from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail @@ -9,11 +10,12 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.caching import DualCache from unittest.mock import MagicMock, AsyncMock, patch + @pytest.mark.asyncio async def test_bedrock_guardrails_pii_masking(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = BedrockGuardrail( guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", @@ -25,14 +27,17 @@ async def test_bedrock_guardrails_pii_masking(): {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, {"role": "assistant", "content": "Hello, how can I help you today?"}, {"role": "user", "content": "I need to cancel my order"}, - {"role": "user", "content": "ok, my credit card number is 1234-5678-9012-3456"}, + { + "role": "user", + "content": "ok, my credit card number is 1234-5678-9012-3456", + }, ], } response = await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - call_type="completion" + call_type="completion", ) print("response after moderation hook", response) @@ -40,14 +45,17 @@ async def test_bedrock_guardrails_pii_masking(): assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}" assert response["messages"][1]["content"] == "Hello, how can I help you today?" assert response["messages"][2]["content"] == "I need to cancel my order" - assert response["messages"][3]["content"] == "ok, my credit card number is {CREDIT_DEBIT_CARD_NUMBER}" + assert ( + response["messages"][3]["content"] + == "ok, my credit card number is {CREDIT_DEBIT_CARD_NUMBER}" + ) @pytest.mark.asyncio async def test_bedrock_guardrails_pii_masking_content_list(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = BedrockGuardrail( guardrailIdentifier="wf0hkdb5x07f", guardrailVersion="DRAFT", @@ -56,34 +64,41 @@ async def test_bedrock_guardrails_pii_masking_content_list(): request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": [ - {"type": "text", "text": "Hello, my phone number is +1 412 555 1212"}, - {"type": "text", "text": "what time is it?"}, - ]}, - {"role": "assistant", "content": "Hello, how can I help you today?"}, { "role": "user", - "content": "who is the president of the united states?" - } + "content": [ + { + "type": "text", + "text": "Hello, my phone number is +1 412 555 1212", + }, + {"type": "text", "text": "what time is it?"}, + ], + }, + {"role": "assistant", "content": "Hello, how can I help you today?"}, + {"role": "user", "content": "who is the president of the united states?"}, ], } response = await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - call_type="completion" + call_type="completion", ) print(response) - + if response: # Only assert if response is not None # Verify that the list content is properly masked assert isinstance(response["messages"][0]["content"], list) - assert response["messages"][0]["content"][0]["text"] == "Hello, my phone number is {PHONE}" + assert ( + response["messages"][0]["content"][0]["text"] + == "Hello, my phone number is {PHONE}" + ) assert response["messages"][0]["content"][1]["text"] == "what time is it?" assert response["messages"][1]["content"] == "Hello, how can I help you today?" - assert response["messages"][2]["content"] == "who is the president of the united states?" - - + assert ( + response["messages"][2]["content"] + == "who is the president of the united states?" + ) @pytest.mark.asyncio @@ -92,10 +107,10 @@ async def test_bedrock_guardrails_block_messages_api(): Test that guardrails block messages API requests containing 'coffee' and raise the expected exception. """ from fastapi import HTTPException - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = BedrockGuardrail( guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", @@ -104,14 +119,17 @@ async def test_bedrock_guardrails_block_messages_api(): request_data = { "model": "claude-3-5-sonnet-20240620", "messages": [ - {"role": "user", "content": [ - {"type": "text", "text": "Hello, my phone number is +1 412 555 1212"}, - {"type": "text", "text": "what time is it?"}, - ]}, { "role": "user", - "content": "tell me about coffee" - } + "content": [ + { + "type": "text", + "text": "Hello, my phone number is +1 412 555 1212", + }, + {"type": "text", "text": "what time is it?"}, + ], + }, + {"role": "user", "content": "tell me about coffee"}, ], } @@ -122,13 +140,17 @@ async def test_bedrock_guardrails_block_messages_api(): call_type="anthropic_messages", cache=MagicMock(spec=DualCache), ) - + exception = exc_info.value assert exception.status_code == 400 detail = exception.detail assert isinstance(detail, dict) assert detail["error"] == "Violated guardrail policy" - assert detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question. coffee guardrail applied " + assert ( + detail["bedrock_guardrail_response"] + == "Sorry, the model cannot answer this question. coffee guardrail applied " + ) + @pytest.mark.asyncio async def test_bedrock_guardrails_block_responses_api(): @@ -136,10 +158,10 @@ async def test_bedrock_guardrails_block_responses_api(): Test that guardrails block responses API requests containing 'coffee' and raise the expected exception. """ from fastapi import HTTPException - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = BedrockGuardrail( guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", @@ -158,14 +180,16 @@ async def test_bedrock_guardrails_block_responses_api(): call_type="responses", cache=MagicMock(spec=DualCache), ) - + exception = exc_info.value assert exception.status_code == 400 detail = exception.detail assert isinstance(detail, dict) assert detail["error"] == "Violated guardrail policy" - assert detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question. coffee guardrail applied " - + assert ( + detail["bedrock_guardrail_response"] + == "Sorry, the model cannot answer this question. coffee guardrail applied " + ) @pytest.mark.asyncio @@ -182,7 +206,7 @@ async def test_bedrock_guardrails_with_streaming(): user_api_key_cache=mock_user_api_key_cache, premium_user=True, ) - + guardrail = BedrockGuardrail( guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", @@ -194,14 +218,9 @@ async def test_bedrock_guardrails_with_streaming(): request_data = { "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": "Hi I like coffee" - } - ], + "messages": [{"role": "user", "content": "Hi I like coffee"}], "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]} + "metadata": {"guardrails": ["bedrock-post-guard"]}, } response = await litellm.acompletion( @@ -213,7 +232,7 @@ async def test_bedrock_guardrails_with_streaming(): response=response, request_data=request_data, ) - + async for chunk in response: print(chunk) @@ -231,7 +250,7 @@ async def test_bedrock_guardrails_with_streaming_no_violation(): user_api_key_cache=mock_user_api_key_cache, premium_user=True, ) - + guardrail = BedrockGuardrail( guardrailIdentifier="ff6ujrregl1q", guardrailVersion="DRAFT", @@ -241,17 +260,11 @@ async def test_bedrock_guardrails_with_streaming_no_violation(): litellm.callbacks.append(guardrail) - request_data = { "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": "hi" - } - ], + "messages": [{"role": "user", "content": "hi"}], "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]} + "metadata": {"guardrails": ["bedrock-post-guard"]}, } response = await litellm.acompletion( @@ -263,11 +276,10 @@ async def test_bedrock_guardrails_with_streaming_no_violation(): response=response, request_data=request_data, ) - - + async for chunk in response: print(chunk) - + @pytest.mark.asyncio async def test_bedrock_guardrails_streaming_request_body_mock(): @@ -277,7 +289,7 @@ async def test_bedrock_guardrails_streaming_request_body_mock(): from litellm.proxy._types import UserAPIKeyAuth from litellm.caching import DualCache from litellm.types.guardrails import GuardrailEventHooks - + # Create mock objects mock_user_api_key_dict = UserAPIKeyAuth() mock_cache = MagicMock(spec=DualCache) @@ -297,79 +309,68 @@ async def test_bedrock_guardrails_streaming_request_body_mock(): litellm.Choices( index=0, message=litellm.Message( - role="assistant", - content="The capital of Spain is Madrid." + role="assistant", content="The capital of Spain is Madrid." ), - finish_reason="stop" + finish_reason="stop", ) ], created=1234567890, model="gpt-4o", - object="chat.completion" + object="chat.completion", ) # Mock Bedrock API response mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "NONE", - "outputs": [] - } + mock_bedrock_response.json.return_value = {"action": "NONE", "outputs": []} # Patch the async_handler.post method to capture the request body - with patch.object(guardrail, 'async_handler') as mock_async_handler: + with patch.object(guardrail, "async_handler") as mock_async_handler: mock_async_handler.post = AsyncMock(return_value=mock_bedrock_response) - + # Test data - simulating request data and assembled response request_data = { "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": "what's the capital of spain?" - } - ], + "messages": [{"role": "user", "content": "what's the capital of spain?"}], "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]} + "metadata": {"guardrails": ["bedrock-post-guard"]}, } # Call the method that should make the Bedrock API request await guardrail.make_bedrock_api_request( - source="OUTPUT", - response=mock_response, - request_data=request_data + source="OUTPUT", response=mock_response, request_data=request_data ) # Verify the API call was made mock_async_handler.post.assert_called_once() - + # Get the request data that was passed call_args = mock_async_handler.post.call_args - + # The data should be in the 'data' parameter of the prepared request # We need to parse the JSON from the prepared request body - prepared_request_body = call_args.kwargs.get('data') - + prepared_request_body = call_args.kwargs.get("data") + # Parse the JSON body if isinstance(prepared_request_body, bytes): - actual_body = json.loads(prepared_request_body.decode('utf-8')) + actual_body = json.loads(prepared_request_body.decode("utf-8")) else: actual_body = json.loads(prepared_request_body) - + # Expected body based on the convert_to_bedrock_format method behavior expected_body = { - 'source': 'OUTPUT', - 'content': [ - {'text': {'text': 'The capital of Spain is Madrid.'}} - ] + "source": "OUTPUT", + "content": [{"text": {"text": "The capital of Spain is Madrid."}}], } - + print("Actual Bedrock request body:", json.dumps(actual_body, indent=2)) print("Expected Bedrock request body:", json.dumps(expected_body, indent=2)) - + # Assert the request body matches exactly - assert actual_body == expected_body, f"Request body mismatch. Expected: {expected_body}, Got: {actual_body}" - + assert ( + actual_body == expected_body + ), f"Request body mismatch. Expected: {expected_body}, Got: {actual_body}" + @pytest.mark.asyncio async def test_bedrock_guardrail_aws_param_persistence(): @@ -387,23 +388,31 @@ async def test_bedrock_guardrail_aws_param_persistence(): guardrail_name="bedrock-post-guard", ) - with patch.object(guardrail, "get_credentials", wraps=guardrail.get_credentials) as mock_get_creds: + with patch.object( + guardrail, "get_credentials", wraps=guardrail.get_credentials + ) as mock_get_creds: for i in range(3): request_data = { "model": "gpt-4o", - "messages": [ - {"role": "user", "content": f"request {i}"} - ], + "messages": [{"role": "user", "content": f"request {i}"}], "stream": False, - "metadata": {"guardrails": ["bedrock-post-guard"]} + "metadata": {"guardrails": ["bedrock-post-guard"]}, } - with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: # Configure the mock response properly mock_response = AsyncMock() mock_response.status_code = 200 - mock_response.json = MagicMock(return_value={"action": "NONE", "outputs": []}) + mock_response.json = MagicMock( + return_value={"action": "NONE", "outputs": []} + ) mock_post.return_value = mock_response - await guardrail.make_bedrock_api_request(source="INPUT", messages=request_data.get("messages"), request_data=request_data) + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, + ) assert mock_get_creds.call_count == 3 for call in mock_get_creds.call_args_list: @@ -413,114 +422,124 @@ async def test_bedrock_guardrail_aws_param_persistence(): assert kwargs["aws_secret_access_key"] == "test-secret-key" assert kwargs["aws_region_name"] == "us-east-1" + @pytest.mark.asyncio async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" from unittest.mock import MagicMock - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrailResponse - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT" + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, ) - + from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrailResponse, + ) + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + # Test 1: ANONYMIZED action should NOT raise exception anonymized_response: BedrockGuardrailResponse = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "Hello, my phone number is {PHONE}" - }], - "assessments": [{ - "sensitiveInformationPolicy": { - "piiEntities": [{ - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED" - }] + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + } } - }] + ], } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(anonymized_response) + + should_raise = guardrail._should_raise_guardrail_blocked_exception( + anonymized_response + ) assert should_raise is False, "ANONYMIZED actions should not raise exceptions" - + # Test 2: BLOCKED action should raise exception blocked_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "I can't provide that information." - }], - "assessments": [{ - "topicPolicy": { - "topics": [{ - "name": "Sensitive Topic", - "type": "DENY", - "action": "BLOCKED" - }] + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "topicPolicy": { + "topics": [ + {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + ] + } } - }] + ], } - + should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response) assert should_raise is True, "BLOCKED actions should raise exceptions" - + # Test 3: Mixed actions - should raise if ANY action is BLOCKED mixed_response: BedrockGuardrailResponse = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "I can't provide that information." - }], - "assessments": [{ - "sensitiveInformationPolicy": { - "piiEntities": [{ - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED" - }] - }, - "topicPolicy": { - "topics": [{ - "name": "Blocked Topic", - "type": "DENY", - "action": "BLOCKED" - }] + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + }, + "topicPolicy": { + "topics": [ + {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} + ] + }, } - }] + ], } - + should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert should_raise is True, "Mixed actions with any BLOCKED should raise exceptions" - + assert ( + should_raise is True + ), "Mixed actions with any BLOCKED should raise exceptions" + # Test 4: NONE action should not raise exception none_response: BedrockGuardrailResponse = { "action": "NONE", "outputs": [], - "assessments": [] + "assessments": [], } - + should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response) assert should_raise is False, "NONE actions should not raise exceptions" - + # Test 5: Test other policy types with BLOCKED actions content_blocked_response: BedrockGuardrailResponse = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "I can't provide that information." - }], - "assessments": [{ - "contentPolicy": { - "filters": [{ - "type": "VIOLENCE", - "confidence": "HIGH", - "action": "BLOCKED" - }] + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "contentPolicy": { + "filters": [ + {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} + ] + } } - }] + ], } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(content_blocked_response) - assert should_raise is True, "Content policy BLOCKED actions should raise exceptions" + + should_raise = guardrail._should_raise_guardrail_blocked_exception( + content_blocked_response + ) + assert ( + should_raise is True + ), "Content policy BLOCKED actions should raise exceptions" @pytest.mark.asyncio @@ -529,10 +548,10 @@ async def test_bedrock_guardrail_masking_with_anonymized_response(): from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from litellm.caching import DualCache - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", @@ -544,18 +563,20 @@ async def test_bedrock_guardrail_masking_with_anonymized_response(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "Hello, my phone number is {PHONE}" - }], - "assessments": [{ - "sensitiveInformationPolicy": { - "piiEntities": [{ - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED" - }] + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + } + ] + } } - }] + ], } request_data = { @@ -566,21 +587,28 @@ async def test_bedrock_guardrail_masking_with_anonymized_response(): } # Patch the async_handler.post method - with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # This should NOT raise an exception since action is ANONYMIZED try: response = await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - call_type="completion" + call_type="completion", ) # Should succeed and return data with masked content assert response is not None - assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}" + assert ( + response["messages"][0]["content"] + == "Hello, my phone number is {PHONE}" + ) except Exception as e: - pytest.fail(f"Should not raise exception for ANONYMIZED actions, but got: {e}") + pytest.fail( + f"Should not raise exception for ANONYMIZED actions, but got: {e}" + ) @pytest.mark.asyncio @@ -588,10 +616,10 @@ async def test_bedrock_guardrail_uses_masked_output_without_masking_flags(): """Test that masked output from guardrails is used even when masking flags are not enabled""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + # Create guardrail WITHOUT masking flags enabled guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", @@ -604,48 +632,56 @@ async def test_bedrock_guardrail_uses_masked_output_without_masking_flags(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "Hello, my phone number is {PHONE} and email is {EMAIL}" - }], - "assessments": [{ - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED" - }, - { - "type": "EMAIL", - "match": "user@example.com", - "action": "ANONYMIZED" - } - ] + "outputs": [{"text": "Hello, my phone number is {PHONE} and email is {EMAIL}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + }, + { + "type": "EMAIL", + "match": "user@example.com", + "action": "ANONYMIZED", + }, + ] + } } - }] + ], } request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": "Hello, my phone number is +1 412 555 1212 and email is user@example.com"}, + { + "role": "user", + "content": "Hello, my phone number is +1 412 555 1212 and email is user@example.com", + }, ], } # Patch the async_handler.post method - with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # This should use the masked output even without masking flags response = await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - call_type="completion" + call_type="completion", ) - + # Should use the masked content from guardrail output assert response is not None - assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE} and email is {EMAIL}" + assert ( + response["messages"][0]["content"] + == "Hello, my phone number is {PHONE} and email is {EMAIL}" + ) print("✅ Masked output was applied even without masking flags enabled") @@ -654,10 +690,10 @@ async def test_bedrock_guardrail_response_pii_masking_non_streaming(): """Test that PII masking is applied to response content in non-streaming scenarios""" from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + # Create guardrail with response masking enabled guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", @@ -669,25 +705,29 @@ async def test_bedrock_guardrail_response_pii_masking_non_streaming(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" - }], - "assessments": [{ - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "CREDIT_DEBIT_CARD_NUMBER", - "match": "1234-5678-9012-3456", - "action": "ANONYMIZED" - }, - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED" - } - ] + "outputs": [ + { + "text": "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" } - }] + ], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "CREDIT_DEBIT_CARD_NUMBER", + "match": "1234-5678-9012-3456", + "action": "ANONYMIZED", + }, + { + "type": "PHONE", + "match": "+1 412 555 1212", + "action": "ANONYMIZED", + }, + ] + } + } + ], } # Create a mock response that contains PII @@ -697,15 +737,15 @@ async def test_bedrock_guardrail_response_pii_masking_non_streaming(): litellm.Choices( index=0, message=litellm.Message( - role="assistant", - content="My credit card number is 1234-5678-9012-3456 and my phone is +1 412 555 1212" + role="assistant", + content="My credit card number is 1234-5678-9012-3456 and my phone is +1 412 555 1212", ), - finish_reason="stop" + finish_reason="stop", ) ], created=1234567890, model="gpt-4o", - object="chat.completion" + object="chat.completion", ) request_data = { @@ -716,18 +756,23 @@ async def test_bedrock_guardrail_response_pii_masking_non_streaming(): } # Patch the async_handler.post method - with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # Call the post-call success hook await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - response=mock_response + response=mock_response, ) - + # Verify that the response content was masked - assert mock_response.choices[0].message.content == "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" + assert ( + mock_response.choices[0].message.content + == "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" + ) print("✓ Non-streaming response PII masking test passed") @@ -737,10 +782,10 @@ async def test_bedrock_guardrail_response_pii_masking_streaming(): from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ModelResponseStream - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + # Create guardrail with response masking enabled guardrail = BedrockGuardrail( guardrailIdentifier="test-guardrail", @@ -752,25 +797,25 @@ async def test_bedrock_guardrail_response_pii_masking_streaming(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "Sure! My email is {EMAIL} and SSN is {US_SSN}" - }], - "assessments": [{ - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "john@example.com", - "action": "ANONYMIZED" - }, - { - "type": "US_SSN", - "match": "123-45-6789", - "action": "ANONYMIZED" - } - ] + "outputs": [{"text": "Sure! My email is {EMAIL} and SSN is {US_SSN}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "john@example.com", + "action": "ANONYMIZED", + }, + { + "type": "US_SSN", + "match": "123-45-6789", + "action": "ANONYMIZED", + }, + ] + } } - }] + ], } # Create mock streaming chunks @@ -782,25 +827,27 @@ async def test_bedrock_guardrail_response_pii_masking_streaming(): litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="Sure! My email is "), - finish_reason=None + finish_reason=None, ) ], created=1234567890, model="gpt-4o", - object="chat.completion.chunk" + object="chat.completion.chunk", ), ModelResponseStream( id="test-id", choices=[ litellm.utils.StreamingChoices( index=0, - delta=litellm.utils.Delta(content="john@example.com and SSN is "), - finish_reason=None + delta=litellm.utils.Delta( + content="john@example.com and SSN is " + ), + finish_reason=None, ) ], created=1234567890, model="gpt-4o", - object="chat.completion.chunk" + object="chat.completion.chunk", ), ModelResponseStream( id="test-id", @@ -808,13 +855,13 @@ async def test_bedrock_guardrail_response_pii_masking_streaming(): litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="123-45-6789"), - finish_reason="stop" + finish_reason="stop", ) ], created=1234567890, model="gpt-4o", - object="chat.completion.chunk" - ) + object="chat.completion.chunk", + ), ] for chunk in chunks: yield chunk @@ -828,32 +875,37 @@ async def test_bedrock_guardrail_response_pii_masking_streaming(): } # Patch the async_handler.post method - with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # Call the streaming hook masked_stream = guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, response=mock_streaming_response(), - request_data=request_data + request_data=request_data, ) - + # Collect all chunks from the masked stream masked_chunks = [] async for chunk in masked_stream: masked_chunks.append(chunk) - + # Verify that we got chunks back assert len(masked_chunks) > 0 - + # Reconstruct the full response from chunks to verify masking full_content = "" for chunk in masked_chunks: - if hasattr(chunk, 'choices') and chunk.choices: - if hasattr(chunk.choices[0], 'delta') and chunk.choices[0].delta: - if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content: + if hasattr(chunk, "choices") and chunk.choices: + if hasattr(chunk.choices[0], "delta") and chunk.choices[0].delta: + if ( + hasattr(chunk.choices[0].delta, "content") + and chunk.choices[0].delta.content + ): full_content += chunk.choices[0].delta.content - + # Verify that the reconstructed content contains the masked PII assert "Sure! My email is {EMAIL} and SSN is {US_SSN}" == full_content print("✓ Streaming response PII masking test passed") @@ -862,64 +914,70 @@ async def test_bedrock_guardrail_response_pii_masking_streaming(): @pytest.mark.asyncio async def test_convert_to_bedrock_format_input_source(): """Test convert_to_bedrock_format with INPUT source and mock messages""" - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockRequest + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockRequest, + ) from unittest.mock import patch - + # Create the guardrail instance guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT" + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - + # Mock messages mock_messages = [ {"role": "user", "content": "Hello, how are you?"}, {"role": "assistant", "content": "I'm doing well, thank you!"}, - {"role": "user", "content": [ - {"type": "text", "text": "What's the weather like?"}, - {"type": "text", "text": "Is it sunny today?"} - ]} + { + "role": "user", + "content": [ + {"type": "text", "text": "What's the weather like?"}, + {"type": "text", "text": "Is it sunny today?"}, + ], + }, ] - + # Call the method - result = guardrail.convert_to_bedrock_format( - source="INPUT", - messages=mock_messages - ) - + result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) + # Verify the result structure assert isinstance(result, dict) assert result.get("source") == "INPUT" assert "content" in result assert isinstance(result.get("content"), list) - + # Verify content items expected_content_items = [ {"text": {"text": "Hello, how are you?"}}, {"text": {"text": "I'm doing well, thank you!"}}, {"text": {"text": "What's the weather like?"}}, - {"text": {"text": "Is it sunny today?"}} + {"text": {"text": "Is it sunny today?"}}, ] - + assert result.get("content") == expected_content_items print("✅ INPUT source test passed - result:", result) -@pytest.mark.asyncio +@pytest.mark.asyncio async def test_convert_to_bedrock_format_output_source(): """Test convert_to_bedrock_format with OUTPUT source and mock ModelResponse""" - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockRequest + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + ) + from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockRequest, + ) import litellm from unittest.mock import patch - - # Create the guardrail instance + + # Create the guardrail instance guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT" + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - + # Mock ModelResponse mock_response = litellm.ModelResponse( id="test-response-id", @@ -927,43 +985,40 @@ async def test_convert_to_bedrock_format_output_source(): litellm.Choices( index=0, message=litellm.Message( - role="assistant", - content="This is a test response from the model." + role="assistant", content="This is a test response from the model." ), - finish_reason="stop" + finish_reason="stop", ), litellm.Choices( - index=1, + index=1, message=litellm.Message( - role="assistant", - content="This is a second choice response." + role="assistant", content="This is a second choice response." ), - finish_reason="stop" - ) + finish_reason="stop", + ), ], created=1234567890, model="gpt-4o", - object="chat.completion" + object="chat.completion", ) - + # Call the method result = guardrail.convert_to_bedrock_format( - source="OUTPUT", - response=mock_response + source="OUTPUT", response=mock_response ) - + # Verify the result structure assert isinstance(result, dict) assert result.get("source") == "OUTPUT" assert "content" in result assert isinstance(result.get("content"), list) - + # Verify content items - should contain both choice contents expected_content_items = [ {"text": {"text": "This is a test response from the model."}}, - {"text": {"text": "This is a second choice response."}} + {"text": {"text": "This is a second choice response."}}, ] - + assert result.get("content") == expected_content_items print("✅ OUTPUT source test passed - result:", result) @@ -975,16 +1030,15 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ModelResponseStream import litellm - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + # Create guardrail instance guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT" + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - + # Mock streaming chunks that contain PII async def mock_streaming_response(): chunks = [ @@ -994,12 +1048,12 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="My email is "), - finish_reason=None + finish_reason=None, ) ], created=1234567890, model="gpt-4o", - object="chat.completion.chunk" + object="chat.completion.chunk", ), ModelResponseStream( id="test-id", @@ -1007,99 +1061,121 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="john@example.com"), - finish_reason="stop" + finish_reason="stop", ) ], created=1234567890, model="gpt-4o", - object="chat.completion.chunk" - ) + object="chat.completion.chunk", + ), ] for chunk in chunks: yield chunk - + # Mock Bedrock API response with PII masking mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "My email is {EMAIL}" - }], - "assessments": [{ - "sensitiveInformationPolicy": { - "piiEntities": [{ - "type": "EMAIL", - "match": "john@example.com", - "action": "ANONYMIZED" - }] + "outputs": [{"text": "My email is {EMAIL}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "john@example.com", + "action": "ANONYMIZED", + } + ] + } } - }] + ], } - + request_data = { "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "What's your email?"} - ], - "stream": True + "messages": [{"role": "user", "content": "What's your email?"}], + "stream": True, } - + # Track which bedrock API calls were made bedrock_calls = [] - + # Mock the make_bedrock_api_request method to track calls - async def mock_make_bedrock_api_request(source, messages=None, response=None, request_data=None): - bedrock_calls.append({ - "source": source, - "messages": messages, - "response": response, - "request_data": request_data - }) + async def mock_make_bedrock_api_request( + source, messages=None, response=None, request_data=None + ): + bedrock_calls.append( + { + "source": source, + "messages": messages, + "response": response, + "request_data": request_data, + } + ) # Return the mock bedrock response - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrailResponse + from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrailResponse, + ) + return BedrockGuardrailResponse(**mock_bedrock_response.json()) - + # Patch the bedrock API request method - with patch.object(guardrail, 'make_bedrock_api_request', side_effect=mock_make_bedrock_api_request): - + with patch.object( + guardrail, "make_bedrock_api_request", side_effect=mock_make_bedrock_api_request + ): + # Call the streaming hook result_generator = guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=mock_user_api_key_dict, response=mock_streaming_response(), - request_data=request_data + request_data=request_data, ) - + # Collect all chunks from the result result_chunks = [] async for chunk in result_generator: result_chunks.append(chunk) - + # Verify bedrock API calls were made - assert len(bedrock_calls) == 2, f"Expected 2 bedrock calls (INPUT and OUTPUT), got {len(bedrock_calls)}" - + assert ( + len(bedrock_calls) == 2 + ), f"Expected 2 bedrock calls (INPUT and OUTPUT), got {len(bedrock_calls)}" + # Find the OUTPUT call output_calls = [call for call in bedrock_calls if call["source"] == "OUTPUT"] - assert len(output_calls) == 1, f"Expected 1 OUTPUT call, got {len(output_calls)}" - + assert ( + len(output_calls) == 1 + ), f"Expected 1 OUTPUT call, got {len(output_calls)}" + output_call = output_calls[0] assert output_call["source"] == "OUTPUT" assert output_call["response"] is not None assert output_call["messages"] is None # OUTPUT calls don't need messages - + # Verify that the response content was masked # The streaming chunks should now contain the masked content full_content = "" for chunk in result_chunks: - if hasattr(chunk, 'choices') and chunk.choices: - if hasattr(chunk.choices[0], 'delta') and chunk.choices[0].delta.content: + if hasattr(chunk, "choices") and chunk.choices: + if ( + hasattr(chunk.choices[0], "delta") + and chunk.choices[0].delta.content + ): full_content += chunk.choices[0].delta.content - + # The content should be masked (contains {EMAIL} instead of john@example.com) - assert "{EMAIL}" in full_content, f"Expected masked content with {{EMAIL}}, got: {full_content}" - assert "john@example.com" not in full_content, f"Original email should be masked, got: {full_content}" - - print("✅ Post-call streaming hook test passed - OUTPUT source used for masking") + assert ( + "{EMAIL}" in full_content + ), f"Expected masked content with {{EMAIL}}, got: {full_content}" + assert ( + "john@example.com" not in full_content + ), f"Original email should be masked, got: {full_content}" + + print( + "✅ Post-call streaming hook test passed - OUTPUT source used for masking" + ) print(f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]}") print(f"✅ Final masked content: {full_content}") @@ -1110,13 +1186,12 @@ async def test_bedrock_guardrail_blocked_action_shows_output_text(): from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from fastapi import HTTPException - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT" + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) # Mock the Bedrock API response with BLOCKED action and output text @@ -1124,20 +1199,16 @@ async def test_bedrock_guardrail_blocked_action_shows_output_text(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "outputs": [ + "outputs": [{"text": "this violates litellm corporate guardrail policy"}], + "assessments": [ { - "text": "this violates litellm corporate guardrail policy" + "topicPolicy": { + "topics": [ + {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + ] + } } ], - "assessments": [{ - "topicPolicy": { - "topics": [{ - "name": "Sensitive Topic", - "type": "DENY", - "action": "BLOCKED" - }] - } - }] } request_data = { @@ -1148,32 +1219,36 @@ async def test_bedrock_guardrail_blocked_action_shows_output_text(): } # Patch the async_handler.post method - with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # This should raise HTTPException due to BLOCKED action with pytest.raises(HTTPException) as exc_info: await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - call_type="completion" + call_type="completion", ) - + # Verify the exception details exception = exc_info.value assert exception.status_code == 400 assert "detail" in exception.__dict__ - + # Check that the detail contains the expected structure detail = exception.detail assert isinstance(detail, dict) assert detail["error"] == "Violated guardrail policy" - + # Verify that the output text from both outputs is included expected_output_text = "this violates litellm corporate guardrail policy" assert detail["bedrock_guardrail_response"] == expected_output_text - - print("✅ BLOCKED action HTTPException test passed - output text properly included") + + print( + "✅ BLOCKED action HTTPException test passed - output text properly included" + ) @pytest.mark.asyncio @@ -1182,13 +1257,12 @@ async def test_bedrock_guardrail_blocked_action_empty_outputs(): from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from fastapi import HTTPException - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT" + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) # Mock the Bedrock API response with BLOCKED action but no outputs @@ -1197,15 +1271,15 @@ async def test_bedrock_guardrail_blocked_action_empty_outputs(): mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [], # Empty outputs - "assessments": [{ - "contentPolicy": { - "filters": [{ - "type": "VIOLENCE", - "confidence": "HIGH", - "action": "BLOCKED" - }] + "assessments": [ + { + "contentPolicy": { + "filters": [ + {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} + ] + } } - }] + ], } request_data = { @@ -1216,27 +1290,29 @@ async def test_bedrock_guardrail_blocked_action_empty_outputs(): } # Patch the async_handler.post method - with patch.object(guardrail.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # This should raise HTTPException due to BLOCKED action with pytest.raises(HTTPException) as exc_info: await guardrail.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - call_type="completion" + call_type="completion", ) - + # Verify the exception details exception = exc_info.value assert exception.status_code == 400 - + # Check that the detail contains the expected structure with empty output text detail = exception.detail assert isinstance(detail, dict) assert detail["error"] == "Violated guardrail policy" assert detail["bedrock_guardrail_response"] == "" # Empty string for no outputs - + print("✅ BLOCKED action with empty outputs test passed") @@ -1246,15 +1322,15 @@ async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): from unittest.mock import AsyncMock, MagicMock, patch from litellm.proxy._types import UserAPIKeyAuth from fastapi import HTTPException - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + # Test 1: disable_exception_on_block=False (default) - should raise exception guardrail_default = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", - disable_exception_on_block=False + disable_exception_on_block=False, ) # Mock the Bedrock API response with BLOCKED action @@ -1262,18 +1338,16 @@ async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "I can't provide that information." - }], - "assessments": [{ - "topicPolicy": { - "topics": [{ - "name": "Sensitive Topic", - "type": "DENY", - "action": "BLOCKED" - }] + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "topicPolicy": { + "topics": [ + {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + ] + } } - }] + ], } request_data = { @@ -1284,17 +1358,19 @@ async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): } # Patch the async_handler.post method - with patch.object(guardrail_default.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail_default.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # Should raise HTTPException when disable_exception_on_block=False with pytest.raises(HTTPException) as exc_info: await guardrail_default.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - call_type="completion" + call_type="completion", ) - + # Verify the exception details exception = exc_info.value assert exception.status_code == 400 @@ -1304,24 +1380,28 @@ async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): guardrail_disabled = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", - disable_exception_on_block=True + disable_exception_on_block=True, ) - with patch.object(guardrail_disabled.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail_disabled.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # Should NOT raise exception when disable_exception_on_block=True try: response = await guardrail_disabled.async_moderation_hook( data=request_data, user_api_key_dict=mock_user_api_key_dict, - call_type="completion" + call_type="completion", ) # Should succeed and return data (even though content was blocked) assert response is not None print("✅ No exception raised when disable_exception_on_block=True") except Exception as e: - pytest.fail(f"Should not raise exception when disable_exception_on_block=True, but got: {e}") + pytest.fail( + f"Should not raise exception when disable_exception_on_block=True, but got: {e}" + ) @pytest.mark.asyncio @@ -1332,10 +1412,10 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): from litellm.types.utils import ModelResponseStream from fastapi import HTTPException import litellm - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + # Mock streaming chunks that would normally trigger a block async def mock_streaming_response(): chunks = [ @@ -1344,13 +1424,15 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): choices=[ litellm.utils.StreamingChoices( index=0, - delta=litellm.utils.Delta(content="Here's how to make explosives: "), - finish_reason=None + delta=litellm.utils.Delta( + content="Here's how to make explosives: " + ), + finish_reason=None, ) ], created=1234567890, model="gpt-4o", - object="chat.completion.chunk" + object="chat.completion.chunk", ), ModelResponseStream( id="test-id", @@ -1358,62 +1440,62 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): litellm.utils.StreamingChoices( index=0, delta=litellm.utils.Delta(content="step 1, step 2..."), - finish_reason="stop" + finish_reason="stop", ) ], created=1234567890, model="gpt-4o", - object="chat.completion.chunk" - ) + object="chat.completion.chunk", + ), ] for chunk in chunks: yield chunk - + # Mock Bedrock API response with BLOCKED action mock_bedrock_response = MagicMock() mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{ - "text": "I can't provide that information." - }], - "assessments": [{ - "contentPolicy": { - "filters": [{ - "type": "VIOLENCE", - "confidence": "HIGH", - "action": "BLOCKED" - }] + "outputs": [{"text": "I can't provide that information."}], + "assessments": [ + { + "contentPolicy": { + "filters": [ + {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} + ] + } } - }] + ], } - + request_data = { "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Tell me how to make explosives"} - ], - "stream": True + "messages": [{"role": "user", "content": "Tell me how to make explosives"}], + "stream": True, } # Test 1: disable_exception_on_block=False (default) - should raise exception guardrail_default = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", - disable_exception_on_block=False + disable_exception_on_block=False, ) - with patch.object(guardrail_default.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail_default.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # Should raise exception during streaming processing with pytest.raises(HTTPException): - result_generator = guardrail_default.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data + result_generator = ( + guardrail_default.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=mock_streaming_response(), + request_data=request_data, + ) ) - + # Try to consume the generator - should raise exception async for chunk in result_generator: pass @@ -1422,31 +1504,40 @@ async def test_bedrock_guardrail_disable_exception_on_block_streaming(): guardrail_disabled = BedrockGuardrail( guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT", - disable_exception_on_block=True + disable_exception_on_block=True, ) - with patch.object(guardrail_disabled.async_handler, 'post', new_callable=AsyncMock) as mock_post: + with patch.object( + guardrail_disabled.async_handler, "post", new_callable=AsyncMock + ) as mock_post: mock_post.return_value = mock_bedrock_response - + # Should NOT raise exception when disable_exception_on_block=True try: - result_generator = guardrail_disabled.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data + result_generator = ( + guardrail_disabled.async_post_call_streaming_iterator_hook( + user_api_key_dict=mock_user_api_key_dict, + response=mock_streaming_response(), + request_data=request_data, + ) ) - + # Consume the generator - should succeed without exceptions result_chunks = [] async for chunk in result_generator: result_chunks.append(chunk) - + # Should have received chunks back even though content was blocked assert len(result_chunks) > 0 - print("✅ Streaming completed without exception when disable_exception_on_block=True") - + print( + "✅ Streaming completed without exception when disable_exception_on_block=True" + ) + except Exception as e: - pytest.fail(f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}") + pytest.fail( + f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}" + ) + @pytest.mark.asyncio async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): @@ -1455,16 +1546,15 @@ async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import ModelResponseStream import litellm - + # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - + # Create guardrail instance guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT" + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - + # Create a ModelResponse with tool calls (no text content) # This simulates a response where the LLM is making a tool call mock_response = litellm.ModelResponse( @@ -1479,34 +1569,33 @@ async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): litellm.utils.ChatCompletionMessageToolCall( id="tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", function=litellm.utils.Function( - name="top_song", - arguments='{"sign": "WZPZ"}' + name="top_song", arguments='{"sign": "WZPZ"}' ), - type="function" + type="function", ) - ] + ], ), - finish_reason="tool_calls" + finish_reason="tool_calls", ) ], created=1234567890, model="gpt-4o", - object="chat.completion" + object="chat.completion", ) - + data = { "model": "gpt-4o", "messages": [ {"role": "user", "content": "Hello"}, ], - } + } mock_user_api_key_dict = UserAPIKeyAuth() result = await guardrail.async_post_call_success_hook( data=data, - response=mock_response, + response=mock_response, user_api_key_dict=mock_user_api_key_dict, ) # If no error is raised and result is None, then the test passes assert result is None - print("✅ No output text in response test passed") \ No newline at end of file + print("✅ No output text in response test passed") diff --git a/tests/litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/litellm/llms/vertex_ai/text_to_speech/test_transformation.py new file mode 100644 index 00000000000..9a5888e8506 --- /dev/null +++ b/tests/litellm/llms/vertex_ai/text_to_speech/test_transformation.py @@ -0,0 +1,190 @@ +import json +import os +import sys +from unittest.mock import MagicMock, Mock, patch + +import httpx +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.llms.vertex_ai.text_to_speech.transformation import ( + VertexAITextToSpeechConfig, +) + + +class TestVertexAITextToSpeechConfig: + """Tests for VertexAITextToSpeechConfig transformation""" + + def test_get_complete_url(self): + """Test that get_complete_url returns the correct Google Cloud TTS API URL""" + config = VertexAITextToSpeechConfig() + + url = config.get_complete_url( + model="vertex_ai/chirp", + api_base=None, + litellm_params={}, + ) + + assert url == "https://texttospeech.googleapis.com/v1/text:synthesize" + + def test_get_complete_url_with_custom_api_base(self): + """Test that get_complete_url uses custom api_base when provided""" + config = VertexAITextToSpeechConfig() + + custom_url = "https://custom-tts-endpoint.example.com/v1/synthesize" + url = config.get_complete_url( + model="vertex_ai/chirp", + api_base=custom_url, + litellm_params={}, + ) + + assert url == custom_url + + @patch.object(VertexAITextToSpeechConfig, "_ensure_access_token") + @patch.object(VertexAITextToSpeechConfig, "_get_token_and_url") + def test_transform_text_to_speech_request_body( + self, mock_get_token, mock_ensure_token + ): + """Test that transform_text_to_speech_request generates correct request body""" + # Mock authentication + mock_ensure_token.return_value = ("mock-token", "test-project") + mock_get_token.return_value = ("mock-token", "mock-url") + + config = VertexAITextToSpeechConfig() + + # Test with voice dict in litellm_params (as set by dispatch) + result = config.transform_text_to_speech_request( + model="vertex_ai/chirp", + input="Hello, this is a test", + voice=None, + optional_params={ + "vertex_voice_dict": { + "languageCode": "en-US", + "name": "en-US-Chirp3-HD-Charon", + } + }, + litellm_params={ + "vertex_credentials": None, + "vertex_project": "test-project", + "vertex_location": "us-central1", + }, + headers={}, + ) + + # Verify request body structure + assert "dict_body" in result + request_body = result["dict_body"] + + assert "input" in request_body + assert request_body["input"] == {"text": "Hello, this is a test"} + + assert "voice" in request_body + assert request_body["voice"]["languageCode"] == "en-US" + assert request_body["voice"]["name"] == "en-US-Chirp3-HD-Charon" + + assert "audioConfig" in request_body + + # Verify headers contain auth + assert "headers" in result + assert "Authorization" in result["headers"] + + def test_voice_mapping_openai_to_vertex(self): + """Test that OpenAI voice names are correctly mapped to Vertex AI voices""" + config = VertexAITextToSpeechConfig() + + # Test the _map_voice_to_vertex_format helper + voice_str, voice_dict = config._map_voice_to_vertex_format("alloy") + + assert voice_str == "alloy" + assert voice_dict is not None + assert voice_dict["name"] == "en-US-Studio-O" + assert voice_dict["languageCode"] == "en-US" + + def test_voice_mapping_vertex_voice_passthrough(self): + """Test that Vertex AI voice names are passed through directly""" + config = VertexAITextToSpeechConfig() + + # Test with a Chirp3 HD voice + voice_str, voice_dict = config._map_voice_to_vertex_format( + "en-US-Chirp3-HD-Charon" + ) + + assert voice_str == "en-US-Chirp3-HD-Charon" + assert voice_dict is not None + assert voice_dict["name"] == "en-US-Chirp3-HD-Charon" + assert voice_dict["languageCode"] == "en-US" + + def test_voice_mapping_dict_passthrough(self): + """Test that voice dict is passed through unchanged""" + config = VertexAITextToSpeechConfig() + + voice_input = { + "languageCode": "de-DE", + "name": "de-DE-Chirp3-HD-Charon", + } + voice_str, voice_dict = config._map_voice_to_vertex_format(voice_input) + + assert voice_str is None + assert voice_dict == voice_input + + +@patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") +@patch.object(VertexAITextToSpeechConfig, "_ensure_access_token") +@patch.object(VertexAITextToSpeechConfig, "_get_token_and_url") +def test_litellm_speech_vertex_ai_chirp(mock_get_token, mock_ensure_token, mock_post): + """ + Test that litellm.speech(model="vertex_ai/chirp") sends the correct URL and request body + """ + # Mock authentication + mock_ensure_token.return_value = ("mock-token", "test-project") + mock_get_token.return_value = ("mock-token", "mock-url") + + # Mock HTTP response + mock_response = Mock(spec=httpx.Response) + mock_response.content = b'{"audioContent": "SGVsbG8gV29ybGQ="}' # base64 encoded "Hello World" + mock_response.status_code = 200 + mock_response.headers = {"content-type": "application/json"} + mock_response.json.return_value = {"audioContent": "SGVsbG8gV29ybGQ="} + mock_post.return_value = mock_response + + litellm.speech( + model="vertex_ai/chirp", + input="Hello, this is a test", + voice="en-US-Chirp3-HD-Charon", + vertex_project="test-project", + vertex_location="us-central1", + ) + + # Verify the HTTP call was made + mock_post.assert_called_once() + call_kwargs = mock_post.call_args.kwargs + + # Verify the URL is the Google Cloud TTS API + assert call_kwargs["url"] == "https://texttospeech.googleapis.com/v1/text:synthesize" + + # Verify request body structure + assert "data" in call_kwargs + request_body = json.loads(call_kwargs["data"]) + + # Verify input + assert "input" in request_body + assert request_body["input"] == {"text": "Hello, this is a test"} + + # Verify voice + assert "voice" in request_body + assert request_body["voice"]["name"] == "en-US-Chirp3-HD-Charon" + assert request_body["voice"]["languageCode"] == "en-US" + + # Verify audioConfig + assert "audioConfig" in request_body + + # Verify headers contain authorization + assert "headers" in call_kwargs + assert "Authorization" in call_kwargs["headers"] + assert call_kwargs["headers"]["Authorization"] == "Bearer mock-token" + + diff --git a/tests/litellm_utils_tests/test_cyberark.py b/tests/litellm_utils_tests/test_cyberark.py index 24c84687705..b94e5949534 100644 --- a/tests/litellm_utils_tests/test_cyberark.py +++ b/tests/litellm_utils_tests/test_cyberark.py @@ -53,13 +53,14 @@ async def test_cyberark_write_and_read_secret(): secret_value = f"test-value-{uuid.uuid4()}" # Mock sync httpx client (for auth, ensure variable exists, sync read) + # The _get_httpx_client returns an HTTPHandler with a .client property mock_sync_client = MagicMock() - # Auth response - mock_sync_client.post.return_value = create_mock_response( + # Auth response - note: the actual client is accessed via .client property + mock_sync_client.client.post.return_value = create_mock_response( status_code=200, text="mock-token" ) # Sync read response - mock_sync_client.get.return_value = create_mock_response( + mock_sync_client.client.get.return_value = create_mock_response( status_code=200, text=secret_value ) @@ -123,9 +124,10 @@ async def test_cyberark_rotate_secret(): current_value = {"value": initial_key_value} # Mock sync httpx client (for auth, ensure variable exists, sync reads) + # The _get_httpx_client returns an HTTPHandler with a .client property mock_sync_client = MagicMock() - # Auth response - mock_sync_client.post.return_value = create_mock_response( + # Auth response - note: the actual client is accessed via .client property + mock_sync_client.client.post.return_value = create_mock_response( status_code=200, text="mock-token" ) @@ -133,7 +135,7 @@ async def test_cyberark_rotate_secret(): def get_mock_sync_read_response(*args, **kwargs): return create_mock_response(status_code=200, text=current_value["value"]) - mock_sync_client.get.side_effect = get_mock_sync_read_response + mock_sync_client.client.get.side_effect = get_mock_sync_read_response # Mock async httpx client (for async writes and reads) mock_async_client = AsyncMock() @@ -228,9 +230,10 @@ async def test_cyberark_rotate_secret_with_new_alias(): secrets_store = {} # Mock sync httpx client (for auth, ensure variable exists, sync reads) + # The _get_httpx_client returns an HTTPHandler with a .client property mock_sync_client = MagicMock() - # Auth response - mock_sync_client.post.return_value = create_mock_response( + # Auth response - note: the actual client is accessed via .client property + mock_sync_client.client.post.return_value = create_mock_response( status_code=200, text="mock-token" ) @@ -243,7 +246,7 @@ async def test_cyberark_rotate_secret_with_new_alias(): return create_mock_response(status_code=200, text=secret_val) return create_mock_response(status_code=404, text="Not found") - mock_sync_client.get.side_effect = get_mock_sync_read + mock_sync_client.client.get.side_effect = get_mock_sync_read # Mock async httpx client (for async writes and reads) mock_async_client = AsyncMock() diff --git a/tests/llm_translation/test_bedrock_completion.py b/tests/llm_translation/test_bedrock_completion.py index f43e939c681..bd08d4444f6 100644 --- a/tests/llm_translation/test_bedrock_completion.py +++ b/tests/llm_translation/test_bedrock_completion.py @@ -3531,3 +3531,336 @@ def test_bedrock_openai_imported_model(): # Check max_tokens and temperature assert request_body["max_tokens"] == 300 assert request_body["temperature"] == 0.5 + +def test_bedrock_openai_provider_detection(): + """ + Test that the OpenAI provider is correctly detected from model strings. + """ + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + # Test various OpenAI model formats + test_cases = [ + "openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/abc123", + "bedrock/openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/xyz789", + ] + + for model in test_cases: + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + assert provider == "openai", f"Failed for model: {model}, got provider: {provider}" + print(f"✓ Provider detection works for: {model}") + + +def test_bedrock_openai_model_id_extraction(): + """ + Test that the model ID (ARN) is correctly extracted and encoded for OpenAI models. + """ + from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM + + model = "openai/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-model-123" + provider = BaseAWSLLM.get_bedrock_invoke_provider(model) + + model_id = BaseAWSLLM.get_bedrock_model_id( + model=model, + provider=provider, + optional_params={} + ) + + # The ARN should be double URL encoded + assert "arn" in model_id + assert "imported-model" in model_id + print(f"✓ Model ID extracted and encoded: {model_id}") + + +def test_bedrock_openai_convert_messages_to_prompt(): + """ + Test that convert_messages_to_prompt returns empty string for OpenAI models. + """ + from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM + + bedrock_llm = BedrockLLM() + messages = [ + {"role": "system", "content": "You are helpful"}, + {"role": "user", "content": "Hello"} + ] + + prompt, chat_history = bedrock_llm.convert_messages_to_prompt( + model="test-model", + messages=messages, + provider="openai", + custom_prompt_dict={} + ) + + # OpenAI models use messages directly, no prompt conversion + assert prompt == "" + assert chat_history is None + print("✓ convert_messages_to_prompt returns empty for OpenAI") + + +def test_bedrock_openai_response_parsing(): + """ + Test that OpenAI responses are correctly parsed. + """ + from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM + from litellm import ModelResponse + from unittest.mock import Mock + import json + + bedrock_llm = BedrockLLM() + + # Mock OpenAI-style response + openai_response = { + "choices": [ + { + "message": { + "content": "The capital of France is Paris.", + "role": "assistant" + }, + "finish_reason": "stop", + "index": 0 + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 8, + "total_tokens": 18 + } + } + + mock_response = Mock() + mock_response.json.return_value = openai_response + mock_response.text = json.dumps(openai_response) + mock_response.status_code = 200 + mock_response.headers = {} + + model_response = ModelResponse() + mock_logging = Mock() + + result = bedrock_llm.process_response( + model="openai/arn:aws:bedrock:us-east-1:123:imported-model/test", + response=mock_response, + model_response=model_response, + stream=False, + logging_obj=mock_logging, + optional_params={}, + api_key="", + data={}, + messages=[{"role": "user", "content": "What is the capital of France?"}], + print_verbose=lambda x: None, + encoding=None + ) + + # Verify response content + assert result.choices[0].message.content == "The capital of France is Paris." + assert result.choices[0].finish_reason == "stop" + + # Verify usage + assert result.usage.prompt_tokens == 10 + assert result.usage.completion_tokens == 8 + assert result.usage.total_tokens == 18 + + print("✓ OpenAI response parsing works correctly") + + +def test_bedrock_openai_request_transformation(): + """ + Test that the request is correctly transformed for OpenAI models. + """ + from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import AmazonInvokeConfig + + config = AmazonInvokeConfig() + + model = "openai/arn:aws:bedrock:us-east-1:123:imported-model/test" + messages = [ + {"role": "system", "content": "You are helpful"}, + {"role": "user", "content": "Hello"} + ] + + optional_params = { + "max_tokens": 100, + "temperature": 0.7, + "top_p": 0.9, + "stream": False + } + + litellm_params = {} + headers = {} + + with patch.object(config, 'get_bedrock_invoke_provider', return_value="openai"): + result = config.transform_request( + model=model, + messages=messages, + optional_params=optional_params.copy(), + litellm_params=litellm_params, + headers=headers + ) + + # Verify the request uses messages format (not prompt) + assert "messages" in result + assert len(result["messages"]) == 2 + assert result["messages"][0]["role"] == "system" + assert result["messages"][1]["role"] == "user" + + # Verify parameters are included + assert "max_tokens" in result + assert "temperature" in result + + print("✓ Request transformation works correctly") + + +def test_bedrock_openai_parameter_filtering(): + """ + Test that only supported OpenAI parameters are included in the request. + """ + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig + + config = AmazonBedrockOpenAIConfig() + model = "test-model" + + supported_params = config.get_supported_openai_params(model=model) + + # Verify common OpenAI parameters are supported + assert "max_tokens" in supported_params + assert "temperature" in supported_params + assert "top_p" in supported_params + assert "stream" in supported_params + assert "stop" in supported_params + + print(f"✓ Parameter filtering supports: {len(supported_params)} parameters") + print(f" Supported params: {supported_params}") + + +def test_bedrock_openai_route_detection(): + """ + Test that the OpenAI route is correctly detected. + """ + from litellm.llms.bedrock.common_utils import BedrockModelInfo + + test_cases = [ + ("openai/arn:aws:bedrock:us-east-1:123:imported-model/test", "openai"), + ("bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test", "openai"), + ] + + for model, expected_route in test_cases: + route = BedrockModelInfo.get_bedrock_route(model) + assert route == expected_route, f"Failed for model: {model}, got route: {route}" + print(f"✓ Route detection works for: {model} -> {route}") + + +def test_bedrock_openai_explicit_route_check(): + """ + Test the explicit OpenAI route checker helper method. + """ + from litellm.llms.bedrock.common_utils import BedrockModelInfo + + # Test with openai/ prefix + assert BedrockModelInfo._explicit_openai_route("openai/arn:aws:bedrock:us-east-1:123:imported-model/test") is True + assert BedrockModelInfo._explicit_openai_route("bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test") is True + + # Test without openai/ prefix + assert BedrockModelInfo._explicit_openai_route("anthropic.claude-3-sonnet") is False + assert BedrockModelInfo._explicit_openai_route("arn:aws:bedrock:us-east-1:123:imported-model/test") is False + + print("✓ Explicit route check works correctly") + + +def test_bedrock_openai_config_initialization(): + """ + Test that AmazonBedrockOpenAIConfig can be properly initialized. + """ + from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import AmazonBedrockOpenAIConfig + + config = AmazonBedrockOpenAIConfig() + + # Verify it has the necessary methods + assert hasattr(config, 'get_supported_openai_params') + assert hasattr(config, 'transform_request') + assert hasattr(config, 'transform_response') + assert hasattr(config, 'map_openai_params') + + print("✓ AmazonBedrockOpenAIConfig initializes correctly") + + +def test_bedrock_openai_multiple_message_types(): + """ + Test that various message content types are handled correctly. + """ + from litellm.llms.custom_httpx.http_handler import HTTPHandler + + client = HTTPHandler() + + # Test with mixed content types + messages = [ + {"role": "system", "content": "You are helpful"}, + {"role": "user", "content": "Simple text message"}, + { + "role": "user", + "content": [ + {"type": "text", "text": "Complex message with text"}, + {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,iVBORw0KGg"}} + ] + } + ] + + with patch.object(client, "post") as mock_post: + try: + response = completion( + model="bedrock/openai/arn:aws:bedrock:us-east-1:123:imported-model/test", + messages=messages, + max_tokens=50, + client=client, + ) + except Exception as e: + pass + + # Verify the request was made + if mock_post.called: + request_body = json.loads(mock_post.call_args.kwargs["data"]) + + # Verify messages are preserved + assert "messages" in request_body + assert len(request_body["messages"]) == 3 + + # Verify mixed content is handled + assert isinstance(request_body["messages"][2]["content"], list) + + print("✓ Multiple message types handled correctly") + + +def test_bedrock_openai_error_handling(): + """ + Test that errors from OpenAI models are properly handled. + """ + from litellm.llms.bedrock.chat.invoke_handler import BedrockLLM + from litellm import ModelResponse + from litellm.llms.bedrock.common_utils import BedrockError + from unittest.mock import Mock + import json + + bedrock_llm = BedrockLLM() + + # Mock error response + mock_response = Mock() + mock_response.json.side_effect = Exception("Invalid JSON") + mock_response.text = "Invalid response" + mock_response.status_code = 422 + + model_response = ModelResponse() + mock_logging = Mock() + + with pytest.raises(BedrockError) as exc_info: + bedrock_llm.process_response( + model="openai/arn:aws:bedrock:us-east-1:123:imported-model/test", + response=mock_response, + model_response=model_response, + stream=False, + logging_obj=mock_logging, + optional_params={}, + api_key="", + data={}, + messages=[], + print_verbose=lambda x: None, + encoding=None + ) + + assert exc_info.value.status_code == 422 + print("✓ Error handling works correctly") diff --git a/tests/llm_translation/test_groq.py b/tests/llm_translation/test_groq.py index fc84580404c..33c1a425870 100644 --- a/tests/llm_translation/test_groq.py +++ b/tests/llm_translation/test_groq.py @@ -9,6 +9,7 @@ import pytest # ) # noqa # ) # Adds the parent directory to the system path +import litellm from base_llm_unit_tests import BaseLLMChatTest from litellm.llms.groq.chat.transformation import GroqChatConfig @@ -30,3 +31,136 @@ class TestGroq(BaseLLMChatTest): """Test that reasoning_effort is in the list of supported parameters for Groq""" supported_params = GroqChatConfig().get_supported_openai_params(model=model) assert "reasoning_effort" in supported_params + + +class TestGroqStructuredOutputs: + """ + Tests for Groq structured outputs handling. + Related issues: + - https://github.com/BerriAI/litellm/issues/11001 + - https://github.com/openai/openai-agents-python/issues/2140 + """ + + def test_structured_output_with_tools_raises_error_for_non_native_models(self): + """ + Test that using structured outputs + tools with models that don't support + native json_schema raises a clear error message. + + Groq does not support structured outputs + tools together. + See: https://console.groq.com/docs/structured-outputs + "Streaming and tool use are not currently supported with Structured Outputs" + """ + config = GroqChatConfig() + + # Model that doesn't support native json_schema + model = "llama-3.3-70b-versatile" + + non_default_params = { + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test", + "schema": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"] + } + } + }, + "tools": [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": {"type": "object", "properties": {}} + } + } + ] + } + + with pytest.raises(litellm.BadRequestError) as exc_info: + config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + drop_params=False, + ) + + assert "does not support native structured outputs" in str(exc_info.value) + assert "incompatible with user-provided tools" in str(exc_info.value) + + def test_structured_output_without_tools_uses_workaround_for_non_native_models(self): + """ + Test that structured outputs without tools works using the json_tool_call workaround + for models that don't support native json_schema. + """ + config = GroqChatConfig() + + model = "llama-3.3-70b-versatile" + + non_default_params = { + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test", + "schema": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"] + } + } + } + } + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + drop_params=False, + ) + + # Should use the workaround (json_tool_call) + assert "tools" in result + assert len(result["tools"]) == 1 + assert result["tools"][0]["function"]["name"] == "json_tool_call" + assert result["tool_choice"]["function"]["name"] == "json_tool_call" + assert result.get("json_mode") is True + + def test_structured_output_passes_through_for_native_models(self): + """ + Test that structured outputs pass through directly for models that + support native json_schema (e.g., gpt-oss-120b). + """ + config = GroqChatConfig() + + # Model that supports native json_schema + model = "openai/gpt-oss-120b" + + non_default_params = { + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "test", + "schema": { + "type": "object", + "properties": {"name": {"type": "string"}}, + "required": ["name"] + } + } + } + } + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params={}, + model=model, + drop_params=False, + ) + + # Should NOT use the workaround - response_format should pass through + # The workaround sets json_mode=True, so if it's not set, we know it passed through + assert result.get("json_mode") is not True + # Should not have the json_tool_call tool + if "tools" in result: + tool_names = [t.get("function", {}).get("name") for t in result["tools"]] + assert "json_tool_call" not in tool_names diff --git a/tests/local_testing/test_gemini_reasoning_content.py b/tests/local_testing/test_gemini_reasoning_content.py new file mode 100644 index 00000000000..7e516ae8439 --- /dev/null +++ b/tests/local_testing/test_gemini_reasoning_content.py @@ -0,0 +1,20 @@ +import json +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexGeminiConfig + + +def test_empty_part_does_not_create_thinking_block(): + parts = [{"text": "", "thoughtSignature": "sig-1"}] + config = VertexGeminiConfig() + thinking_blocks = config._extract_thinking_blocks_from_parts(parts) + assert thinking_blocks == [] + + +def test_non_empty_part_creates_thinking_block(): + parts = [{"text": "Some thinking", "thoughtSignature": "sig-2"}] + config = VertexGeminiConfig() + thinking_blocks = config._extract_thinking_blocks_from_parts(parts) + assert len(thinking_blocks) == 1 + block = thinking_blocks[0] + # thinking should be valid JSON containing the text + parsed = json.loads(block["thinking"]) if isinstance(block["thinking"], str) else None + assert parsed is not None and parsed.get("text") == "Some thinking" diff --git a/tests/local_testing/test_get_llm_provider.py b/tests/local_testing/test_get_llm_provider.py index 261f9bd0eac..4df910bc73e 100644 --- a/tests/local_testing/test_get_llm_provider.py +++ b/tests/local_testing/test_get_llm_provider.py @@ -137,6 +137,9 @@ def test_default_api_base(): # Get the API base for the given provider if provider == "github_copilot": continue + # Skip ragflow as it requires specific model format: ragflow/chat/{id}/{model} or ragflow/agent/{id}/{model} + if provider == "ragflow": + continue _, _, _, api_base = _get_openai_compatible_provider_info( model=f"{provider}/*", api_base=None, api_key=None, dynamic_api_key=None ) diff --git a/tests/logging_callback_tests/test_alerting.py b/tests/logging_callback_tests/test_alerting.py index b9ecfaeb3f0..ac7f5cd6aa1 100644 --- a/tests/logging_callback_tests/test_alerting.py +++ b/tests/logging_callback_tests/test_alerting.py @@ -477,6 +477,7 @@ async def test_send_daily_reports_all_zero_or_none(): "token_budget", "user_budget", "team_budget", + "organization_budget", "proxy_budget", "projected_limit_exceeded", ], @@ -514,6 +515,7 @@ async def test_send_token_budget_crossed_alerts(alerting_type): "token_budget", "user_budget", "team_budget", + "organization_budget", "proxy_budget", "projected_limit_exceeded", ], diff --git a/tests/logging_callback_tests/test_datadog.py b/tests/logging_callback_tests/test_datadog.py index 13125aa4952..c877f34ac03 100644 --- a/tests/logging_callback_tests/test_datadog.py +++ b/tests/logging_callback_tests/test_datadog.py @@ -633,11 +633,14 @@ async def test_datadog_message_redaction(): def test_datadog_agent_configuration(): """ - Test that DataDog logger correctly configures agent endpoint when DD_AGENT_HOST is set + Test that DataDog logger correctly configures agent endpoint when LITELLM_DD_AGENT_HOST is set. + + Note: We use LITELLM_DD_AGENT_HOST instead of DD_AGENT_HOST to avoid conflicts + with ddtrace which automatically sets DD_AGENT_HOST for APM tracing. """ test_env = { - "DD_AGENT_HOST": "localhost", - "DD_AGENT_PORT": "10518", + "LITELLM_DD_AGENT_HOST": "localhost", + "LITELLM_DD_AGENT_PORT": "10518", } # Remove DD_SITE and DD_API_KEY to verify they're not required for agent mode @@ -654,4 +657,40 @@ def test_datadog_agent_configuration(): assert dd_logger.intake_url == "http://localhost:10518/api/v2/logs", f"Expected agent URL, got {dd_logger.intake_url}" # Verify DD_API_KEY is optional (can be None) - assert dd_logger.DD_API_KEY is None or isinstance(dd_logger.DD_API_KEY, str) \ No newline at end of file + assert dd_logger.DD_API_KEY is None or isinstance(dd_logger.DD_API_KEY, str) + + +def test_datadog_ignores_ddtrace_agent_host(): + """ + Regression test: Ensure DD_AGENT_HOST set by ddtrace doesn't interfere with LiteLLM logging. + + When users have ddtrace installed for APM tracing, it automatically sets DD_AGENT_HOST. + LiteLLM should ignore DD_AGENT_HOST and only use LITELLM_DD_AGENT_HOST for agent mode. + + This prevents the 404 error when ddtrace's DD_AGENT_HOST points to an APM endpoint + that doesn't support /api/v2/logs. + + Regression test for: https://github.com/BerriAI/litellm/issues/16379 + """ + test_env = { + # User's explicit config for LiteLLM logging (direct API) + "DD_API_KEY": "fake-api-key", + "DD_SITE": "us5.datadoghq.com", + # ddtrace automatically sets these for APM tracing + "DD_AGENT_HOST": "10.176.100.40", + "DD_AGENT_PORT": "8126", + } + + with patch.dict(os.environ, test_env, clear=False): + with patch("asyncio.create_task"): + dd_logger = DataDogLogger() + + # Verify direct API endpoint is used (DD_AGENT_HOST should be ignored) + expected_url = "https://http-intake.logs.us5.datadoghq.com/api/v2/logs" + assert dd_logger.intake_url == expected_url, ( + f"Expected direct API URL '{expected_url}', got '{dd_logger.intake_url}'. " + "DD_AGENT_HOST (set by ddtrace) should be ignored - only LITELLM_DD_AGENT_HOST should trigger agent mode." + ) + + # Verify API key is set correctly + assert dd_logger.DD_API_KEY == "fake-api-key" \ No newline at end of file diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/login_to_ui.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/login_to_ui.spec.ts index f691389de53..e5a397a6a66 100644 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/login_to_ui.spec.ts +++ b/tests/proxy_admin_ui_tests/e2e_ui_tests/login_to_ui.spec.ts @@ -11,15 +11,20 @@ import { test, expect } from "@playwright/test"; test("admin login test", async ({ page }) => { // Go to the specified URL await page.goto("http://localhost:4000/ui"); + await page.waitForLoadState("networkidle"); + + await page.screenshot({ path: "test-results/login_before.png" }); // Enter "admin" in the username input field - await page.fill('input[name="username"]', "admin"); + await page.fill('input[placeholder="Enter your username"]', "admin"); // Enter "gm" in the password input field - await page.fill('input[name="password"]', "gm"); + await page.fill('input[placeholder="Enter your password"]', "gm"); + + page.screenshot({ path: "test-results/login_after_inputs.png" }); // Optionally, you can add an assertion to verify the login button is enabled - const loginButton = page.locator('input[type="submit"]'); + const loginButton = page.getByRole("button", { name: "Login" }); await expect(loginButton).toBeEnabled(); // Optionally, you can click the login button to submit the form diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/require_auth_for_dashboard.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/require_auth_for_dashboard.spec.ts index 18bc91c0977..4e4bd2fcd93 100644 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/require_auth_for_dashboard.spec.ts +++ b/tests/proxy_admin_ui_tests/e2e_ui_tests/require_auth_for_dashboard.spec.ts @@ -2,15 +2,19 @@ import { test, expect } from "@playwright/test"; test.describe("Authentication Checks", () => { - test("should redirect unauthenticated user from a protected page", async ({ page }) => { + test("should redirect unauthenticated user from a protected page", async ({ + page, + }) => { test.setTimeout(30000); page.on("console", (msg) => console.log("PAGE LOG:", msg.text())); const protectedPageUrl = "http://localhost:4000/ui?page=llm-playground"; - const expectedRedirectUrl = "http://localhost:4000/sso/key/generate"; + const expectedRedirectUrl = "http://localhost:4000/ui/login/"; - console.log(`Attempting to navigate to protected page: ${protectedPageUrl}`); + console.log( + `Attempting to navigate to protected page: ${protectedPageUrl}` + ); await page.goto(protectedPageUrl); @@ -20,7 +24,9 @@ test.describe("Authentication Checks", () => { await page.waitForURL(expectedRedirectUrl, { timeout: 10000 }); console.log(`Waited for URL. Current URL is now: ${page.url()}`); } catch (error) { - console.error(`Timeout waiting for URL: ${expectedRedirectUrl}. Current URL: ${page.url()}`); + console.error( + `Timeout waiting for URL: ${expectedRedirectUrl}. Current URL: ${page.url()}` + ); await page.screenshot({ path: "redirect-fail-screenshot.png" }); throw error; } diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/search_users.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/search_users.spec.ts index 7b9da6a27db..d72c44ab8cc 100644 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/search_users.spec.ts +++ b/tests/proxy_admin_ui_tests/e2e_ui_tests/search_users.spec.ts @@ -21,17 +21,22 @@ test("user search test", async ({ page }) => { // Login first await page.goto("http://localhost:4000/ui"); + await page.waitForLoadState("networkidle"); console.log("Navigated to login page"); + page.screenshot({ path: "test-results/search_users_before_login.png" }); + // Wait for login form to be visible - await page.waitForSelector('input[name="username"]', { timeout: 10000 }); + await page.waitForSelector('input[placeholder="Enter your username"]', { + timeout: 10000, + }); console.log("Login form is visible"); - await page.fill('input[name="username"]', "admin"); - await page.fill('input[name="password"]', "gm"); + await page.fill('input[placeholder="Enter your username"]', "admin"); + await page.fill('input[placeholder="Enter your password"]', "gm"); console.log("Filled login credentials"); - const loginButton = page.locator('input[type="submit"]'); + const loginButton = page.getByRole("button", { name: "Login" }); await expect(loginButton).toBeEnabled(); await loginButton.click(); console.log("Clicked login button"); @@ -128,17 +133,20 @@ test("user filter test", async ({ page }) => { // Login first await page.goto("http://localhost:4000/ui"); + await page.waitForLoadState("networkidle"); console.log("Navigated to login page"); // Wait for login form to be visible - await page.waitForSelector('input[name="username"]', { timeout: 10000 }); + await page.waitForSelector('input[placeholder="Enter your username"]', { + timeout: 10000, + }); console.log("Login form is visible"); - await page.fill('input[name="username"]', "admin"); - await page.fill('input[name="password"]', "gm"); + await page.fill('input[placeholder="Enter your username"]', "admin"); + await page.fill('input[placeholder="Enter your password"]', "gm"); console.log("Filled login credentials"); - const loginButton = page.locator('input[type="submit"]'); + const loginButton = page.getByRole("button", { name: "Login" }); await expect(loginButton).toBeEnabled(); await loginButton.click(); console.log("Clicked login button"); diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts index 1d263e50511..2aae9e2bb54 100644 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts +++ b/tests/proxy_admin_ui_tests/e2e_ui_tests/view_internal_user.spec.ts @@ -7,15 +7,18 @@ import { test, expect } from "@playwright/test"; test("view internal user page", async ({ page }) => { // Go to the specified URL await page.goto("http://localhost:4000/ui"); + await page.waitForLoadState("networkidle"); + + page.screenshot({ path: "test-results/view_internal_user_before_login.png" }); // Enter "admin" in the username input field - await page.fill('input[name="username"]', "admin"); + await page.fill('input[placeholder="Enter your username"]', "admin"); // Enter "gm" in the password input field - await page.fill('input[name="password"]', "gm"); + await page.fill('input[placeholder="Enter your password"]', "gm"); // Click the login button - const loginButton = page.locator('input[type="submit"]'); + const loginButton = page.getByRole("button", { name: "Login" }); await expect(loginButton).toBeEnabled(); await loginButton.click(); @@ -34,10 +37,10 @@ test("view internal user page", async ({ page }) => { // The UI renders badges in each row - we just verify the column structure exists const rowCount = await page.locator("tbody tr").count(); expect(rowCount).toBeGreaterThan(0); - - // Verify table headers are present (including API Keys column) - const apiKeysHeader = page.locator("th", { hasText: "API Keys" }); - await expect(apiKeysHeader).toBeVisible(); + + const userIdHeader = page.locator("th", { hasText: "User ID" }); + page.screenshot({ path: "user_id_header.png" }); + await expect(userIdHeader).toBeVisible(); // test pagination // Wait for pagination controls to be visible diff --git a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_user_info.spec.ts b/tests/proxy_admin_ui_tests/e2e_ui_tests/view_user_info.spec.ts index 01eadc9ad1a..adda3088f12 100644 --- a/tests/proxy_admin_ui_tests/e2e_ui_tests/view_user_info.spec.ts +++ b/tests/proxy_admin_ui_tests/e2e_ui_tests/view_user_info.spec.ts @@ -2,22 +2,56 @@ import { test, expect } from "@playwright/test"; import { loginToUI } from "../utils/login"; test.describe("User Info View", () => { - test.beforeEach(async ({ page }) => { - await loginToUI(page); - // Navigate to users page - await page.goto("http://localhost:4000/ui?page=users"); - }); - test("should display user info when clicking on user ID", async ({ page, }) => { + await page.goto("http://localhost:4000/ui"); + await page.waitForLoadState("networkidle"); + + page.screenshot({ + path: "test-results/view_user_info_before_login.png", + }); + + // Enter "admin" in the username input field + await page.fill('input[placeholder="Enter your username"]', "admin"); + page.screenshot({ + path: "test-results/view_user_info_after_username_input.png", + }); + + // Enter "gm" in the password input field + await page.fill('input[placeholder="Enter your password"]', "gm"); + page.screenshot({ + path: "test-results/view_user_info_after_password_input.png", + }); + + // Click the login button + const loginButton = page.getByRole("button", { name: "Login" }); + await expect(loginButton).toBeEnabled(); + await loginButton.click(); + page.screenshot({ + path: "test-results/view_user_info_after_login_button_click.png", + }); + + // Wait for navigation to complete and dashboard to load + await page.waitForLoadState("networkidle"); + const tabElement = page.locator("span.ant-menu-title-content", { + hasText: "Internal User", + }); + await tabElement.click(); + page.screenshot({ + path: "test-results/view_user_info_after_internal_user_tab_click.png", + }); // Wait for loading state to disappear await page.waitForSelector('text="🚅 Loading users..."', { state: "hidden", + timeout: 10000, }); + page.screenshot({ path: "test-results/view_user_info_after_loading.png" }); // Wait for users table to load await page.waitForSelector("table"); - + page.screenshot({ + path: "test-results/view_user_info_after_table_load.png", + }); // Get the first user ID cell const firstUserIdCell = page.locator( "table tbody tr:first-child td:first-child" @@ -27,6 +61,7 @@ test.describe("User Info View", () => { // Click on the user ID await firstUserIdCell.click(); + await page.waitForLoadState("networkidle"); // Check for tabs await expect(page.locator('button:has-text("Overview")')).toBeVisible({ diff --git a/tests/proxy_admin_ui_tests/utils/login.ts b/tests/proxy_admin_ui_tests/utils/login.ts index e8755089976..25858d9f570 100644 --- a/tests/proxy_admin_ui_tests/utils/login.ts +++ b/tests/proxy_admin_ui_tests/utils/login.ts @@ -3,17 +3,21 @@ import { Page, expect } from "@playwright/test"; export async function loginToUI(page: Page) { // Login first await page.goto("http://localhost:4000/ui"); + await page.waitForLoadState("networkidle"); console.log("Navigated to login page"); + page.screenshot({ path: "test-results/login_utils_before.png" }); // Wait for login form to be visible - await page.waitForSelector('input[name="username"]', { timeout: 10000 }); + await page.waitForSelector('input[placeholder="Enter your username"]', { + timeout: 10000, + }); console.log("Login form is visible"); - await page.fill('input[name="username"]', "admin"); - await page.fill('input[name="password"]', "gm"); + await page.fill('input[placeholder="Enter your username"]', "admin"); + await page.fill('input[placeholder="Enter your password"]', "gm"); console.log("Filled login credentials"); - const loginButton = page.locator('input[type="submit"]'); + const loginButton = page.getByRole("button", { name: "Login" }); await expect(loginButton).toBeEnabled(); await loginButton.click(); console.log("Clicked login button"); diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 6dad7cb08d0..dc34a50f87e 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -175,6 +175,50 @@ def test_chat_completion(mock_acompletion, client_no_auth): pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") +def test_chat_completion_malformed_messages_returns_400(client_no_auth): + """ + Test that malformed messages (strings instead of dicts) return 400 instead of 500. + + This test verifies that when a client sends messages as raw strings instead of + {role, content} objects, LiteLLM returns a 400 invalid_request_error instead + of a 500 Internal Server Error. + """ + global headers + try: + # Test data with malformed messages (string instead of dict) + test_data = { + "model": "gpt-3.5-turbo", + "messages": ["hi how are you"], # Invalid: should be [{"role": "user", "content": "hi how are you"}] + } + + print("testing proxy server with malformed messages") + response = client_no_auth.post("/v1/chat/completions", json=test_data, headers=headers) + + print(f"response status: {response.status_code}") + print(f"response text: {response.text}") + + # Should return 400, not 500 + assert response.status_code == 400, f"Expected 400, got {response.status_code}. Response: {response.text}" + + # Verify error format + result = response.json() + assert "error" in result, "Response should contain 'error' key" + error = result["error"] + + # Verify error type and message + assert error.get("type") == "invalid_request_error" or error.get("type") is None, \ + f"Expected invalid_request_error or None, got {error.get('type')}" + assert error.get("code") == "400" or error.get("code") == 400, \ + f"Expected code 400, got {error.get('code')}" + + # Error message should indicate invalid request format + error_message = error.get("message", "") + assert len(error_message) > 0, "Error message should not be empty" + + except Exception as e: + pytest.fail(f"LiteLLM Proxy test failed. Exception - {str(e)}") + + def test_get_settings_request_timeout(client_no_auth): """ When no timeout is set, it should use the litellm.request_timeout value diff --git a/tests/test_litellm/integrations/test_weave_otel.py b/tests/test_litellm/integrations/test_weave_otel.py new file mode 100644 index 00000000000..440c0888512 --- /dev/null +++ b/tests/test_litellm/integrations/test_weave_otel.py @@ -0,0 +1,147 @@ +import os +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.integrations.weave.weave_otel import ( + _set_weave_specific_attributes, + get_weave_otel_config, +) +from litellm.types.integrations.weave_otel import WeaveOtelConfig, WeaveSpanAttributes + + +def test_get_weave_otel_config(): + """Test config creation with required env vars and error cases for missing vars.""" + # Test successful config creation with required environment variables + with patch.dict( + os.environ, + { + "WANDB_API_KEY": "test_api_key", + "WANDB_PROJECT_ID": "test-entity/test-project", + }, + clear=True, + ): + config = get_weave_otel_config() + + assert isinstance(config, WeaveOtelConfig) + assert config.protocol == "otlp_http" + assert config.project_id == "test-entity/test-project" + assert config.otlp_auth_headers is not None + assert "Authorization=" in config.otlp_auth_headers + assert "project_id=test-entity/test-project" in config.otlp_auth_headers + assert config.endpoint == "https://trace.wandb.ai/otel/v1/traces" + + # Verify environment variables were set + assert os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] == "https://trace.wandb.ai/otel/v1/traces" + assert os.environ["OTEL_EXPORTER_OTLP_HEADERS"] == config.otlp_auth_headers + + # Test ValueError when WANDB_API_KEY is missing + with patch.dict(os.environ, {"WANDB_PROJECT_ID": "test-entity/test-project"}, clear=True): + with pytest.raises(ValueError, match="WANDB_API_KEY must be set"): + get_weave_otel_config() + + # Test ValueError when WANDB_PROJECT_ID is missing + with patch.dict(os.environ, {"WANDB_API_KEY": "test_api_key"}, clear=True): + with pytest.raises(ValueError, match="WANDB_PROJECT_ID must be set"): + get_weave_otel_config() + + +def test_get_weave_otel_config_with_custom_host(): + """Test config creation with custom WANDB_HOST.""" + # Test with host that already has https:// + with patch.dict( + os.environ, + { + "WANDB_API_KEY": "test_api_key", + "WANDB_PROJECT_ID": "test-entity/test-project", + "WANDB_HOST": "https://custom.wandb.io", + }, + clear=True, + ): + config = get_weave_otel_config() + assert config.endpoint == "https://custom.wandb.io/otel/v1/traces" + assert os.environ["OTEL_EXPORTER_OTLP_ENDPOINT"] == "https://custom.wandb.io/otel/v1/traces" + + # Test with host without http:// or https:// + with patch.dict( + os.environ, + { + "WANDB_API_KEY": "test_api_key", + "WANDB_PROJECT_ID": "test-entity/test-project", + "WANDB_HOST": "custom.wandb.io", + }, + clear=True, + ): + config = get_weave_otel_config() + assert config.endpoint == "https://custom.wandb.io/otel/v1/traces" + + # Test with host with trailing slash + with patch.dict( + os.environ, + { + "WANDB_API_KEY": "test_api_key", + "WANDB_PROJECT_ID": "test-entity/test-project", + "WANDB_HOST": "https://custom.wandb.io/", + }, + clear=True, + ): + config = get_weave_otel_config() + assert config.endpoint == "https://custom.wandb.io/otel/v1/traces" + + + + + +def test_set_weave_specific_attributes_display_name_from_metadata(): + """Test _set_weave_specific_attributes sets display_name from metadata.""" + mock_span = MagicMock() + kwargs = { + "metadata": {"display_name": "custom-display-name"}, + "model": "gpt-4", + } + + with patch("litellm.integrations.weave.weave_otel.safe_set_attribute") as mock_safe_set: + _set_weave_specific_attributes(mock_span, kwargs, None) + + # Should set display_name from metadata + mock_safe_set.assert_any_call( + mock_span, WeaveSpanAttributes.DISPLAY_NAME.value, "custom-display-name" + ) + + +def test_set_weave_specific_attributes_display_name_from_model(): + """Test _set_weave_specific_attributes sets display_name from model when not in metadata.""" + mock_span = MagicMock() + kwargs = { + "model": "openai/gpt-4o-mini", + "metadata": {}, + } + + with patch("litellm.integrations.weave.weave_otel.safe_set_attribute") as mock_safe_set: + _set_weave_specific_attributes(mock_span, kwargs, None) + + # Should set display_name from model + mock_safe_set.assert_any_call( + mock_span, WeaveSpanAttributes.DISPLAY_NAME.value, "openai__gpt-4o-mini" + ) + + + +def test_set_weave_specific_attributes_thread_id_and_is_turn(): + """Test _set_weave_specific_attributes sets thread_id and is_turn from session_id.""" + mock_span = MagicMock() + kwargs = { + "metadata": {"session_id": "session-123"}, + } + + with patch("litellm.integrations.weave.weave_otel.safe_set_attribute") as mock_safe_set: + _set_weave_specific_attributes(mock_span, kwargs, None) + + # Should set thread_id and is_turn + mock_safe_set.assert_any_call( + mock_span, WeaveSpanAttributes.THREAD_ID.value, "session-123" + ) + mock_safe_set.assert_any_call( + mock_span, WeaveSpanAttributes.IS_TURN.value, True + ) + diff --git a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py index 09c16add77a..6e05ca564fc 100644 --- a/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/test_litellm/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -1290,10 +1290,95 @@ def test_effort_with_other_features(): litellm_params={}, headers={} ) - + # Verify all features are present assert "output_config" in result assert result["output_config"]["effort"] == "low" assert "tools" in result assert len(result["tools"]) > 0 assert "thinking" in result + + +def test_translate_system_message_skips_empty_string_content(): + """ + Test that translate_system_message skips system messages with empty string content. + + Fixes: Vertex AI Anthropic API error "messages: text content blocks must be non-empty" + """ + config = AnthropicConfig() + + # Test empty string content - should not produce any anthropic system message content + messages = [ + {"role": "system", "content": ""}, + {"role": "user", "content": "Hello"}, + ] + + result = config.translate_system_message(messages) + + # Empty system message should produce no anthropic content blocks + assert len(result) == 0 + + +def test_translate_system_message_skips_empty_list_content(): + """ + Test that translate_system_message skips empty text blocks in list content. + + Fixes: Vertex AI Anthropic API error "messages: text content blocks must be non-empty" + """ + config = AnthropicConfig() + + # Test list content with empty text block + messages = [ + {"role": "system", "content": [ + {"type": "text", "text": ""}, + {"type": "text", "text": "Valid content"}, + {"type": "text", "text": ""}, + ]}, + {"role": "user", "content": "Hello"}, + ] + + result = config.translate_system_message(messages) + + # Only non-empty text blocks should be included + assert len(result) == 1 + assert result[0]["text"] == "Valid content" + + +def test_translate_system_message_preserves_valid_content(): + """ + Test that translate_system_message preserves valid system message content. + """ + config = AnthropicConfig() + + # Test valid string content + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello"}, + ] + + result = config.translate_system_message(messages) + + assert len(result) == 1 + assert result[0]["type"] == "text" + assert result[0]["text"] == "You are a helpful assistant." + + +def test_translate_system_message_preserves_cache_control(): + """ + Test that translate_system_message preserves cache_control on valid content. + """ + config = AnthropicConfig() + + # Test list content with cache_control + messages = [ + {"role": "system", "content": [ + {"type": "text", "text": "Cached content", "cache_control": {"type": "ephemeral"}}, + ]}, + {"role": "user", "content": "Hello"}, + ] + + result = config.translate_system_message(messages) + + assert len(result) == 1 + assert result[0]["text"] == "Cached content" + assert result[0]["cache_control"] == {"type": "ephemeral"} diff --git a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py index 3095ff87f5a..91d664c3216 100644 --- a/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py +++ b/tests/test_litellm/llms/azure/chat/test_azure_gpt5_transformation.py @@ -104,34 +104,33 @@ def test_azure_gpt5_codex_series_transform_request(config: AzureOpenAIGPT5Config # GPT-5.1 temperature handling tests for Azure def test_azure_gpt5_1_temperature_with_reasoning_effort_none(config: AzureOpenAIGPT5Config): - """Test that Azure GPT-5.1 supports any temperature when reasoning_effort='none' and drop_params=True. - - Note: Azure OpenAI doesn't support reasoning_effort='none', so it's dropped from the params - when drop_params=True. The temperature logic still works correctly because the parent treats - missing reasoning_effort the same as 'none' for gpt-5.1. + """Test that Azure GPT-5.1 supports any temperature when reasoning_effort='none'. + + Azure OpenAI supports reasoning_effort='none' for gpt-5.1 models. + See: https://learn.microsoft.com/en-us/azure/ai-foundry/openai/how-to/reasoning """ params = config.map_openai_params( non_default_params={"temperature": 0.5, "reasoning_effort": "none"}, optional_params={}, model="azure/gpt-5.1", - drop_params=True, + drop_params=False, api_version="2024-05-01-preview", ) assert params["temperature"] == 0.5 - # Azure doesn't support reasoning_effort="none", so it should be dropped - assert "reasoning_effort" not in params or params.get("reasoning_effort") != "none" + # Azure supports reasoning_effort="none" for gpt-5.1 + assert params.get("reasoning_effort") == "none" -def test_azure_gpt5_1_reasoning_effort_none_error_when_drop_params_false(config: AzureOpenAIGPT5Config): - """Test that Azure GPT-5.1 raises error for reasoning_effort='none' when drop_params=False.""" - with pytest.raises(litellm.utils.UnsupportedParamsError): - config.map_openai_params( - non_default_params={"reasoning_effort": "none"}, - optional_params={}, - model="azure/gpt-5.1", - drop_params=False, - api_version="2024-05-01-preview", - ) +def test_azure_gpt5_1_reasoning_effort_none_supported(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5.1 supports reasoning_effort='none' without error.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": "none"}, + optional_params={}, + model="azure/gpt-5.1", + drop_params=False, + api_version="2024-05-01-preview", + ) + assert params.get("reasoning_effort") == "none" def test_azure_gpt5_1_temperature_without_reasoning_effort(config: AzureOpenAIGPT5Config): @@ -181,3 +180,27 @@ def test_azure_gpt5_1_series_temperature_handling(config: AzureOpenAIGPT5Config) ) assert params["temperature"] == 0.6 + +def test_azure_gpt5_reasoning_effort_none_error(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5 (non-5.1) raises error for reasoning_effort='none' when drop_params=False.""" + with pytest.raises(litellm.utils.UnsupportedParamsError): + config.map_openai_params( + non_default_params={"reasoning_effort": "none"}, + optional_params={}, + model="azure/gpt-5", + drop_params=False, + api_version="2024-05-01-preview", + ) + + +def test_azure_gpt5_reasoning_effort_none_dropped(config: AzureOpenAIGPT5Config): + """Test that Azure GPT-5 (non-5.1) drops reasoning_effort='none' when drop_params=True.""" + params = config.map_openai_params( + non_default_params={"reasoning_effort": "none"}, + optional_params={}, + model="azure/gpt-5", + drop_params=True, + api_version="2024-05-01-preview", + ) + assert "reasoning_effort" not in params or params.get("reasoning_effort") != "none" + diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py new file mode 100644 index 00000000000..737e1279e65 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py @@ -0,0 +1,292 @@ +import asyncio +import json +import os +import sys +from unittest.mock import Mock + +import pytest + +# Ensure the project root is on the import path so `litellm` can be imported when +# tests are executed from any working directory. +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.llms.bedrock.chat.invoke_transformations.amazon_qwen2_transformation import ( + AmazonQwen2Config, +) +from litellm.types.utils import ModelResponse + + +def test_qwen2_get_supported_params(): + """Test that Qwen2 config returns correct supported parameters""" + config = AmazonQwen2Config() + params = config.get_supported_openai_params(model="qwen2/test-model") + + expected_params = ["max_tokens", "temperature", "top_p", "top_k", "stop", "stream"] + for param in expected_params: + assert param in params + + +def test_qwen2_map_openai_params(): + """Test that OpenAI parameters are correctly mapped to Qwen2 format""" + config = AmazonQwen2Config() + non_default_params = { + "max_tokens": 100, + "temperature": 0.7, + "top_p": 0.9, + "top_k": 40, + "stop": ["", "<|im_end|>"], + "stream": True + } + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model="qwen2/test-model", + drop_params=False + ) + + assert result["max_tokens"] == 100 + assert result["temperature"] == 0.7 + assert result["top_p"] == 0.9 + assert result["top_k"] == 40 + assert result["stop"] == ["", "<|im_end|>"] + assert result["stream"] is True + + +def test_qwen2_convert_messages_to_prompt(): + """Test that messages are correctly converted to Qwen2 prompt format""" + config = AmazonQwen2Config() + + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello, how are you?"}, + {"role": "assistant", "content": "I'm doing well, thank you!"}, + {"role": "user", "content": "What's the weather like?"} + ] + + prompt = config._convert_messages_to_prompt(messages) + + expected_prompt = """<|im_start|>system +You are a helpful assistant.<|im_end|> +<|im_start|>user +Hello, how are you?<|im_end|> +<|im_start|>assistant +I'm doing well, thank you!<|im_end|> +<|im_start|>user +What's the weather like?<|im_end|> +<|im_start|>assistant +""" + + assert prompt == expected_prompt + + +def test_qwen2_transform_request(): + """Test that the request is correctly transformed to Qwen2 format""" + config = AmazonQwen2Config() + + messages = [ + {"role": "user", "content": "Hello, world!"} + ] + + optional_params = { + "max_tokens": 50, + "temperature": 0.8, + "top_p": 0.9 + } + + request_body = config.transform_request( + model="qwen2/test-model", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={} + ) + + assert "prompt" in request_body + assert request_body["max_gen_len"] == 50 + assert request_body["temperature"] == 0.8 + assert request_body["top_p"] == 0.9 + + # Check that the prompt contains the expected format + assert "<|im_start|>user" in request_body["prompt"] + assert "Hello, world!" in request_body["prompt"] + assert "<|im_end|>" in request_body["prompt"] + + +def test_qwen2_transform_response_with_text_field(): + """Test that Qwen2 response with 'text' field is correctly transformed to OpenAI format""" + config = AmazonQwen2Config() + + # Mock response data with 'text' field (Qwen2 format) + mock_response_data = { + "text": "<|im_start|>assistant\nHello! How can I help you today?<|im_end|>", + "usage": { + "prompt_tokens": 10, + "completion_tokens": 15, + "total_tokens": 25 + } + } + + # Mock the raw response + mock_raw_response = Mock() + mock_raw_response.json.return_value = mock_response_data + + model_response = ModelResponse() + messages = [{"role": "user", "content": "Hello!"}] + + result = config.transform_response( + model="qwen2/test-model", + messages=messages, + raw_response=mock_raw_response, + model_response=model_response, + logging_obj=Mock(), + optional_params={}, + litellm_params={}, + api_key="test-key", + request_data={}, + encoding=None + ) + + # Check that the response is correctly formatted + assert len(result.choices) == 1 + assert result.choices[0]["message"]["role"] == "assistant" + assert result.choices[0]["message"]["content"] == "Hello! How can I help you today?" + assert result.choices[0]["finish_reason"] == "stop" + + # Check usage information + assert result.usage["prompt_tokens"] == 10 + assert result.usage["completion_tokens"] == 15 + assert result.usage["total_tokens"] == 25 + + +def test_qwen2_transform_response_with_generation_field(): + """Test that Qwen2 response also supports 'generation' field for compatibility""" + config = AmazonQwen2Config() + + # Mock response data with 'generation' field (Qwen3 format, but Qwen2 should handle it) + mock_response_data = { + "generation": "<|im_start|>assistant\nHello! How can I help you today?<|im_end|>", + "usage": { + "prompt_tokens": 10, + "completion_tokens": 15, + "total_tokens": 25 + } + } + + # Mock the raw response + mock_raw_response = Mock() + mock_raw_response.json.return_value = mock_response_data + + model_response = ModelResponse() + messages = [{"role": "user", "content": "Hello!"}] + + result = config.transform_response( + model="qwen2/test-model", + messages=messages, + raw_response=mock_raw_response, + model_response=model_response, + logging_obj=Mock(), + optional_params={}, + litellm_params={}, + api_key="test-key", + request_data={}, + encoding=None + ) + + # Check that the response is correctly formatted + assert len(result.choices) == 1 + assert result.choices[0]["message"]["role"] == "assistant" + assert result.choices[0]["message"]["content"] == "Hello! How can I help you today?" + assert result.choices[0]["finish_reason"] == "stop" + + +def test_qwen2_transform_response_prefers_generation_over_text(): + """Test that Qwen2 prefers 'generation' field over 'text' when both are present""" + config = AmazonQwen2Config() + + # Mock response data with both fields + mock_response_data = { + "generation": "This is from generation field", + "text": "This is from text field", + "usage": { + "prompt_tokens": 10, + "completion_tokens": 15, + "total_tokens": 25 + } + } + + # Mock the raw response + mock_raw_response = Mock() + mock_raw_response.json.return_value = mock_response_data + + model_response = ModelResponse() + messages = [{"role": "user", "content": "Hello!"}] + + result = config.transform_response( + model="qwen2/test-model", + messages=messages, + raw_response=mock_raw_response, + model_response=model_response, + logging_obj=Mock(), + optional_params={}, + litellm_params={}, + api_key="test-key", + request_data={}, + encoding=None + ) + + # Should prefer 'generation' field + assert result.choices[0]["message"]["content"] == "This is from generation field" + + +def test_qwen2_transform_response_without_usage(): + """Test response transformation when usage information is not provided""" + config = AmazonQwen2Config() + + # Mock response data without usage + mock_response_data = { + "text": "Hello! How can I help you today?" + } + + # Mock the raw response + mock_raw_response = Mock() + mock_raw_response.json.return_value = mock_response_data + + model_response = ModelResponse() + messages = [{"role": "user", "content": "Hello!"}] + + result = config.transform_response( + model="qwen2/test-model", + messages=messages, + raw_response=mock_raw_response, + model_response=model_response, + logging_obj=Mock(), + optional_params={}, + litellm_params={}, + api_key="test-key", + request_data={}, + encoding=None + ) + + # Check that the response is correctly formatted + assert len(result.choices) == 1 + assert result.choices[0]["message"]["role"] == "assistant" + assert result.choices[0]["message"]["content"] == "Hello! How can I help you today?" + assert result.choices[0]["finish_reason"] == "stop" + + +def test_qwen2_provider_detection(): + """Test that Qwen2 provider is correctly detected from model names""" + from litellm.utils import ProviderConfigManager + from litellm.types.utils import LlmProviders + + # Test with qwen2/ prefix + config = ProviderConfigManager.get_provider_chat_config( + model="qwen2/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-qwen2", + provider=LlmProviders.BEDROCK + ) + + assert config is not None + assert isinstance(config, AmazonQwen2Config) + diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 37c95be72ce..e603f94ab87 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -2702,3 +2702,36 @@ def test_empty_assistant_message_handling(): finally: # Restore original modify_params setting litellm.modify_params = original_modify_params + + +def test_is_nova_lite_2_model(): + """Test the _is_nova_lite_2_model() method for detecting Nova 2 models.""" + config = AmazonConverseConfig() + + # Test with amazon.nova-2-lite-v1:0 + assert config._is_nova_lite_2_model("amazon.nova-2-lite-v1:0") is True + + # Test with regional variants + assert config._is_nova_lite_2_model("us.amazon.nova-2-lite-v1:0") is True + assert config._is_nova_lite_2_model("eu.amazon.nova-2-lite-v1:0") is True + assert config._is_nova_lite_2_model("apac.amazon.nova-2-lite-v1:0") is True + + # Test with other Nova 2 variants (pro, micro) + assert config._is_nova_lite_2_model("amazon.nova-pro-1-5-v1:0") is False + assert config._is_nova_lite_2_model("amazon.nova-micro-1-5-v1:0") is False + assert config._is_nova_lite_2_model("us.amazon.nova-pro-1-5-v1:0") is False + assert config._is_nova_lite_2_model("eu.amazon.nova-micro-1-5-v1:0") is False + + # Test with non-Nova-1.5 lite models (should return False) + assert config._is_nova_lite_2_model("amazon.nova-lite-v1:0") is False + assert config._is_nova_lite_2_model("amazon.nova-pro-v1:0") is False + assert config._is_nova_lite_2_model("amazon.nova-micro-v1:0") is False + + # Test with Nova v1:0 models (should return False) + assert config._is_nova_lite_2_model("us.amazon.nova-lite-v1:0") is False + assert config._is_nova_lite_2_model("eu.amazon.nova-pro-v1:0") is False + + # Test with completely different models (should return False) + assert config._is_nova_lite_2_model("anthropic.claude-3-5-sonnet-20240620-v1:0") is False + assert config._is_nova_lite_2_model("meta.llama3-70b-instruct-v1:0") is False + assert config._is_nova_lite_2_model("mistral.mistral-7b-instruct-v0:2") is False diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py new file mode 100644 index 00000000000..23243dac201 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation_nova_2.py @@ -0,0 +1,794 @@ +""" +Unit tests for Amazon Nova 2 reasoning configuration transformation. + +Tests the _transform_reasoning_effort_to_reasoning_config method in AmazonConverseConfig. +""" + +import pytest +import sys +import os + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + +import litellm +from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig + + +class TestNova15ReasoningTransformation: + """Test suite for Nova 2 reasoning effort transformation.""" + + def test_reasoning_effort_low_transformation(self): + """Test that reasoning_effort='low' is transformed to correct reasoningConfig structure.""" + config = AmazonConverseConfig() + + result = config._transform_reasoning_effort_to_reasoning_config("low") + + # Verify the structure + assert "reasoningConfig" in result + assert result["reasoningConfig"]["type"] == "enabled" + assert result["reasoningConfig"]["maxReasoningEffort"] == "low" + + def test_reasoning_effort_high_transformation(self): + """Test that reasoning_effort='high' is transformed to correct reasoningConfig structure.""" + config = AmazonConverseConfig() + + result = config._transform_reasoning_effort_to_reasoning_config("high") + + # Verify the structure + assert "reasoningConfig" in result + assert result["reasoningConfig"]["type"] == "enabled" + assert result["reasoningConfig"]["maxReasoningEffort"] == "high" + + def test_invalid_reasoning_effort_value(self): + """Test that invalid reasoning_effort values raise BadRequestError.""" + config = AmazonConverseConfig() + + # Test with invalid value "invalid" + with pytest.raises(litellm.exceptions.BadRequestError) as exc_info: + config._transform_reasoning_effort_to_reasoning_config("invalid") + + # Verify error message contains the invalid value and valid values + error_message = str(exc_info.value) + assert "invalid" in error_message + assert "low" in error_message + assert "high" in error_message + assert "Nova 2" in error_message + + def test_invalid_reasoning_effort_empty_string(self): + """Test that empty string raises BadRequestError.""" + config = AmazonConverseConfig() + + with pytest.raises(litellm.exceptions.BadRequestError) as exc_info: + config._transform_reasoning_effort_to_reasoning_config("") + + # Verify error message + error_message = str(exc_info.value) + assert "low" in error_message + assert "high" in error_message + + def test_invalid_reasoning_effort_wrong_case(self): + """Test that case-sensitive values are rejected (e.g., 'Low' instead of 'low').""" + config = AmazonConverseConfig() + + with pytest.raises(litellm.exceptions.BadRequestError): + config._transform_reasoning_effort_to_reasoning_config("Low") + + with pytest.raises(litellm.exceptions.BadRequestError): + config._transform_reasoning_effort_to_reasoning_config("HIGH") + + +class TestNova2ParameterMapping: + """Test suite for Nova 2 parameter mapping integration.""" + + def test_nova_2_reasoning_effort_low_mapping(self): + """Test that reasoning_effort='low' is correctly mapped to reasoningConfig for Nova 2.""" + config = AmazonConverseConfig() + + model = "amazon.nova-2-lite-v1:0" + non_default_params = {"reasoning_effort": "low"} + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + # Verify reasoningConfig is in result + assert "reasoningConfig" in result + assert result["reasoningConfig"]["type"] == "enabled" + assert result["reasoningConfig"]["maxReasoningEffort"] == "low" + + # Verify thinking is NOT in result + assert "thinking" not in result + + # Verify reasoning_effort is NOT kept as-is (should be transformed) + assert "reasoning_effort" not in result + + def test_nova_2_reasoning_effort_high_mapping(self): + """Test that reasoning_effort='high' is correctly mapped to reasoningConfig for Nova 2.""" + config = AmazonConverseConfig() + + model = "amazon.nova-2-lite-v1:0" + non_default_params = {"reasoning_effort": "high"} + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + # Verify reasoningConfig is in result + assert "reasoningConfig" in result + assert result["reasoningConfig"]["type"] == "enabled" + assert result["reasoningConfig"]["maxReasoningEffort"] == "high" + + # Verify thinking is NOT in result + assert "thinking" not in result + + # Verify reasoning_effort is NOT kept as-is (should be transformed) + assert "reasoning_effort" not in result + + def test_nova_2_without_reasoning_effort(self): + """Test that Nova 2 without reasoning_effort has no reasoningConfig in result.""" + config = AmazonConverseConfig() + + model = "amazon.nova-2-lite-v1:0" + non_default_params = {"temperature": 0.7} + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + # Verify reasoningConfig is NOT in result + assert "reasoningConfig" not in result + + # Verify thinking is NOT in result + assert "thinking" not in result + + # Verify reasoning_effort is NOT in result + assert "reasoning_effort" not in result + + def test_nova_2_regional_variant_us(self): + """Test that US regional variant of Nova 2 works correctly.""" + config = AmazonConverseConfig() + + model = "us.amazon.nova-2-lite-v1:0" + non_default_params = {"reasoning_effort": "high"} + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + # Verify reasoningConfig is in result + assert "reasoningConfig" in result + assert result["reasoningConfig"]["type"] == "enabled" + assert result["reasoningConfig"]["maxReasoningEffort"] == "high" + + def test_nova_2_regional_variant_eu(self): + """Test that EU regional variant of Nova 2 works correctly.""" + config = AmazonConverseConfig() + + model = "eu.amazon.nova-2-lite-v1:0" + non_default_params = {"reasoning_effort": "low"} + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + # Verify reasoningConfig is in result + assert "reasoningConfig" in result + assert result["reasoningConfig"]["type"] == "enabled" + assert result["reasoningConfig"]["maxReasoningEffort"] == "low" + + def test_nova_2_regional_variant_apac(self): + """Test that APAC regional variant of Nova 2 works correctly.""" + config = AmazonConverseConfig() + + model = "apac.amazon.nova-2-lite-v1:0" + non_default_params = {"reasoning_effort": "high"} + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + # Verify reasoningConfig is in result + assert "reasoningConfig" in result + assert result["reasoningConfig"]["type"] == "enabled" + assert result["reasoningConfig"]["maxReasoningEffort"] == "high" + + def test_nova_2_with_other_params(self): + """Test that Nova 2 reasoning works alongside other parameters.""" + config = AmazonConverseConfig() + + model = "amazon.nova-2-lite-v1:0" + non_default_params = { + "reasoning_effort": "high", + "temperature": 0.8, + "max_tokens": 1000, + "top_p": 0.9, + } + optional_params = {} + + result = config.map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=False, + ) + + # Verify reasoningConfig is in result + assert "reasoningConfig" in result + assert result["reasoningConfig"]["type"] == "enabled" + assert result["reasoningConfig"]["maxReasoningEffort"] == "high" + + # Verify other params are also present + assert result["temperature"] == 0.8 + assert result["maxTokens"] == 1000 + assert result["topP"] == 0.9 + + +class TestNova15SupportedParameters: + """Test suite for Nova 2 supported parameters.""" + + def test_nova_2_supports_reasoning_effort(self): + """Test that Nova 2 model reports reasoning_effort in supported params.""" + config = AmazonConverseConfig() + + model = "amazon.nova-2-lite-v1:0" + supported_params = config.get_supported_openai_params(model) + + # Verify reasoning_effort is in supported params + assert "reasoning_effort" in supported_params + + # Verify thinking is NOT in supported params (Nova 2 uses reasoningConfig, not thinking) + assert "thinking" not in supported_params + + def test_nova_2_regional_variant_us_supported_params(self): + """Test that US regional variant returns same supported params.""" + config = AmazonConverseConfig() + + model = "us.amazon.nova-2-lite-v1:0" + supported_params = config.get_supported_openai_params(model) + + # Verify reasoning_effort is in supported params + assert "reasoning_effort" in supported_params + + # Verify thinking is NOT in supported params + assert "thinking" not in supported_params + + def test_nova_2_regional_variant_eu_supported_params(self): + """Test that EU regional variant returns same supported params.""" + config = AmazonConverseConfig() + + model = "eu.amazon.nova-2-lite-v1:0" + supported_params = config.get_supported_openai_params(model) + + # Verify reasoning_effort is in supported params + assert "reasoning_effort" in supported_params + + # Verify thinking is NOT in supported params + assert "thinking" not in supported_params + + def test_nova_2_regional_variant_apac_supported_params(self): + """Test that APAC regional variant returns same supported params.""" + config = AmazonConverseConfig() + + model = "apac.amazon.nova-2-lite-v1:0" + supported_params = config.get_supported_openai_params(model) + + # Verify reasoning_effort is in supported params + assert "reasoning_effort" in supported_params + + # Verify thinking is NOT in supported params + assert "thinking" not in supported_params + + def test_nova_2_has_standard_params(self): + """Test that Nova 2 still has all standard supported params.""" + config = AmazonConverseConfig() + + model = "amazon.nova-2-lite-v1:0" + supported_params = config.get_supported_openai_params(model) + + # Verify standard params are present + assert "max_tokens" in supported_params + assert "max_completion_tokens" in supported_params + assert "stream" in supported_params + assert "stream_options" in supported_params + assert "stop" in supported_params + assert "temperature" in supported_params + assert "top_p" in supported_params + assert "tools" in supported_params + assert "response_format" in supported_params + + +class TestNova15ResponseParsing: + """Test suite for Nova 2 response parsing.""" + + def test_transform_reasoning_content_single_block(self): + """Test that reasoning content is extracted correctly from a single block.""" + config = AmazonConverseConfig() + + reasoning_blocks = [ + {"reasoningText": {"text": "Let me think through this step by step..."}} + ] + + result = config._transform_reasoning_content(reasoning_blocks) + + assert result == "Let me think through this step by step..." + + def test_transform_reasoning_content_multiple_blocks(self): + """Test that reasoning content is concatenated from multiple blocks.""" + config = AmazonConverseConfig() + + reasoning_blocks = [ + {"reasoningText": {"text": "First, I need to analyze the problem. "}}, + {"reasoningText": {"text": "Then, I'll consider the solution."}}, + ] + + result = config._transform_reasoning_content(reasoning_blocks) + + assert ( + result + == "First, I need to analyze the problem. Then, I'll consider the solution." + ) + + def test_transform_reasoning_content_empty_blocks(self): + """Test that empty reasoning blocks return empty string.""" + config = AmazonConverseConfig() + + reasoning_blocks = [] + + result = config._transform_reasoning_content(reasoning_blocks) + + assert result == "" + + def test_transform_thinking_blocks_with_text(self): + """Test that thinking blocks are populated correctly with text.""" + config = AmazonConverseConfig() + + reasoning_blocks = [{"reasoningText": {"text": "My reasoning process..."}}] + + result = config._transform_thinking_blocks(reasoning_blocks) + + assert len(result) == 1 + assert result[0]["type"] == "thinking" + assert result[0]["thinking"] == "My reasoning process..." + assert "signature" not in result[0] + + def test_transform_thinking_blocks_with_signature(self): + """Test that signature field is preserved when present.""" + config = AmazonConverseConfig() + + reasoning_blocks = [ + { + "reasoningText": { + "text": "My reasoning...", + "signature": "signature-hash-12345", + } + } + ] + + result = config._transform_thinking_blocks(reasoning_blocks) + + assert len(result) == 1 + assert result[0]["type"] == "thinking" + assert result[0]["thinking"] == "My reasoning..." + assert result[0]["signature"] == "signature-hash-12345" + + def test_transform_thinking_blocks_with_redacted_content(self): + """Test that redacted content blocks are handled correctly.""" + config = AmazonConverseConfig() + + reasoning_blocks = [ + {"reasoningText": {"text": "First part of reasoning..."}}, + {"redactedContent": {}}, + {"reasoningText": {"text": "Second part after redaction..."}}, + ] + + result = config._transform_thinking_blocks(reasoning_blocks) + + assert len(result) == 3 + assert result[0]["type"] == "thinking" + assert result[0]["thinking"] == "First part of reasoning..." + assert result[1]["type"] == "redacted_thinking" + assert result[2]["type"] == "thinking" + assert result[2]["thinking"] == "Second part after redaction..." + + def test_transform_thinking_blocks_multiple_blocks(self): + """Test that multiple thinking blocks are all transformed.""" + config = AmazonConverseConfig() + + reasoning_blocks = [ + {"reasoningText": {"text": "Step 1: Analyze the problem"}}, + { + "reasoningText": { + "text": "Step 2: Consider solutions", + "signature": "sig-abc", + } + }, + {"reasoningText": {"text": "Step 3: Choose best approach"}}, + ] + + result = config._transform_thinking_blocks(reasoning_blocks) + + assert len(result) == 3 + assert all(block["type"] == "thinking" for block in result) + assert result[0]["thinking"] == "Step 1: Analyze the problem" + assert result[1]["thinking"] == "Step 2: Consider solutions" + assert result[1]["signature"] == "sig-abc" + assert result[2]["thinking"] == "Step 3: Choose best approach" + + def test_transform_thinking_blocks_empty_list(self): + """Test that empty thinking blocks list returns empty list.""" + config = AmazonConverseConfig() + + reasoning_blocks = [] + + result = config._transform_thinking_blocks(reasoning_blocks) + + assert result == [] + + def test_response_parsing_integration(self): + """Test that response parsing works end-to-end with Nova 2 structure.""" + config = AmazonConverseConfig() + + # Simulate a Nova 2 response with reasoning content + reasoning_blocks = [ + { + "reasoningText": { + "text": "Let me analyze this carefully. ", + "signature": "test-signature", + } + }, + {"reasoningText": {"text": "Based on my analysis, the answer is clear."}}, + ] + + # Test reasoning content extraction + reasoning_content = config._transform_reasoning_content(reasoning_blocks) + assert ( + reasoning_content + == "Let me analyze this carefully. Based on my analysis, the answer is clear." + ) + + # Test thinking blocks transformation + thinking_blocks = config._transform_thinking_blocks(reasoning_blocks) + assert len(thinking_blocks) == 2 + assert thinking_blocks[0]["thinking"] == "Let me analyze this carefully. " + assert thinking_blocks[0]["signature"] == "test-signature" + assert ( + thinking_blocks[1]["thinking"] + == "Based on my analysis, the answer is clear." + ) + + +class TestNova15StreamingResponseParsing: + """Test suite for Nova 2 streaming response parsing.""" + + def test_streaming_reasoning_content_start_event(self): + """Test that streaming start event with reasoningContent is handled correctly.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + # Simulate a start event with redacted reasoning content + chunk_data = { + "start": {"reasoningContent": {"redactedContent": {}}}, + "contentBlockIndex": 0, + } + + result = handler.converse_chunk_parser(chunk_data) + + # Verify thinking blocks are populated + assert result.choices[0].delta.thinking_blocks is not None + assert len(result.choices[0].delta.thinking_blocks) == 1 + assert result.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking" + + def test_streaming_reasoning_content_delta_text(self): + """Test that streaming delta event with reasoning text is handled correctly.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + # Simulate a delta event with reasoning text + chunk_data = { + "delta": {"reasoningContent": {"text": "Let me think about this..."}}, + "contentBlockIndex": 0, + } + + result = handler.converse_chunk_parser(chunk_data) + + # Verify reasoning content is extracted + assert result.choices[0].delta.reasoning_content == "Let me think about this..." + + # Verify thinking blocks are populated + assert result.choices[0].delta.thinking_blocks is not None + assert len(result.choices[0].delta.thinking_blocks) == 1 + assert result.choices[0].delta.thinking_blocks[0]["type"] == "thinking" + assert ( + result.choices[0].delta.thinking_blocks[0]["thinking"] + == "Let me think about this..." + ) + + def test_streaming_reasoning_content_delta_signature(self): + """Test that streaming delta event with signature is handled correctly.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + # Simulate a delta event with signature + chunk_data = { + "delta": {"reasoningContent": {"signature": "signature-hash-xyz"}}, + "contentBlockIndex": 0, + } + + result = handler.converse_chunk_parser(chunk_data) + + # Verify reasoning content is set to empty string for consistency + assert result.choices[0].delta.reasoning_content == "" + + # Verify thinking blocks are populated with signature + assert result.choices[0].delta.thinking_blocks is not None + assert len(result.choices[0].delta.thinking_blocks) == 1 + assert result.choices[0].delta.thinking_blocks[0]["type"] == "thinking" + assert ( + result.choices[0].delta.thinking_blocks[0]["signature"] + == "signature-hash-xyz" + ) + assert result.choices[0].delta.thinking_blocks[0]["thinking"] == "" + + def test_streaming_reasoning_content_multiple_deltas(self): + """Test that multiple reasoning content deltas are accumulated correctly.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + # Simulate multiple delta events + chunks = [ + { + "delta": {"reasoningContent": {"text": "First, "}}, + "contentBlockIndex": 0, + }, + { + "delta": {"reasoningContent": {"text": "I need to analyze "}}, + "contentBlockIndex": 0, + }, + { + "delta": {"reasoningContent": {"text": "the problem."}}, + "contentBlockIndex": 0, + }, + ] + + results = [] + for chunk_data in chunks: + result = handler.converse_chunk_parser(chunk_data) + results.append(result) + + # Verify each delta has the correct reasoning content + assert results[0].choices[0].delta.reasoning_content == "First, " + assert results[1].choices[0].delta.reasoning_content == "I need to analyze " + assert results[2].choices[0].delta.reasoning_content == "the problem." + + # Verify thinking blocks are populated for each delta + for result in results: + assert result.choices[0].delta.thinking_blocks is not None + assert len(result.choices[0].delta.thinking_blocks) == 1 + assert result.choices[0].delta.thinking_blocks[0]["type"] == "thinking" + + def test_streaming_reasoning_then_text_content(self): + """Test that reasoning content followed by text content is handled correctly.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + # Simulate reasoning content followed by text content + chunks = [ + { + "delta": {"reasoningContent": {"text": "Let me think..."}}, + "contentBlockIndex": 0, + }, + {"delta": {"text": "Based on my reasoning, "}, "contentBlockIndex": 1}, + {"delta": {"text": "the answer is 42."}, "contentBlockIndex": 1}, + ] + + results = [] + for chunk_data in chunks: + result = handler.converse_chunk_parser(chunk_data) + results.append(result) + + # Verify first chunk has reasoning content + assert results[0].choices[0].delta.reasoning_content == "Let me think..." + assert results[0].choices[0].delta.thinking_blocks is not None + + # Verify subsequent chunks have text content + assert results[1].choices[0].delta.content == "Based on my reasoning, " + assert results[2].choices[0].delta.content == "the answer is 42." + + def test_streaming_redacted_content_delta(self): + """Test that streaming delta with redacted content is handled correctly.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + # Simulate a delta event with redacted content + chunk_data = { + "delta": {"reasoningContent": {"redactedContent": {}}}, + "contentBlockIndex": 0, + } + + result = handler.converse_chunk_parser(chunk_data) + + # Verify reasoning content is set to empty string for consistency + assert result.choices[0].delta.reasoning_content == "" + + # Verify thinking blocks contain redacted block + assert result.choices[0].delta.thinking_blocks is not None + assert len(result.choices[0].delta.thinking_blocks) == 1 + assert result.choices[0].delta.thinking_blocks[0]["type"] == "redacted_thinking" + + def test_streaming_provider_specific_fields(self): + """Test that provider_specific_fields are populated in streaming responses.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + # Simulate a delta event with reasoning content + chunk_data = { + "delta": {"reasoningContent": {"text": "Reasoning text"}}, + "contentBlockIndex": 0, + } + + result = handler.converse_chunk_parser(chunk_data) + + # Verify provider_specific_fields are populated + assert result.choices[0].delta.provider_specific_fields is not None + assert "reasoningContent" in result.choices[0].delta.provider_specific_fields + assert ( + result.choices[0].delta.provider_specific_fields["reasoningContent"]["text"] + == "Reasoning text" + ) + + def test_streaming_mixed_content_blocks(self): + """Test streaming with mixed content blocks (reasoning, text, tool calls).""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + # Simulate a complex streaming scenario + chunks = [ + # Start with reasoning + { + "delta": { + "reasoningContent": { + "text": "I need to call a tool to get information." + } + }, + "contentBlockIndex": 0, + }, + # Tool use start + { + "start": {"toolUse": {"toolUseId": "tool-123", "name": "get_weather"}}, + "contentBlockIndex": 1, + }, + # Tool use delta + { + "delta": {"toolUse": {"input": '{"location": "NYC"}'}}, + "contentBlockIndex": 1, + }, + # Text response + {"delta": {"text": "The weather is sunny."}, "contentBlockIndex": 2}, + ] + + results = [] + for chunk_data in chunks: + result = handler.converse_chunk_parser(chunk_data) + results.append(result) + + # Verify reasoning content in first chunk + assert ( + results[0].choices[0].delta.reasoning_content + == "I need to call a tool to get information." + ) + + # Verify tool call in second and third chunks + assert results[1].choices[0].delta.tool_calls is not None + assert ( + results[1].choices[0].delta.tool_calls[0]["function"]["name"] + == "get_weather" + ) + assert results[2].choices[0].delta.tool_calls is not None + + # Verify text content in fourth chunk + assert results[3].choices[0].delta.content == "The weather is sunny." + + def test_extract_reasoning_content_str_with_text(self): + """Test extract_reasoning_content_str method with text.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + reasoning_block = {"text": "This is reasoning text"} + + result = handler.extract_reasoning_content_str(reasoning_block) + + assert result == "This is reasoning text" + + def test_extract_reasoning_content_str_without_text(self): + """Test extract_reasoning_content_str method without text (e.g., signature only).""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + reasoning_block = {"signature": "sig-123"} + + result = handler.extract_reasoning_content_str(reasoning_block) + + assert result is None + + def test_translate_thinking_blocks_streaming_text(self): + """Test translate_thinking_blocks method with text.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + thinking_block = {"text": "Thinking content"} + + result = handler.translate_thinking_blocks(thinking_block) + + assert result is not None + assert len(result) == 1 + assert result[0]["type"] == "thinking" + assert result[0]["thinking"] == "Thinking content" + + def test_translate_thinking_blocks_streaming_signature(self): + """Test translate_thinking_blocks method with signature.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + thinking_block = {"signature": "sig-abc"} + + result = handler.translate_thinking_blocks(thinking_block) + + assert result is not None + assert len(result) == 1 + assert result[0]["type"] == "thinking" + assert result[0]["signature"] == "sig-abc" + assert ( + result[0]["thinking"] == "" + ) # Empty string for consistency with Anthropic + + def test_translate_thinking_blocks_streaming_redacted(self): + """Test translate_thinking_blocks method with redacted content.""" + from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder + + handler = AWSEventStreamDecoder(model="amazon.nova-2-lite-v1:0") + + thinking_block = {"redactedContent": {}} + + result = handler.translate_thinking_blocks(thinking_block) + + assert result is not None + assert len(result) == 1 + assert result[0]["type"] == "redacted_thinking" diff --git a/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py new file mode 100644 index 00000000000..37a0daa1d50 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/files/test_bedrock_files_integration.py @@ -0,0 +1,110 @@ +""" +Test Bedrock files integration with main files API +""" + +import base64 +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.utils import SpecialEnums + + +class TestBedrockFilesIntegration: + """Test integration of Bedrock files with main litellm API""" + + @pytest.mark.asyncio + async def test_litellm_afile_content_bedrock_provider_with_s3_uri(self): + """Test litellm.afile_content with bedrock provider using direct S3 URI""" + file_id = "s3://test-bucket/test-file.jsonl" + expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' + + # Mock the bedrock_files_instance.file_content method + with patch( + "litellm.files.main.bedrock_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx + + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request( + method="GET", url="s3://test-bucket/test-file.jsonl" + ), + ) + mock_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) + + # Call litellm.afile_content + result = await litellm.afile_content( + file_id=file_id, + custom_llm_provider="bedrock", + aws_region_name="us-west-2", + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 + + # Verify the mock was called with correct parameters + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is True + assert call_kwargs["file_content_request"]["file_id"] == file_id + + @pytest.mark.asyncio + async def test_litellm_afile_content_bedrock_provider_with_unified_file_id(self): + """Test litellm.afile_content with bedrock provider using unified file ID""" + # Create a unified file ID + s3_uri = "s3://test-bucket/batch-outputs/output.jsonl" + unified_id = "test-unified-id-123" + model_id = "test-model-id-456" + + unified_file_id_str = f"litellm_proxy:application/json;unified_id,{unified_id};target_model_names,;llm_output_file_id,{s3_uri};llm_output_file_model_id,{model_id}" + encoded_file_id = base64.urlsafe_b64encode(unified_file_id_str.encode()).decode().rstrip("=") + + expected_content = b'{"recordId": "request-1", "modelInput": {}, "modelOutput": {}}' + + # Mock the bedrock_files_instance.file_content method + with patch( + "litellm.files.main.bedrock_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx + + mock_response = httpx.Response( + status_code=200, + content=expected_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request(method="GET", url=s3_uri), + ) + mock_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) + + # Call litellm.afile_content with unified file ID + result = await litellm.afile_content( + file_id=encoded_file_id, + custom_llm_provider="bedrock", + aws_region_name="us-west-2", + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 + + # Verify the mock was called - the handler should extract S3 URI from unified file ID + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is True + # The handler extracts S3 URI from the unified file ID + assert call_kwargs["file_content_request"]["file_id"] == encoded_file_id diff --git a/tests/test_litellm/llms/cohere/rerank/test_rerank_guardrail_handler.py b/tests/test_litellm/llms/cohere/rerank/test_rerank_guardrail_handler.py index 8fdc09fc752..9c2bbeb7a68 100644 --- a/tests/test_litellm/llms/cohere/rerank/test_rerank_guardrail_handler.py +++ b/tests/test_litellm/llms/cohere/rerank/test_rerank_guardrail_handler.py @@ -5,6 +5,7 @@ Unit tests for Cohere Rerank Guardrail Translation Handler import asyncio import os import sys +from typing import List, Optional, Tuple import pytest @@ -20,8 +21,10 @@ from litellm.types.utils import CallTypes class MockGuardrail(CustomGuardrail): """Mock guardrail for testing""" - async def apply_guardrail(self, text: str, language=None, entities=None) -> str: - return f"{text} [GUARDRAILED]" + async def apply_guardrail( + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: + return ([f"{text} [GUARDRAILED]" for text in texts], None) class TestHandlerDiscovery: @@ -183,17 +186,20 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: import re - masked = re.sub( - r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", - "[EMAIL_REDACTED]", - text, - ) - masked = masked.replace("John Doe", "[NAME_REDACTED]") - return masked + masked_texts = [] + for text in texts: + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + text, + ) + masked = masked.replace("John Doe", "[NAME_REDACTED]") + masked_texts.append(masked) + return (masked_texts, None) handler = CohereRerankHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") @@ -231,21 +237,24 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: import re - # Mask emails - masked = re.sub( - r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", - "[EMAIL_REDACTED]", - text, - ) - # Mask phone numbers - masked = re.sub(r"\d{3}-\d{3}-\d{4}", "[PHONE_REDACTED]", masked) - # Mask names - masked = masked.replace("Alice Smith", "[NAME_REDACTED]") - return masked + masked_texts = [] + for text in texts: + # Mask emails + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + text, + ) + # Mask phone numbers + masked = re.sub(r"\d{3}-\d{3}-\d{4}", "[PHONE_REDACTED]", masked) + # Mask names + masked = masked.replace("Alice Smith", "[NAME_REDACTED]") + masked_texts.append(masked) + return (masked_texts, None) handler = CohereRerankHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") @@ -340,13 +349,16 @@ class TestContentFilteringScenario: """Mock content filter guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: bad_words = ["inappropriate", "offensive"] - filtered = text - for word in bad_words: - filtered = filtered.replace(word, "[FILTERED]") - return filtered + filtered_texts = [] + for text in texts: + filtered = text + for word in bad_words: + filtered = filtered.replace(word, "[FILTERED]") + filtered_texts.append(filtered) + return (filtered_texts, None) handler = CohereRerankHandler() guardrail = ContentFilterGuardrail(guardrail_name="content_filter") diff --git a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py index 1f1a36fd7ab..f0dac113645 100644 --- a/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py +++ b/tests/test_litellm/llms/custom_httpx/test_aiohttp_transport.py @@ -1,6 +1,6 @@ +import asyncio import os import sys -from unittest.mock import AsyncMock, MagicMock, patch import aiohttp import aiohttp.client_exceptions @@ -8,14 +8,11 @@ import aiohttp.http_exceptions import httpx import pytest -sys.path.insert( - 0, os.path.abspath("../../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../../..")) # Adds the parent directory to the system path from litellm.llms.custom_httpx.aiohttp_transport import ( AiohttpResponseStream, LiteLLMAiohttpTransport, - map_aiohttp_exceptions, ) @@ -32,9 +29,7 @@ class MockAiohttpResponse: ): self.status = status self.headers = headers or {} - self.content = MockContent( - content_chunks, exception_to_raise, exception_at_chunk - ) + self.content = MockContent(content_chunks, exception_to_raise, exception_at_chunk) async def __aexit__(self, exc_type, exc_val, exc_tb): pass @@ -74,7 +69,6 @@ async def test_aiohttp_response_stream_normal_flow(): @pytest.mark.asyncio async def test_transfer_encoding_error_no_httpx_read_error(): """Test that TransferEncodingError doesn't get converted to httpx.ReadError""" - import logging # Create a TransferEncodingError wrapped in ClientPayloadError (like in real scenarios) transfer_error = aiohttp.http_exceptions.TransferEncodingError( @@ -82,9 +76,7 @@ async def test_transfer_encoding_error_no_httpx_read_error(): ) # Wrap it in ClientPayloadError as aiohttp does - client_payload_error = aiohttp.ClientPayloadError( - "Response payload is not completed" - ) + client_payload_error = aiohttp.ClientPayloadError("Response payload is not completed") client_payload_error.__cause__ = transfer_error mock_response = MockAiohttpResponse( @@ -111,9 +103,7 @@ async def test_transfer_encoding_error_no_httpx_read_error(): async def test_client_payload_error_graceful_handling(): """Test that ClientPayloadError is handled gracefully without stacktrace""" # Create a ClientPayloadError directly - client_error = aiohttp.client_exceptions.ClientPayloadError( - "Response payload is not completed" - ) + client_error = aiohttp.client_exceptions.ClientPayloadError("Response payload is not completed") mock_response = MockAiohttpResponse( content_chunks=[b"data1", b"data2", b"data3"], @@ -181,7 +171,6 @@ async def test_timeout_exception_gets_mapped(): @pytest.mark.asyncio async def test_handle_async_request_uses_env_proxy(monkeypatch): """Aiohttp transport should honor HTTP(S)_PROXY env vars""" - import asyncio proxy_url = "http://proxy.local:3128" monkeypatch.setenv("HTTP_PROXY", proxy_url) monkeypatch.setenv("http_proxy", proxy_url) @@ -200,7 +189,7 @@ async def test_handle_async_request_uses_env_proxy(monkeypatch): self._loop = asyncio.get_running_loop() except RuntimeError: self._loop = None - + def request(self, *args, **kwargs): captured["proxy"] = kwargs.get("proxy") @@ -231,30 +220,118 @@ async def test_handle_async_request_uses_env_proxy(monkeypatch): assert captured["proxy"] == proxy_url +@pytest.mark.asyncio +async def test_handle_async_request_uses_env_proxy_per_url(monkeypatch): + """Aiohttp transport should honor HTTP(S)_PROXY env vars unless NO_PROXY matches""" + proxy_url = "http://proxy.local:3128" + monkeypatch.setenv("NO_PROXY", "example.com") + monkeypatch.setenv("HTTP_PROXY", proxy_url) + monkeypatch.setenv("http_proxy", proxy_url) + monkeypatch.setenv("HTTPS_PROXY", proxy_url) + monkeypatch.setenv("https_proxy", proxy_url) + monkeypatch.delenv("DISABLE_AIOHTTP_TRUST_ENV", raising=False) + + request_count = 0 + proxied_count = 0 + + class FakeSession: + def __init__(self): + self.closed = False + try: + self._loop = asyncio.get_running_loop() + except RuntimeError: + self._loop = None + + def request(self, *args, **kwargs): + nonlocal request_count + nonlocal proxied_count + request_count += 1 + + if kwargs.get("proxy") is not None: + proxied_count += 1 + + class Resp: + status = 200 + headers = {} + + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + pass + + @property + def content(self): + class C: + async def iter_chunked(self, size): + yield b"" + + return C() + + return Resp() + + transport = LiteLLMAiohttpTransport(client=lambda: FakeSession()) # type: ignore + request = httpx.Request("GET", "http://example.com") + await transport.handle_async_request(request) + + request = httpx.Request("GET", "http://foo.com") + await transport.handle_async_request(request) + + assert request_count == 2 + assert proxied_count == 1 + + +@pytest.mark.asyncio +async def test_handle_async_request_proxy_cache_per_host(monkeypatch): + """Aiohttp transport should only cache a proxy per host rather than full URL""" + proxy_url = "http://proxy.local:3128" + monkeypatch.setenv("NO_PROXY", "example.com") + monkeypatch.setenv("HTTP_PROXY", proxy_url) + monkeypatch.setenv("http_proxy", proxy_url) + monkeypatch.setenv("HTTPS_PROXY", proxy_url) + monkeypatch.setenv("https_proxy", proxy_url) + monkeypatch.delenv("DISABLE_AIOHTTP_TRUST_ENV", raising=False) + + def factory(): + return _make_mock_session() + + transport = LiteLLMAiohttpTransport(client=factory) # type: ignore + request = httpx.Request("GET", "http://foo.com/path1") + await transport.handle_async_request(request) + + request = httpx.Request("GET", "http://foo.com/path2") + await transport.handle_async_request(request) + + assert len(transport.proxy_cache) == 1 + + def _make_mock_response(should_fail=False, fail_count={"count": 0}): """Helper to create a mock aiohttp response""" + class MockResp: status = 200 headers = {} - + async def __aenter__(self): if should_fail and fail_count["count"] < 1: fail_count["count"] += 1 raise RuntimeError("Session is closed") return self - + async def __aexit__(self, *args): pass - + @property def content(self): class C: async def iter_chunked(self, size): yield b"test" + return C() - + return MockResp() + @pytest.mark.asyncio async def test_handle_async_request_total_timeout_triggers(): """ @@ -298,10 +375,10 @@ async def test_handle_async_request_total_timeout_triggers(): await transport.aclose() await runner.cleanup() + def _make_mock_session(closed=False): """Helper to create a mock aiohttp session""" - import asyncio - + class MockSession: def __init__(self): self.closed = closed @@ -309,10 +386,10 @@ def _make_mock_session(closed=False): self._loop = asyncio.get_running_loop() except RuntimeError: self._loop = None - + def request(self, *args, **kwargs): return _make_mock_response() - + return MockSession() @@ -320,14 +397,14 @@ def _make_mock_session(closed=False): async def test_handle_closed_session_before_request(): """Test that closed sessions are detected and recreated""" counts = {"sessions": 0} - + def factory(): counts["sessions"] += 1 return _make_mock_session(closed=counts["sessions"] == 1) - + transport = LiteLLMAiohttpTransport(client=factory) # type: ignore response = await transport.handle_async_request(httpx.Request("GET", "http://example.com")) - + assert counts["sessions"] == 2 # Created 2 sessions: closed one, then open one assert response.status_code == 200 @@ -337,7 +414,7 @@ async def test_handle_session_closed_during_request(): """Test that sessions closed during request are handled with retry""" counts = {"sessions": 0, "requests": 0} fail_count = {"count": 0} - + class MockSession: def __init__(self): self.closed = False @@ -345,18 +422,18 @@ async def test_handle_session_closed_during_request(): self._loop = __import__("asyncio").get_running_loop() except RuntimeError: self._loop = None - + def request(self, *args, **kwargs): counts["requests"] += 1 return _make_mock_response(should_fail=True, fail_count=fail_count) - + def factory(): counts["sessions"] += 1 return MockSession() - + transport = LiteLLMAiohttpTransport(client=factory) # type: ignore response = await transport.handle_async_request(httpx.Request("GET", "http://example.com")) - + assert counts["requests"] == 2 # First request failed, second succeeded assert counts["sessions"] == 2 # Created 2 sessions for retry assert response.status_code == 200 diff --git a/tests/test_litellm/llms/openai/chat/__init__.py b/tests/test_litellm/llms/openai/chat/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py new file mode 100644 index 00000000000..09f3f8dcf0f --- /dev/null +++ b/tests/test_litellm/llms/openai/chat/guardrail_translation/test_handler.py @@ -0,0 +1,518 @@ +""" +Unit tests for OpenAI Chat Completions Guardrail Translation Handler + +Tests the handler's ability to process input/output for Chat Completions API +with guardrail transformations, including tool calls. +""" + +import json +import os +import sys +from typing import Any, List, Literal, Optional, Tuple +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../../../../..") +) # Adds the parent directory to the system path + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.llms.openai.chat.guardrail_translation.handler import ( + OpenAIChatCompletionsHandler, +) +from litellm.types.guardrails import GenericGuardrailAPIInputs +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Choices, + Function, + Message, + ModelResponse, +) + + +class MockGuardrail(CustomGuardrail): + """Mock guardrail for testing that transforms text and tool calls""" + + def __init__(self, guardrail_name: str = "test"): + super().__init__(guardrail_name=guardrail_name) + self.last_inputs = None + self.last_request_data = None + self.tool_calls_modified = False + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> Tuple[List[str], Optional[List[str]]]: + """Mock apply_guardrail that uppercases text and modifies tool calls""" + self.last_inputs = inputs + self.last_request_data = request_data + + # Return modified texts (uppercase for testing) + texts = inputs.get("texts", []) + modified_texts = [text.upper() for text in texts] + + # Modify tool calls in place if present + tool_calls = inputs.get("tool_calls", []) + if tool_calls: + self.tool_calls_modified = True + for tool_call in tool_calls: + if isinstance(tool_call, dict) and "function" in tool_call: + function = tool_call["function"] + if "arguments" in function: + # Modify arguments to uppercase JSON string + try: + args_dict = json.loads(function["arguments"]) + # Uppercase all string values + for key, value in args_dict.items(): + if isinstance(value, str): + args_dict[key] = value.upper() + function["arguments"] = json.dumps(args_dict) + except json.JSONDecodeError: + # If not JSON, just uppercase the string + function["arguments"] = function["arguments"].upper() + + return modified_texts, [] + + +class TestOpenAIChatCompletionsHandlerToolCallsInput: + """Test input processing with tool calls""" + + @pytest.mark.asyncio + async def test_extract_tool_calls_from_input_messages(self): + """Test that tool calls are extracted from input messages""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + # Create input data with tool calls (assistant message) + data = { + "messages": [ + {"role": "user", "content": "What's the weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps( + {"location": "San Francisco", "unit": "celsius"} + ), + }, + } + ], + }, + ] + } + + # Process the input + await handler.process_input_messages(data, guardrail) + + # Verify tool calls were extracted and passed to guardrail + assert guardrail.last_inputs is not None + assert "tool_calls" in guardrail.last_inputs + assert len(guardrail.last_inputs["tool_calls"]) == 1 + + tool_call = guardrail.last_inputs["tool_calls"][0] + assert tool_call["id"] == "call_123" + assert tool_call["function"]["name"] == "get_weather" + # Note: tool call arguments may already be modified by guardrail + # Check that it contains location parameter + assert "location" in tool_call["function"]["arguments"] + + # Verify tool call was modified by guardrail + assert guardrail.tool_calls_modified is True + + # Verify the message was updated with modified tool call + modified_tool_call = data["messages"][1]["tool_calls"][0] + args = json.loads(modified_tool_call["function"]["arguments"]) + assert args["location"] == "SAN FRANCISCO" # Should be uppercased + assert args["unit"] == "CELSIUS" # Should be uppercased + + @pytest.mark.asyncio + async def test_extract_tool_calls_and_text_from_input_messages(self): + """Test that both tool calls and text content are extracted""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + # Create input data with both text and tool calls + data = { + "messages": [ + {"role": "user", "content": "What's the weather?"}, + { + "role": "assistant", + "content": "Let me check that for you.", + "tool_calls": [ + { + "id": "call_456", + "type": "function", + "function": { + "name": "get_current_weather", + "arguments": json.dumps({"location": "Boston"}), + }, + } + ], + }, + ] + } + + # Process the input + await handler.process_input_messages(data, guardrail) + + # Verify both texts and tool calls were extracted + assert guardrail.last_inputs is not None + assert "texts" in guardrail.last_inputs + assert "tool_calls" in guardrail.last_inputs + + # Should have 2 texts (user message + assistant message) + assert len(guardrail.last_inputs["texts"]) == 2 + assert "What's the weather?" in guardrail.last_inputs["texts"] + assert "Let me check that for you." in guardrail.last_inputs["texts"] + + # Should have 1 tool call + assert len(guardrail.last_inputs["tool_calls"]) == 1 + assert ( + guardrail.last_inputs["tool_calls"][0]["function"]["name"] + == "get_current_weather" + ) + + # Verify text content was modified + assert data["messages"][0]["content"] == "WHAT'S THE WEATHER?" + assert data["messages"][1]["content"] == "LET ME CHECK THAT FOR YOU." + + # Verify tool call was modified + modified_tool_call = data["messages"][1]["tool_calls"][0] + args = json.loads(modified_tool_call["function"]["arguments"]) + assert args["location"] == "BOSTON" + + @pytest.mark.asyncio + async def test_extract_multiple_tool_calls_from_input(self): + """Test extraction of multiple tool calls""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + # Create input data with multiple tool calls + data = { + "messages": [ + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"location": "NYC"}), + }, + }, + { + "id": "call_2", + "type": "function", + "function": { + "name": "get_time", + "arguments": json.dumps({"timezone": "EST"}), + }, + }, + ], + } + ] + } + + # Process the input + await handler.process_input_messages(data, guardrail) + + # Verify multiple tool calls were extracted + assert guardrail.last_inputs is not None + assert "tool_calls" in guardrail.last_inputs + assert len(guardrail.last_inputs["tool_calls"]) == 2 + + # Verify both tool calls + tool_calls = guardrail.last_inputs["tool_calls"] + assert tool_calls[0]["function"]["name"] == "get_weather" + assert tool_calls[1]["function"]["name"] == "get_time" + + # Verify both were modified + modified_tool_calls = data["messages"][0]["tool_calls"] + args1 = json.loads(modified_tool_calls[0]["function"]["arguments"]) + args2 = json.loads(modified_tool_calls[1]["function"]["arguments"]) + assert args1["location"] == "NYC" + assert args2["timezone"] == "EST" + + @pytest.mark.asyncio + async def test_tool_calls_separate_from_texts(self): + """Test that tool calls are passed as a separate parameter, not mixed with texts""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + data = { + "messages": [ + {"role": "user", "content": "Get weather for LA"}, + { + "role": "assistant", + "content": "Sure!", + "tool_calls": [ + { + "id": "call_xyz", + "type": "function", + "function": { + "name": "get_weather", + "arguments": json.dumps({"city": "Los Angeles"}), + }, + } + ], + }, + ] + } + + # Process the input + await handler.process_input_messages(data, guardrail) + + # Verify tool calls and texts are separate + assert guardrail.last_inputs is not None + texts = guardrail.last_inputs.get("texts", []) + tool_calls = guardrail.last_inputs.get("tool_calls", []) + + # Texts should only contain the content strings + assert len(texts) == 2 + assert "Get weather for LA" in texts + assert "Sure!" in texts + + # Tool call arguments should NOT be in texts + assert not any("Los Angeles" in text for text in texts) + + # Tool calls should be separate + assert len(tool_calls) == 1 + assert tool_calls[0]["function"]["name"] == "get_weather" + # Check that it contains city parameter (may be modified by guardrail) + assert "city" in tool_calls[0]["function"]["arguments"] + + @pytest.mark.asyncio + async def test_no_tool_calls_in_input(self): + """Test that messages without tool calls work correctly""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + # Create input data without tool calls + data = { + "messages": [ + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi there!"}, + ] + } + + # Process the input + await handler.process_input_messages(data, guardrail) + + # Verify no tool calls were passed to guardrail + assert guardrail.last_inputs is not None + tool_calls = guardrail.last_inputs.get("tool_calls", []) + assert len(tool_calls) == 0 + + # Verify text was still processed + assert len(guardrail.last_inputs["texts"]) == 2 + assert data["messages"][0]["content"] == "HELLO" + assert data["messages"][1]["content"] == "HI THERE!" + + @pytest.mark.asyncio + async def test_empty_tool_calls_list(self): + """Test that empty tool_calls list is handled correctly""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + data = { + "messages": [ + {"role": "assistant", "content": "Hello", "tool_calls": []}, + ] + } + + # Process the input + await handler.process_input_messages(data, guardrail) + + # Verify empty tool_calls doesn't cause issues + assert guardrail.last_inputs is not None + tool_calls = guardrail.last_inputs.get("tool_calls", []) + assert len(tool_calls) == 0 + + +class TestOpenAIChatCompletionsHandlerToolCallsOutput: + """Test output processing with tool calls""" + + @pytest.mark.asyncio + async def test_extract_tool_calls_from_output_response(self): + """Test that tool calls are extracted from output responses""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + # Create a mock response with tool calls + response = ModelResponse( + id="chatcmpl-123", + created=1234567890, + model="gpt-4", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_789", + type="function", + function=Function( + name="search_database", + arguments=json.dumps({"query": "python tutorials"}), + ), + ) + ], + ), + ) + ], + ) + + # Process the output + await handler.process_output_response(response, guardrail) + + # Verify tool calls were extracted and passed to guardrail + assert guardrail.last_inputs is not None + assert "tool_calls" in guardrail.last_inputs + assert len(guardrail.last_inputs["tool_calls"]) == 1 + + tool_call = guardrail.last_inputs["tool_calls"][0] + assert tool_call["function"]["name"] == "search_database" + # Check that it contains query parameter (may be modified by guardrail) + assert "query" in tool_call["function"]["arguments"] + + # Verify tool call was modified in response + response_tool_call = response.choices[0].message.tool_calls[0] + args = json.loads(response_tool_call.function.arguments) + assert args["query"] == "PYTHON TUTORIALS" # Should be uppercased + + @pytest.mark.asyncio + async def test_extract_tool_calls_and_content_from_output(self): + """Test extraction of both content and tool calls from output""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + response = ModelResponse( + id="chatcmpl-456", + created=1234567890, + model="gpt-4", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content="I'll search for that information.", + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_999", + type="function", + function=Function( + name="web_search", + arguments=json.dumps( + {"keywords": "litellm documentation"} + ), + ), + ) + ], + ), + ) + ], + ) + + # Process the output + await handler.process_output_response(response, guardrail) + + # Verify both texts and tool calls were extracted + assert guardrail.last_inputs is not None + assert "texts" in guardrail.last_inputs + assert "tool_calls" in guardrail.last_inputs + + assert len(guardrail.last_inputs["texts"]) == 1 + assert "I'll search for that information." in guardrail.last_inputs["texts"] + + assert len(guardrail.last_inputs["tool_calls"]) == 1 + + # Verify both were modified + assert ( + response.choices[0].message.content == "I'LL SEARCH FOR THAT INFORMATION." + ) + response_tool_call = response.choices[0].message.tool_calls[0] + args = json.loads(response_tool_call.function.arguments) + assert args["keywords"] == "LITELLM DOCUMENTATION" + + @pytest.mark.asyncio + async def test_extract_multiple_tool_calls_from_output(self): + """Test extraction of multiple tool calls from output""" + handler = OpenAIChatCompletionsHandler() + guardrail = MockGuardrail() + + response = ModelResponse( + id="chatcmpl-789", + created=1234567890, + model="gpt-4", + object="chat.completion", + choices=[ + Choices( + finish_reason="tool_calls", + index=0, + message=Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + id="call_1", + type="function", + function=Function( + name="get_weather", + arguments=json.dumps({"location": "Tokyo"}), + ), + ), + ChatCompletionMessageToolCall( + id="call_2", + type="function", + function=Function( + name="get_news", + arguments=json.dumps({"topic": "technology"}), + ), + ), + ], + ), + ) + ], + ) + + # Process the output + await handler.process_output_response(response, guardrail) + + # Verify multiple tool calls were extracted + assert guardrail.last_inputs is not None + assert "tool_calls" in guardrail.last_inputs + assert len(guardrail.last_inputs["tool_calls"]) == 2 + + # Verify both tool calls + tool_calls = guardrail.last_inputs["tool_calls"] + assert tool_calls[0]["function"]["name"] == "get_weather" + assert tool_calls[1]["function"]["name"] == "get_news" + + # Verify both were modified + response_tool_calls = response.choices[0].message.tool_calls + args1 = json.loads(response_tool_calls[0].function.arguments) + args2 = json.loads(response_tool_calls[1].function.arguments) + assert args1["location"] == "TOKYO" + assert args2["topic"] == "TECHNOLOGY" + + +if __name__ == "__main__": + # Run the tests + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/llms/openai/completion/test_text_completion_guardrail_handler.py b/tests/test_litellm/llms/openai/completion/test_text_completion_guardrail_handler.py index 97ca423f773..c861e48ad48 100644 --- a/tests/test_litellm/llms/openai/completion/test_text_completion_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/completion/test_text_completion_guardrail_handler.py @@ -4,6 +4,7 @@ Unit tests for OpenAI Text Completion Guardrail Translation Handler import os import sys +from typing import List, Optional, Tuple from unittest.mock import MagicMock import pytest @@ -21,8 +22,10 @@ from litellm.types.utils import CallTypes, TextChoices, TextCompletionResponse class MockGuardrail(CustomGuardrail): """Mock guardrail for testing""" - async def apply_guardrail(self, text: str, language=None, entities=None) -> str: - return f"{text} [GUARDRAILED]" + async def apply_guardrail( + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: + return ([f"{text} [GUARDRAILED]" for text in texts], None) class TestHandlerDiscovery: @@ -243,19 +246,22 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: # Simple mock: replace email-like patterns import re - masked = re.sub( - r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", - "[EMAIL_REDACTED]", - text, - ) - # Replace names (simple mock) - masked = masked.replace("John Doe", "[NAME_REDACTED]") - return masked + masked_texts = [] + for text in texts: + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + text, + ) + # Replace names (simple mock) + masked = masked.replace("John Doe", "[NAME_REDACTED]") + masked_texts.append(masked) + return (masked_texts, None) handler = OpenAITextCompletionHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") @@ -303,15 +309,19 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: import re - return re.sub( - r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", - "[EMAIL_REDACTED]", - text, - ) + masked_texts = [] + for text in texts: + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + text, + ) + masked_texts.append(masked) + return (masked_texts, None) handler = OpenAITextCompletionHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") diff --git a/tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py b/tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py index 529d74f63d9..5e183e32208 100644 --- a/tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/image_generation/test_image_generation_guardrail_handler.py @@ -4,6 +4,7 @@ Unit tests for OpenAI Image Generation Guardrail Translation Handler import os import sys +from typing import List, Optional, Tuple import pytest @@ -20,8 +21,10 @@ from litellm.types.utils import CallTypes, ImageObject, ImageResponse class MockGuardrail(CustomGuardrail): """Mock guardrail for testing""" - async def apply_guardrail(self, text: str, language=None, entities=None) -> str: - return f"{text} [GUARDRAILED]" + async def apply_guardrail( + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: + return ([f"{text} [GUARDRAILED]" for text in texts], None) class TestHandlerDiscovery: @@ -141,19 +144,22 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: # Simple mock: replace email-like patterns import re - masked = re.sub( - r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", - "[EMAIL_REDACTED]", - text, - ) - # Replace names (simple mock) - masked = masked.replace("John Doe", "[NAME_REDACTED]") - return masked + masked_texts = [] + for text in texts: + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + text, + ) + # Replace names (simple mock) + masked = masked.replace("John Doe", "[NAME_REDACTED]") + masked_texts.append(masked) + return (masked_texts, None) handler = OpenAIImageGenerationHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index b447c281aa3..cc76be08178 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -7,7 +7,7 @@ with guardrail transformations. import os import sys -from typing import Any +from typing import Any, List, Literal, Optional, Tuple from unittest.mock import AsyncMock, MagicMock import pytest @@ -16,6 +16,8 @@ sys.path.insert( 0, os.path.abspath("../../../../../..") ) # Adds the parent directory to the system path +from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.llms import get_guardrail_translation_mapping from litellm.llms.openai.responses.guardrail_translation.handler import ( @@ -27,11 +29,28 @@ from litellm.types.utils import CallTypes class MockGuardrail(CustomGuardrail): - """Mock guardrail for testing that transforms text""" + """Mock guardrail for testing that transforms text for requests and blocks responses""" - async def apply_guardrail(self, text: str) -> str: - """Append [GUARDRAILED] to text""" - return f"{text} [GUARDRAILED]" + async def apply_guardrail( + self, + texts: List[str], + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + images: Optional[List[str]] = None, + ) -> Tuple[List[str], Optional[List[str]]]: + """ + For requests: Append [GUARDRAILED] to text + For responses: Block by raising HTTPException (masking responses is no longer supported) + """ + if input_type == "response": + # Responses should be blocked, not masked + raise HTTPException( + status_code=400, + detail={"error": "Response blocked by guardrail", "texts": texts}, + ) + # For requests, we can still mask/transform + return ([f"{text} [GUARDRAILED]" for text in texts], None) class TestOpenAIResponsesHandlerDiscovery: @@ -167,11 +186,15 @@ class TestOpenAIResponsesHandlerOutputProcessing: @pytest.mark.asyncio async def test_process_output_response_simple(self): - """Test processing simple output response""" + """Test processing simple output response - should block, not mask + + After unified_guardrail.py changes, responses can only be blocked/rejected, not masked. + This test verifies that the guardrail properly blocks responses. + """ handler = OpenAIResponsesHandler() guardrail = MockGuardrail(guardrail_name="test") - # Create a mock response + # Create a mock response with dict format (works with current handler) response = ResponsesAPIResponse( id="resp_123", created_at=1234567890, @@ -179,30 +202,36 @@ class TestOpenAIResponsesHandlerOutputProcessing: object="response", status="completed", output=[ - GenericResponseOutputItem( - type="message", - id="msg_123", - status="completed", - role="assistant", - content=[ - OutputText( - type="output_text", text="Hello user", annotations=None - ), + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Hello user"}, ], - ) + } ], ) - result = await handler.process_output_response(response, guardrail) + # Response should be blocked, not masked + with pytest.raises(HTTPException) as exc_info: + await handler.process_output_response(response, guardrail) - assert result.output[0].content[0].text == "Hello user [GUARDRAILED]" + assert exc_info.value.status_code == 400 + assert "Response blocked by guardrail" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_process_output_response_multiple_items(self): - """Test processing output response with multiple output items""" + """Test processing output response with multiple output items - should block, not mask + + After unified_guardrail.py changes, responses can only be blocked/rejected, not masked. + This test verifies that the guardrail properly blocks responses with multiple items. + """ handler = OpenAIResponsesHandler() guardrail = MockGuardrail(guardrail_name="test") + # Use dict format (works with current handler) response = ResponsesAPIResponse( id="resp_123", created_at=1234567890, @@ -210,46 +239,45 @@ class TestOpenAIResponsesHandlerOutputProcessing: object="response", status="completed", output=[ - GenericResponseOutputItem( - type="message", - id="msg_123", - status="completed", - role="assistant", - content=[ - OutputText( - type="output_text", - text="First message", - annotations=None, - ), + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "First message"}, ], - ), - GenericResponseOutputItem( - type="message", - id="msg_124", - status="completed", - role="assistant", - content=[ - OutputText( - type="output_text", - text="Second message", - annotations=None, - ), + }, + { + "type": "message", + "id": "msg_124", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Second message"}, ], - ), + }, ], ) - result = await handler.process_output_response(response, guardrail) + # Response should be blocked, not masked + with pytest.raises(HTTPException) as exc_info: + await handler.process_output_response(response, guardrail) - assert result.output[0].content[0].text == "First message [GUARDRAILED]" - assert result.output[1].content[0].text == "Second message [GUARDRAILED]" + assert exc_info.value.status_code == 400 + assert "Response blocked by guardrail" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_process_output_response_multiple_content_items(self): - """Test processing output response with multiple content items in one output""" + """Test processing output response with multiple content items - should block, not mask + + After unified_guardrail.py changes, responses can only be blocked/rejected, not masked. + This test verifies that the guardrail properly blocks responses with multiple content items. + """ handler = OpenAIResponsesHandler() guardrail = MockGuardrail(guardrail_name="test") + # Use dict format (works with current handler) response = ResponsesAPIResponse( id="resp_123", created_at=1234567890, @@ -257,27 +285,33 @@ class TestOpenAIResponsesHandlerOutputProcessing: object="response", status="completed", output=[ - GenericResponseOutputItem( - type="message", - id="msg_123", - status="completed", - role="assistant", - content=[ - OutputText(type="output_text", text="Part 1", annotations=None), - OutputText(type="output_text", text="Part 2", annotations=None), + { + "type": "message", + "id": "msg_123", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "Part 1"}, + {"type": "output_text", "text": "Part 2"}, ], - ) + } ], ) - result = await handler.process_output_response(response, guardrail) + # Response should be blocked, not masked + with pytest.raises(HTTPException) as exc_info: + await handler.process_output_response(response, guardrail) - assert result.output[0].content[0].text == "Part 1 [GUARDRAILED]" - assert result.output[0].content[1].text == "Part 2 [GUARDRAILED]" + assert exc_info.value.status_code == 400 + assert "Response blocked by guardrail" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_process_output_response_with_dict_format(self): - """Test processing output response where content items are dicts instead of OutputText objects""" + """Test processing output response with dict format - should block, not mask + + After unified_guardrail.py changes, responses can only be blocked/rejected, not masked. + This test verifies blocking works even when content items are dicts instead of OutputText objects. + """ handler = OpenAIResponsesHandler() guardrail = MockGuardrail(guardrail_name="test") @@ -301,9 +335,12 @@ class TestOpenAIResponsesHandlerOutputProcessing: ], ) - result = await handler.process_output_response(response, guardrail) + # Response should be blocked, not masked + with pytest.raises(HTTPException) as exc_info: + await handler.process_output_response(response, guardrail) - assert result.output[0]["content"][0]["text"] == "Hello from dict [GUARDRAILED]" + assert exc_info.value.status_code == 400 + assert "Response blocked by guardrail" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_process_output_response_no_text_content(self): @@ -450,7 +487,10 @@ class TestOpenAIResponsesHandlerEdgeCases: "role": "user", "content": [ {"type": "text", "text": "List content"}, - {"type": "image_url", "image_url": {"url": "http://example.com"}}, + { + "type": "image_url", + "image_url": {"url": "http://example.com"}, + }, ], "type": "message", }, @@ -492,4 +532,3 @@ class TestOpenAIResponsesHandlerEdgeCases: # Should skip processing and return unchanged assert result == response - diff --git a/tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py b/tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py index b4064a22c17..dfd96beb2f4 100644 --- a/tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/speech/test_text_to_speech_guardrail_handler.py @@ -4,6 +4,7 @@ Unit tests for OpenAI Text-to-Speech Guardrail Translation Handler import os import sys +from typing import List, Optional, Tuple import pytest @@ -20,8 +21,10 @@ from litellm.types.utils import CallTypes class MockGuardrail(CustomGuardrail): """Mock guardrail for testing""" - async def apply_guardrail(self, text: str, language=None, entities=None) -> str: - return f"{text} [GUARDRAILED]" + async def apply_guardrail( + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: + return ([f"{text} [GUARDRAILED]" for text in texts], None) class MockBinaryResponse: @@ -169,20 +172,23 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: # Simple mock: replace email-like patterns import re - masked = re.sub( - r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", - "[EMAIL_REDACTED]", - text, - ) - # Replace names (simple mock) - masked = masked.replace("John Doe", "[NAME_REDACTED]") - masked = masked.replace("555-1234", "[PHONE_REDACTED]") - return masked + masked_texts = [] + for text in texts: + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + text, + ) + # Replace names (simple mock) + masked = masked.replace("John Doe", "[NAME_REDACTED]") + masked = masked.replace("555-1234", "[PHONE_REDACTED]") + masked_texts.append(masked) + return (masked_texts, None) handler = OpenAITextToSpeechHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") @@ -211,17 +217,24 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: import re - # Mask account numbers - masked = re.sub(r"account number \d{8,12}", "account number [REDACTED]", text) - # Mask SSNs - masked = re.sub(r"\d{3}-\d{2}-\d{4}", "[SSN_REDACTED]", masked) - # Mask credit cards - masked = re.sub(r"\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}", "[CC_REDACTED]", masked) - return masked + masked_texts = [] + for text in texts: + # Mask account numbers + masked = re.sub( + r"account number \d{8,12}", "account number [REDACTED]", text + ) + # Mask SSNs + masked = re.sub(r"\d{3}-\d{2}-\d{4}", "[SSN_REDACTED]", masked) + # Mask credit cards + masked = re.sub( + r"\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}", "[CC_REDACTED]", masked + ) + masked_texts.append(masked) + return (masked_texts, None) handler = OpenAITextToSpeechHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") @@ -256,14 +269,17 @@ class TestContentModerationScenario: """Mock content filter guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: # Simple mock: filter inappropriate words bad_words = ["badword", "inappropriate", "offensive"] - filtered = text - for word in bad_words: - filtered = filtered.replace(word, "[FILTERED]") - return filtered + filtered_texts = [] + for text in texts: + filtered = text + for word in bad_words: + filtered = filtered.replace(word, "[FILTERED]") + filtered_texts.append(filtered) + return (filtered_texts, None) handler = OpenAITextToSpeechHandler() guardrail = ContentFilterGuardrail(guardrail_name="content_filter") @@ -322,4 +338,3 @@ class TestMultilingualTTS: assert f"Testing with {voice} voice [GUARDRAILED]" == result["input"] assert result["voice"] == voice - diff --git a/tests/test_litellm/llms/openai/test_openai_empty_response.py b/tests/test_litellm/llms/openai/test_openai_empty_response.py new file mode 100644 index 00000000000..fb42918f381 --- /dev/null +++ b/tests/test_litellm/llms/openai/test_openai_empty_response.py @@ -0,0 +1,97 @@ +""" +Test for issue #17209: Clearer error when LLM endpoint returns empty response +""" + +import os +import sys +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +from litellm.llms.openai.openai import OpenAIChatCompletion +from litellm.llms.openai.common_utils import OpenAIError + +class TestEmptyResponseHandling: + """Test that empty/invalid responses from LLM endpoints produce clear error messages""" + + def test_sync_empty_string_response_raises_clear_error(self): + """ + Test that when an OpenAI-compatible endpoint returns an empty string, + we get a clear error instead of "'str' object has no attribute 'model_dump'" + """ + openai_chat = OpenAIChatCompletion() + + # Mock the raw response to return an empty string from parse() + mock_raw_response = MagicMock() + mock_raw_response.headers = {} + mock_raw_response.parse.return_value = "" # Empty string response + + mock_client = MagicMock() + mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_raw_response + ) + + with pytest.raises(OpenAIError) as exc_info: + openai_chat.make_sync_openai_chat_completion_request( + openai_client=mock_client, + data={"messages": [{"role": "user", "content": "test"}]}, + timeout=30, + logging_obj=MagicMock(), + ) + + assert "Empty or invalid response from LLM endpoint" in str(exc_info.value) + assert "Check the reverse proxy or model server configuration" in str( + exc_info.value + ) + + def test_sync_none_response_raises_clear_error(self): + """Test that None response also produces a clear error""" + openai_chat = OpenAIChatCompletion() + + mock_raw_response = MagicMock() + mock_raw_response.headers = {} + mock_raw_response.parse.return_value = None + + mock_client = MagicMock() + mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_raw_response + ) + + with pytest.raises(OpenAIError) as exc_info: + openai_chat.make_sync_openai_chat_completion_request( + openai_client=mock_client, + data={"messages": [{"role": "user", "content": "test"}]}, + timeout=30, + logging_obj=MagicMock(), + ) + + assert "Empty or invalid response from LLM endpoint" in str(exc_info.value) + + def test_valid_response_passes_through(self): + """Test that a valid response with model_dump passes through correctly""" + openai_chat = OpenAIChatCompletion() + + # Create a mock response that has model_dump (like a real Pydantic model) + mock_response = MagicMock() + mock_response.model_dump.return_value = {"choices": []} + + mock_raw_response = MagicMock() + mock_raw_response.headers = {"x-request-id": "123"} + mock_raw_response.parse.return_value = mock_response + + mock_client = MagicMock() + mock_client.chat.completions.with_raw_response.create.return_value = ( + mock_raw_response + ) + + headers, response = openai_chat.make_sync_openai_chat_completion_request( + openai_client=mock_client, + data={"messages": [{"role": "user", "content": "test"}]}, + timeout=30, + logging_obj=MagicMock(), + ) + + assert response == mock_response + assert headers == {"x-request-id": "123"} diff --git a/tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py b/tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py index faa425eb714..4d2cb142b35 100644 --- a/tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/transcriptions/test_audio_transcription_guardrail_handler.py @@ -4,6 +4,7 @@ Unit tests for OpenAI Audio Transcription Guardrail Translation Handler import os import sys +from typing import List, Optional, Tuple import pytest @@ -21,8 +22,10 @@ from litellm.utils import TranscriptionResponse class MockGuardrail(CustomGuardrail): """Mock guardrail for testing""" - async def apply_guardrail(self, text: str, language=None, entities=None) -> str: - return f"{text} [GUARDRAILED]" + async def apply_guardrail( + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: + return ([f"{text} [GUARDRAILED]" for text in texts], None) class TestHandlerDiscovery: @@ -140,20 +143,23 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: # Simple mock: replace email-like patterns import re - masked = re.sub( - r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", - "[EMAIL_REDACTED]", - text, - ) - # Replace names (simple mock) - masked = masked.replace("John Doe", "[NAME_REDACTED]") - masked = masked.replace("555-1234", "[PHONE_REDACTED]") - return masked + masked_texts = [] + for text in texts: + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + text, + ) + # Replace names (simple mock) + masked = masked.replace("John Doe", "[NAME_REDACTED]") + masked = masked.replace("555-1234", "[PHONE_REDACTED]") + masked_texts.append(masked) + return (masked_texts, None) handler = OpenAIAudioTranscriptionHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") @@ -181,23 +187,26 @@ class TestPIIMaskingScenario: """Mock PII masking guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: import re - # Mask credit card numbers - masked = re.sub( - r"\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}", "[CC_REDACTED]", text - ) - # Mask SSNs - masked = re.sub(r"\d{3}-\d{2}-\d{4}", "[SSN_REDACTED]", masked) - # Mask emails - masked = re.sub( - r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", - "[EMAIL_REDACTED]", - masked, - ) - return masked + masked_texts = [] + for text in texts: + # Mask credit card numbers + masked = re.sub( + r"\d{4}[- ]?\d{4}[- ]?\d{4}[- ]?\d{4}", "[CC_REDACTED]", text + ) + # Mask SSNs + masked = re.sub(r"\d{3}-\d{2}-\d{4}", "[SSN_REDACTED]", masked) + # Mask emails + masked = re.sub( + r"\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b", + "[EMAIL_REDACTED]", + masked, + ) + masked_texts.append(masked) + return (masked_texts, None) handler = OpenAIAudioTranscriptionHandler() guardrail = PIIMaskingGuardrail(guardrail_name="mask_pii") @@ -231,14 +240,17 @@ class TestContentModerationScenario: """Mock profanity filter guardrail""" async def apply_guardrail( - self, text: str, language=None, entities=None - ) -> str: + self, texts: List[str], request_data: dict, input_type: str, **kwargs + ) -> Tuple[List[str], Optional[List[str]]]: # Simple mock: replace common profanity bad_words = ["badword1", "badword2", "inappropriate"] - filtered = text - for word in bad_words: - filtered = filtered.replace(word, "[FILTERED]") - return filtered + filtered_texts = [] + for text in texts: + filtered = text + for word in bad_words: + filtered = filtered.replace(word, "[FILTERED]") + filtered_texts.append(filtered) + return (filtered_texts, None) handler = OpenAIAudioTranscriptionHandler() guardrail = ProfanityFilterGuardrail(guardrail_name="content_filter") diff --git a/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py index 64ac299fd79..d5a73b3fd12 100644 --- a/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/test_litellm/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -489,3 +489,30 @@ def test_openrouter_cost_tracking_streaming(): # Verify cost field is preserved in the Usage object - this is the key data for cost tracking # The chunk_parser converts the dict to a Usage Pydantic model which includes the cost field assert result2.usage.cost == 0.0001 + + +def test_openrouter_reasoning_models_allow_reasoning_effort_param(): + """ + OpenRouter reasoning-capable models should accept the reasoning_effort param. + """ + config = OpenrouterConfig() + + supported_params = config.get_supported_openai_params( + model="openrouter/deepseek/deepseek-v3.2" + ) + + assert "reasoning_effort" in supported_params + assert supported_params.count("reasoning_effort") == 1 + + +def test_openrouter_non_reasoning_models_do_not_add_reasoning_effort(): + """ + Models without reasoning support should not gain reasoning-specific params. + """ + config = OpenrouterConfig() + + supported_params = config.get_supported_openai_params( + model="openrouter/anthropic/claude-3-5-haiku" + ) + + assert "reasoning_effort" not in supported_params diff --git a/tests/test_litellm/llms/ragflow/chat/__init__.py b/tests/test_litellm/llms/ragflow/chat/__init__.py new file mode 100644 index 00000000000..4e074b84150 --- /dev/null +++ b/tests/test_litellm/llms/ragflow/chat/__init__.py @@ -0,0 +1,4 @@ +""" +RAGFlow chat transformation tests. +""" + diff --git a/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py b/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py new file mode 100644 index 00000000000..90f2504f94c --- /dev/null +++ b/tests/test_litellm/llms/ragflow/chat/test_ragflow_chat_transformation.py @@ -0,0 +1,376 @@ +""" +Test file for RAGFlow chat transformation functionality. + +Tests the model name parsing, URL construction, and request transformation +for RAGFlow's OpenAI-compatible API with custom path structures. +""" + +import os +import sys +from unittest.mock import Mock, patch + +import pytest + +# Add the project root to Python path +sys.path.insert(0, os.path.abspath("../../../../..")) + +import litellm +from litellm.llms.ragflow.chat.transformation import RAGFlowConfig +from litellm.types.llms.openai import AllMessageValues + + +class TestRAGFlowChatTransformation: + """Test suite for RAGFlow chat transformation functionality.""" + + def test_parse_ragflow_model_chat(self): + """Test parsing of chat model format.""" + config = RAGFlowConfig() + + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + endpoint_type, entity_id, model_name = config._parse_ragflow_model(model) + + assert endpoint_type == "chat" + assert entity_id == "my-chat-id" + assert model_name == "gpt-4o-mini" + + def test_parse_ragflow_model_agent(self): + """Test parsing of agent model format.""" + config = RAGFlowConfig() + + model = "ragflow/agent/my-agent-id/gpt-4o-mini" + endpoint_type, entity_id, model_name = config._parse_ragflow_model(model) + + assert endpoint_type == "agent" + assert entity_id == "my-agent-id" + assert model_name == "gpt-4o-mini" + + def test_parse_ragflow_model_with_slashes_in_model_name(self): + """Test parsing when model name contains slashes.""" + config = RAGFlowConfig() + + model = "ragflow/chat/my-chat-id/openai/gpt-4o-mini" + endpoint_type, entity_id, model_name = config._parse_ragflow_model(model) + + assert endpoint_type == "chat" + assert entity_id == "my-chat-id" + assert model_name == "openai/gpt-4o-mini" + + def test_parse_ragflow_model_invalid_format(self): + """Test parsing with invalid model format.""" + config = RAGFlowConfig() + + with pytest.raises(ValueError, match="Invalid RAGFlow model format"): + config._parse_ragflow_model("ragflow/chat/model-name") + + with pytest.raises(ValueError, match="Invalid RAGFlow model format"): + config._parse_ragflow_model("invalid/chat/id/model") + + with pytest.raises(ValueError, match="Must start with 'ragflow/'"): + config._parse_ragflow_model("not-ragflow/chat/id/model") + + def test_parse_ragflow_model_invalid_endpoint_type(self): + """Test parsing with invalid endpoint type.""" + config = RAGFlowConfig() + + with pytest.raises(ValueError, match="Invalid RAGFlow endpoint type"): + config._parse_ragflow_model("ragflow/invalid/my-id/model") + + def test_get_complete_url_chat(self): + """Test URL construction for chat endpoint.""" + config = RAGFlowConfig() + + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + api_base = "http://localhost:9380" + + url = config.get_complete_url( + api_base=api_base, + api_key=None, + model=model, + optional_params={}, + litellm_params={}, + stream=False, + ) + + assert url == "http://localhost:9380/api/v1/chats_openai/my-chat-id/chat/completions" + + def test_get_complete_url_agent(self): + """Test URL construction for agent endpoint.""" + config = RAGFlowConfig() + + model = "ragflow/agent/my-agent-id/gpt-4o-mini" + api_base = "http://localhost:9380" + + url = config.get_complete_url( + api_base=api_base, + api_key=None, + model=model, + optional_params={}, + litellm_params={}, + stream=False, + ) + + assert url == "http://localhost:9380/api/v1/agents_openai/my-agent-id/chat/completions" + + def test_get_complete_url_strips_v1(self): + """Test URL construction when api_base ends with /v1.""" + config = RAGFlowConfig() + + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + api_base = "http://localhost:9380/v1" + + url = config.get_complete_url( + api_base=api_base, + api_key=None, + model=model, + optional_params={}, + litellm_params={}, + stream=False, + ) + + assert url == "http://localhost:9380/api/v1/chats_openai/my-chat-id/chat/completions" + + def test_get_complete_url_strips_api_v1(self): + """Test URL construction when api_base ends with /api/v1.""" + config = RAGFlowConfig() + + model = "ragflow/agent/my-agent-id/gpt-4o-mini" + api_base = "http://localhost:9380/api/v1" + + url = config.get_complete_url( + api_base=api_base, + api_key=None, + model=model, + optional_params={}, + litellm_params={}, + stream=False, + ) + + assert url == "http://localhost:9380/api/v1/agents_openai/my-agent-id/chat/completions" + + def test_get_complete_url_from_litellm_params(self): + """Test URL construction with api_base from litellm_params.""" + config = RAGFlowConfig() + + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + # Create a simple dict-like object for litellm_params + class LiteLLMParams: + def __init__(self): + self.api_base = "http://ragflow-server:9380" + + litellm_params = LiteLLMParams() + + url = config.get_complete_url( + api_base=None, + api_key=None, + model=model, + optional_params={}, + litellm_params=litellm_params, + stream=False, + ) + + assert url == "http://ragflow-server:9380/api/v1/chats_openai/my-chat-id/chat/completions" + + def test_get_complete_url_missing_api_base(self): + """Test URL construction when api_base is missing.""" + config = RAGFlowConfig() + + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + + with pytest.raises(ValueError, match="api_base is required"): + config.get_complete_url( + api_base=None, + api_key=None, + model=model, + optional_params={}, + litellm_params={}, + stream=False, + ) + + @patch.dict(os.environ, {"RAGFLOW_API_BASE": "http://env-ragflow:9380"}) + def test_get_complete_url_from_environment(self): + """Test URL construction with api_base from environment variable.""" + config = RAGFlowConfig() + + model = "ragflow/agent/my-agent-id/gpt-4o-mini" + + url = config.get_complete_url( + api_base=None, + api_key=None, + model=model, + optional_params={}, + litellm_params={}, + stream=False, + ) + + assert url == "http://env-ragflow:9380/api/v1/agents_openai/my-agent-id/chat/completions" + + def test_validate_environment_sets_headers(self): + """Test that validate_environment sets proper headers.""" + config = RAGFlowConfig() + + headers = {} + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + messages = [{"role": "user", "content": "Hello"}] + api_key = "test-api-key" + + result_headers = config.validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params={}, + litellm_params={}, + api_key=api_key, + api_base="http://localhost:9380", + ) + + assert result_headers["Authorization"] == "Bearer test-api-key" + assert result_headers["Content-Type"] == "application/json" + + def test_validate_environment_stores_actual_model(self): + """Test that validate_environment stores actual model name.""" + config = RAGFlowConfig() + + headers = {} + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + messages = [{"role": "user", "content": "Hello"}] + litellm_params = {} + + config.validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params={}, + litellm_params=litellm_params, + api_key="test-key", + api_base="http://localhost:9380", + ) + + assert litellm_params["_ragflow_actual_model"] == "gpt-4o-mini" + + @patch.dict(os.environ, {"RAGFLOW_API_KEY": "env-api-key"}) + def test_validate_environment_from_environment(self): + """Test that validate_environment gets api_key from environment.""" + config = RAGFlowConfig() + + headers = {} + model = "ragflow/agent/my-agent-id/gpt-4o-mini" + messages = [{"role": "user", "content": "Hello"}] + + result_headers = config.validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params={}, + litellm_params={}, + api_key=None, + api_base="http://localhost:9380", + ) + + assert result_headers["Authorization"] == "Bearer env-api-key" + + def test_validate_environment_from_litellm_params(self): + """Test that validate_environment gets api_key from litellm_params.""" + config = RAGFlowConfig() + + headers = {} + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + messages = [{"role": "user", "content": "Hello"}] + # Create a simple object for litellm_params with api_key attribute + class LiteLLMParams: + def __init__(self): + self.api_key = "litellm-params-key" + def __setitem__(self, key, value): + setattr(self, key, value) + + litellm_params = LiteLLMParams() + + result_headers = config.validate_environment( + headers=headers, + model=model, + messages=messages, + optional_params={}, + litellm_params=litellm_params, + api_key=None, + api_base="http://localhost:9380", + ) + + assert result_headers["Authorization"] == "Bearer litellm-params-key" + + def test_transform_request_uses_actual_model(self): + """Test that transform_request uses the actual model name.""" + config = RAGFlowConfig() + + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + messages = [{"role": "user", "content": "Hello"}] + litellm_params = {"_ragflow_actual_model": "gpt-4o-mini"} + + # Test the actual behavior by checking the model in the result + result = config.transform_request( + model=model, + messages=messages, + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + + # The result should contain the actual model name, not the full ragflow path + assert result["model"] == "gpt-4o-mini" + assert result["messages"] == messages + + def test_transform_request_fallback_parsing(self): + """Test that transform_request falls back to parsing if _ragflow_actual_model is missing.""" + config = RAGFlowConfig() + + model = "ragflow/agent/my-agent-id/gpt-4o-mini" + messages = [{"role": "user", "content": "Hello"}] + litellm_params = {} # Missing _ragflow_actual_model + + result = config.transform_request( + model=model, + messages=messages, + optional_params={}, + litellm_params=litellm_params, + headers={}, + ) + + # Should parse and use the actual model name + assert result["model"] == "gpt-4o-mini" + assert result["messages"] == messages + + def test_get_openai_compatible_provider_info(self): + """Test _get_openai_compatible_provider_info returns correct values.""" + config = RAGFlowConfig() + + model = "ragflow/chat/my-chat-id/gpt-4o-mini" + api_base = "http://localhost:9380" + api_key = "test-key" + + result_api_base, result_api_key, result_provider = config._get_openai_compatible_provider_info( + model=model, + api_base=api_base, + api_key=api_key, + custom_llm_provider="ragflow", + ) + + assert result_api_base == api_base + assert result_api_key == api_key + assert result_provider == "ragflow" + + @patch.dict(os.environ, {"RAGFLOW_API_BASE": "http://env-base:9380", "RAGFLOW_API_KEY": "env-key"}) + def test_get_openai_compatible_provider_info_from_env(self): + """Test _get_openai_compatible_provider_info gets values from environment.""" + config = RAGFlowConfig() + + model = "ragflow/agent/my-agent-id/gpt-4o-mini" + + result_api_base, result_api_key, result_provider = config._get_openai_compatible_provider_info( + model=model, + api_base=None, + api_key=None, + custom_llm_provider="ragflow", + ) + + assert result_api_base == "http://env-base:9380" + assert result_api_key == "env-key" + assert result_provider == "ragflow" + diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py b/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py index 68c0f3bdbfc..46bb8930a7a 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_thought_signature_in_tool_call_id.py @@ -4,9 +4,13 @@ Tests for embedding thought signatures in tool call IDs for OpenAI client compat When using OpenAI clients (instead of LiteLLM SDK), provider_specific_fields are not preserved. This test suite validates that thought signatures can be embedded in tool call IDs and extracted when converting back to Gemini format. + +Note: Embedding signatures in tool call IDs is a beta feature that requires +enable_preview_features=True to be enabled. """ import pytest +import litellm from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) @@ -62,36 +66,57 @@ def test_encode_tool_call_id_without_signature(): assert decoded_signature is None -def test_tool_call_id_includes_signature_in_response(): - """Test that tool call IDs in responses include embedded thought signatures""" +@pytest.mark.parametrize("enable_preview_features", [True, False]) +def test_tool_call_id_includes_signature_in_response(enable_preview_features): + """Test that tool call IDs in responses include embedded thought signatures only when preview features are enabled""" test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - parts_with_signature = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, + # Save original state + original_flag = litellm.enable_preview_features + litellm.enable_preview_features = enable_preview_features + + try: + parts_with_signature = [ + HttpxPartType( + functionCall={ + "name": "get_current_temperature", + "args": {"location": "Paris"}, + }, + thoughtSignature=test_signature, + ) + ] + + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=parts_with_signature, + cumulative_tool_call_idx=0, + is_function_call=False, ) - ] - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=parts_with_signature, - cumulative_tool_call_idx=0, - is_function_call=False, - ) + # Verify tool call exists + assert tools is not None + assert len(tools) == 1 + tool_call_id = tools[0]["id"] + + # Verify signature is always in provider_specific_fields + assert tools[0].get("provider_specific_fields", {}).get("thought_signature") == test_signature - # Verify tool call ID includes thought signature - assert tools is not None - assert len(tools) == 1 - tool_call_id = tools[0]["id"] - assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id - - # Verify we can decode it using the factory function - tool_obj = {"id": tool_call_id, "type": "function"} - decoded_sig = _get_thought_signature_from_tool(tool_obj) - assert decoded_sig == test_signature + if enable_preview_features: + # When preview features enabled, signature should be embedded in ID + assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id + # Verify we can decode it using the factory function + tool_obj = {"id": tool_call_id, "type": "function"} + decoded_sig = _get_thought_signature_from_tool(tool_obj) + assert decoded_sig == test_signature + else: + # When preview features disabled, signature should NOT be embedded in ID + assert THOUGHT_SIGNATURE_SEPARATOR not in tool_call_id + # But we can still extract from provider_specific_fields + tool_obj = {"id": tool_call_id, "type": "function", "provider_specific_fields": {"thought_signature": test_signature}} + decoded_sig = _get_thought_signature_from_tool(tool_obj) + assert decoded_sig == test_signature + finally: + # Restore original state + litellm.enable_preview_features = original_flag def test_get_thought_signature_backward_compatibility(): @@ -168,97 +193,157 @@ def test_convert_to_gemini_with_embedded_signature(): assert gemini_parts[0]["thoughtSignature"] == test_signature -def test_openai_client_e2e_flow(): +@pytest.mark.parametrize("enable_preview_features", [True, False]) +def test_openai_client_e2e_flow(enable_preview_features): """ End-to-end test simulating OpenAI client usage: 1. LiteLLM receives response from Gemini with thought signature - 2. LiteLLM embeds signature in tool call ID + 2. LiteLLM embeds signature in tool call ID (if preview features enabled) 3. OpenAI client sends message back with same tool call ID - 4. LiteLLM extracts signature from ID and sends to Gemini + 4. LiteLLM extracts signature from ID/provider_specific_fields and sends to Gemini """ test_signature = "Co4CAdHtim/rWgXbz2Ghp4tShzLeMASrPw6JJyYIC3cbVyZnKzU3uv8/wVzyS2sKRPL2m8QQHHXbNQhEEz500G7n/4ZMmksdTtfQcJMoT76S1DGwhnAiLwTgWCNXs3lEb4M19EVYoWFxhrH5Lr9YMIquoU9U4paydGwvZyIyigamIg4B6WnxrRsf0KZV12gJed0DZuKczvOFtHz3zUnmZRlOiTzd5gBVyQM+5jv1VI8m4WUKd6cN/5a5ZvaA0ggiO6kdVhlpIVs7GczSEVJD8KH4u02X7VSnb7CvykqDntZzV0y8rZFBEFGKrChmeHlWXP4D1IB3F9KQyhuLgWImMzg4BajKVxxMU737JGnNISy5" - # Step 1: Gemini returns function call with thought signature - gemini_parts = [ - HttpxPartType( - functionCall={ - "name": "get_current_temperature", - "args": {"location": "Paris"}, - }, - thoughtSignature=test_signature, - ) - ] + # Save original state + original_flag = litellm.enable_preview_features + litellm.enable_preview_features = enable_preview_features - # Step 2: LiteLLM transforms to OpenAI format with embedded signature - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=gemini_parts, - cumulative_tool_call_idx=0, - is_function_call=False, - ) - - assert tools is not None - assert len(tools) == 1 - tool_call_id = tools[0]["id"] - assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id - - # Step 3: OpenAI client sends back assistant message (preserves tool_call_id) - openai_assistant_message = { - "role": "assistant", - "content": "", - "tool_calls": [ - { - "id": tool_call_id, # Preserved from response - "type": "function", - "function": { + try: + # Step 1: Gemini returns function call with thought signature + gemini_parts = [ + HttpxPartType( + functionCall={ "name": "get_current_temperature", - "arguments": '{"location": "Paris"}', + "args": {"location": "Paris"}, }, + thoughtSignature=test_signature, + ) + ] + + # Step 2: LiteLLM transforms to OpenAI format + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=gemini_parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) + + assert tools is not None + assert len(tools) == 1 + tool_call_id = tools[0]["id"] + + if enable_preview_features: + # When preview features enabled, signature should be embedded in ID + assert THOUGHT_SIGNATURE_SEPARATOR in tool_call_id + else: + # When preview features disabled, signature should NOT be embedded in ID + assert THOUGHT_SIGNATURE_SEPARATOR not in tool_call_id + + # Step 3: OpenAI client sends back assistant message + # For the disabled case, we simulate that the client might have provider_specific_fields + # or we use the embedded ID if preview features were enabled + if enable_preview_features: + openai_assistant_message = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": tool_call_id, # Preserved from response (with embedded signature) + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + }, + } + ], + } + else: + # When preview features disabled, simulate that provider_specific_fields might be preserved + # (though in real OpenAI client usage, this might not happen) + # For this test, we'll use provider_specific_fields to show extraction still works + openai_assistant_message = { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": tool_call_id, # ID without embedded signature + "type": "function", + "function": { + "name": "get_current_temperature", + "arguments": '{"location": "Paris"}', + }, + "provider_specific_fields": {"thought_signature": test_signature}, + } + ], } - ], - } - # Step 4: LiteLLM converts back to Gemini format, extracting signature - gemini_parts_converted = convert_to_gemini_tool_call_invoke( - openai_assistant_message - ) + # Step 4: LiteLLM converts back to Gemini format, extracting signature + gemini_parts_converted = convert_to_gemini_tool_call_invoke( + openai_assistant_message + ) - # Verify signature is preserved through the round trip - assert len(gemini_parts_converted) == 1 - assert "thoughtSignature" in gemini_parts_converted[0] - assert gemini_parts_converted[0]["thoughtSignature"] == test_signature + # Verify signature is preserved through the round trip + assert len(gemini_parts_converted) == 1 + assert "thoughtSignature" in gemini_parts_converted[0] + assert gemini_parts_converted[0]["thoughtSignature"] == test_signature + finally: + # Restore original state + litellm.enable_preview_features = original_flag -def test_parallel_tool_calls_with_signatures(): +@pytest.mark.parametrize("enable_preview_features", [True, False]) +def test_parallel_tool_calls_with_signatures(enable_preview_features): """Test that parallel tool calls preserve signatures correctly""" signature1 = "signature_for_first_call" # Only first call has signature (Gemini behavior for parallel calls) - gemini_parts = [ - HttpxPartType( - functionCall={"name": "get_temperature", "args": {"location": "Paris"}}, - thoughtSignature=signature1, - ), - HttpxPartType( - functionCall={"name": "get_temperature", "args": {"location": "London"}}, - # No signature for second parallel call - ), - ] + # Save original state + original_flag = litellm.enable_preview_features + litellm.enable_preview_features = enable_preview_features - function, tools, _ = VertexGeminiConfig._transform_parts( - parts=gemini_parts, - cumulative_tool_call_idx=0, - is_function_call=False, - ) + try: + gemini_parts = [ + HttpxPartType( + functionCall={"name": "get_temperature", "args": {"location": "Paris"}}, + thoughtSignature=signature1, + ), + HttpxPartType( + functionCall={"name": "get_temperature", "args": {"location": "London"}}, + # No signature for second parallel call + ), + ] - assert tools is not None - assert len(tools) == 2 + function, tools, _ = VertexGeminiConfig._transform_parts( + parts=gemini_parts, + cumulative_tool_call_idx=0, + is_function_call=False, + ) - # First tool call has signature in ID - assert THOUGHT_SIGNATURE_SEPARATOR in tools[0]["id"] - sig1 = _get_thought_signature_from_tool({"id": tools[0]["id"], "type": "function"}) - assert sig1 == signature1 + assert tools is not None + assert len(tools) == 2 - # Second tool call has no signature in ID - assert THOUGHT_SIGNATURE_SEPARATOR not in tools[1]["id"] - sig2 = _get_thought_signature_from_tool({"id": tools[1]["id"], "type": "function"}) - assert sig2 is None + # First tool call should have signature in provider_specific_fields + assert tools[0].get("provider_specific_fields", {}).get("thought_signature") == signature1 + + if enable_preview_features: + # When preview features enabled, first tool call has signature in ID + assert THOUGHT_SIGNATURE_SEPARATOR in tools[0]["id"] + sig1 = _get_thought_signature_from_tool({"id": tools[0]["id"], "type": "function"}) + assert sig1 == signature1 + else: + # When preview features disabled, signature should NOT be in ID + assert THOUGHT_SIGNATURE_SEPARATOR not in tools[0]["id"] + # But we can extract from provider_specific_fields + sig1 = _get_thought_signature_from_tool({ + "id": tools[0]["id"], + "type": "function", + "provider_specific_fields": {"thought_signature": signature1} + }) + assert sig1 == signature1 + + # Second tool call has no signature in ID (regardless of flag) + assert THOUGHT_SIGNATURE_SEPARATOR not in tools[1]["id"] + sig2 = _get_thought_signature_from_tool({"id": tools[1]["id"], "type": "function"}) + assert sig2 is None + finally: + # Restore original state + litellm.enable_preview_features = original_flag 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 940dcf1a7a1..86e25c46c8e 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 @@ -9,6 +9,8 @@ from fastapi import HTTPException sys.path.insert(0, "../../../../../") import httpx +from mcp import ReadResourceResult, Resource +from mcp.types import GetPromptResult, Prompt, ResourceTemplate, TextResourceContents from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( MCPServerManager, @@ -17,8 +19,6 @@ from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( from litellm.proxy._types import LiteLLM_MCPServerTable, MCPTransport from litellm.types.mcp import MCPAuth from litellm.types.mcp_server.mcp_server_manager import MCPOAuthMetadata, MCPServer -from mcp import ReadResourceResult, Resource -from mcp.types import GetPromptResult, Prompt, ResourceTemplate, TextResourceContents class TestMCPServerManager: @@ -1606,6 +1606,104 @@ class TestMCPServerManager: # Verify the MCP client call was awaited exactly once assert mock_client.call_tool.await_count == 1 + @pytest.mark.asyncio + async def test_get_allowed_mcp_servers_with_user_api_key_auth(self): + """ + Test that get_allowed_mcp_servers properly receives and uses user_api_key_auth + when called. This verifies the fix where user_api_key_auth is passed through + litellm_metadata from responses API. + """ + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) + from litellm.proxy._types import LiteLLM_ObjectPermissionTable, UserAPIKeyAuth + + manager = MCPServerManager() + + # Create a mock user_api_key_auth with object_permission + object_permission = LiteLLM_ObjectPermissionTable( + object_permission_id="perm_123", + mcp_servers=["test_server_1", "test_server_2"], + mcp_access_groups=[], + ) + + user_api_key_auth = UserAPIKeyAuth( + api_key="sk-test", + user_id="user-123", + object_permission=object_permission, + object_permission_id="perm_123", + ) + + # Mock MCPRequestHandler.get_allowed_mcp_servers to verify it receives user_api_key_auth + with patch.object( + MCPRequestHandler, + "get_allowed_mcp_servers", + new_callable=AsyncMock, + ) as mock_get_allowed: + # Configure mock to return servers from object_permission + mock_get_allowed.return_value = ["test_server_1", "test_server_2"] + + # Call get_allowed_mcp_servers with user_api_key_auth + result = await manager.get_allowed_mcp_servers(user_api_key_auth) + + # Verify MCPRequestHandler.get_allowed_mcp_servers was called with user_api_key_auth + mock_get_allowed.assert_called_once() + call_args = mock_get_allowed.call_args + assert call_args[0][0] is user_api_key_auth # First positional arg should be user_api_key_auth + assert call_args[0][0].user_id == "user-123" + assert call_args[0][0].object_permission_id == "perm_123" + assert call_args[0][0].object_permission is not None + assert call_args[0][0].object_permission.mcp_servers == ["test_server_1", "test_server_2"] + + # Verify result contains the expected servers + assert "test_server_1" in result + assert "test_server_2" in result + + def test_get_mcp_server_from_tool_name_uses_server_name_not_name(self): + """ + Test that _get_mcp_server_from_tool_name uses server.server_name instead of server.name + when extracting server name from prefixed tool name (second case). + This ensures the fix for using server_name instead of name works correctly. + """ + from litellm.proxy._experimental.mcp_server.utils import ( + add_server_prefix_to_name, + ) + + manager = MCPServerManager() + + # Create a server where server_name differs from name + # This tests the scenario where server.name != server.server_name + server = MCPServer( + server_id="test-server-id", + name="Test Server Name", # Different from server_name + server_name="test_server", # This is what should be used + alias="test_server", + transport=MCPTransport.http, + ) + + # Register the server + manager.registry = {server.server_id: server} + + # Create a tool with prefixed name + tool_name = "test_tool" + prefixed_tool_name = add_server_prefix_to_name(tool_name, "test_server") + + # Populate the mapping with the original tool name + manager.tool_name_to_mcp_server_name_mapping[tool_name] = "test_server" + manager.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = "test_server" + + # Test: _get_mcp_server_from_tool_name should find the server using server.server_name + # even when server.name is different + resolved_server = manager._get_mcp_server_from_tool_name(prefixed_tool_name) + + # Verify the server was found correctly + assert resolved_server is not None + assert resolved_server.server_id == server.server_id + assert resolved_server.server_name == "test_server" + # Verify it matched using server_name, not name + assert resolved_server.name == "Test Server Name" # name is different + assert resolved_server.server_name == "test_server" # server_name matches + if __name__ == "__main__": pytest.main([__file__]) diff --git a/tests/test_litellm/proxy/agent_endpoints/__init__.py b/tests/test_litellm/proxy/agent_endpoints/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py new file mode 100644 index 00000000000..c7257073b33 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -0,0 +1,114 @@ +""" +Mock tests for A2A endpoints. + +Tests that invoke_agent_a2a properly integrates with add_litellm_data_to_request. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +@pytest.mark.asyncio +async def test_invoke_agent_a2a_adds_litellm_data(): + """ + Test that invoke_agent_a2a calls add_litellm_data_to_request + and the resulting data includes proxy_server_request. + """ + from litellm.proxy._types import UserAPIKeyAuth + + # Track the data passed to add_litellm_data_to_request + captured_data = {} + + async def mock_add_litellm_data(data, **kwargs): + # Simulate what add_litellm_data_to_request does + data["proxy_server_request"] = { + "url": "http://localhost:4000/a2a/test-agent", + "method": "POST", + "headers": {}, + "body": dict(data), + } + captured_data.update(data) + return data + + # Mock response from asend_message + mock_response = MagicMock() + mock_response.model_dump.return_value = { + "jsonrpc": "2.0", + "id": "test-id", + "result": {"status": "success"}, + } + + # Mock agent + mock_agent = MagicMock() + mock_agent.agent_card_params = { + "url": "http://backend-agent:10001", + "name": "Test Agent", + } + + # Mock request + mock_request = MagicMock() + mock_request.json = AsyncMock(return_value={ + "jsonrpc": "2.0", + "id": "test-id", + "method": "message/send", + "params": { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-123", + } + }, + }) + + mock_user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key", + user_id="test-user", + team_id="test-team", + ) + + # Patch at the source modules + with patch( + "litellm.proxy.agent_endpoints.a2a_endpoints._get_agent", + return_value=mock_agent, + ), patch( + "litellm.proxy.litellm_pre_call_utils.add_litellm_data_to_request", + side_effect=mock_add_litellm_data, + ) as mock_add_data, patch( + "litellm.a2a_protocol.create_a2a_client", + new_callable=AsyncMock, + ), patch( + "litellm.a2a_protocol.asend_message", + new_callable=AsyncMock, + return_value=mock_response, + ), patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ), patch( + "litellm.proxy.proxy_server.proxy_config", + MagicMock(), + ), patch( + "litellm.proxy.proxy_server.version", + "1.0.0", + ): + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + mock_fastapi_response = MagicMock() + + result = await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify add_litellm_data_to_request was called + mock_add_data.assert_called_once() + + # Verify model and custom_llm_provider were set + assert captured_data.get("model") == "a2a_agent/Test Agent" + assert captured_data.get("custom_llm_provider") == "a2a_agent" + + # Verify proxy_server_request was added + assert "proxy_server_request" in captured_data + assert captured_data["proxy_server_request"]["method"] == "POST" diff --git a/tests/test_litellm/proxy/auth/test_login_utils.py b/tests/test_litellm/proxy/auth/test_login_utils.py new file mode 100644 index 00000000000..201461dc8b5 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_login_utils.py @@ -0,0 +1,284 @@ +""" +Tests for login_utils module. + +This module tests the refactored login logic that was moved from proxy_server.py +to login_utils.py for better reusability. +""" + +import os +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.proxy._types import ( + LiteLLM_UserTable, + LitellmUserRoles, + ProxyErrorTypes, + ProxyException, + hash_token, +) +from litellm.proxy.auth.login_utils import ( + LoginResult, + authenticate_user, + get_ui_credentials, +) + + +def test_get_ui_credentials_prefers_explicit_password(): + """The configured UI password should be returned when available.""" + with patch.dict( + os.environ, + {"UI_USERNAME": "test-admin", "UI_PASSWORD": "secure-pass"}, + clear=True, + ): + username, password = get_ui_credentials(master_key="sk-123") + + assert username == "test-admin" + assert password == "secure-pass" + + +def test_get_ui_credentials_can_use_master_key(): + """Master key should be used as password when UI_PASSWORD is missing.""" + with patch.dict(os.environ, {"UI_USERNAME": "fallback-admin"}, clear=True): + username, password = get_ui_credentials(master_key="fallback-key") + + assert username == "fallback-admin" + assert password == "fallback-key" + + +def test_get_ui_credentials_requires_password(): + """Missing UI password and master key results in error.""" + with patch.dict(os.environ, {}, clear=True): + with pytest.raises(ProxyException) as exc_info: + get_ui_credentials(master_key=None) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "500" + + +@pytest.mark.asyncio +async def test_authenticate_user_admin_login_with_ui_credentials(): + """Test admin login using UI_USERNAME and UI_PASSWORD""" + master_key = "sk-1234" + ui_username = "admin" + ui_password = "sk-1234" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict( + os.environ, + { + "UI_USERNAME": ui_username, + "UI_PASSWORD": ui_password, + "DATABASE_URL": "postgresql://test:test@localhost/test", + }, + ): + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.return_value = { + "token": "test-token-123", + "user_id": LITELLM_PROXY_ADMIN_NAME, + } + + with patch( + "litellm.proxy.auth.login_utils.user_update", + new_callable=AsyncMock, + return_value=None, + ) as mock_user_update: + with patch( + "litellm.proxy.auth.login_utils.get_secret_bool", + return_value=False, + ): + result = await authenticate_user( + username=ui_username, + password=ui_password, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + assert result.key == "test-token-123" + assert result.user_email is None + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + assert result.login_method == "username_password" + + +@pytest.mark.asyncio +async def test_authenticate_user_admin_login_with_master_key_as_password(): + """Test admin login when UI_PASSWORD is not set, should use master_key""" + master_key = "sk-1234" + ui_username = "admin" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + env_vars = {"UI_USERNAME": ui_username, "DATABASE_URL": "postgresql://test:test@localhost/test"} + # Remove UI_PASSWORD to test fallback to master_key + if "UI_PASSWORD" in os.environ: + # Keep other env vars but don't set UI_PASSWORD + pass + else: + # Ensure UI_PASSWORD is not in the patched env + pass + + with patch.dict(os.environ, env_vars, clear=False): + # Explicitly remove UI_PASSWORD if it exists + original_ui_password = os.environ.pop("UI_PASSWORD", None) + try: + with patch( + "litellm.proxy.auth.login_utils.generate_key_helper_fn", + new_callable=AsyncMock, + ) as mock_generate_key: + mock_generate_key.return_value = { + "token": "test-token-123", + "user_id": LITELLM_PROXY_ADMIN_NAME, + } + + with patch( + "litellm.proxy.auth.login_utils.user_update", + new_callable=AsyncMock, + return_value=None, + ) as mock_user_update: + with patch( + "litellm.proxy.auth.login_utils.get_secret_bool", + return_value=False, + ): + result = await authenticate_user( + username=ui_username, + password=master_key, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert isinstance(result, LoginResult) + assert result.user_id == LITELLM_PROXY_ADMIN_NAME + assert result.user_role == LitellmUserRoles.PROXY_ADMIN + finally: + if original_ui_password: + os.environ["UI_PASSWORD"] = original_ui_password + +@pytest.mark.asyncio +async def test_authenticate_user_invalid_credentials(): + """Test authentication failure with invalid credentials""" + master_key = "sk-1234" + ui_username = "admin" + wrong_password = "wrong-password" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": "correct-password"}): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=ui_username, + password=wrong_password, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "401" + assert "Invalid credentials" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_authenticate_user_missing_master_key(): + """Test authentication failure when master_key is None""" + mock_prisma_client = MagicMock() + + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username="admin", + password="password", + master_key=None, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "500" + assert "Master Key not set" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_authenticate_user_wrong_password(): + """Test authentication failure with wrong password for database user""" + master_key = "sk-1234" + user_email = "test@example.com" + correct_password = "correct-password" + wrong_password = "wrong-password" + hashed_password = hash_token(token=correct_password) + + mock_user = LiteLLM_UserTable( + user_id="test-user-123", + user_email=user_email, + password=hashed_password, + user_role=LitellmUserRoles.INTERNAL_USER, + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock( + return_value=mock_user + ) + + with patch.dict( + os.environ, + { + "DATABASE_URL": "postgresql://test:test@localhost/test", + "UI_USERNAME": "admin", + "UI_PASSWORD": "admin-password", + }, + ): + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=user_email, + password=wrong_password, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "401" + assert "Invalid credentials" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_authenticate_user_database_required_for_admin(): + """Test that database is required for admin login""" + master_key = "sk-1234" + ui_username = "admin" + ui_password = "sk-1234" + + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_usertable.find_first = AsyncMock(return_value=None) + + with patch.dict(os.environ, {"UI_USERNAME": ui_username, "UI_PASSWORD": ui_password}): + with patch( + "litellm.proxy.auth.login_utils.user_update", + new_callable=AsyncMock, + return_value=None, + ): + # Remove DATABASE_URL to simulate no database + original_db_url = os.environ.get("DATABASE_URL") + if "DATABASE_URL" in os.environ: + del os.environ["DATABASE_URL"] + + try: + with pytest.raises(ProxyException) as exc_info: + await authenticate_user( + username=ui_username, + password=ui_password, + master_key=master_key, + prisma_client=mock_prisma_client, + ) + + assert exc_info.value.type == ProxyErrorTypes.auth_error + assert exc_info.value.code == "500" + assert "No Database connected" in exc_info.value.message + finally: + if original_db_url: + os.environ["DATABASE_URL"] = original_db_url diff --git a/tests/test_litellm/proxy/auth/test_organization_budget_enforcement.py b/tests/test_litellm/proxy/auth/test_organization_budget_enforcement.py new file mode 100644 index 00000000000..9c2adca9cd3 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_organization_budget_enforcement.py @@ -0,0 +1,344 @@ +""" +Tests for organization budget enforcement. + +These tests verify that organization-level budgets are properly enforced during +request authentication. When an organization's spend exceeds its max_budget, +requests should fail with BudgetExceededError. + +This prevents teams within an organization from collectively exceeding the +organization's budget limit. +""" + +import asyncio +import os +import sys +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../../")) + +import litellm +from litellm.proxy._types import ( + LiteLLM_BudgetTable, + LiteLLM_OrganizationTable, + LiteLLM_TeamTable, + UserAPIKeyAuth, +) +from litellm.proxy.auth.auth_checks import common_checks +from litellm.proxy.utils import ProxyLogging + + +@pytest.mark.asyncio +async def test_organization_budget_exceeded_blocks_request(): + """ + Bug: Organization budget is retrieved but NEVER enforced. + + When organization spend >= organization_max_budget, requests should fail + with BudgetExceededError. Currently this passes because no check exists. + """ + org_id = "test-org-budget-exceeded" + + # Organization with max_budget of 100, but spend is 150 + org_object = LiteLLM_OrganizationTable( + organization_id=org_id, + budget_id="org-budget-1", + spend=150.0, # Over budget! + models=["gpt-4"], + created_by="test", + updated_by="test", + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=100.0, # Budget is 100 + ), + ) + + # Team within the organization (team itself is under budget) + team_object = LiteLLM_TeamTable( + team_id="test-team-1", + organization_id=org_id, + max_budget=50.0, # Team budget is 50 + spend=10.0, # Team spend is only 10 - under budget + models=["gpt-4"], + ) + + # Valid token with organization info + valid_token = UserAPIKeyAuth( + token="sk-test-123", + team_id="test-team-1", + org_id=org_id, + organization_max_budget=100.0, # This is set but never checked! + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = AsyncMock() + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache: + with patch("litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock) as mock_get_org: + mock_get_org.return_value = org_object + + # BUG: This should raise BudgetExceededError but currently passes + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body={"model": "gpt-4"}, + team_object=team_object, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token, + request=mock_request, + ) + + assert "Organization" in str(exc_info.value.message) + assert exc_info.value.current_cost == 150.0 + assert exc_info.value.max_budget == 100.0 + + +@pytest.mark.asyncio +async def test_multiple_teams_exceed_organization_budget(): + """ + Test that organization budget is enforced even when individual teams are under budget. + + Scenario: + - Organization max_budget = $5000, spend = $5000 (at limit) + - Team A spend = $1500 (under team budget of $2000) + - Request via Team A should FAIL because org is at budget limit + + Expected: Request fails with BudgetExceededError + """ + org_id = "multi-team-org" + + # Organization at budget limit + org_object = LiteLLM_OrganizationTable( + organization_id=org_id, + budget_id="org-budget-2", + spend=5000.0, # At $5000 limit + models=["gpt-4"], + created_by="test", + updated_by="test", + litellm_budget_table=LiteLLM_BudgetTable( + max_budget=5000.0, # Org budget is $5000 + ), + ) + + # Team A - under its own budget, but org is almost at limit + team_a = LiteLLM_TeamTable( + team_id="team-a", + organization_id=org_id, + max_budget=2000.0, + spend=1500.0, # Team A has spent $1500 of its $2000 budget + models=["gpt-4"], + ) + + valid_token = UserAPIKeyAuth( + token="sk-team-a-key", + team_id="team-a", + org_id=org_id, + organization_max_budget=5000.0, # Set but never enforced + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = AsyncMock() + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache: + with patch("litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock) as mock_get_org: + mock_get_org.return_value = org_object + + # Org is at budget limit, should raise BudgetExceededError + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body={"model": "gpt-4"}, + team_object=team_a, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token, + request=mock_request, + ) + + # Verify the error message mentions organization + assert "Organization" in str(exc_info.value.message) + assert exc_info.value.current_cost == 5000.0 + assert exc_info.value.max_budget == 5000.0 + + +@pytest.mark.asyncio +async def test_organization_budget_fields_are_checked(): + """ + Verify that organization_max_budget is populated in UserAPIKeyAuth + and BudgetExceededError is raised when organization is over budget. + """ + # Token has org budget info + valid_token = UserAPIKeyAuth( + token="sk-test", + team_id="test-team", + org_id="test-org", + organization_max_budget=100.0, # Budget is $100 + ) + + # Verify the field exists and is set + assert valid_token.organization_max_budget == 100.0 + assert valid_token.org_id == "test-org" + + team_object = LiteLLM_TeamTable( + team_id="test-team", + organization_id="test-org", + max_budget=None, + spend=0.0, + models=["gpt-4"], + ) + + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = AsyncMock() + + # Organization is over budget + org_over_budget = LiteLLM_OrganizationTable( + organization_id="test-org", + budget_id="budget-1", + spend=150.0, # Over $100 budget + models=["gpt-4"], + created_by="test", + updated_by="test", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache: + with patch("litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock) as mock_get_org: + mock_get_org.return_value = org_over_budget + + # Organization is over budget, should raise BudgetExceededError + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body={"model": "gpt-4"}, + team_object=team_object, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token, + request=mock_request, + ) + + assert exc_info.value.current_cost == 150.0 + assert exc_info.value.max_budget == 100.0 + + +@pytest.mark.asyncio +async def test_both_team_and_org_budget_enforced(): + """ + Verify that both team budget and organization budget are enforced consistently. + + This test verifies: + 1. Team over budget raises BudgetExceededError + 2. Organization over budget also raises BudgetExceededError + """ + mock_request = MagicMock() + mock_request.url.path = "/v1/chat/completions" + + mock_proxy_logging = MagicMock(spec=ProxyLogging) + mock_proxy_logging.budget_alerts = AsyncMock() + + # Scenario A: Team over budget - should raise BudgetExceededError + team_over_budget = LiteLLM_TeamTable( + team_id="team-over", + max_budget=100.0, + spend=150.0, # Over budget + models=["gpt-4"], + ) + + valid_token_team = UserAPIKeyAuth( + token="sk-team-test", + team_id="team-over", + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache: + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body={"model": "gpt-4"}, + team_object=team_over_budget, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token_team, + request=mock_request, + ) + assert "Team" in str(exc_info.value.message) + + # Scenario B: Org over budget - should also raise BudgetExceededError + org_over_budget = LiteLLM_OrganizationTable( + organization_id="org-over", + budget_id="budget-1", + spend=150.0, # Over $100 budget + models=["gpt-4"], + created_by="test", + updated_by="test", + litellm_budget_table=LiteLLM_BudgetTable(max_budget=100.0), + ) + + team_under_budget = LiteLLM_TeamTable( + team_id="team-under", + organization_id="org-over", + max_budget=50.0, + spend=10.0, # Team is fine + models=["gpt-4"], + ) + + valid_token_org = UserAPIKeyAuth( + token="sk-org-test", + team_id="team-under", + org_id="org-over", + organization_max_budget=100.0, + ) + + with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: + with patch("litellm.proxy.proxy_server.user_api_key_cache") as mock_cache: + with patch("litellm.proxy.auth.auth_checks.get_org_object", new_callable=AsyncMock) as mock_get_org: + mock_get_org.return_value = org_over_budget + + # Organization is over budget, should raise BudgetExceededError + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await common_checks( + request_body={"model": "gpt-4"}, + team_object=team_under_budget, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings={}, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=mock_proxy_logging, + valid_token=valid_token_org, + request=mock_request, + ) + + assert "Organization" in str(exc_info.value.message) + assert exc_info.value.current_cost == 150.0 + assert exc_info.value.max_budget == 100.0 diff --git a/tests/test_litellm/proxy/common_utils/test_custom_openapi_spec.py b/tests/test_litellm/proxy/common_utils/test_custom_openapi_spec.py index 7549b4259af..5fef35eb821 100644 --- a/tests/test_litellm/proxy/common_utils/test_custom_openapi_spec.py +++ b/tests/test_litellm/proxy/common_utils/test_custom_openapi_spec.py @@ -90,6 +90,35 @@ class TestCustomOpenAPISpec: ) assert result == base_openapi_schema +def test_defs_rewritten_in_add_schema_to_components(): + """ + Test that defs are rewritten to components/schemas in add_schema_to_components. + """ + + openapi_schema = {} + schema_name = "SchemaName" + schema_def = { + "type": "object", + "properties": { + "messages": { + "type": "array", + "items": { + "anyOf": [ + {"$ref": "#/$defs/UserMessage"}, + {"$ref": "#/$defs/AssistantMessage"} + ] + } + } + }, + "$defs": { + "UserMessage": {"type": "object"}, + "AssistantMessage": {"type": "object"} + } + } + CustomOpenAPISpec.add_schema_to_components(openapi_schema=openapi_schema, schema_name=schema_name, schema_def=schema_def) + assert "$defs" not in openapi_schema + assert openapi_schema["components"]["schemas"]["SchemaName"]["properties"]["messages"]["items"]["anyOf"][0]["$ref"] == "#/components/schemas/UserMessage" + assert openapi_schema["components"]["schemas"]["SchemaName"]["properties"]["messages"]["items"]["anyOf"][1]["$ref"] == "#/components/schemas/AssistantMessage" def test_move_defs_to_components(): """ diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py new file mode 100644 index 00000000000..599d5437589 --- /dev/null +++ b/tests/test_litellm/proxy/discovery_endpoints/test_ui_discovery_endpoints.py @@ -0,0 +1,146 @@ +import os +import sys +from unittest.mock import patch + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +sys.path.insert( + 0, os.path.abspath("../../..") +) + +from litellm.proxy.discovery_endpoints.ui_discovery_endpoints import router + + +def test_ui_discovery_endpoints_with_defaults(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \ + patch.dict(os.environ, {}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["server_root_path"] == "/" + assert data["proxy_base_url"] is None + assert data["auto_redirect_to_sso"] is False + + +def test_ui_discovery_endpoints_with_custom_server_root_path(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \ + patch.dict(os.environ, {}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["server_root_path"] == "/litellm" + assert data["proxy_base_url"] is None + assert data["auto_redirect_to_sso"] is False + + +def test_ui_discovery_endpoints_with_proxy_base_url_when_set(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com"), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \ + patch.dict(os.environ, {}, clear=False): + + response = client.get("/litellm/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["server_root_path"] == "/" + assert data["proxy_base_url"] == "https://proxy.example.com" + assert data["auto_redirect_to_sso"] is False + + +def test_ui_discovery_endpoints_with_sso_configured_and_auto_redirect_enabled(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com"), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), \ + patch.dict(os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "true"}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["server_root_path"] == "/litellm" + assert data["proxy_base_url"] == "https://proxy.example.com" + assert data["auto_redirect_to_sso"] is True + + +def test_ui_discovery_endpoints_with_sso_configured_but_auto_redirect_disabled(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com"), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), \ + patch.dict(os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "false"}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["server_root_path"] == "/litellm" + assert data["proxy_base_url"] == "https://proxy.example.com" + assert data["auto_redirect_to_sso"] is False + + +def test_ui_discovery_endpoints_with_sso_not_configured_but_auto_redirect_enabled(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value=None), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=False), \ + patch.dict(os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "true"}, clear=False): + + response = client.get("/.well-known/litellm-ui-config") + + assert response.status_code == 200 + data = response.json() + assert data["server_root_path"] == "/" + assert data["proxy_base_url"] is None + assert data["auto_redirect_to_sso"] is False + + +def test_ui_discovery_endpoints_both_routes_return_same_data(): + app = FastAPI() + app.include_router(router) + client = TestClient(app) + + with patch("litellm.proxy.utils.get_server_root_path", return_value="/litellm"), \ + patch("litellm.proxy.utils.get_proxy_base_url", return_value="https://proxy.example.com"), \ + patch("litellm.proxy.auth.auth_utils._has_user_setup_sso", return_value=True), \ + patch.dict(os.environ, {"AUTO_REDIRECT_UI_LOGIN_TO_SSO": "true"}, clear=False): + + response1 = client.get("/.well-known/litellm-ui-config") + response2 = client.get("/litellm/.well-known/litellm-ui-config") + + assert response1.status_code == 200 + assert response2.status_code == 200 + assert response1.json() == response2.json() + diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index e756dd6bd5d..b5c093fdf06 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -40,12 +40,12 @@ class TestContentFilterGuardrail: action=ContentFilterAction.BLOCK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-content-filter", patterns=patterns, ) - + assert guardrail.guardrail_name == "test-content-filter" assert len(guardrail.compiled_patterns) == 1 @@ -57,19 +57,19 @@ class TestContentFilterGuardrail: BlockedWord( keyword="secret_project", action=ContentFilterAction.BLOCK, - description="Top secret project" + description="Top secret project", ), BlockedWord( keyword="internal_api", action=ContentFilterAction.MASK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-content-filter", blocked_words=blocked_words, ) - + assert len(guardrail.blocked_words) == 2 assert "secret_project" in guardrail.blocked_words assert guardrail.blocked_words["secret_project"][0] == ContentFilterAction.BLOCK @@ -85,18 +85,18 @@ class TestContentFilterGuardrail: action=ContentFilterAction.BLOCK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-ssn", patterns=patterns, ) - + # Test with SSN result = guardrail._check_patterns("My SSN is 123-45-6789") assert result is not None assert result[1] == "us_ssn" assert result[2] == ContentFilterAction.BLOCK - + # Test without SSN result = guardrail._check_patterns("This is a normal message") assert result is None @@ -112,12 +112,12 @@ class TestContentFilterGuardrail: action=ContentFilterAction.MASK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-email", patterns=patterns, ) - + result = guardrail._check_patterns("Contact me at test@example.com") assert result is not None assert result[1] == "email" @@ -135,12 +135,12 @@ class TestContentFilterGuardrail: action=ContentFilterAction.BLOCK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-custom", patterns=patterns, ) - + result = guardrail._check_patterns("My ID is ABC-1234") assert result is not None assert result[1] == "custom_id" @@ -155,18 +155,18 @@ class TestContentFilterGuardrail: action=ContentFilterAction.BLOCK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-words", blocked_words=blocked_words, ) - + # Test with blocked word result = guardrail._check_blocked_words("This is CONFIDENTIAL information") assert result is not None assert result[0] == "confidential" assert result[1] == ContentFilterAction.BLOCK - + # Test without blocked word result = guardrail._check_blocked_words("This is normal information") assert result is None @@ -183,15 +183,19 @@ class TestContentFilterGuardrail: action=ContentFilterAction.BLOCK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-block", patterns=patterns, ) - + with pytest.raises(HTTPException) as exc_info: - await guardrail.apply_guardrail(text="My SSN is 123-45-6789") - + await guardrail.apply_guardrail( + inputs={"texts": ["My SSN is 123-45-6789"]}, + request_data={}, + input_type="request", + ) + assert exc_info.value.status_code == 400 assert "us_ssn" in str(exc_info.value.detail) @@ -207,17 +211,23 @@ class TestContentFilterGuardrail: action=ContentFilterAction.MASK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-mask", patterns=patterns, ) - - result = await guardrail.apply_guardrail(text="Contact me at test@example.com") - + + guardrailed_inputs = await guardrail.apply_guardrail( + inputs={"texts": ["Contact me at test@example.com"]}, + request_data={}, + input_type="request", + ) + result = guardrailed_inputs.get("texts", []) + assert result is not None - assert "[EMAIL_REDACTED]" in result - assert "test@example.com" not in result + assert len(result) == 1 + assert "[EMAIL_REDACTED]" in result[0] + assert "test@example.com" not in result[0] @pytest.mark.asyncio async def test_apply_guardrail_blocked_word_mask(self): @@ -230,17 +240,23 @@ class TestContentFilterGuardrail: action=ContentFilterAction.MASK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-word-mask", blocked_words=blocked_words, ) - - result = await guardrail.apply_guardrail(text="This is PROPRIETARY information") - + + guardrailed_inputs = await guardrail.apply_guardrail( + inputs={"texts": ["This is PROPRIETARY information"]}, + request_data={}, + input_type="request", + ) + result = guardrailed_inputs.get("texts", []) + assert result is not None - assert "[KEYWORD_REDACTED]" in result - assert "PROPRIETARY" not in result + assert len(result) == 1 + assert "[KEYWORD_REDACTED]" in result[0] + assert "PROPRIETARY" not in result[0] @pytest.mark.asyncio async def test_apply_guardrail_multiple_patterns(self): @@ -259,19 +275,23 @@ class TestContentFilterGuardrail: action=ContentFilterAction.MASK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-multiple", patterns=patterns, ) - - result = await guardrail.apply_guardrail( - text="Contact user@test.com or SSN: 123-45-6789" + + guardrailed_inputs = await guardrail.apply_guardrail( + inputs={"texts": ["Contact user@test.com or SSN: 123-45-6789"]}, + request_data={}, + input_type="request", ) - + result = guardrailed_inputs.get("texts", []) + assert result is not None + assert len(result) == 1 # At least one pattern should be redacted (first match wins) - assert "[EMAIL_REDACTED]" in result or "[US_SSN_REDACTED]" in result + assert "[EMAIL_REDACTED]" in result[0] or "[US_SSN_REDACTED]" in result[0] def test_mask_content(self): """ @@ -280,7 +300,7 @@ class TestContentFilterGuardrail: guardrail = ContentFilterGuardrail( guardrail_name="test-mask", ) - + masked = guardrail._mask_content("sensitive text", "us_ssn") assert masked == "[US_SSN_REDACTED]" @@ -291,28 +311,34 @@ class TestContentFilterGuardrail: import tempfile # Create a temporary blocked words file - with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f: - f.write("""blocked_words: + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write( + """blocked_words: - keyword: "test_keyword" action: "BLOCK" description: "Test keyword" - keyword: "another_word" action: "MASK" -""") +""" + ) temp_file = f.name - + try: guardrail = ContentFilterGuardrail( guardrail_name="test-file-load", blocked_words_file=temp_file, ) - + assert len(guardrail.blocked_words) == 2 assert "test_keyword" in guardrail.blocked_words - assert guardrail.blocked_words["test_keyword"][0] == ContentFilterAction.BLOCK + assert ( + guardrail.blocked_words["test_keyword"][0] == ContentFilterAction.BLOCK + ) assert guardrail.blocked_words["test_keyword"][1] == "Test keyword" assert "another_word" in guardrail.blocked_words - assert guardrail.blocked_words["another_word"][0] == ContentFilterAction.MASK + assert ( + guardrail.blocked_words["another_word"][0] == ContentFilterAction.MASK + ) finally: os.unlink(temp_file) @@ -327,17 +353,17 @@ class TestContentFilterGuardrail: action=ContentFilterAction.BLOCK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-cc", patterns=patterns, ) - + # Test Visa card result = guardrail._check_patterns("My card is 4532-1234-5678-9010") assert result is not None assert result[1] == "visa" - + def test_api_key_patterns(self): """ Test API key pattern detection @@ -349,26 +375,33 @@ class TestContentFilterGuardrail: action=ContentFilterAction.BLOCK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-api-key", patterns=patterns, ) - + # Test AWS Access Key result = guardrail._check_patterns("My key is AKIAIOSFODNN7EXAMPLE") assert result is not None assert result[1] == "aws_access_key" + @pytest.mark.skip( + reason="Masking in streaming responses is no longer supported after unified_guardrail.py changes. Only blocking/rejecting is supported for responses." + ) @pytest.mark.asyncio async def test_streaming_hook_mask(self): """ Test streaming hook with MASK action + + Note: After changes to unified_guardrail.py, masking responses to users + is no longer supported. This test is skipped as the feature is deprecated. + Only BLOCK actions (test_streaming_hook_block) are supported for streaming responses. """ from unittest.mock import AsyncMock from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - + patterns = [ ContentFilterPattern( pattern_type="prebuilt", @@ -376,35 +409,41 @@ class TestContentFilterGuardrail: action=ContentFilterAction.MASK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-streaming-mask", patterns=patterns, event_hook=GuardrailEventHooks.during_call, ) - + # Create mock streaming chunks async def mock_stream(): # Chunk 1: contains email chunk1 = ModelResponseStream( id="chunk1", - choices=[StreamingChoices(delta=Delta(content="Contact me at test@example.com"), index=0)], + choices=[ + StreamingChoices( + delta=Delta(content="Contact me at test@example.com"), index=0 + ) + ], model="gpt-4", ) yield chunk1 - + # Chunk 2: normal content chunk2 = ModelResponseStream( id="chunk2", - choices=[StreamingChoices(delta=Delta(content=" for more info"), index=0)], + choices=[ + StreamingChoices(delta=Delta(content=" for more info"), index=0) + ], model="gpt-4", ) yield chunk2 - + user_api_key_dict = MagicMock() request_data = {} - - # Process streaming response + + # Process streaming response - no masking expected result_chunks = [] async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, @@ -412,13 +451,9 @@ class TestContentFilterGuardrail: request_data=request_data, ): result_chunks.append(chunk) - + + # Chunks should pass through unchanged since masking is no longer supported assert len(result_chunks) == 2 - # First chunk should have email masked - assert "[EMAIL_REDACTED]" in result_chunks[0].choices[0].delta.content - assert "test@example.com" not in result_chunks[0].choices[0].delta.content - # Second chunk should be unchanged - assert result_chunks[1].choices[0].delta.content == " for more info" @pytest.mark.asyncio async def test_streaming_hook_block(self): @@ -428,7 +463,7 @@ class TestContentFilterGuardrail: from unittest.mock import AsyncMock from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices - + patterns = [ ContentFilterPattern( pattern_type="prebuilt", @@ -436,25 +471,27 @@ class TestContentFilterGuardrail: action=ContentFilterAction.BLOCK, ), ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-streaming-block", patterns=patterns, event_hook=GuardrailEventHooks.during_call, ) - + # Create mock streaming chunks with SSN async def mock_stream(): chunk = ModelResponseStream( id="chunk1", - choices=[StreamingChoices(delta=Delta(content="SSN: 123-45-6789"), index=0)], + choices=[ + StreamingChoices(delta=Delta(content="SSN: 123-45-6789"), index=0) + ], model="gpt-4", ) yield chunk - + user_api_key_dict = MagicMock() request_data = {} - + # Should raise HTTPException when SSN is detected with pytest.raises(HTTPException) as exc_info: async for chunk in guardrail.async_post_call_streaming_iterator_hook( @@ -463,7 +500,7 @@ class TestContentFilterGuardrail: request_data=request_data, ): pass - + assert exc_info.value.status_code == 400 assert "us_ssn" in str(exc_info.value.detail) @@ -487,9 +524,9 @@ class TestContentFilterGuardrail: "action": "MASK", "name": "email", "pattern": None, - } + }, ] - + blocked_words = [ { "keyword": "langchain", @@ -500,19 +537,19 @@ class TestContentFilterGuardrail: "keyword": "openai", "action": "MASK", "description": "Competitor name", - } + }, ] - + guardrail = ContentFilterGuardrail( guardrail_name="test-db-format", patterns=patterns, blocked_words=blocked_words, ) - + assert guardrail.guardrail_name == "test-db-format" assert len(guardrail.compiled_patterns) == 2 assert len(guardrail.blocked_words) == 2 - + # Verify blocked_words are stored as dict assert "langchain" in guardrail.blocked_words assert guardrail.blocked_words["langchain"] == ("BLOCK", None) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index e6ebada4d2e..f3578fbcefa 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -4,6 +4,7 @@ Tests for Generic Guardrail API integration This test file tests the Generic Guardrail API implementation, specifically focusing on metadata extraction and passing. """ + import os from unittest.mock import AsyncMock, MagicMock, patch @@ -108,9 +109,14 @@ class TestGenericGuardrailAPIConfiguration: headers={"Authorization": "Bearer test-key"}, additional_provider_specific_params={"custom_param": "value"}, ) - assert guardrail.api_base == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api" + assert ( + guardrail.api_base + == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api" + ) assert guardrail.headers == {"Authorization": "Bearer test-key"} - assert guardrail.additional_provider_specific_params == {"custom_param": "value"} + assert guardrail.additional_provider_specific_params == { + "custom_param": "value" + } def test_init_with_env_vars(self): """Test initialization with environment variables""" @@ -121,7 +127,10 @@ class TestGenericGuardrailAPIConfiguration: }, ): guardrail = GenericGuardrailAPI() - assert guardrail.api_base == "https://env.api.guardrail.com/beta/litellm_basic_guardrail_api" + assert ( + guardrail.api_base + == "https://env.api.guardrail.com/beta/litellm_basic_guardrail_api" + ) def test_init_without_api_base_raises_error(self): """Test that initialization without API base raises ValueError""" @@ -134,14 +143,20 @@ class TestGenericGuardrailAPIConfiguration: guardrail = GenericGuardrailAPI( api_base="https://api.test.guardrail.com/v1", ) - assert guardrail.api_base == "https://api.test.guardrail.com/v1/beta/litellm_basic_guardrail_api" + assert ( + guardrail.api_base + == "https://api.test.guardrail.com/v1/beta/litellm_basic_guardrail_api" + ) def test_api_base_not_duplicated(self): """Test that endpoint path is not duplicated if already present""" guardrail = GenericGuardrailAPI( api_base="https://api.test.guardrail.com/beta/litellm_basic_guardrail_api", ) - assert guardrail.api_base == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api" + assert ( + guardrail.api_base + == "https://api.test.guardrail.com/beta/litellm_basic_guardrail_api" + ) class TestMetadataExtraction: @@ -164,7 +179,7 @@ class TestMetadataExtraction: generic_guardrail.async_handler, "post", return_value=mock_response ) as mock_post: await generic_guardrail.apply_guardrail( - texts=["Who is Ishaan?"], + inputs=GenericGuardrailAPIInputs(texts=["Who is Ishaan?"]), request_data=mock_request_data_input, input_type="request", ) @@ -180,7 +195,10 @@ class TestMetadataExtraction: request_metadata = json_payload["request_data"] # Verify metadata was extracted from request_data["metadata"] - assert request_metadata["user_api_key_hash"] == "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" + assert ( + request_metadata["user_api_key_hash"] + == "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b" + ) assert request_metadata["user_api_key_user_id"] == "default_user_id" assert request_metadata["user_api_key_user_email"] == "test@example.com" assert request_metadata["user_api_key_team_id"] == "test-team" @@ -192,7 +210,7 @@ class TestMetadataExtraction: """Test extracting metadata from output response (litellm_metadata field)""" # Create request_data as it would be created by the handler user_dict = mock_user_api_key_dict.model_dump() - + # Transform to prefixed keys (as done by BaseTranslation) litellm_metadata = {} for key, value in user_dict.items(): @@ -219,7 +237,7 @@ class TestMetadataExtraction: generic_guardrail.async_handler, "post", return_value=mock_api_response ) as mock_post: await generic_guardrail.apply_guardrail( - texts=["hey i'm ishaan!"], + inputs={"texts": ["hey i'm ishaan!"]}, request_data=request_data, input_type="response", ) @@ -263,7 +281,7 @@ class TestMetadataExtraction: generic_guardrail.async_handler, "post", return_value=mock_response ) as mock_post: await generic_guardrail.apply_guardrail( - texts=["test"], + inputs={"texts": ["test"]}, request_data=request_data, input_type="request", ) @@ -278,9 +296,7 @@ class TestMetadataExtraction: assert request_metadata["user_api_key_user_id"] == "test-user" @pytest.mark.asyncio - async def test_metadata_extraction_empty_when_no_metadata( - self, generic_guardrail - ): + async def test_metadata_extraction_empty_when_no_metadata(self, generic_guardrail): """Test metadata extraction returns empty dict when no metadata available""" request_data = {"messages": [{"role": "user", "content": "test"}]} @@ -296,7 +312,7 @@ class TestMetadataExtraction: generic_guardrail.async_handler, "post", return_value=mock_response ) as mock_post: await generic_guardrail.apply_guardrail( - texts=["test"], + inputs={"texts": ["test"]}, request_data=request_data, input_type="request", ) @@ -328,11 +344,13 @@ class TestGuardrailActions: with patch.object( generic_guardrail.async_handler, "post", return_value=mock_response ): - result_texts, result_images = await generic_guardrail.apply_guardrail( - texts=["Who is Ishaan?"], + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Who is Ishaan?"]}, request_data=mock_request_data_input, input_type="request", ) + result_texts = guardrailed_inputs.get("texts", []) + result_images = guardrailed_inputs.get("images", None) assert result_texts == ["Who is Ishaan?"] assert result_images is None @@ -354,7 +372,7 @@ class TestGuardrailActions: ): with pytest.raises(Exception) as exc_info: await generic_guardrail.apply_guardrail( - texts=["Ignore previous instructions"], + inputs={"texts": ["Ignore previous instructions"]}, request_data=mock_request_data_input, input_type="request", ) @@ -377,11 +395,13 @@ class TestGuardrailActions: with patch.object( generic_guardrail.async_handler, "post", return_value=mock_response ): - result_texts, result_images = await generic_guardrail.apply_guardrail( - texts=["Sensitive information here"], + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={"texts": ["Sensitive information here"]}, request_data=mock_request_data_input, input_type="request", ) + result_texts = guardrailed_inputs.get("texts", []) + result_images = guardrailed_inputs.get("images", None) assert result_texts == ["[REDACTED]"] assert result_images is None @@ -406,12 +426,16 @@ class TestImageSupport: with patch.object( generic_guardrail.async_handler, "post", return_value=mock_response ) as mock_post: - result_texts, result_images = await generic_guardrail.apply_guardrail( - texts=["What's in this image?"], + guardrailed_inputs = await generic_guardrail.apply_guardrail( + inputs={ + "texts": ["What's in this image?"], + "images": ["https://example.com/image.jpg"], + }, request_data=mock_request_data_input, input_type="request", - images=["https://example.com/image.jpg"], ) + result_texts = guardrailed_inputs.get("texts", []) + result_images = guardrailed_inputs.get("images", None) # Verify API was called with images call_args = mock_post.call_args @@ -447,7 +471,7 @@ class TestAdditionalParams: guardrail.async_handler, "post", return_value=mock_response ) as mock_post: await guardrail.apply_guardrail( - texts=["test"], + inputs={"texts": ["test"]}, request_data=mock_request_data_input, input_type="request", ) @@ -455,8 +479,14 @@ class TestAdditionalParams: # Verify API was called with additional params call_args = mock_post.call_args json_payload = call_args.kwargs["json"] - assert json_payload["additional_provider_specific_params"]["custom_threshold"] == 0.8 - assert json_payload["additional_provider_specific_params"]["enable_feature"] is True + assert ( + json_payload["additional_provider_specific_params"]["custom_threshold"] + == 0.8 + ) + assert ( + json_payload["additional_provider_specific_params"]["enable_feature"] + is True + ) class TestErrorHandling: @@ -476,7 +506,7 @@ class TestErrorHandling: ): with pytest.raises(Exception) as exc_info: await generic_guardrail.apply_guardrail( - texts=["test"], + inputs={"texts": ["test"]}, request_data=mock_request_data_input, input_type="request", ) @@ -495,10 +525,9 @@ class TestErrorHandling: ): with pytest.raises(Exception) as exc_info: await generic_guardrail.apply_guardrail( - texts=["test"], + inputs={"texts": ["test"]}, request_data=mock_request_data_input, input_type="request", ) assert "Generic Guardrail API failed" in str(exc_info.value) - diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 9543b61ef69..a6b2ae5b3a0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -537,24 +537,25 @@ async def test_logging_hook_multiple_content_items(presidio_guardrail): async def test_presidio_sets_guardrail_information_in_request_data(): """ Test that Presidio populates guardrail information into request_data metadata. - + This validates that add_standard_logging_guardrail_information_to_request_data correctly sets the guardrail information that will be used for logging. """ presidio = _OPTIONAL_PresidioPIIMasking( guardrail_name="test_presidio", output_parse_pii=True, + mock_testing=True, ) - + request_data = { "messages": [{"role": "user", "content": "Test"}], "model": "gpt-4o", "metadata": {}, } - + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): assert request_data is not None - + presidio.add_standard_logging_guardrail_information_to_request_data( guardrail_provider="presidio", guardrail_json_response=[], @@ -565,27 +566,30 @@ async def test_presidio_sets_guardrail_information_in_request_data(): duration=1.0, masked_entity_count={"EMAIL_ADDRESS": 1, "PERSON": 1}, ) - + return text - - with patch.object(presidio, 'check_pii', mock_check_pii): + + with patch.object(presidio, "check_pii", mock_check_pii): await presidio.apply_guardrail( - text="Test message", + inputs={"texts": ["Test message"]}, request_data=request_data, + input_type="request", ) - + assert "metadata" in request_data assert "standard_logging_guardrail_information" in request_data["metadata"] - - guardrail_info_list = request_data["metadata"]["standard_logging_guardrail_information"] + + guardrail_info_list = request_data["metadata"][ + "standard_logging_guardrail_information" + ] assert isinstance(guardrail_info_list, list) assert len(guardrail_info_list) > 0 - + guardrail_info = guardrail_info_list[0] assert "masked_entity_count" in guardrail_info assert guardrail_info["masked_entity_count"]["EMAIL_ADDRESS"] == 1 assert guardrail_info["masked_entity_count"]["PERSON"] == 1 - + print("✓ Presidio sets guardrail_information in request_data") @@ -593,7 +597,7 @@ async def test_presidio_sets_guardrail_information_in_request_data(): async def test_request_data_flows_to_apply_guardrail(): """ Test that request_data is correctly passed to apply_guardrail method. - + This validates the fix where guardrail translation handler passes data as request_data to apply_guardrail so guardrails can store metadata for logging. """ @@ -601,31 +605,32 @@ async def test_request_data_flows_to_apply_guardrail(): guardrail_name="test_presidio", output_parse_pii=True, ) - + request_data = { "messages": [{"role": "user", "content": "Test message"}], "model": "gpt-4o", "metadata": {}, } - + async def mock_check_pii(text, output_parse_pii, presidio_config, request_data): assert request_data is not None, "request_data should be passed to check_pii" assert "metadata" in request_data, "request_data should have metadata" - + request_data.setdefault("metadata", {}) request_data["metadata"]["test_flag"] = "passed_correctly" - + return text - - with patch.object(presidio, 'check_pii', mock_check_pii): + + with patch.object(presidio, "check_pii", mock_check_pii): result = await presidio.apply_guardrail( - text="Test message", + inputs={"texts": ["Test message"]}, request_data=request_data, + input_type="request", ) - + assert "metadata" in request_data assert request_data["metadata"].get("test_flag") == "passed_correctly" - + print("✓ request_data correctly passed to apply_guardrail") diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 6574c500fcf..c8c30d41b5e 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -22,6 +22,7 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types.utils import ModelResponse, Usage + class TimeController: def __init__(self): self._current = datetime.utcnow() @@ -461,10 +462,11 @@ async def test_token_rate_limit_type_respected_v3(monkeypatch, token_rate_limit_ ) # Create mock kwargs for the success event + # Use standard_logging_object which is the canonical source for metadata mock_kwargs = { - "litellm_params": { + "standard_logging_object": { "metadata": { - "user_api_key": _api_key, + "user_api_key_hash": _api_key, "user_api_key_user_id": None, "user_api_key_team_id": None, "user_api_key_end_user_id": None, @@ -532,8 +534,8 @@ async def test_async_log_failure_event_v3(): internal_usage_cache=InternalUsageCache(local_cache) ) - # Mock kwargs with user_api_key - mock_kwargs = {"litellm_params": {"metadata": {"user_api_key": _api_key}}} + # Mock kwargs with user_api_key via standard_logging_object + mock_kwargs = {"standard_logging_object": {"metadata": {"user_api_key_hash": _api_key}}} # Capture pipeline operations captured_ops = [] diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index de7a847a918..d30cce067a0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -315,7 +315,32 @@ async def test_user_info_url_encoding_plus_character(mocker): # Mock the prisma client mock_prisma_client = mocker.MagicMock() - mock_prisma_client.get_data = mocker.AsyncMock() + + # Create a real LiteLLM_UserTable instance (BaseModel) so isinstance check passes + mock_user = LiteLLM_UserTable( + user_id="machine-user+alp-air-admin-b58-b@tempus.com", + user_email="machine-user+alp-air-admin-b58-b@tempus.com", + teams=[], + ) + + # Mock get_data to return user when called with user_id, empty list for keys + async def mock_get_data(*args, **kwargs): + if kwargs.get("table_name") == "key": + return [] + elif kwargs.get("table_name") == "team": + return [] + elif kwargs.get("user_id") is not None: + return mock_user + return None + + mock_prisma_client.get_data = mocker.AsyncMock(side_effect=mock_get_data) + + # Mock list_team to return None (patch it from where it's imported) + mock_list_team = mocker.AsyncMock(return_value=None) + mocker.patch( + "litellm.proxy.management_endpoints.team_endpoints.list_team", + mock_list_team, + ) # Patch the prisma client import in the endpoint mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) @@ -335,20 +360,73 @@ async def test_user_info_url_encoding_plus_character(mocker): "machine-user alp-air-admin-b58-b@tempus.com" # What FastAPI gives us ) expected_user_id = "machine-user+alp-air-admin-b58-b@tempus.com" - try: - response = await user_info( - user_id=decoded_user_id, + + response = await user_info( + user_id=decoded_user_id, + user_api_key_dict=mock_user_api_key_dict, + request=mock_request, + ) + + # Verify that the response contains the correct user data + # Check that get_data was called with the correct user_id (first call should be for user) + user_call = None + for call in mock_prisma_client.get_data.call_args_list: + if call.kwargs.get("user_id") and not call.kwargs.get("table_name"): + user_call = call + break + + assert user_call is not None, "get_data should be called with user_id" + assert user_call.kwargs["user_id"] == expected_user_id + + +@pytest.mark.asyncio +async def test_user_info_nonexistent_user(mocker): + """ + Test that /user/info endpoint returns 404 when a non-existent user_id is provided. + """ + from fastapi import Request + + from litellm.proxy._types import ProxyException, UserAPIKeyAuth + from litellm.proxy.management_endpoints.internal_user_endpoints import user_info + + # Mock the prisma client + mock_prisma_client = mocker.MagicMock() + + # Mock get_data to return None (user doesn't exist) + async def mock_get_data(*args, **kwargs): + if kwargs.get("table_name") == "key": + return [] + elif kwargs.get("user_id") is not None: + return None # User not found + return None + + mock_prisma_client.get_data = mocker.AsyncMock(side_effect=mock_get_data) + + # Patch the prisma client import in the endpoint + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Create a mock request + mock_request = mocker.MagicMock(spec=Request) + + # Mock user_api_key_dict + mock_user_api_key_dict = UserAPIKeyAuth( + user_id="test_admin", user_role="proxy_admin" + ) + + # Call user_info function with a non-existent user_id + nonexistent_user_id = "nonexistent-user@example.com" + + # Should raise ProxyException with 404 status code (HTTPException is converted by decorator) + with pytest.raises(ProxyException) as exc_info: + await user_info( + user_id=nonexistent_user_id, user_api_key_dict=mock_user_api_key_dict, request=mock_request, ) - except Exception as e: - print(f"Error in user_info: {e}") - # Verify that the response contains the correct user data - print( - f"mock_prisma_client.get_data.call_args: {mock_prisma_client.get_data.call_args.kwargs}" - ) - assert mock_prisma_client.get_data.call_args.kwargs["user_id"] == expected_user_id + # Verify the exception details + assert exc_info.value.code == "404" # ProxyException.code is a string + assert f"User {nonexistent_user_id} not found" in str(exc_info.value.message) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/test_enforce_user_param.py b/tests/test_litellm/proxy/test_enforce_user_param.py new file mode 100644 index 00000000000..5d9369d61bb --- /dev/null +++ b/tests/test_litellm/proxy/test_enforce_user_param.py @@ -0,0 +1,438 @@ +""" +Tests for enforce_user_param feature with POST/GET method filtering and MCP route exclusion. + +Tests verify that: +1. enforce_user_param only applies to POST requests +2. GET requests like /v1/models are not affected +3. MCP routes are excluded from enforcement +4. POST requests to completion endpoints still require user param when enforced +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +import pytest +from fastapi import Request + +from litellm.proxy._types import LiteLLMRoutes, UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import common_checks +from litellm.proxy.auth.route_checks import RouteChecks + + +class MockRequest: + """Mock FastAPI Request object""" + def __init__(self, method: str = "POST"): + self.method = method + + +def get_mock_user_token(): + """Create a mock UserAPIKeyAuth token for testing""" + return UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + team_id="test-team", + org_id="test-org", + models=["*"], + metadata={} + ) + + +class TestEnforceUserParamPostGetFiltering: + """Test POST/GET method filtering for enforce_user_param""" + + @pytest.mark.asyncio + async def test_post_completion_without_user_param_should_fail(self): + """POST to /v1/chat/completions without user param should raise error""" + request = MockRequest(method="POST") + general_settings = {"enforce_user_param": True} + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}] + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + with pytest.raises(Exception) as exc_info: + await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert "user" in str(exc_info.value).lower() + + @pytest.mark.asyncio + async def test_post_completion_with_user_param_should_pass(self): + """POST to /v1/chat/completions with user param should pass""" + request = MockRequest(method="POST") + general_settings = {"enforce_user_param": True} + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "user": "user123" + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + # Should not raise exception + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + @pytest.mark.asyncio + async def test_get_models_without_user_param_should_pass(self): + """GET to /v1/models without user param should NOT raise error""" + request = MockRequest(method="GET") + general_settings = {"enforce_user_param": True} + request_body = {} # GET requests typically don't have body + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + # Should not raise exception + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/models", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + @pytest.mark.asyncio + async def test_get_files_without_user_param_should_pass(self): + """GET to /v1/files without user param should NOT raise error""" + request = MockRequest(method="GET") + general_settings = {"enforce_user_param": True} + request_body = {} + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/files", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + @pytest.mark.asyncio + async def test_post_embeddings_without_user_param_should_fail(self): + """POST to /v1/embeddings without user param should raise error""" + request = MockRequest(method="POST") + general_settings = {"enforce_user_param": True} + request_body = { + "model": "text-embedding-ada-002", + "input": "test" + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + with pytest.raises(Exception) as exc_info: + await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/embeddings", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert "user" in str(exc_info.value).lower() + + @pytest.mark.asyncio + async def test_post_embeddings_with_user_param_should_pass(self): + """POST to /v1/embeddings with user param should pass""" + request = MockRequest(method="POST") + general_settings = {"enforce_user_param": True} + request_body = { + "model": "text-embedding-ada-002", + "input": "test", + "user": "user123" + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/embeddings", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + +class TestEnforceUserParamMCPExclusion: + """Test MCP route exclusion from enforce_user_param""" + + @pytest.mark.asyncio + async def test_mcp_route_without_user_param_should_pass(self): + """POST to MCP route without user param should NOT raise error""" + request = MockRequest(method="POST") + general_settings = {"enforce_user_param": True} + request_body = {"action": "list_tools"} + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + # Should not raise exception for MCP routes + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/mcp/tools/list", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + @pytest.mark.asyncio + async def test_mcp_root_route_without_user_param_should_pass(self): + """POST to /mcp without user param should NOT raise error""" + request = MockRequest(method="POST") + general_settings = {"enforce_user_param": True} + request_body = {"data": "test"} + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/mcp", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + +class TestEnforceUserParamDisabled: + """Test behavior when enforce_user_param is disabled""" + + @pytest.mark.asyncio + async def test_post_without_user_param_when_disabled_should_pass(self): + """POST without user param when enforce_user_param=False should pass""" + request = MockRequest(method="POST") + general_settings = {"enforce_user_param": False} + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}] + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + @pytest.mark.asyncio + async def test_post_without_user_param_when_not_set_should_pass(self): + """POST without user param when enforce_user_param not set should pass""" + request = MockRequest(method="POST") + general_settings = {} # enforce_user_param not set + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}] + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + +class TestEnforceUserParamEdgeCases: + """Test edge cases for enforce_user_param""" + + @pytest.mark.asyncio + async def test_request_without_method_attribute_should_pass(self): + """Request without method attribute should not raise error""" + request = MagicMock() + del request.method # Remove method attribute + request.__hasattr__ = MagicMock(return_value=False) + + general_settings = {"enforce_user_param": True} + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}] + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + # Should not raise error even without method + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + @pytest.mark.asyncio + async def test_case_insensitive_http_method(self): + """HTTP method comparison should be case-insensitive""" + request = MockRequest(method="post") # lowercase + general_settings = {"enforce_user_param": True} + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}] + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + with pytest.raises(Exception) as exc_info: + await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert "user" in str(exc_info.value).lower() + + @pytest.mark.asyncio + async def test_put_method_should_not_enforce_user_param(self): + """PUT method should not enforce user param (only POST)""" + request = MockRequest(method="PUT") + general_settings = {"enforce_user_param": True} + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}] + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + # Should not raise for PUT method + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + @pytest.mark.asyncio + async def test_patch_method_should_not_enforce_user_param(self): + """PATCH method should not enforce user param (only POST)""" + request = MockRequest(method="PATCH") + general_settings = {"enforce_user_param": True} + request_body = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}] + } + + with patch('litellm.proxy.auth.auth_checks._is_api_route_allowed', new_callable=AsyncMock, return_value=True): + # Should not raise for PATCH method + result = await common_checks( + request_body=request_body, + team_object=None, + user_object=None, + end_user_object=None, + global_proxy_spend=None, + general_settings=general_settings, + route="/v1/chat/completions", + llm_router=None, + proxy_logging_obj=MagicMock(), + valid_token=get_mock_user_token(), + request=request, + ) + + assert result is True + + +if __name__ == "__main__": + # Run tests with: pytest tests/test_litellm/proxy/test_enforce_user_param.py -v + pytest.main([__file__, "-v"]) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8e87b67933f..7b3157f8ee1 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -71,6 +71,58 @@ def client_no_auth(): return TestClient(app) +def test_login_v2_returns_redirect_url_and_sets_cookie(monkeypatch): + mock_login_result = {"user_id": "test-user"} + mock_prisma_client = MagicMock() + mock_authenticate_user = AsyncMock(return_value=mock_login_result) + mock_create_ui_token_object = MagicMock(return_value={"user_id": "test-user"}) + mock_jwt_encode = MagicMock(return_value="signed-token") + + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.authenticate_user", + mock_authenticate_user, + ) + monkeypatch.setattr( + "litellm.proxy.auth.login_utils.create_ui_token_object", + mock_create_ui_token_object, + ) + monkeypatch.setattr("jwt.encode", mock_jwt_encode) + monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "test-master-key") + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {}) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + client = TestClient(app) + response = client.post( + "/v2/login", + json={"username": "alice", "password": "secret"}, + ) + + assert response.status_code == 200 + assert ( + response.json() + == {"redirect_url": "http://testserver/ui/?login=success"} + ) + assert response.cookies.get("token") == "signed-token" + + mock_authenticate_user.assert_awaited_once_with( + username="alice", + password="secret", + master_key="test-master-key", + prisma_client=mock_prisma_client, + ) + mock_create_ui_token_object.assert_called_once_with( + login_result=mock_login_result, + general_settings={}, + premium_user=False, + ) + mock_jwt_encode.assert_called_once_with( + {"user_id": "test-user"}, + "test-master-key", + algorithm="HS256", + ) + + @pytest.mark.asyncio async def test_initialize_scheduled_jobs_credentials(monkeypatch): """ diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 71ede0958f5..d2416f8db8c 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -1461,3 +1461,43 @@ async def test_async_mock_completion_stream_with_model_response(): accumulated_content += chunk.choices[0].delta.content assert "This is an async test response" in accumulated_content or len(chunks) > 0 + + +class TestCallTypesOCR: + """Test that OCR call types are properly defined in CallTypes enum. + + Fixes https://github.com/BerriAI/litellm/issues/17381 + """ + + def test_ocr_call_type_exists(self): + """Test that CallTypes.ocr exists and has correct value.""" + from litellm.types.utils import CallTypes + + assert hasattr(CallTypes, "ocr") + assert CallTypes.ocr.value == "ocr" + + def test_aocr_call_type_exists(self): + """Test that CallTypes.aocr exists and has correct value.""" + from litellm.types.utils import CallTypes + + assert hasattr(CallTypes, "aocr") + assert CallTypes.aocr.value == "aocr" + + def test_ocr_call_type_from_string(self): + """Test that CallTypes can be constructed from 'ocr' string.""" + from litellm.types.utils import CallTypes + + call_type = CallTypes("ocr") + assert call_type == CallTypes.ocr + + def test_aocr_call_type_from_string(self): + """Test that CallTypes can be constructed from 'aocr' string. + + This is the actual use case that was failing - the OCR endpoint + uses route_type='aocr' and guardrails try to instantiate + CallTypes('aocr'). + """ + from litellm.types.utils import CallTypes + + call_type = CallTypes("aocr") + assert call_type == CallTypes.aocr diff --git a/tests/test_litellm/vector_stores/test_vector_store_create_provider_logic.py b/tests/test_litellm/vector_stores/test_vector_store_create_provider_logic.py index 501a5ff0389..cca20847f12 100644 --- a/tests/test_litellm/vector_stores/test_vector_store_create_provider_logic.py +++ b/tests/test_litellm/vector_stores/test_vector_store_create_provider_logic.py @@ -9,9 +9,12 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm -from litellm.utils import ProviderConfigManager from litellm.llms.openai.vector_stores.transformation import OpenAIVectorStoreConfig -from litellm.llms.vertex_ai.vector_stores.rag_api.transformation import VertexVectorStoreConfig +from litellm.llms.ragflow.vector_stores.transformation import RAGFlowVectorStoreConfig +from litellm.llms.vertex_ai.vector_stores.rag_api.transformation import ( + VertexVectorStoreConfig, +) +from litellm.utils import ProviderConfigManager def test_vector_store_create_with_simple_provider_name(): @@ -100,3 +103,40 @@ def test_vector_store_create_with_provider_api_type(): print("✅ Test passed: Provider with api_type 'vertex_ai/rag_api' handled correctly") + +def test_vector_store_create_with_ragflow_provider(): + """ + Test that vector store create correctly handles RAGFlow provider. + + This should: + - Return correct RAGFlowVectorStoreConfig + - Support dataset management operations + """ + custom_llm_provider = "ragflow" + + # Simulate the logic from vector_stores/main.py create function + if "/" in custom_llm_provider: + pytest.fail("Should not enter this branch for RAGFlow provider") + else: + api_type = None + custom_llm_provider = custom_llm_provider # Keep as-is + + # Verify api_type is None + assert api_type is None, "api_type should be None for RAGFlow provider" + + # Verify custom_llm_provider is unchanged + assert custom_llm_provider == "ragflow", "custom_llm_provider should remain 'ragflow'" + + # Verify ProviderConfigManager returns correct config + vector_store_provider_config = ProviderConfigManager.get_provider_vector_stores_config( + provider=litellm.LlmProviders(custom_llm_provider), + api_type=api_type, + ) + + assert vector_store_provider_config is not None, "Should return a config for RAGFlow" + assert isinstance( + vector_store_provider_config, RAGFlowVectorStoreConfig + ), "Should return RAGFlowVectorStoreConfig for RAGFlow provider" + + print("✅ Test passed: RAGFlow provider handled correctly") + diff --git a/tests/vector_store_tests/test_ragflow_vector_store.py b/tests/vector_store_tests/test_ragflow_vector_store.py new file mode 100644 index 00000000000..0839bd7153b --- /dev/null +++ b/tests/vector_store_tests/test_ragflow_vector_store.py @@ -0,0 +1,359 @@ +""" +Test RAGFlow Vector Store helper functions and transformation. +""" +import os +import sys +import json +import pytest +from unittest.mock import Mock, patch, MagicMock +import httpx + +sys.path.insert(0, os.path.abspath("../..")) +import litellm + +from tests.vector_store_tests.base_vector_store_test import BaseVectorStoreTest +from litellm.llms.ragflow.vector_stores.transformation import RAGFlowVectorStoreConfig +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.types.vector_stores import VectorStoreCreateOptionalRequestParams + + +class TestRAGFlowVectorStore(BaseVectorStoreTest): + """ + Test the RAGFlow vector store transformation functionality. + """ + + def get_base_create_vector_store_args(self) -> dict: + """Must return the base create vector store args""" + return { + "custom_llm_provider": "ragflow", + "api_key": os.getenv("RAGFLOW_API_KEY", "test-api-key"), + "api_base": os.getenv("RAGFLOW_API_BASE", "http://localhost:9380") + } + + def get_base_request_args(self): + # RAGFlow doesn't support search, so we'll skip search tests + return { + "vector_store_id": "test-dataset-id", + "custom_llm_provider": "ragflow", + "query": "test query" + } + + def test_get_auth_credentials(self): + """Test that auth credentials are correctly extracted.""" + config = RAGFlowVectorStoreConfig() + + # Test with api_key in params + litellm_params = {"api_key": "test-api-key-123"} + credentials = config.get_auth_credentials(litellm_params) + assert "headers" in credentials + assert credentials["headers"]["Authorization"] == "Bearer test-api-key-123" + + # Test with missing api_key (should raise ValueError) + with pytest.raises(ValueError, match="api_key is required"): + config.get_auth_credentials({}) + + def test_get_complete_url(self): + """Test that complete URL is correctly constructed.""" + config = RAGFlowVectorStoreConfig() + + # Test with api_base in params + litellm_params = {"api_base": "http://custom-host:9999"} + url = config.get_complete_url(api_base=None, litellm_params=litellm_params) + assert url == "http://custom-host:9999/api/v1/datasets" + + # Test with api_base parameter + url = config.get_complete_url(api_base="http://test-host:8888", litellm_params={}) + assert url == "http://test-host:8888/api/v1/datasets" + + # Test with default (no api_base provided) + with patch.dict(os.environ, {}, clear=True): + url = config.get_complete_url(api_base=None, litellm_params={}) + assert url == "http://localhost:9380/api/v1/datasets" + + # Test with trailing slash removal + url = config.get_complete_url(api_base="http://test-host:8888/", litellm_params={}) + assert url == "http://test-host:8888/api/v1/datasets" + + def test_validate_environment(self): + """Test environment validation and header setting.""" + config = RAGFlowVectorStoreConfig() + from litellm.types.router import GenericLiteLLMParams + + # Test with api_key in litellm_params + litellm_params = GenericLiteLLMParams(api_key="test-key") + headers = config.validate_environment({}, litellm_params) + assert headers["Authorization"] == "Bearer test-key" + assert headers["Content-Type"] == "application/json" + + # Test with missing api_key + with pytest.raises(ValueError, match="RAGFLOW_API_KEY"): + config.validate_environment({}, GenericLiteLLMParams()) + + def test_get_vector_store_endpoints_by_type(self): + """Test that endpoints are correctly configured (empty for management only).""" + config = RAGFlowVectorStoreConfig() + endpoints = config.get_vector_store_endpoints_by_type() + assert endpoints["read"] == [] + assert endpoints["write"] == [] + + def test_transform_create_vector_store_request_basic(self): + """Test basic dataset creation request transformation.""" + config = RAGFlowVectorStoreConfig() + + params: VectorStoreCreateOptionalRequestParams = { + "name": "test-dataset" + } + + url, body = config.transform_create_vector_store_request( + params, "http://localhost:9380/api/v1/datasets" + ) + + assert url == "http://localhost:9380/api/v1/datasets" + assert body["name"] == "test-dataset" + assert body["chunk_method"] == "naive" # Default chunk method + + def test_transform_create_vector_store_request_with_metadata(self): + """Test dataset creation with RAGFlow-specific metadata.""" + config = RAGFlowVectorStoreConfig() + + params: VectorStoreCreateOptionalRequestParams = { + "name": "test-dataset-advanced", + "metadata": { + "description": "Test dataset", + "embedding_model": "BAAI/bge-large-zh-v1.5@BAAI", + "permission": "me", + "chunk_method": "naive", + "parser_config": { + "chunk_token_num": 512, + "delimiter": "\n" + } + } + } + + url, body = config.transform_create_vector_store_request( + params, "http://localhost:9380/api/v1/datasets" + ) + + assert body["name"] == "test-dataset-advanced" + assert body["description"] == "Test dataset" + assert body["embedding_model"] == "BAAI/bge-large-zh-v1.5@BAAI" + assert body["permission"] == "me" + assert body["chunk_method"] == "naive" + assert "parser_config" in body + assert body["parser_config"]["chunk_token_num"] == 512 + + def test_transform_create_vector_store_request_missing_name(self): + """Test that missing name raises ValueError.""" + config = RAGFlowVectorStoreConfig() + + params: VectorStoreCreateOptionalRequestParams = {} + + with pytest.raises(ValueError, match="name is required"): + config.transform_create_vector_store_request( + params, "http://localhost:9380/api/v1/datasets" + ) + + def test_transform_create_vector_store_request_mutually_exclusive(self): + """Test that chunk_method and pipeline_id are mutually exclusive.""" + config = RAGFlowVectorStoreConfig() + + params: VectorStoreCreateOptionalRequestParams = { + "name": "test-dataset", + "metadata": { + "chunk_method": "naive", + "pipeline_id": "d0bebe30ae2211f0970942010a8e0005" + } + } + + with pytest.raises(ValueError, match="mutually exclusive"): + config.transform_create_vector_store_request( + params, "http://localhost:9380/api/v1/datasets" + ) + + def test_transform_create_vector_store_request_with_pipeline(self): + """Test dataset creation with ingestion pipeline.""" + config = RAGFlowVectorStoreConfig() + + params: VectorStoreCreateOptionalRequestParams = { + "name": "test-pipeline-dataset", + "metadata": { + "parse_type": 2, + "pipeline_id": "d0bebe30ae2211f0970942010a8e0005" + } + } + + url, body = config.transform_create_vector_store_request( + params, "http://localhost:9380/api/v1/datasets" + ) + + assert body["name"] == "test-pipeline-dataset" + assert body["parse_type"] == 2 + assert body["pipeline_id"] == "d0bebe30ae2211f0970942010a8e0005" + assert "chunk_method" not in body + + def test_transform_create_vector_store_response_success(self): + """Test successful response transformation.""" + config = RAGFlowVectorStoreConfig() + + # Mock RAGFlow response + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = { + "code": 0, + "data": { + "id": "3b4de7d4241d11f0a6a79f24fc270c7f", + "name": "test-dataset", + "create_time": 1745836841611, + "chunk_method": "naive", + "embedding_model": "BAAI/bge-large-zh-v1.5@BAAI" + } + } + + response = config.transform_create_vector_store_response(mock_response) + + assert response["id"] == "3b4de7d4241d11f0a6a79f24fc270c7f" + assert response["name"] == "test-dataset" + assert response["object"] == "vector_store" + assert response["status"] == "completed" + assert response["created_at"] == 1745836841 # Converted from milliseconds + assert response["bytes"] == 0 + assert "file_counts" in response + + def test_transform_create_vector_store_response_error(self): + """Test error response transformation.""" + config = RAGFlowVectorStoreConfig() + + # Mock RAGFlow error response + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 400 + mock_response.headers = {} + mock_response.json.return_value = { + "code": 101, + "message": "Dataset name 'test-dataset' already exists" + } + + with pytest.raises(Exception): # Should raise BaseLLMException + config.transform_create_vector_store_response(mock_response) + + def test_transform_create_vector_store_response_missing_id(self): + """Test response with missing dataset ID.""" + config = RAGFlowVectorStoreConfig() + + mock_response = Mock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.headers = {} + mock_response.json.return_value = { + "code": 0, + "data": { + "name": "test-dataset" + # Missing "id" + } + } + + with pytest.raises(ValueError, match="missing dataset id"): + config.transform_create_vector_store_response(mock_response) + + def test_transform_search_vector_store_request_not_implemented(self): + """Test that search operations raise NotImplementedError.""" + config = RAGFlowVectorStoreConfig() + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + + with pytest.raises(NotImplementedError, match="management only"): + config.transform_search_vector_store_request( + vector_store_id="test-id", + query="test query", + vector_store_search_optional_params={}, + api_base="http://localhost:9380", + litellm_logging_obj=logging_obj, + litellm_params={} + ) + + def test_transform_search_vector_store_response_not_implemented(self): + """Test that search response transformation raises NotImplementedError.""" + config = RAGFlowVectorStoreConfig() + logging_obj = MagicMock(spec=LiteLLMLoggingObj) + mock_response = Mock(spec=httpx.Response) + + with pytest.raises(NotImplementedError, match="management only"): + config.transform_search_vector_store_response(mock_response, logging_obj) + + def _validate_vector_store_create_response(self, response): + """Override to handle RAGFlow-specific response format.""" + # RAGFlow IDs are hex strings (not OpenAI-style vs_* format) + # So we override the base validation to not check for vs_ prefix + assert isinstance(response, dict), f"Response should be a dict, got {type(response)}" + assert "id" in response, "Missing required field 'id' in create response" + assert "object" in response, "Missing required field 'object' in create response" + assert "created_at" in response, "Missing required field 'created_at' in create response" + + assert response["object"] == "vector_store", \ + f"Expected object to be 'vector_store', got '{response['object']}'" + + assert isinstance(response["id"], str), \ + f"id should be a string, got {type(response['id'])}" + assert len(response["id"]) > 0, "id should not be empty" + # RAGFlow IDs are hex strings, not OpenAI-style vs_* format + + assert isinstance(response["created_at"], int), \ + f"created_at should be an integer, got {type(response['created_at'])}" + assert response["created_at"] > 0, "created_at should be a positive timestamp" + + print(f"✅ RAGFlow create response validation passed: Dataset '{response['id']}' created successfully") + + @pytest.mark.parametrize("sync_mode", [True, False]) + @pytest.mark.asyncio + async def test_basic_create_vector_store(self, sync_mode): + """Override to handle RAGFlow-specific connection errors.""" + litellm._turn_on_debug() + litellm.set_verbose = True + base_request_args = self.get_base_create_vector_store_args() + + # Skip if no API key is set + if not os.getenv("RAGFLOW_API_KEY") and not base_request_args.get("api_key"): + pytest.skip("RAGFLOW_API_KEY not set, skipping integration test") + + # Extract custom_llm_provider from base args if present + create_args = base_request_args + try: + if sync_mode: + response = litellm.vector_stores.create( + name=f"test-ragflow-{int(__import__('time').time())}", + **create_args + ) + else: + response = await litellm.vector_stores.acreate( + name=f"test-ragflow-{int(__import__('time').time())}", + **create_args + ) + except litellm.InternalServerError: + pytest.skip("Skipping test due to litellm.InternalServerError") + except Exception as e: + error_str = str(e).lower() + error_type = type(e).__name__ + + # Check if it's a connection error + if (isinstance(e, (ConnectionError, OSError)) or + "connection" in error_str or + "connect" in error_str or + "APIConnectionError" in error_type): + pytest.skip(f"Skipping test due to connection error (RAGFlow instance may not be running): {e}") + + # If this is an authentication or permission error, skip the test + if "authentication" in error_str or "permission" in error_str or "unauthorized" in error_str: + pytest.skip(f"Skipping test due to authentication/permission error: {e}") + + # Re-raise if it's not a handled error + raise + + print("litellm create response=", json.dumps(response, indent=4, default=str)) + + # Validate response structure + self._validate_vector_store_create_response(response) + + @pytest.mark.parametrize("sync_mode", [True, False]) + @pytest.mark.asyncio + async def test_basic_search_vector_store(self, sync_mode): + """Override search test - RAGFlow doesn't support search.""" + pytest.skip("RAGFlow vector stores support dataset management only, not search") + diff --git a/ui/litellm-dashboard/public/assets/logos/a2a_agent.png b/ui/litellm-dashboard/public/assets/logos/a2a_agent.png new file mode 100644 index 00000000000..305ae1acaf4 Binary files /dev/null and b/ui/litellm-dashboard/public/assets/logos/a2a_agent.png differ diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.ts new file mode 100644 index 00000000000..aa0a6c2c9fb --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/credentials/useCredentials.ts @@ -0,0 +1,13 @@ +import { credentialListCall, CredentialsResponse } from "@/components/networking"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const credentialsKeys = createQueryKeys("credentials"); + +export const useCredentials = (accessToken: string | null) => { + return useQuery({ + queryKey: credentialsKeys.list({}), + queryFn: async () => await credentialListCall(accessToken!), + enabled: Boolean(accessToken), + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts new file mode 100644 index 00000000000..a15b4a06d13 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/login/useLogin.ts @@ -0,0 +1,11 @@ +import { useMutation } from "@tanstack/react-query"; +import { loginCall, LoginRequest } from "@/components/networking"; + +export const useLogin = () => { + return useMutation({ + mutationFn: async ({ username, password }: LoginRequest) => { + const result = await loginCall(username, password); + return result; + }, + }); +}; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.ts new file mode 100644 index 00000000000..147fdc241b0 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/uiConfig/useUIConfig.ts @@ -0,0 +1,14 @@ +import { getUiConfig, LiteLLMWellKnownUiConfig } from "@/components/networking"; +import { useQuery } from "@tanstack/react-query"; +import { createQueryKeys } from "../common/queryKeysFactory"; + +const uiConfigKeys = createQueryKeys("uiConfig"); + +export const useUIConfig = () => { + return useQuery({ + queryKey: uiConfigKeys.list({}), + queryFn: async () => await getUiConfig(), + staleTime: 24 * 60 * 60 * 1000, // 24 hours - data rarely changes + gcTime: 24 * 60 * 60 * 1000, // 24 hours - keep in cache for 24 hours + }); +}; 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 4f33ba6585d..0ef3f6302a9 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,47 +1,45 @@ -import React, { useState, useEffect, useRef } from "react"; -import { Text, Grid, Col } from "@tremor/react"; import { useQueryClient } from "@tanstack/react-query"; -import { CredentialItem, credentialListCall, CredentialsResponse } from "@/components/networking"; +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 { useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels"; +import { Team } from "@/components/key_team_helpers/key_list"; import CredentialsPanel from "@/components/model_add/credentials"; -import { getDisplayModelName } from "@/components/view_model/model_name_display"; -import { TabPanel, TabPanels, TabGroup, TabList, Tab, Icon } from "@tremor/react"; -import { DateRangePickerValue } from "@tremor/react"; import { - modelCostMap, - modelMetricsCall, - streamingModelMetricsCall, - modelExceptionsCall, - modelMetricsSlowResponsesCall, - getCallbacksCall, - setCallbacksCall, - modelSettingsCall, adminGlobalActivityExceptions, adminGlobalActivityExceptionsPerDeployment, allEndUsersCall, + getCallbacksCall, + modelCostMap, + modelExceptionsCall, + modelMetricsCall, + modelMetricsSlowResponsesCall, + modelSettingsCall, + setCallbacksCall, + streamingModelMetricsCall, } from "@/components/networking"; -import { useModelsInfo } from "@/app/(dashboard)/hooks/models/useModels"; -import { Form } from "antd"; -import { Typography } from "antd"; -import { RefreshIcon } from "@heroicons/react/outline"; -import type { UploadProps } from "antd"; -import { Team } from "@/components/key_team_helpers/key_list"; -import TeamInfoView from "../../../components/team/team_info"; import { Providers, getPlaceholder, getProviderModels } from "@/components/provider_info_helpers"; -import ModelInfoView from "../../../components/model_info_view"; +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 HealthCheckComponent from "../../../components/model_dashboard/HealthCheckComponent"; -import PassThroughSettings from "../../../components/pass_through_settings"; -import ModelGroupAliasSettings from "../../../components/model_group_alias_settings"; -import { all_admin_roles } from "@/utils/roles"; -import NotificationsManager from "../../../components/molecules/notifications_manager"; import AllModelsTab from "@/app/(dashboard)/models-and-endpoints/components/AllModelsTab"; -import PriceDataManagementTab from "@/app/(dashboard)/models-and-endpoints/components/PriceDataManagementTab"; -import ModelRetrySettingsTab from "@/app/(dashboard)/models-and-endpoints/components/ModelRetrySettingsTab"; 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 { all_admin_roles } from "@/utils/roles"; +import HealthCheckComponent from "../../../components/model_dashboard/HealthCheckComponent"; +import ModelGroupAliasSettings from "../../../components/model_group_alias_settings"; +import NotificationsManager from "../../../components/molecules/notifications_manager"; +import PassThroughSettings from "../../../components/pass_through_settings"; interface ModelDashboardProps { accessToken: string | null; @@ -100,7 +98,7 @@ const ModelsAndEndpointsView: React.FC = ({ const [providerModels, setProviderModels] = useState>([]); // Explicitly typing providerModels as a string array const [providerSettings, setProviderSettings] = useState([]); - const [selectedProvider, setSelectedProvider] = useState(Providers.OpenAI); + const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); const [editModalVisible, setEditModalVisible] = useState(false); const [selectedModel, setSelectedModel] = useState(null); @@ -134,8 +132,6 @@ const ModelsAndEndpointsView: React.FC = ({ const [allEndUsers, setAllEndUsers] = useState([]); - const [credentialsList, setCredentialsList] = useState([]); - // Model Group Alias state const [modelGroupAlias, setModelGroupAlias] = useState<{ [key: string]: string }>({}); @@ -160,21 +156,14 @@ const ModelsAndEndpointsView: React.FC = ({ isLoading: isLoadingModels, refetch: refetchModels, } = useModelsInfo(accessToken, userID, userRole); + const { data: credentialsResponse } = useCredentials(accessToken); + const credentialsList = credentialsResponse?.credentials || []; const setProviderModelsFn = (provider: Providers) => { const _providerModels = getProviderModels(provider, modelMap); setProviderModels(_providerModels); }; - const fetchCredentials = async (accessToken: string) => { - try { - const response: CredentialsResponse = await credentialListCall(accessToken); - setCredentialsList(response.credentials); - } catch (error) { - console.error("Error fetching credentials:", error); - } - }; - useEffect(() => { const handleClickOutside = (event: MouseEvent) => { if (dropdownRef.current && !dropdownRef.current.contains(event.target as Node)) { @@ -686,12 +675,7 @@ const ModelsAndEndpointsView: React.FC = ({ /> - + ({ + useRouter: vi.fn(() => ({ + push: mockPush, + replace: mockReplace, + })), +})); + +vi.mock("@/app/(dashboard)/hooks/uiConfig/useUIConfig", () => ({ + useUIConfig: vi.fn(), +})); + +vi.mock("@/utils/cookieUtils", () => ({ + getCookie: vi.fn(), +})); + +vi.mock("@/utils/jwtUtils", () => ({ + isJwtExpired: vi.fn(), +})); + +vi.mock("@/components/networking", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + getProxyBaseUrl: vi.fn().mockReturnValue("http://localhost:4000"), + }; +}); + +vi.mock("@/app/(dashboard)/hooks/login/useLogin", () => ({ + useLogin: vi.fn(() => ({ + mutate: vi.fn(), + isPending: false, + error: null, + })), +})); + +import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig"; +import { getCookie } from "@/utils/cookieUtils"; +import { isJwtExpired } from "@/utils/jwtUtils"; + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +describe("LoginPage", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockPush.mockClear(); + mockReplace.mockClear(); + }); + + it("should render", async () => { + (useUIConfig as ReturnType).mockReturnValue({ + data: { auto_redirect_to_sso: false, server_root_path: "/", proxy_base_url: null }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(null); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(screen.getByRole("heading", { name: "Login" })).toBeInTheDocument(); + }); + }); + + it("should call router.replace to dashboard when jwt is valid", async () => { + const validToken = "valid-token"; + (useUIConfig as ReturnType).mockReturnValue({ + data: { auto_redirect_to_sso: false, server_root_path: "/", proxy_base_url: null }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(validToken); + (isJwtExpired as ReturnType).mockReturnValue(false); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui"); + }); + }); + + it("should call router.push to SSO when jwt is invalid and auto_redirect_to_sso is true", async () => { + const invalidToken = "invalid-token"; + (useUIConfig as ReturnType).mockReturnValue({ + data: { auto_redirect_to_sso: true, server_root_path: "/", proxy_base_url: null }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(invalidToken); + (isJwtExpired as ReturnType).mockReturnValue(true); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(mockPush).toHaveBeenCalledWith("http://localhost:4000/sso/key/generate"); + }); + }); + + it("should not call router when jwt is invalid and auto_redirect_to_sso is false", async () => { + const invalidToken = "invalid-token"; + (useUIConfig as ReturnType).mockReturnValue({ + data: { auto_redirect_to_sso: false, server_root_path: "/", proxy_base_url: null }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(invalidToken); + (isJwtExpired as ReturnType).mockReturnValue(true); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(screen.getByRole("heading", { name: "Login" })).toBeInTheDocument(); + }); + + expect(mockPush).not.toHaveBeenCalled(); + expect(mockReplace).not.toHaveBeenCalled(); + }); + + it("should send user to dashboard when jwt is valid even if auto_redirect_to_sso is true", async () => { + const validToken = "valid-token"; + (useUIConfig as ReturnType).mockReturnValue({ + data: { auto_redirect_to_sso: true, server_root_path: "/", proxy_base_url: null }, + isLoading: false, + }); + (getCookie as ReturnType).mockReturnValue(validToken); + (isJwtExpired as ReturnType).mockReturnValue(false); + + const queryClient = createQueryClient(); + render( + + + , + ); + + await waitFor(() => { + expect(mockReplace).toHaveBeenCalledWith("http://localhost:4000/ui"); + }); + + expect(mockPush).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx new file mode 100644 index 00000000000..85f2c6dd870 --- /dev/null +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -0,0 +1,159 @@ +"use client"; + +import { useLogin } from "@/app/(dashboard)/hooks/login/useLogin"; +import { useUIConfig } from "@/app/(dashboard)/hooks/uiConfig/useUIConfig"; +import LoadingScreen from "@/components/common_components/LoadingScreen"; +import { getProxyBaseUrl } from "@/components/networking"; +import { getCookie } from "@/utils/cookieUtils"; +import { isJwtExpired } from "@/utils/jwtUtils"; +import { InfoCircleOutlined } from "@ant-design/icons"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { Alert, Button, Card, Form, Input, Space, Typography } from "antd"; +import { useRouter } from "next/navigation"; +import { useEffect, useState } from "react"; + +function LoginPageContent() { + const [username, setUsername] = useState(""); + const [password, setPassword] = useState(""); + const [isLoading, setIsLoading] = useState(true); + const { data: uiConfig, isLoading: isConfigLoading } = useUIConfig(); + const loginMutation = useLogin(); + const router = useRouter(); + + useEffect(() => { + if (isConfigLoading) { + return; + } + + const rawToken = getCookie("token"); + if (rawToken && !isJwtExpired(rawToken)) { + router.replace(`${getProxyBaseUrl()}/ui`); + return; + } + + if (uiConfig && uiConfig.auto_redirect_to_sso) { + router.push(`${getProxyBaseUrl()}/sso/key/generate`); + return; + } + + setIsLoading(false); + }, [isConfigLoading, router, uiConfig]); + + const handleSubmit = () => { + loginMutation.mutate( + { username, password }, + { + onSuccess: (data) => { + router.push(data.redirect_url); + }, + }, + ); + }; + + const error = loginMutation.error instanceof Error ? loginMutation.error.message : null; + const isLoginLoading = loginMutation.isPending; + + const { Title, Text, Paragraph } = Typography; + + if (isConfigLoading || isLoading) { + return ; + } + + return ( +
+ + +
+ 🚅 LiteLLM +
+ +
+ Login + Access your LiteLLM Admin UI. +
+ + + + By default, Username is admin and + Password is your set LiteLLM Proxy + MASTER_KEY. + + + Need to set UI credentials or SSO?{" "} + + Check the documentation + + . + + + } + type="info" + icon={} + showIcon + /> + + {error && } + +
+ + setUsername(e.target.value)} + disabled={isLoginLoading} + size="large" + className="rounded-md border-gray-300" + /> + + + + setPassword(e.target.value)} + disabled={isLoginLoading} + size="large" + /> + + + + + +
+ + +
+ ); +} + +export default function LoginPage() { + const queryClient = new QueryClient(); + + return ( + + + + ); +} diff --git a/ui/litellm-dashboard/src/app/login/page.tsx b/ui/litellm-dashboard/src/app/login/page.tsx new file mode 100644 index 00000000000..a539f87ee5e --- /dev/null +++ b/ui/litellm-dashboard/src/app/login/page.tsx @@ -0,0 +1,5 @@ +"use client"; + +import LoginPage from "./LoginPage"; + +export default LoginPage; diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 422d140e07c..20f5480c970 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -1,49 +1,47 @@ "use client"; -import React, { Suspense, useEffect, useState } from "react"; -import { useSearchParams } from "next/navigation"; -import { jwtDecode } from "jwt-decode"; -import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import { Team } from "@/components/key_team_helpers/key_list"; -import Navbar from "@/components/navbar"; -import { ThemeProvider } from "@/contexts/ThemeContext"; -import UserDashboard from "@/components/user_dashboard"; -import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; -import ViewUserDashboard from "@/components/view_users"; -import Organizations from "@/components/organizations"; -import { fetchOrganizations } from "@/components/organizations"; -import AdminPanel from "@/components/admins"; -import Settings from "@/components/settings"; -import GeneralSettings from "@/components/general_settings"; -import PassThroughSettings from "@/components/pass_through_settings"; -import BudgetPanel from "@/components/budgets/budget_panel"; -import SpendLogsTable from "@/components/view_logs"; -import ModelHubTable from "@/components/model_hub_table"; -import PublicModelHub from "@/components/public_model_hub"; -import NewUsagePage from "@/components/new_usage"; import APIReferenceView from "@/app/(dashboard)/api-reference/APIReferenceView"; -import PlaygroundPage from "@/app/(dashboard)/playground/page"; -import Usage from "@/components/usage"; -import CacheDashboard from "@/components/cache_dashboard"; -import { getUiConfig, proxyBaseUrl, setGlobalLitellmHeaderName } from "@/components/networking"; -import { Organization } from "@/components/networking"; -import GuardrailsPanel from "@/components/guardrails"; -import AgentsPanel from "@/components/agents"; -import PromptsPanel from "@/components/prompts"; -import TransformRequestPanel from "@/components/transform_request"; -import { fetchUserModels } from "@/components/organisms/create_key_button"; -import { fetchTeams } from "@/components/common_components/fetch_teams"; -import { MCPServers } from "@/components/mcp_tools"; -import TagManagement from "@/components/tag_management"; -import VectorStoreManagement from "@/components/vector_store_management"; -import UIThemeSettings from "@/components/ui_theme_settings"; -import { CostTrackingSettings } from "@/components/CostTrackingSettings"; -import { UiLoadingSpinner } from "@/components/ui/ui-loading-spinner"; -import { cx } from "@/lib/cva.config"; import SidebarProvider from "@/app/(dashboard)/components/SidebarProvider"; +import OldModelDashboard from "@/app/(dashboard)/models-and-endpoints/ModelsAndEndpointsView"; +import PlaygroundPage from "@/app/(dashboard)/playground/page"; +import AdminPanel from "@/components/admins"; +import AgentsPanel from "@/components/agents"; +import BudgetPanel from "@/components/budgets/budget_panel"; +import CacheDashboard from "@/components/cache_dashboard"; +import { fetchTeams } from "@/components/common_components/fetch_teams"; +import LoadingScreen from "@/components/common_components/LoadingScreen"; +import { CostTrackingSettings } from "@/components/CostTrackingSettings"; +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 Navbar from "@/components/navbar"; +import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName } from "@/components/networking"; +import NewUsagePage from "@/components/new_usage"; import OldTeams from "@/components/OldTeams"; +import { fetchUserModels } from "@/components/organisms/create_key_button"; +import Organizations, { fetchOrganizations } from "@/components/organizations"; +import PassThroughSettings from "@/components/pass_through_settings"; +import PromptsPanel from "@/components/prompts"; +import PublicModelHub from "@/components/public_model_hub"; import { SearchTools } from "@/components/search_tools"; +import Settings from "@/components/settings"; +import TagManagement from "@/components/tag_management"; +import TransformRequestPanel from "@/components/transform_request"; +import UIThemeSettings from "@/components/ui_theme_settings"; +import Usage from "@/components/usage"; +import UserDashboard from "@/components/user_dashboard"; +import VectorStoreManagement from "@/components/vector_store_management"; +import SpendLogsTable from "@/components/view_logs"; +import ViewUserDashboard from "@/components/view_users"; +import { ThemeProvider } from "@/contexts/ThemeContext"; +import { isJwtExpired } from "@/utils/jwtUtils"; import { isAdminRole } from "@/utils/roles"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { jwtDecode } from "jwt-decode"; +import { useSearchParams } from "next/navigation"; +import { Suspense, useEffect, useState } from "react"; function getCookie(name: string) { // Safer cookie read + decoding; handles '=' inside values @@ -62,19 +60,6 @@ function deleteCookie(name: string, path = "/") { document.cookie = `${name}=; Max-Age=0; Path=${path}`; } -function isJwtExpired(token: string): boolean { - try { - const decoded: any = jwtDecode(token); - if (decoded && typeof decoded.exp === "number") { - return decoded.exp * 1000 <= Date.now(); - } - return false; - } catch { - // If we can't decode, treat as invalid/expired - return true; - } -} - function formatUserRole(userRole: string) { if (!userRole) { return "Undefined Role"; @@ -112,19 +97,6 @@ interface ProxySettings { const queryClient = new QueryClient(); -function LoadingScreen() { - return ( -
-
🚅 LiteLLM
- -
- - Loading... -
-
- ); -} - export default function CreateKeyPage() { const [userRole, setUserRole] = useState(""); const [premiumUser, setPremiumUser] = useState(false); @@ -214,7 +186,7 @@ export default function CreateKeyPage() { useEffect(() => { if (redirectToLogin) { // Replace instead of assigning to avoid back-button loops - const dest = (proxyBaseUrl || "") + "/sso/key/generate"; + const dest = (proxyBaseUrl || "") + "/ui/login"; window.location.replace(dest); } }, [redirectToLogin]); @@ -328,7 +300,7 @@ export default function CreateKeyPage() { sidebarCollapsed={sidebarCollapsed} onToggleSidebar={toggleSidebar} /> -
+
diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index 1878e1364d0..a791ece9bb7 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -38,10 +38,9 @@ const ModelSection = ({ modelName, metrics }: { modelName: string; metrics: Mode - {/* Top API Keys Section */} {metrics.top_api_keys && metrics.top_api_keys.length > 0 && ( - Top API Keys by Spend + Top Virtual Keys by Spend
{metrics.top_api_keys.map((keyData, index) => ( @@ -384,12 +383,12 @@ export const processActivityData = ( }); }); - // Process API key breakdowns for each metric (skip if key is 'api_keys' to avoid duplication) + // Process Virtual Key breakdowns for each metric (skip if key is 'api_keys' to avoid duplication) if (key !== "api_keys") { Object.entries(modelMetrics).forEach(([model, _]) => { const apiKeyBreakdown: Record = {}; - // Aggregate API key data across all days + // Aggregate Virtual Key data across all days dailyActivity.results.forEach((day) => { const modelData = day.breakdown[key]?.[model]; if (modelData && "api_key_breakdown" in modelData) { diff --git a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx index d34b14ceabf..a8046d146a8 100644 --- a/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx +++ b/ui/litellm-dashboard/src/components/bulk_create_users_button.tsx @@ -569,7 +569,7 @@ const BulkCreateUsersButton: React.FC = ({
  • Download our CSV template
  • Add your users' information to the spreadsheet
  • Save the file and upload it here
    • -
  • After creation, download the results file containing the API keys for each user
    • +
  • After creation, download the results file containing the Virtual Keys for each user
    • @@ -809,9 +809,9 @@ const BulkCreateUsersButton: React.FC = ({
    User creation complete - Next step: Download the credentials file containing API - keys and invitation links. Users will need these API keys to make LLM requests through - LiteLLM. + Next step: Download the credentials file containing + Virtual Keys and invitation links. Users will need these Virtual Keys to make LLM requests + through LiteLLM.
    diff --git a/ui/litellm-dashboard/src/components/cache_dashboard.tsx b/ui/litellm-dashboard/src/components/cache_dashboard.tsx index a1c0cb0a664..38c0f1a8f41 100644 --- a/ui/litellm-dashboard/src/components/cache_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/cache_dashboard.tsx @@ -293,7 +293,11 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole - + {uniqueApiKeys.map((key) => ( {key} @@ -388,11 +392,7 @@ const CacheDashboard: React.FC = ({ accessToken, token, userRole /> - + diff --git a/ui/litellm-dashboard/src/components/common_components/LoadingScreen.test.tsx b/ui/litellm-dashboard/src/components/common_components/LoadingScreen.test.tsx new file mode 100644 index 00000000000..5601c53d574 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/LoadingScreen.test.tsx @@ -0,0 +1,10 @@ +import { render, screen } from "@testing-library/react"; +import { describe, it, expect } from "vitest"; +import LoadingScreen from "./LoadingScreen"; + +describe("LoadingScreen", () => { + it("should render", () => { + render(); + expect(screen.getByText("Loading...")).toBeInTheDocument(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/common_components/LoadingScreen.tsx b/ui/litellm-dashboard/src/components/common_components/LoadingScreen.tsx new file mode 100644 index 00000000000..6144fa6aad1 --- /dev/null +++ b/ui/litellm-dashboard/src/components/common_components/LoadingScreen.tsx @@ -0,0 +1,15 @@ +import { cx } from "@/lib/cva.config"; +import { UiLoadingSpinner } from "../ui/ui-loading-spinner"; + +export default function LoadingScreen() { + return ( +
    +
    🚅 LiteLLM
    + +
    + + Loading... +
    +
    + ); +} diff --git a/ui/litellm-dashboard/src/components/common_components/PassThroughSecuritySection.tsx b/ui/litellm-dashboard/src/components/common_components/PassThroughSecuritySection.tsx index c42094abb55..c63770d3c85 100644 --- a/ui/litellm-dashboard/src/components/common_components/PassThroughSecuritySection.tsx +++ b/ui/litellm-dashboard/src/components/common_components/PassThroughSecuritySection.tsx @@ -21,7 +21,7 @@ const PassThroughSecuritySection: React.FC = ({ Security - When enabled, requests to this endpoint will require a valid LiteLLM API key + When enabled, requests to this endpoint will require a valid LiteLLM Virtual Key {premiumUser ? ( @@ -35,22 +35,13 @@ const PassThroughSecuritySection: React.FC = ({ ) : (
    - + Authentication (Premium)
    Setting authentication for pass-through endpoints is a LiteLLM Enterprise feature. Get a trial key{" "} - + here . @@ -63,4 +54,3 @@ const PassThroughSecuritySection: React.FC = ({ }; export default PassThroughSecuritySection; - diff --git a/ui/litellm-dashboard/src/components/dashboard_default_team.tsx b/ui/litellm-dashboard/src/components/dashboard_default_team.tsx index 36805b8912a..6506e1a60a5 100644 --- a/ui/litellm-dashboard/src/components/dashboard_default_team.tsx +++ b/ui/litellm-dashboard/src/components/dashboard_default_team.tsx @@ -69,7 +69,8 @@ const DashboardTeam: React.FC = ({ Select Team - If you belong to multiple teams, this setting controls which team is used by default when creating new API Keys. + If you belong to multiple teams, this setting controls which team is used by default when creating new Virtual + Keys. Default Team: If no team_id is set for a key, it will be grouped under here. diff --git a/ui/litellm-dashboard/src/components/entity_usage.tsx b/ui/litellm-dashboard/src/components/entity_usage.tsx index a5789b7dbac..501eac7124b 100644 --- a/ui/litellm-dashboard/src/components/entity_usage.tsx +++ b/ui/litellm-dashboard/src/components/entity_usage.tsx @@ -550,7 +550,7 @@ const EntityUsage: React.FC = ({ {/* Top API Keys */} - Top API Keys + Top Virtual Keys = ({ setLoading(true); try { const agentIdsToMakePublic = Array.from(selectedAgents); - + // Make batch API call for all agents await makeAgentsPublicCall(accessToken, agentIdsToMakePublic); @@ -127,8 +127,8 @@ const MakeAgentPublicForm: React.FC = ({
    - Select the agents you want to be visible on the public model hub. Users will still require a valid API key to - use these agents. + Select the agents you want to be visible on the public model hub. Users will still require a valid Virtual Key + to use these agents.
    @@ -141,10 +141,7 @@ const MakeAgentPublicForm: React.FC = ({ agentHubData.map((agent) => { const agentId = agent.agent_id || agent.name; return ( -
    +
    handleAgentSelection(agentId, e.target.checked)} @@ -217,9 +214,7 @@ const MakeAgentPublicForm: React.FC = ({ )}
    - {agent?.description && ( - {agent.description} - )} + {agent?.description && {agent.description}}
    ); @@ -296,4 +291,3 @@ const MakeAgentPublicForm: React.FC = ({ }; export default MakeAgentPublicForm; - diff --git a/ui/litellm-dashboard/src/components/make_mcp_public_form.tsx b/ui/litellm-dashboard/src/components/make_mcp_public_form.tsx index 29f866f8bc6..f7bba175800 100644 --- a/ui/litellm-dashboard/src/components/make_mcp_public_form.tsx +++ b/ui/litellm-dashboard/src/components/make_mcp_public_form.tsx @@ -76,7 +76,7 @@ const MakeMCPPublicForm: React.FC = ({ const publicServerIds = mcpHubData .filter((server) => server.mcp_info?.is_public === true) .map((server) => server.server_id); - + // Preselect servers that are already public setSelectedServers(new Set(publicServerIds)); } @@ -91,7 +91,7 @@ const MakeMCPPublicForm: React.FC = ({ setLoading(true); try { const serverIdsToMakePublic = Array.from(selectedServers); - + // Make batch API call for all servers await makeMCPPublicCall(accessToken, serverIdsToMakePublic); @@ -128,8 +128,8 @@ const MakeMCPPublicForm: React.FC = ({
    - Select the MCP servers you want to be visible on the public model hub. Users will still require a valid API key to - use these servers. + Select the MCP servers you want to be visible on the public model hub. Users will still require a valid + Virtual Key to use these servers.
    @@ -161,22 +161,20 @@ const MakeMCPPublicForm: React.FC = ({ {server.transport} - {server.status || "unknown"}
    - - {server.description || server.url} - + {server.description || server.url} {server.allowed_tools && server.allowed_tools.length > 0 && (
    {server.allowed_tools.slice(0, 3).map((tool, idx) => ( @@ -236,14 +234,14 @@ const MakeMCPPublicForm: React.FC = ({ {server.transport} - {server.status || "unknown"} @@ -251,12 +249,8 @@ const MakeMCPPublicForm: React.FC = ({ )}
    - {server?.description && ( - {server.description} - )} - {server?.url && ( - {server.url} - )} + {server?.description && {server.description}} + {server?.url && {server.url}}
    ); @@ -267,8 +261,8 @@ const MakeMCPPublicForm: React.FC = ({
    - Total: {selectedServers.size} MCP server{selectedServers.size !== 1 ? "s" : ""} will be made - public + Total: {selectedServers.size} MCP server{selectedServers.size !== 1 ? "s" : ""} will be + made public
    @@ -333,4 +327,3 @@ const MakeMCPPublicForm: React.FC = ({ }; export default MakeMCPPublicForm; - diff --git a/ui/litellm-dashboard/src/components/make_model_public_form.tsx b/ui/litellm-dashboard/src/components/make_model_public_form.tsx index e67d60fb33b..750bdc24eeb 100644 --- a/ui/litellm-dashboard/src/components/make_model_public_form.tsx +++ b/ui/litellm-dashboard/src/components/make_model_public_form.tsx @@ -152,8 +152,8 @@ const MakeModelPublicForm: React.FC = ({
    - Select the models you want to be visible on the public model hub. Users will still require a valid API key to - use these models. + Select the models you want to be visible on the public model hub. Users will still require a valid Virtual Key + to use these models. {/* Filters */} diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx index 4b4f1ab676b..5a012c1fc5c 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx @@ -220,12 +220,12 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] } - title="API Key Setup" - description="Configure your LiteLLM Proxy API key for authentication" + title="Virtual Key Setup" + description="Configure your LiteLLM Proxy Virtual Key for authentication" >
    - Get your API key from your LiteLLM Proxy dashboard or contact your administrator + Get your Virtual Key from your LiteLLM Proxy dashboard or contact your administrator
    @@ -249,7 +249,7 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] = ({ currentServerAccessGroups = [] "server_url": "${proxyBaseUrl}/mcp", "require_approval": "never", "headers": { - "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY", + "x-litellm-api-key": "Bearer YOUR_LITELLM_VIRTUAL_KEY", "x-mcp-servers": ["Zapier_MCP,dev"] } } diff --git a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.test.tsx b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.test.tsx new file mode 100644 index 00000000000..aee7a0cdd1d --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.test.tsx @@ -0,0 +1,108 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, waitFor } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import { Providers } from "../provider_info_helpers"; +import AddCredentialModal from "./AddCredentialModal"; + +vi.mock("../networking", async () => { + const actual = await vi.importActual("../networking"); + return { + ...actual, + getProviderCreateMetadata: vi.fn().mockResolvedValue([ + { + provider: "OpenAI", + provider_display_name: Providers.OpenAI, + litellm_provider: "openai", + default_model_placeholder: "gpt-3.5-turbo", + credential_fields: [ + { + key: "api_key", + label: "OpenAI API Key", + field_type: "password", + required: true, + }, + { + key: "api_base", + label: "API Base", + field_type: "text", + placeholder: "https://api.openai.com/v1", + }, + ], + }, + { + provider: "Anthropic", + provider_display_name: Providers.Anthropic, + litellm_provider: "anthropic", + default_model_placeholder: "claude-3-opus-20240229", + credential_fields: [ + { + key: "api_key", + label: "Anthropic API Key", + field_type: "password", + required: true, + }, + ], + }, + ]), + }; +}); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +const mockUploadProps = { + beforeUpload: vi.fn(), + onChange: vi.fn(), +}; + +describe("AddCredentialModal", () => { + it("should render", () => { + const queryClient = createQueryClient(); + const onCancel = vi.fn(); + const onAddCredential = vi.fn(); + + render( + + + , + ); + + expect(screen.getByText("Add New Credential")).toBeInTheDocument(); + expect(screen.getByLabelText("Credential Name:")).toBeInTheDocument(); + expect(screen.getByLabelText("Provider:")).toBeInTheDocument(); + }); + + it("should show the correct provider fields", async () => { + const queryClient = createQueryClient(); + const onCancel = vi.fn(); + const onAddCredential = vi.fn(); + + render( + + + , + ); + + await waitFor(() => { + expect(screen.getByLabelText("OpenAI API Key")).toBeInTheDocument(); + expect(screen.getByPlaceholderText("https://api.openai.com/v1")).toBeInTheDocument(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx new file mode 100644 index 00000000000..694a98201c6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_add/AddCredentialModal.tsx @@ -0,0 +1,118 @@ +import { TextInput } from "@tremor/react"; +import { Select as AntdSelect, Button, Form, Modal, Tooltip, Typography } from "antd"; +import type { UploadProps } from "antd/es/upload"; +import React, { useState } from "react"; +import ProviderSpecificFields from "../add_model/provider_specific_fields"; +import { Providers, providerLogoMap } from "../provider_info_helpers"; +const { Link } = Typography; + +interface AddCredentialsModalProps { + open: boolean; + onCancel: () => void; + onAddCredential: (values: any) => void; + uploadProps: UploadProps; +} + +const AddCredentialsModal: React.FC = ({ open, onCancel, onAddCredential, uploadProps }) => { + const [form] = Form.useForm(); + const [selectedProvider, setSelectedProvider] = useState(Providers.OpenAI); + + const handleSubmit = (values: any) => { + const filteredValues = Object.entries(values).reduce((acc, [key, value]) => { + if (value !== "" && value !== undefined && value !== null) { + acc[key] = value; + } + return acc; + }, {} as any); + onAddCredential(filteredValues); + form.resetFields(); + }; + + return ( + { + onCancel(); + form.resetFields(); + }} + footer={null} + width={600} + > +
    + {/* Credential Name */} + + + + + {/* Provider Selection */} + + { + setSelectedProvider(value as Providers); + form.setFieldValue("custom_llm_provider", value); + }} + > + {Object.entries(Providers).map(([providerEnum, providerDisplayName]) => ( + +
    + {`${providerEnum} { + const target = e.target as HTMLImageElement; + const parent = target.parentElement; + if (parent) { + const fallbackDiv = document.createElement("div"); + fallbackDiv.className = + "w-5 h-5 rounded-full bg-gray-200 flex items-center justify-center text-xs"; + fallbackDiv.textContent = providerDisplayName.charAt(0); + parent.replaceChild(fallbackDiv, target); + } + }} + /> + {providerDisplayName} +
    +     + ))} +     +     + + + + {/* Modal Footer */} +
    + + Need Help? + + +
    + + +
    +
    + +
    + ); +}; + +export default AddCredentialsModal; diff --git a/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.test.tsx b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.test.tsx new file mode 100644 index 00000000000..def3b4f6cd7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.test.tsx @@ -0,0 +1,123 @@ +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { render, screen, waitFor } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; +import { Providers } from "../provider_info_helpers"; +import { CredentialItem } from "../networking"; +import EditCredentialModal from "./EditCredentialModal"; + +vi.mock("../networking", async () => { + const actual = await vi.importActual("../networking"); + return { + ...actual, + getProviderCreateMetadata: vi.fn().mockResolvedValue([ + { + provider: "OpenAI", + provider_display_name: Providers.OpenAI, + litellm_provider: "openai", + default_model_placeholder: "gpt-3.5-turbo", + credential_fields: [ + { + key: "api_key", + label: "OpenAI API Key", + field_type: "password", + required: true, + }, + { + key: "api_base", + label: "API Base", + field_type: "text", + placeholder: "https://api.openai.com/v1", + }, + ], + }, + { + provider: "Anthropic", + provider_display_name: Providers.Anthropic, + litellm_provider: "anthropic", + default_model_placeholder: "claude-3-opus-20240229", + credential_fields: [ + { + key: "api_key", + label: "Anthropic API Key", + field_type: "password", + required: true, + }, + ], + }, + ]), + }; +}); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, + }); + +const mockUploadProps = { + beforeUpload: vi.fn(), + onChange: vi.fn(), +}; + +const mockCredential: CredentialItem = { + credential_name: "test-credential", + credential_values: { + api_key: "test-api-key", + api_base: "https://api.test.com", + }, + credential_info: { + custom_llm_provider: Providers.OpenAI, + }, +}; + +describe("EditCredentialModal", () => { + it("should render", () => { + const queryClient = createQueryClient(); + const onCancel = vi.fn(); + const onUpdateCredential = vi.fn(); + + render( + + + , + ); + + expect(screen.getByText("Edit Credential")).toBeInTheDocument(); + expect(screen.getByLabelText("Credential Name:")).toBeInTheDocument(); + expect(screen.getByLabelText("Provider:")).toBeInTheDocument(); + }); + + it("should render initial values", async () => { + const queryClient = createQueryClient(); + const onCancel = vi.fn(); + const onUpdateCredential = vi.fn(); + + render( + + + , + ); + + await waitFor(() => { + const credentialNameInput = screen.getByLabelText("Credential Name:") as HTMLInputElement; + expect(credentialNameInput.value).toBe("test-credential"); + expect(credentialNameInput.disabled).toBe(true); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/components/model_add/add_credentials_tab.tsx b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx similarity index 78% rename from ui/litellm-dashboard/src/components/model_add/add_credentials_tab.tsx rename to ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx index 9c061eb121a..b206ed6c91d 100644 --- a/ui/litellm-dashboard/src/components/model_add/add_credentials_tab.tsx +++ b/ui/litellm-dashboard/src/components/model_add/EditCredentialModal.tsx @@ -1,34 +1,29 @@ -import React, { useEffect, useState } from "react"; -import { Form, Button, Tooltip, Typography, Select as AntdSelect, Modal } from "antd"; -import type { UploadProps } from "antd/es/upload"; -import { Providers, providerLogoMap } from "../provider_info_helpers"; -import ProviderSpecificFields from "../add_model/provider_specific_fields"; import { TextInput } from "@tremor/react"; +import { Select as AntdSelect, Button, Form, Modal, Tooltip, Typography } from "antd"; +import type { UploadProps } from "antd/es/upload"; +import { useEffect, useState } from "react"; +import ProviderSpecificFields from "../add_model/provider_specific_fields"; import { CredentialItem } from "../networking"; -const { Title, Link } = Typography; +import { Providers, providerLogoMap } from "../provider_info_helpers"; +const { Link } = Typography; -interface AddCredentialsModalProps { - isVisible: boolean; +interface EditCredentialsModalProps { + open: boolean; onCancel: () => void; - onAddCredential: (values: any) => void; onUpdateCredential: (values: any) => void; uploadProps: UploadProps; - addOrEdit: "add" | "edit"; existingCredential: CredentialItem | null; } -const AddCredentialsModal: React.FC = ({ - isVisible, +export default function EditCredentialsModal({ + open, onCancel, - onAddCredential, onUpdateCredential, uploadProps, - addOrEdit, existingCredential, -}) => { +}: EditCredentialsModalProps) { const [form] = Form.useForm(); - const [selectedProvider, setSelectedProvider] = useState(Providers.OpenAI); - const [showAdvancedSettings, setShowAdvancedSettings] = useState(false); + const [selectedProvider, setSelectedProvider] = useState(Providers.Anthropic); const handleSubmit = (values: any) => { const filteredValues = Object.entries(values).reduce((acc, [key, value]) => { @@ -37,23 +32,25 @@ const AddCredentialsModal: React.FC = ({ } return acc; }, {} as any); - if (addOrEdit === "add") { - onAddCredential(filteredValues); - } else { - onUpdateCredential(filteredValues); - } + onUpdateCredential(filteredValues); form.resetFields(); }; useEffect(() => { if (existingCredential) { + // Spread all credential_values dynamically, converting undefined/null to null for form compatibility + const credentialValues = Object.entries(existingCredential.credential_values || {}).reduce( + (acc, [key, value]) => { + acc[key] = value ?? null; + return acc; + }, + {} as Record, + ); + form.setFieldsValue({ credential_name: existingCredential.credential_name, custom_llm_provider: existingCredential.credential_info.custom_llm_provider, - api_base: existingCredential.credential_values.api_base, - api_version: existingCredential.credential_values.api_version, - base_model: existingCredential.credential_values.base_model, - api_key: existingCredential.credential_values.api_key, + ...credentialValues, }); setSelectedProvider(existingCredential.credential_info.custom_llm_provider as Providers); } @@ -61,14 +58,15 @@ const AddCredentialsModal: React.FC = ({ return ( { onCancel(); form.resetFields(); }} footer={null} width={600} + destroyOnHidden={true} >
    {/* Credential Name */} @@ -142,12 +140,10 @@ const AddCredentialsModal: React.FC = ({ > Cancel - +
    ); -}; - -export default AddCredentialsModal; +} diff --git a/ui/litellm-dashboard/src/components/model_add/credentials.test.tsx b/ui/litellm-dashboard/src/components/model_add/credentials.test.tsx index 8e356baa12b..2504b2a9789 100644 --- a/ui/litellm-dashboard/src/components/model_add/credentials.test.tsx +++ b/ui/litellm-dashboard/src/components/model_add/credentials.test.tsx @@ -1,33 +1,51 @@ import { CredentialItem } from "@/components/networking"; -import { render, waitFor } from "@testing-library/react"; +import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; +import { act, fireEvent, render, screen, waitFor } from "@testing-library/react"; import { UploadProps } from "antd/es/upload"; import { describe, expect, it, vi } from "vitest"; import CredentialsPanel from "./credentials"; const DEFAULT_UPLOAD_PROPS = {} as UploadProps; -describe("CredentialsPanel", () => { - it("renders without crashing and fetches credentials when token exists", async () => { - const fetchCredentials = vi.fn(() => Promise.resolve()); +const mockUseAuthorized = vi.fn(); +const mockUseCredentials = vi.fn(); - const { getByRole, getByText } = render( - , - ); +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => mockUseAuthorized(), +})); - await waitFor(() => { - expect(getByRole("button", { name: /add credential/i })).toBeInTheDocument(); - expect(getByText("Credential Name")).toBeInTheDocument(); - expect(getByText("Provider")).toBeInTheDocument(); - }); +vi.mock("@/app/(dashboard)/hooks/credentials/useCredentials", () => ({ + useCredentials: () => mockUseCredentials(), +})); + +const createQueryClient = () => + new QueryClient({ + defaultOptions: { + queries: { + retry: false, + gcTime: 0, + }, + }, }); - it("displays provided credentials and still calls the fetch helper", async () => { - const fetchCredentials = vi.fn(() => Promise.resolve()); +describe("CredentialsPanel", () => { + it("should render", () => { + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" }); + mockUseCredentials.mockReturnValue({ + data: { credentials: [] }, + refetch: vi.fn(), + }); + + render( + + + , + ); + + expect(screen.getByRole("button", { name: /add credential/i })).toBeInTheDocument(); + }); + + it("should display provided credentials", () => { const credentials: CredentialItem[] = [ { credential_name: "openai-key", @@ -36,15 +54,58 @@ describe("CredentialsPanel", () => { }, ]; - const { getByText } = render( - , + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" }); + mockUseCredentials.mockReturnValue({ + data: { credentials }, + refetch: vi.fn(), + }); + + render( + + + , ); - await waitFor(() => expect(getByText("openai-key")).toBeInTheDocument()); + expect(screen.getByText("openai-key")).toBeInTheDocument(); + }); + + it("should display empty state when no credentials are provided", () => { + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" }); + mockUseCredentials.mockReturnValue({ + data: { credentials: [] }, + refetch: vi.fn(), + }); + + render( + + + , + ); + + expect(screen.getByText("No credentials configured")).toBeInTheDocument(); + }); + + it("should open add modal when add button is clicked", async () => { + mockUseAuthorized.mockReturnValue({ accessToken: "test-token" }); + mockUseCredentials.mockReturnValue({ + data: { credentials: [] }, + refetch: vi.fn(), + }); + + render( + + + , + ); + + const addButton = screen.getByRole("button", { name: /add credential/i }); + + act(() => { + fireEvent.click(addButton); + }); + + await waitFor(() => { + expect(screen.getByText("Add New Credential")).toBeInTheDocument(); + }); }); }); diff --git a/ui/litellm-dashboard/src/components/model_add/credentials.tsx b/ui/litellm-dashboard/src/components/model_add/credentials.tsx index e36a759294b..3887e340daa 100644 --- a/ui/litellm-dashboard/src/components/model_add/credentials.tsx +++ b/ui/litellm-dashboard/src/components/model_add/credentials.tsx @@ -1,45 +1,46 @@ -import React, { useState, useEffect } from "react"; import { + credentialCreateCall, + credentialDeleteCall, + CredentialItem, + credentialUpdateCall, +} from "@/components/networking"; // Assume this is your networking function +import { PencilAltIcon, TrashIcon } from "@heroicons/react/outline"; +import { + Badge, + Button, + Card, Table, TableBody, TableCell, TableHead, TableHeaderCell, TableRow, - Card, Text, - Badge, - Button, } from "@tremor/react"; -import { PencilAltIcon, TrashIcon } from "@heroicons/react/outline"; -import { UploadProps } from "antd/es/upload"; -import { - credentialCreateCall, - credentialDeleteCall, - credentialUpdateCall, - CredentialItem, -} from "@/components/networking"; // Assume this is your networking function -import AddCredentialsTab from "./add_credentials_tab"; -import CredentialDeleteModal from "./CredentialDeleteModal"; import { Form } from "antd"; +import { UploadProps } from "antd/es/upload"; +import { useState } from "react"; +import DeleteResourceModal from "../common_components/DeleteResourceModal"; import NotificationsManager from "../molecules/notifications_manager"; +import AddCredentialsTab from "./AddCredentialModal"; +import EditCredentialsModal from "./EditCredentialModal"; +import { useCredentials } from "@/app/(dashboard)/hooks/credentials/useCredentials"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; interface CredentialsPanelProps { - accessToken: string | null; uploadProps: UploadProps; - credentialList: CredentialItem[]; - fetchCredentials: (accessToken: string) => Promise; } -const CredentialsPanel: React.FC = ({ - accessToken, - uploadProps, - credentialList, - fetchCredentials, -}) => { +const CredentialsPanel: React.FC = ({ uploadProps }) => { + const { accessToken } = useAuthorized(); + const { data: credentialsResponse, refetch: refetchCredentials } = useCredentials(accessToken); + const credentialList = credentialsResponse?.credentials || []; + const [isAddModalOpen, setIsAddModalOpen] = useState(false); const [isUpdateModalOpen, setIsUpdateModalOpen] = useState(false); const [selectedCredential, setSelectedCredential] = useState(null); - const [credentialToDelete, setCredentialToDelete] = useState(null); + const [credentialToDelete, setCredentialToDelete] = useState(null); + const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); + const [isCredentialDeleting, setIsCredentialDeleting] = useState(false); const [form] = Form.useForm(); const restrictedFields = ["credential_name", "custom_llm_provider"]; @@ -60,10 +61,10 @@ const CredentialsPanel: React.FC = ({ }, }; - const response = await credentialUpdateCall(accessToken, values.credential_name, newCredential); + await credentialUpdateCall(accessToken, values.credential_name, newCredential); NotificationsManager.success("Credential updated successfully"); setIsUpdateModalOpen(false); - fetchCredentials(accessToken); + await refetchCredentials(); }; const handleAddCredential = async (values: any) => { @@ -84,19 +85,12 @@ const CredentialsPanel: React.FC = ({ }; // Add to list and close modal - const response = await credentialCreateCall(accessToken, newCredential); + await credentialCreateCall(accessToken, newCredential); NotificationsManager.success("Credential added successfully"); setIsAddModalOpen(false); - fetchCredentials(accessToken); + await refetchCredentials(); }; - useEffect(() => { - if (!accessToken) { - return; - } - fetchCredentials(accessToken); - }, [accessToken]); - const renderProviderBadge = (provider: string) => { const providerColors: Record = { openai: "blue", @@ -113,27 +107,38 @@ const CredentialsPanel: React.FC = ({ ); }; - const handleDeleteCredential = async (credentialName: string) => { - if (!accessToken) { + const handleDeleteCredential = async () => { + if (!accessToken || !credentialToDelete) { return; } - const response = await credentialDeleteCall(accessToken, credentialName); - NotificationsManager.success("Credential deleted successfully"); - setCredentialToDelete(null); - fetchCredentials(accessToken); + setIsCredentialDeleting(true); + try { + await credentialDeleteCall(accessToken, credentialToDelete.credential_name); + NotificationsManager.success("Credential deleted successfully"); + await refetchCredentials(); + } catch (error) { + NotificationsManager.error("Failed to delete credential"); + } finally { + setCredentialToDelete(null); + setIsDeleteModalOpen(false); + setIsCredentialDeleting(false); + } }; - const openDeleteModal = (credentialName: string) => { - setCredentialToDelete(credentialName); + const openDeleteModal = (credential: CredentialItem) => { + setCredentialToDelete(credential); + setIsDeleteModalOpen(true); }; const closeDeleteModal = () => { setCredentialToDelete(null); + setIsDeleteModalOpen(false); }; return ( -
    -
    +
    + +
    Configured credentials for different AI providers. Add and manage your API credentials.
    @@ -143,6 +148,7 @@ const CredentialsPanel: React.FC = ({ Credential Name Provider + Actions @@ -173,7 +179,8 @@ const CredentialsPanel: React.FC = ({ icon={TrashIcon} variant="light" size="sm" - onClick={() => openDeleteModal(credential.credential_name)} + onClick={() => openDeleteModal(credential)} + className="ml-2" /> @@ -182,41 +189,39 @@ const CredentialsPanel: React.FC = ({ - {isAddModalOpen && ( setIsAddModalOpen(false)} uploadProps={uploadProps} - addOrEdit="add" - onUpdateCredential={handleUpdateCredential} - existingCredential={null} /> )} {isUpdateModalOpen && ( - setIsUpdateModalOpen(false)} - addOrEdit="edit" /> )} - {credentialToDelete && ( - handleDeleteCredential(credentialToDelete)} - credentialName={credentialToDelete} - /> - )} +
    ); }; diff --git a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx index b74060db6c5..344ff2e94f2 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/table.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/table.tsx @@ -100,7 +100,7 @@ export function ModelDataTable({ key={header.id} className={`py-1 h-8 relative ${ header.id === "actions" - ? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)] z-20 w-[120px] ml-8" + ? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)] w-[120px] ml-8" : "" } ${header.column.columnDef.meta?.className || ""}`} style={{ @@ -160,7 +160,7 @@ export function ModelDataTable({ key={cell.id} className={`py-0.5 ${ cell.column.id === "actions" - ? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)] z-20 w-[120px] ml-8" + ? "sticky right-0 bg-white shadow-[-4px_0_8px_-6px_rgba(0,0,0,0.1)] w-[120px] ml-8" : "" } ${cell.column.columnDef.meta?.className || ""}`} style={{ diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 14763930aaf..0e45c0f3a91 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -124,7 +124,7 @@ export interface PromptSpec { prompt_info: PromptInfo; created_at?: string; updated_at?: string; - version?: number; // Explicit version number for version history + version?: number; // Explicit version number for version history } export interface PromptTemplateBase { @@ -205,6 +205,7 @@ export interface PublicModelHubInfo { export interface LiteLLMWellKnownUiConfig { server_root_path: string; proxy_base_url: string | null; + auto_redirect_to_sso: boolean; } export interface CredentialsResponse { @@ -7414,7 +7415,11 @@ interface RegisterMcpOAuthClientPayload { token_endpoint_auth_method?: string; } -export const registerMcpOAuthClient = async (accessToken: string, serverId: string, payload: RegisterMcpOAuthClientPayload) => { +export const registerMcpOAuthClient = async ( + accessToken: string, + serverId: string, + payload: RegisterMcpOAuthClientPayload, +) => { const base = getProxyBaseUrl(); const normalizedServerId = encodeURIComponent(serverId.trim()); const url = `${base}/v1/mcp/server/oauth/${normalizedServerId}/register`; @@ -7424,7 +7429,7 @@ export const registerMcpOAuthClient = async (accessToken: string, serverId: stri headers: { [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", - "Accept": "application/json, text/event-stream", + Accept: "application/json, text/event-stream", }, body: JSON.stringify(payload), }); @@ -7969,3 +7974,40 @@ const deriveErrorMessage = (errorData: any): string => { JSON.stringify(errorData) ); }; + +export interface LoginRequest { + username: string; + password: string; +} + +export interface LoginResponse { + redirect_url: string; +} + +export const loginCall = async (username: string, password: string): Promise => { + const proxyBaseUrl = getProxyBaseUrl(); + const loginUrl = proxyBaseUrl ? `${proxyBaseUrl}/v2/login` : "/v2/login"; + + const body = JSON.stringify({ + username, + password, + }); + + const response = await fetch(loginUrl, { + method: "POST", + body, + credentials: "include", + headers: { + "Content-Type": "application/json", + }, + }); + + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + throw new Error(errorMessage); + } + + const data = await response.json(); + return data; +}; diff --git a/ui/litellm-dashboard/src/components/new_usage.test.tsx b/ui/litellm-dashboard/src/components/new_usage.test.tsx index a4969124f1f..a06045137d7 100644 --- a/ui/litellm-dashboard/src/components/new_usage.test.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.test.tsx @@ -239,7 +239,7 @@ describe("NewUsage", () => { // Check for chart titles expect(screen.getByText("Daily Spend")).toBeInTheDocument(); - expect(screen.getByText("Top API Keys")).toBeInTheDocument(); + expect(screen.getByText("Top Virtual Keys")).toBeInTheDocument(); }); it("should switch between tabs correctly", async () => { diff --git a/ui/litellm-dashboard/src/components/new_usage.tsx b/ui/litellm-dashboard/src/components/new_usage.tsx index 4794a7f091d..a8d30885493 100644 --- a/ui/litellm-dashboard/src/components/new_usage.tsx +++ b/ui/litellm-dashboard/src/components/new_usage.tsx @@ -580,7 +580,7 @@ const NewUsagePage: React.FC = ({ {/* Top API Keys */} - Top API Keys + Top Virtual Keys = ({ setApiKey(response["key"]); setSoftBudget(response["soft_budget"]); - NotificationsManager.success("API Key Created"); + NotificationsManager.success("Virtual Key Created"); form.resetFields(); localStorage.removeItem("userData" + userID); } catch (error) { @@ -415,7 +415,7 @@ const CreateKey: React.FC = ({ }; const handleCopy = () => { - NotificationsManager.success("API Key copied to clipboard"); + NotificationsManager.success("Virtual Key copied to clipboard"); }; useEffect(() => { @@ -505,7 +505,7 @@ const CreateKey: React.FC = ({ label={ Owned By{" "} - + @@ -594,8 +594,8 @@ const CreateKey: React.FC = ({ {isFormDisabled && (
    - Please select a team to continue configuring your API key. If you do not see any teams, please contact - your Proxy Admin to either provide you with access to models or to add you to a team. + Please select a team to continue configuring your Virtual Key. If you do not see any teams, please + contact your Proxy Admin to either provide you with access to models or to add you to a team.
    )} @@ -1277,7 +1277,7 @@ const CreateKey: React.FC = ({ {apiKey != null ? (
    - API Key: + Virtual Key:
    = ({
    - + {/*
    - New API Key: + New Virtual Key: 
    {regeneratedKey}
    NotificationManager.success("API Key copied to clipboard")} + onCopy={() => NotificationManager.success("Virtual Key copied to clipboard")} > - + diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/A2AMetrics.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/A2AMetrics.tsx new file mode 100644 index 00000000000..3b2ea0eb50b --- /dev/null +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/A2AMetrics.tsx @@ -0,0 +1,231 @@ +import React, { useState } from "react"; +import { Tooltip, Button } from "antd"; +import { + CheckCircleOutlined, + ClockCircleOutlined, + LoadingOutlined, + ExclamationCircleOutlined, + CopyOutlined, + DownOutlined, + RightOutlined, + LinkOutlined, + FileTextOutlined, + RobotOutlined, +} from "@ant-design/icons"; + +export interface A2ATaskMetadata { + taskId?: string; + contextId?: string; + status?: { + state?: string; + timestamp?: string; + message?: string; + }; + metadata?: Record; +} + +interface A2AMetricsProps { + a2aMetadata?: A2ATaskMetadata; + timeToFirstToken?: number; + totalLatency?: number; +} + +const getStatusIcon = (state?: string) => { + switch (state) { + case "completed": + return ; + case "working": + case "submitted": + return ; + case "failed": + case "canceled": + return ; + default: + return ; + } +}; + +const getStatusColor = (state?: string) => { + switch (state) { + case "completed": + return "bg-green-100 text-green-700"; + case "working": + case "submitted": + return "bg-blue-100 text-blue-700"; + case "failed": + case "canceled": + return "bg-red-100 text-red-700"; + default: + return "bg-gray-100 text-gray-700"; + } +}; + +const formatTimestamp = (timestamp?: string) => { + if (!timestamp) return null; + try { + const date = new Date(timestamp); + return date.toLocaleTimeString([], { hour: "2-digit", minute: "2-digit" }); + } catch { + return timestamp; + } +}; + +const truncateId = (id?: string, length = 8) => { + if (!id) return null; + return id.length > length ? `${id.substring(0, length)}…` : id; +}; + +const copyToClipboard = (text: string) => { + navigator.clipboard.writeText(text); +}; + +const A2AMetrics: React.FC = ({ a2aMetadata, timeToFirstToken, totalLatency }) => { + const [showDetails, setShowDetails] = useState(false); + + if (!a2aMetadata && !timeToFirstToken && !totalLatency) return null; + + const { taskId, contextId, status, metadata } = a2aMetadata || {}; + const formattedTime = formatTimestamp(status?.timestamp); + + return ( +
    + {/* A2A Metadata Header */} +
    + + A2A Metadata +
    + + {/* Main metrics row */} +
    + {/* Status badge */} + {status?.state && ( + + {getStatusIcon(status.state)} + {status.state} + + )} + + {/* Timestamp */} + {formattedTime && ( + + + + {formattedTime} + + + )} + + {/* Latency */} + {totalLatency !== undefined && ( + + + + {(totalLatency / 1000).toFixed(2)}s + + + )} + + {/* Time to first token */} + {timeToFirstToken !== undefined && ( + + + TTFT: {(timeToFirstToken / 1000).toFixed(2)}s + + + )} +
    + + {/* IDs row */} +
    + {/* Task ID */} + {taskId && ( + + copyToClipboard(taskId)} + > + + Task: {truncateId(taskId)} + + + + )} + + {/* Context/Session ID */} + {contextId && ( + + copyToClipboard(contextId)} + > + + Session: {truncateId(contextId)} + + + + )} + + {/* Details toggle */} + {(metadata || status?.message) && ( + + )} +
    + + {/* Expandable details panel */} + {showDetails && ( +
    + {/* Status message */} + {status?.message && ( +
    + Status Message: + {status.message} +
    + )} + + {/* Full IDs */} + {taskId && ( +
    + Task ID: + {taskId} + copyToClipboard(taskId)} + /> +
    + )} + + {contextId && ( +
    + Session ID: + {contextId} + copyToClipboard(contextId)} + /> +
    + )} + + {/* Metadata fields */} + {metadata && Object.keys(metadata).length > 0 && ( +
    + Custom Metadata: + 
    +                {JSON.stringify(metadata, null, 2)}
+
    +
    + )} +
    + )} +
    + ); +}; + +export default A2AMetrics; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx index 90c43df0dc4..18b5bb82dd4 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/ChatUI.tsx @@ -51,6 +51,10 @@ import { fetchAvailableModels, ModelGroup } from "../llm_calls/fetch_models"; import { makeOpenAIImageEditsRequest } from "../llm_calls/image_edits"; import { makeOpenAIImageGenerationRequest } from "../llm_calls/image_generation"; import { makeOpenAIResponsesRequest } from "../llm_calls/responses_api"; +import { Agent, fetchAvailableAgents } from "../llm_calls/fetch_agents"; +import { makeA2AStreamMessageRequest } from "../llm_calls/a2a_send_message"; +import A2AMetrics from "./A2AMetrics"; +import { A2ATaskMetadata } from "./types"; import MCPEventsDisplay, { MCPEvent } from "./MCPEventsDisplay"; import { EndpointType, getEndpointType } from "./mode_endpoint_mapping"; import ReasoningContent from "./ReasoningContent"; @@ -124,6 +128,8 @@ const ChatUI: React.FC = ({ const [selectedModel, setSelectedModel] = useState(undefined); const [showCustomModelInput, setShowCustomModelInput] = useState(false); const [modelInfo, setModelInfo] = useState([]); + const [agentInfo, setAgentInfo] = useState([]); + const [selectedAgent, setSelectedAgent] = useState(undefined); const customModelTimeout = useRef(null); const [endpointType, setEndpointType] = useState( () => sessionStorage.getItem("endpointType") || EndpointType.CHAT, @@ -340,6 +346,29 @@ const ChatUI: React.FC = ({ loadMCPTools(); }, [accessToken, userID, userRole, apiKeySource, apiKey, token]); + // Fetch agents when A2A endpoint is selected + useEffect(() => { + const userApiKey = apiKeySource === "session" ? accessToken : apiKey; + if (!userApiKey || endpointType !== EndpointType.A2A_AGENTS) { + return; + } + + const loadAgents = async () => { + try { + const agents = await fetchAvailableAgents(userApiKey); + setAgentInfo(agents); + // Clear selection if current agent not in list + if (selectedAgent && !agents.some((a) => a.agent_name === selectedAgent)) { + setSelectedAgent(undefined); + } + } catch (error) { + console.error("Error fetching agents:", error); + } + }; + + loadAgents(); + }, [accessToken, apiKeySource, apiKey, endpointType]); + useEffect(() => { // Scroll to the bottom of the chat whenever chatHistory updates if (chatEndRef.current) { @@ -469,6 +498,23 @@ const ChatUI: React.FC = ({ }); }; + const updateA2AMetadata = (a2aMetadata: A2ATaskMetadata) => { + console.log("Received A2A metadata:", a2aMetadata); + setChatHistory((prevHistory) => { + const lastMessage = prevHistory[prevHistory.length - 1]; + + if (lastMessage && lastMessage.role === "assistant") { + const updatedMessage = { + ...lastMessage, + a2aMetadata, + }; + return [...prevHistory.slice(0, prevHistory.length - 1), updatedMessage]; + } + + return prevHistory; + }); + }; + const updateTotalLatency = (totalLatency: number) => { setChatHistory((prevHistory) => { const lastMessage = prevHistory[prevHistory.length - 1]; @@ -684,6 +730,12 @@ const ChatUI: React.FC = ({ return; } + // For A2A agents, require agent selection + if (endpointType === EndpointType.A2A_AGENTS && !selectedAgent) { + NotificationsManager.fromBackend("Please select an agent to send a message"); + return; + } + if (!token || !userRole || !userID) { return; } @@ -691,7 +743,7 @@ const ChatUI: React.FC = ({ const effectiveApiKey = apiKeySource === "session" ? accessToken : apiKey; if (!effectiveApiKey) { - NotificationsManager.fromBackend("Please provide an API key or select Current UI Session"); + NotificationsManager.fromBackend("Please provide a Virtual Key or select Current UI Session"); return; } @@ -908,6 +960,20 @@ const ChatUI: React.FC = ({ } } } + + // Handle A2A agent calls (separate from model-based calls) - use streaming + if (endpointType === EndpointType.A2A_AGENTS && selectedAgent) { + await makeA2AStreamMessageRequest( + selectedAgent, + inputMessage, + (chunk, model) => updateTextUI("assistant", chunk, model), + effectiveApiKey, + signal, + updateTimingData, + updateTotalLatency, + updateA2AMetadata, + ); + } } catch (error) { if (signal.aborted) { console.log("Request was cancelled"); @@ -1003,7 +1069,7 @@ const ChatUI: React.FC = ({
    - API Key Source + Virtual Key Source { - if (!option.mode) { - //If no mode, show all models - return true; - } - const optionEndpoint = getEndpointType(option.mode); - // Show chat models for responses/anthropic_messages endpoints as they are compatible - if ( - endpointType === EndpointType.RESPONSES || - endpointType === EndpointType.ANTHROPIC_MESSAGES - ) { - return optionEndpoint === endpointType || optionEndpoint === EndpointType.CHAT; - } - // Show image models for image_edits endpoint as they are compatible - if (endpointType === EndpointType.IMAGE_EDITS) { - return optionEndpoint === endpointType || optionEndpoint === EndpointType.IMAGE; - } - return optionEndpoint === endpointType; - }) - .map((option) => option.model_group), - ), - ).map((model_group, index) => ({ - value: model_group, - label: model_group, - key: index, - })), - { value: "custom", label: "Enter custom model", key: "custom" }, - ]} - style={{ width: "100%" }} - showSearch={true} - className="rounded-md" - /> - {showCustomModelInput && ( - { - // Using setTimeout to create a simple debounce effect - if (customModelTimeout.current) { - clearTimeout(customModelTimeout.current); - } - - customModelTimeout.current = setTimeout(() => { - setSelectedModel(value); - }, 500); // 500ms delay after typing stops - }} + + ) : ( + +
    + )}
    @@ -1440,7 +1550,8 @@ const ChatUI: React.FC = ({ )} {message.role === "assistant" && - (message.timeToFirstToken || message.totalLatency || message.usage) && ( + (message.timeToFirstToken || message.totalLatency || message.usage) && + !message.a2aMetadata && ( = ({ toolName={message.toolName} /> )} + + {/* A2A Metrics - show for A2A agent responses */} + {message.role === "assistant" && message.a2aMetadata && ( + + )}
    @@ -1685,13 +1805,15 @@ const ChatUI: React.FC = ({ endpointType === EndpointType.RESPONSES || endpointType === EndpointType.ANTHROPIC_MESSAGES ? "Type your message... (Shift+Enter for new line)" - : endpointType === EndpointType.IMAGE_EDITS - ? "Describe how you want to edit the image..." - : endpointType === EndpointType.SPEECH - ? "Enter text to convert to speech..." - : endpointType === EndpointType.TRANSCRIPTION - ? "Optional: Add context or prompt for transcription..." - : "Describe the image you want to generate..." + : endpointType === EndpointType.A2A_AGENTS + ? "Send a message to the A2A agent..." + : endpointType === EndpointType.IMAGE_EDITS + ? "Describe how you want to edit the image..." + : endpointType === EndpointType.SPEECH + ? "Enter text to convert to speech..." + : endpointType === EndpointType.TRANSCRIPTION + ? "Optional: Add context or prompt for transcription..." + : "Describe the image you want to generate..." } disabled={isLoading} className="flex-1" diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/chatConstants.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/chatConstants.ts index a64c0264f4c..e9b44836455 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/chatConstants.ts +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/chatConstants.ts @@ -42,4 +42,5 @@ export const ENDPOINT_OPTIONS = [ { value: EndpointType.EMBEDDINGS, label: "/v1/embeddings" }, { value: EndpointType.SPEECH, label: "/v1/audio/speech" }, { value: EndpointType.TRANSCRIPTION, label: "/v1/audio/transcriptions" }, + { value: EndpointType.A2A_AGENTS, label: "/v1/a2a/message/send" }, ]; diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/mode_endpoint_mapping.tsx b/ui/litellm-dashboard/src/components/playground/chat_ui/mode_endpoint_mapping.tsx index 479d3cb8940..8ad0ae04a6d 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/mode_endpoint_mapping.tsx +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/mode_endpoint_mapping.tsx @@ -25,6 +25,7 @@ export enum EndpointType { EMBEDDINGS = "embeddings", SPEECH = "speech", TRANSCRIPTION = "transcription", + A2A_AGENTS = "a2a_agents", // add additional endpoint types if required } diff --git a/ui/litellm-dashboard/src/components/playground/chat_ui/types.ts b/ui/litellm-dashboard/src/components/playground/chat_ui/types.ts index 15c90221fc0..4e43e6ec77d 100644 --- a/ui/litellm-dashboard/src/components/playground/chat_ui/types.ts +++ b/ui/litellm-dashboard/src/components/playground/chat_ui/types.ts @@ -70,12 +70,24 @@ export interface VectorStoreSearchResponse { data: VectorStoreSearchResult[]; } +export interface A2ATaskMetadata { + taskId?: string; + contextId?: string; + status?: { + state?: string; + timestamp?: string; + message?: string; + }; + metadata?: Record; +} + export interface MessageType { role: string; content: string | MultimodalContent[]; model?: string; isImage?: boolean; isAudio?: boolean; + isEmbeddings?: boolean; reasoningContent?: string; timeToFirstToken?: number; totalLatency?: number; @@ -93,6 +105,7 @@ export interface MessageType { detail: string; }; searchResults?: VectorStoreSearchResponse[]; + a2aMetadata?: A2ATaskMetadata; } export interface MultimodalContent { diff --git a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx b/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx index 24b4008a1bd..ef4d65d4095 100644 --- a/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx +++ b/ui/litellm-dashboard/src/components/playground/compareUI/CompareUI.tsx @@ -395,7 +395,7 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }: return; } if (!effectiveApiKey) { - NotificationsManager.fromBackend("Please provide an API key or select Current UI Session"); + NotificationsManager.fromBackend("Please provide a Virtual Key or select Current UI Session"); return; } const targetComparisons = comparisons; @@ -551,7 +551,7 @@ export default function CompareUI({ accessToken, disabledPersonalKeyCreation }:
    - API Key Source + Virtual Key Source - {(is_proxy_admin || userModels.includes("all-proxy-models")) && ( - - All Proxy Models - - )} - - No Default Models - - {Array.from(new Set(userModels)).map((model, idx) => ( + {(() => { + let shouldShowAllProxyModels = false; + + if (organization) { + // Team is in an organization + if (organization.models.length === 0 || organization.models.includes("all-proxy-models")) { + // Organization has empty array [] or "all-proxy-models" + shouldShowAllProxyModels = true; + } + // Otherwise (organization has specific models), don't show "all-proxy-models" + } else { + // Team is not in an organization + shouldShowAllProxyModels = is_proxy_admin || userModels.includes("all-proxy-models"); + } + + return shouldShowAllProxyModels ? ( + + All Proxy Models + + ) : null; + })()} + {(() => { + // Show "no-default-models" option if: + // 1. Team is not in an organization, OR + // 2. Team is in an organization and organization's models include "no-default-models" + const shouldShowNoDefaultModels = + !organization || organization.models.includes("no-default-models"); + + return shouldShowNoDefaultModels ? ( + + No Default Models + + ) : null; + })()} + {Array.from(new Set(modelsToPick)).map((model, idx) => ( {getModelDisplayName(model)} @@ -758,7 +835,7 @@ const TeamInfoView: React.FC = ({ - + diff --git a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index f58e392a58f..9897bb4d47a 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -245,7 +245,7 @@ import KeyInfoView from "./key_info_view"; const baseKeyData = { token_id: "tok_123", token: "tok_123", - key_alias: "My API Key", + key_alias: "My Virtual Key", key_name: "sk-xxxx", created_at: new Date().toISOString(), updated_at: new Date().toISOString(), diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index ec2b294d9de..dbbb195a1f8 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -297,7 +297,7 @@ export default function KeyInfoView({ - {currentKeyData.key_alias || "API Key"} + {currentKeyData.key_alias || "Virtual Key"}
    @@ -381,7 +381,7 @@ export default function KeyInfoView({ {/* Delete Confirmation Modal */} {isDeleteModalOpen && (() => { - const keyName = currentKeyData?.key_alias || currentKeyData?.token_id || "API Key"; + const keyName = currentKeyData?.key_alias || currentKeyData?.token_id || "Virtual Key"; const isValid = deleteConfirmInput === keyName; return (
    @@ -415,7 +415,7 @@ export default function KeyInfoView({

    - Warning: You are about to delete this API key. + Warning: You are about to delete this Virtual Key.

    This action is irreversible and will immediately revoke access for any applications using this @@ -423,7 +423,7 @@ export default function KeyInfoView({

    -

    Are you sure you want to delete this API key?

    +

    Are you sure you want to delete this Virtual Key?

    - Warning: You are about to delete this API key. + Warning: You are about to delete this Virtual Key.

    This action is irreversible and will immediately revoke access for any applications using this @@ -374,7 +374,7 @@ const ViewKeyTable: React.FC = ({

    -

    Are you sure you want to delete this API key?

    +

    Are you sure you want to delete this Virtual Key?

    @@ -417,7 +417,7 @@ const ViewKeyTable: React.FC = ({ {/* Regenerate Key Form Modal */} { setRegenerateDialogVisible(false); @@ -516,7 +516,7 @@ const ViewKeyTable: React.FC = ({ {selectedToken?.key_alias || "No alias set"}
    - New API Key: + New Virtual Key:
    = ({
    NotificationManager.success({ description: "API Key copied to clipboard" })} + onCopy={() => NotificationManager.success({ description: "Virtual Key copied to clipboard" })} > - + diff --git a/ui/litellm-dashboard/src/components/usage.tsx b/ui/litellm-dashboard/src/components/usage.tsx index 88b1e3c3fb7..0900a0a9cc1 100644 --- a/ui/litellm-dashboard/src/components/usage.tsx +++ b/ui/litellm-dashboard/src/components/usage.tsx @@ -615,7 +615,7 @@ const UsagePage: React.FC = ({ accessToken, token, userRole, use - Top API Keys + Top Virtual Keys {possibleUIRoles?.[row.original.user_role]?.ui_label || "-"}, }, + { + header: "User Alias", + accessorKey: "user_alias", + enableSorting: false, + cell: ({ row }) => {row.original.user_alias || "-"}, + }, { header: "Spend (USD)", accessorKey: "spend", @@ -78,14 +84,14 @@ export const columns = ( ), }, { - header: "API Keys", + header: "Virtual Keys", accessorKey: "key_count", enableSorting: false, cell: ({ row }) => ( {row.original.key_count > 0 ? ( - {row.original.key_count} Keys + {row.original.key_count} {row.original.key_count === 1 ? "Key" : "Keys"} ) : ( diff --git a/ui/litellm-dashboard/src/components/view_users/table.test.tsx b/ui/litellm-dashboard/src/components/view_users/table.test.tsx index 8ef887932cc..82f49d9618e 100644 --- a/ui/litellm-dashboard/src/components/view_users/table.test.tsx +++ b/ui/litellm-dashboard/src/components/view_users/table.test.tsx @@ -1,63 +1,52 @@ import { act, fireEvent, render, screen } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; - import { UserDataTable } from "./table"; +const defaultFilters = { + email: "", + user_id: "", + user_role: "", + sso_user_id: "", + team: "", + model: "", + min_spend: null, + max_spend: null, + sort_by: "", + sort_order: "asc" as const, +}; + +const getDefaultProps = () => ({ + data: [] as any[], + columns: [] as any[], + accessToken: null, + userRole: "Admin", + possibleUIRoles: null as Record> | null, + filters: defaultFilters, + updateFilters: vi.fn(), + initialFilters: defaultFilters, + teams: [] as any[], + handleEdit: vi.fn(), + handleDelete: vi.fn(), + handleResetPassword: vi.fn(), + userListResponse: { users: [], total: 0, page: 1, page_size: 25, total_pages: 1 }, + currentPage: 1, + handlePageChange: vi.fn(), +}); + describe("UserDataTable", () => { it("should render the UserDataTable component", () => { - const filters = { - email: "", - user_id: "", - user_role: "", - sso_user_id: "", - team: "", - model: "", - min_spend: null, - max_spend: null, - sort_by: "", - sort_order: "asc" as const, - }; - - const updateFilters = vi.fn(); - - render( - , - ); + render(); expect(screen.getByText("Filters")).toBeInTheDocument(); }); it("should call onSortChange when clicking a sortable header", () => { const filters = { - email: "", - user_id: "", - user_role: "", - sso_user_id: "", - team: "", - model: "", - min_spend: null, - max_spend: null, + ...defaultFilters, sort_by: "created_at", sort_order: "desc" as const, }; - const updateFilters = vi.fn(); const onSortChange = vi.fn(); const possibleUIRoles = { @@ -67,21 +56,10 @@ describe("UserDataTable", () => { render( , @@ -96,41 +74,7 @@ describe("UserDataTable", () => { }); it("should show skeleton loaders when isLoading is true", () => { - const filters = { - email: "", - user_id: "", - user_role: "", - sso_user_id: "", - team: "", - model: "", - min_spend: null, - max_spend: null, - sort_by: "", - sort_order: "asc" as const, - }; - - const updateFilters = vi.fn(); - - render( - , - ); + render(); expect(screen.queryByText(/Showing/i)).not.toBeInTheDocument(); expect(screen.queryByRole("button", { name: /Previous/i })).not.toBeInTheDocument(); @@ -138,44 +82,35 @@ describe("UserDataTable", () => { }); it("should show actual content when isLoading is false", () => { - const filters = { - email: "", - user_id: "", - user_role: "", - sso_user_id: "", - team: "", - model: "", - min_spend: null, - max_spend: null, - sort_by: "", - sort_order: "asc" as const, - }; - - const updateFilters = vi.fn(); - - render( - , - ); + render(); expect(screen.getByText(/Showing/i)).toBeInTheDocument(); expect(screen.getByRole("button", { name: /Previous/i })).toBeInTheDocument(); expect(screen.getByRole("button", { name: /Next/i })).toBeInTheDocument(); }); + + it("should render all column headers", () => { + const possibleUIRoles = { + admin: { ui_label: "Admin" }, + user: { ui_label: "User" }, + }; + + render(); + + [ + "User ID", + "Email", + "Global Proxy Role", + "User Alias", + "Spend (USD)", + "Budget (USD)", + "SSO ID", + "Virtual Keys", + "Created At", + "Updated At", + "Actions", + ].forEach((header) => { + expect(screen.getByRole("columnheader", { name: header })).toBeInTheDocument(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/view_users/types.ts b/ui/litellm-dashboard/src/components/view_users/types.ts index d976d46ebc1..d674db5c7db 100644 --- a/ui/litellm-dashboard/src/components/view_users/types.ts +++ b/ui/litellm-dashboard/src/components/view_users/types.ts @@ -1,6 +1,7 @@ export interface UserInfo { user_id: string; user_email: string; + user_alias: string | null; user_role: string; spend: number; max_budget: number | null; diff --git a/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx index 456f07d1882..2caae7d861f 100644 --- a/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/components/view_users/user_info_view.tsx @@ -320,9 +320,11 @@ export default function UserInfoView({ - API Keys + Virtual Keys
    - {userData.keys?.length || 0} keys + + {userData.keys?.length || 0} {userData.keys?.length === 1 ? "Key" : "Keys"} +
    @@ -467,7 +469,7 @@ export default function UserInfoView({
    - API Keys + Virtual Keys
    {userData.keys?.length && userData.keys?.length > 0 ? ( userData.keys.map((key, index) => ( @@ -476,7 +478,7 @@ export default function UserInfoView({ )) ) : ( - No API keys + No Virtual Keys )}
    diff --git a/ui/litellm-dashboard/src/utils/jwtUtils.test.ts b/ui/litellm-dashboard/src/utils/jwtUtils.test.ts new file mode 100644 index 00000000000..d695a3c37bb --- /dev/null +++ b/ui/litellm-dashboard/src/utils/jwtUtils.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; +import { isJwtExpired } from "./jwtUtils"; +import { jwtDecode } from "jwt-decode"; + +vi.mock("jwt-decode"); + +describe("jwtUtils", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("should return true if the token is expired", () => { + const mockDateNow = 1716838401000; + vi.spyOn(Date, "now").mockReturnValue(mockDateNow); + vi.mocked(jwtDecode).mockReturnValue({ + exp: Math.floor(mockDateNow / 1000) - 1, + user_id: "test", + }); + + expect(isJwtExpired("any-token")).toBe(true); + }); + + it("should return false if the token is not expired", () => { + const mockDateNow = 1716838401000; + vi.spyOn(Date, "now").mockReturnValue(mockDateNow); + vi.mocked(jwtDecode).mockReturnValue({ + exp: Math.floor(mockDateNow / 1000) + 1000, + user_id: "test", + }); + + expect(isJwtExpired("any-token")).toBe(false); + }); + + it("should return false if the token does not have an exp field", () => { + vi.mocked(jwtDecode).mockReturnValue({ + user_id: "test", + }); + + expect(isJwtExpired("any-token")).toBe(false); + }); + + it("should return true if jwtDecode throws an error", () => { + vi.mocked(jwtDecode).mockImplementation(() => { + throw new Error("Invalid token"); + }); + + expect(isJwtExpired("invalid-token")).toBe(true); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/jwtUtils.ts b/ui/litellm-dashboard/src/utils/jwtUtils.ts new file mode 100644 index 00000000000..d3db41a411a --- /dev/null +++ b/ui/litellm-dashboard/src/utils/jwtUtils.ts @@ -0,0 +1,14 @@ +import { jwtDecode } from "jwt-decode"; + +export function isJwtExpired(token: string): boolean { + try { + const decoded: any = jwtDecode(token); + if (decoded && typeof decoded.exp === "number") { + return decoded.exp * 1000 <= Date.now(); + } + return false; + } catch { + // If we can't decode, treat as invalid/expired + return true; + } +} diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index 9287e203182..c3c5ae59237 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -193,7 +193,7 @@ describe("CreateKeyPage auth behavior", () => { // Assert: we eventually redirect to SSO login (single replace, not assign/href) await waitFor(() => { - expect(window.location.replace).toHaveBeenCalledWith("https://example.com/sso/key/generate"); + expect(window.location.replace).toHaveBeenCalledWith("https://example.com/ui/login"); }); // And we attempted to clear the cookie (defensive deletion)