Merge branch 'main' into litellm_add_stability.stable-image-ultra-v1

deploy/charts/litellm-helm/templates/migrations-job.yaml

@@ -13,18 +13,18 @@ spec:
     spec:
       containers:
         - name: prisma-migrations
-          image: "ghcr.io/berriai/litellm:main-stable"
+          image: ghcr.io/berriai/litellm-database:main-latest
           command: ["python", "litellm/proxy/prisma_migration.py"]
           workingDir: "/app"
           env:
-            {{- if .Values.db.deployStandalone }}
+            {{- if .Values.db.useExisting }}
-            - name: DATABASE_URL
-              value: postgresql://{{ .Values.postgresql.auth.username }}:{{ .Values.postgresql.auth.password }}@{{ .Release.Name }}-postgresql/{{ .Values.postgresql.auth.database }}
-            {{- else if .Values.db.useExisting }}
             - name: DATABASE_URL
               value: {{ .Values.db.url | quote }}
+            {{- else }}
+            - name: DATABASE_URL
+              value: postgresql://{{ .Values.postgresql.auth.username }}:{{ .Values.postgresql.auth.password }}@{{ .Release.Name }}-postgresql/{{ .Values.postgresql.auth.database }}
             {{- end }}
             - name: DISABLE_SCHEMA_UPDATE
-              value: "{{ .Values.migrationJob.disableSchemaUpdate }}"
+              value: "false" # always run the migration from the Helm PreSync hook, override the value set
       restartPolicy: OnFailure
   backoffLimit: {{ .Values.migrationJob.backoffLimit }}

docs/my-website/docs/proxy/configs.md

@@ -791,9 +791,9 @@ general_settings:
 | store_model_in_db | boolean | If true, allows `/model/new` endpoint to store model information in db. Endpoint disabled by default. [Doc on `/model/new` endpoint](./model_management.md#create-a-new-model) |
 | max_request_size_mb | int | The maximum size for requests in MB. Requests above this size will be rejected. |
 | 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. |
+| 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. |
+| 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. |
+| proxy_batch_write_at | int | Time (in seconds) to wait before batch writing spend logs to the db. **Default is 10 seconds** |
 | 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) |
 | allowed_ips | List[str] | List of IPs allowed to access the proxy. If not set, all IPs are allowed. |

docs/my-website/docs/proxy/prod.md

@@ -1,5 +1,6 @@
 import Tabs from '@theme/Tabs';
 import TabItem from '@theme/TabItem';
+import Image from '@theme/IdealImage';
 
 # ⚡ Best Practices for Production
 
@@ -112,7 +113,35 @@ general_settings:
   disable_spend_logs: True
 ```
 
-## 7. Set LiteLLM Salt Key 
+## 7. Use Helm PreSync Hook for Database Migrations
+
+To ensure only one service manages database migrations, use our [Helm PreSync hook for Database Migrations](https://github.com/BerriAI/litellm/blob/main/deploy/charts/litellm-helm/templates/migrations-job.yaml). This ensures migrations are handled during `helm upgrade` or `helm install`, while LiteLLM pods explicitly disable migrations.
+
+
+1. **Helm PreSync Hook**:
+   - The Helm PreSync hook is configured in the chart to run database migrations during deployments.
+   - The hook always sets `DISABLE_SCHEMA_UPDATE=false`, ensuring migrations are executed reliably.
+  
+  Reference Settings to set on ArgoCD for `values.yaml`
+
+  ```yaml
+  db:
+    useExisting: true # use existing Postgres DB
+    url: postgresql://ishaanjaffer0324:3rnwpOBau6hT@ep-withered-mud-a5dkdpke.us-east-2.aws.neon.tech/test-argo-cd?sslmode=require # url of existing Postgres DB
+  ```
+
+2. **LiteLLM Pods**:
+   - Set `DISABLE_SCHEMA_UPDATE=true` in LiteLLM pod configurations to prevent them from running migrations.
+   
+   Example configuration for LiteLLM pod:
+   ```yaml
+   env:
+     - name: DISABLE_SCHEMA_UPDATE
+       value: "true"
+   ```
+
+
+## 8. Set LiteLLM Salt Key 
 
 If you plan on using the DB, set a salt key for encrypting/decrypting variables in the DB. 
 

litellm/proxy/_new_secret_config.yaml

@@ -68,7 +68,7 @@ model_list:
 
 litellm_settings:
   fallbacks: [{ "claude-3-5-sonnet-20240620": ["claude-3-5-sonnet-aihubmix"] }]
-  callbacks: ["otel", "prometheus"]
+  # callbacks: ["otel", "prometheus"]
   default_redis_batch_cache_expiry: 10
   # default_team_settings: 
   #   - team_id: "dbe2f686-a686-4896-864a-4c3924458709"

litellm/proxy/management_endpoints/key_management_endpoints.py

@@ -303,21 +303,17 @@ async def generate_key_fn(  # noqa: PLR0915
         )
 
 
-async def prepare_key_update_data(
+def prepare_key_update_data(
     data: Union[UpdateKeyRequest, RegenerateKeyRequest], existing_key_row
 ):
-    data_json: dict = data.dict(exclude_unset=True)
+    data_json: dict = data.model_dump(exclude_unset=True)
     data_json.pop("key", None)
     _metadata_fields = ["model_rpm_limit", "model_tpm_limit", "guardrails"]
     non_default_values = {}
     for k, v in data_json.items():
         if k in _metadata_fields:
             continue
-        if v is not None:
+        non_default_values[k] = v
-            if not isinstance(v, bool) and v in ([], {}, 0):
-                pass
-            else:
-                non_default_values[k] = v
 
     if "duration" in non_default_values:
         duration = non_default_values.pop("duration")
@@ -379,7 +375,7 @@ async def update_key_fn(
     )
 
     try:
-        data_json: dict = data.json()
+        data_json: dict = data.model_dump(exclude_unset=True)
         key = data_json.pop("key")
         # get the row from db
         if prisma_client is None:
@@ -395,7 +391,7 @@ async def update_key_fn(
                 detail={"error": f"Team not found, passed team_id={data.team_id}"},
             )
 
-        non_default_values = await prepare_key_update_data(
+        non_default_values = prepare_key_update_data(
             data=data, existing_key_row=existing_key_row
         )
 
@@ -1144,7 +1140,7 @@ async def regenerate_key_fn(
     non_default_values = {}
     if data is not None:
         # Update with any provided parameters from GenerateKeyRequest
-        non_default_values = await prepare_key_update_data(
+        non_default_values = prepare_key_update_data(
             data=data, existing_key_row=_key_in_db
         )
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "1.52.6"
+version = "1.52.7"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT"
@@ -91,7 +91,7 @@ requires = ["poetry-core", "wheel"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "1.52.6"
+version = "1.52.7"
 version_files = [
     "pyproject.toml:^version"
 ]

tests/local_testing/test_router_fallbacks.py

@@ -1138,9 +1138,9 @@ async def test_router_content_policy_fallbacks(
     router = Router(
         model_list=[
             {
-                "model_name": "claude-2",
+                "model_name": "claude-2.1",
                 "litellm_params": {
-                    "model": "claude-2",
+                    "model": "claude-2.1",
                     "api_key": "",
                     "mock_response": mock_response,
                 },
@@ -1164,7 +1164,7 @@ async def test_router_content_policy_fallbacks(
             {
                 "model_name": "my-general-model",
                 "litellm_params": {
-                    "model": "claude-2",
+                    "model": "claude-2.1",
                     "api_key": "",
                     "mock_response": Exception("Should not have called this."),
                 },
@@ -1172,14 +1172,14 @@ async def test_router_content_policy_fallbacks(
             {
                 "model_name": "my-context-window-model",
                 "litellm_params": {
-                    "model": "claude-2",
+                    "model": "claude-2.1",
                     "api_key": "",
                     "mock_response": Exception("Should not have called this."),
                 },
             },
         ],
         content_policy_fallbacks=(
-            [{"claude-2": ["my-fallback-model"]}]
+            [{"claude-2.1": ["my-fallback-model"]}]
             if fallback_type == "model-specific"
             else None
         ),
@@ -1190,12 +1190,12 @@ async def test_router_content_policy_fallbacks(
 
     if sync_mode is True:
         response = router.completion(
-            model="claude-2",
+            model="claude-2.1",
             messages=[{"role": "user", "content": "Hey, how's it going?"}],
         )
     else:
         response = await router.acompletion(
-            model="claude-2",
+            model="claude-2.1",
             messages=[{"role": "user", "content": "Hey, how's it going?"}],
         )
 

tests/local_testing/test_streaming.py

@@ -718,7 +718,7 @@ async def test_acompletion_claude_2_stream():
     try:
         litellm.set_verbose = True
         response = await litellm.acompletion(
-            model="claude-2",
+            model="claude-2.1",
             messages=[{"role": "user", "content": "hello from litellm"}],
             stream=True,
         )

tests/proxy_unit_tests/test_proxy_utils.py

@@ -510,3 +510,23 @@ async def test_proxy_config_update_from_db():
                 "success_callback": "langfuse",
             }
         }
+
+
+def test_prepare_key_update_data():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        prepare_key_update_data,
+    )
+    from litellm.proxy._types import UpdateKeyRequest
+
+    existing_key_row = MagicMock()
+    data = UpdateKeyRequest(key="test_key", models=["gpt-4"], duration="120s")
+    updated_data = prepare_key_update_data(data, existing_key_row)
+    assert "expires" in updated_data
+
+    data = UpdateKeyRequest(key="test_key", metadata={})
+    updated_data = prepare_key_update_data(data, existing_key_row)
+    assert updated_data["metadata"] == {}
+
+    data = UpdateKeyRequest(key="test_key", metadata=None)
+    updated_data = prepare_key_update_data(data, existing_key_row)
+    assert updated_data["metadata"] == None

tests/test_keys.py

@@ -66,6 +66,7 @@ async def generate_key(
     max_parallel_requests: Optional[int] = None,
     user_id: Optional[str] = None,
     team_id: Optional[str] = None,
+    metadata: Optional[dict] = None,
     calling_key="sk-1234",
 ):
     url = "http://0.0.0.0:4000/key/generate"
@@ -82,6 +83,7 @@ async def generate_key(
         "max_parallel_requests": max_parallel_requests,
         "user_id": user_id,
         "team_id": team_id,
+        "metadata": metadata,
     }
 
     print(f"data: {data}")
@@ -136,16 +138,21 @@ async def test_key_gen_bad_key():
             pass
 
 
-async def update_key(session, get_key):
+async def update_key(session, get_key, metadata: Optional[dict] = None):
     """
     Make sure only models user has access to are returned
     """
     url = "http://0.0.0.0:4000/key/update"
     headers = {
-        "Authorization": f"Bearer sk-1234",
+        "Authorization": "Bearer sk-1234",
         "Content-Type": "application/json",
     }
-    data = {"key": get_key, "models": ["gpt-4"], "duration": "120s"}
+    data = {"key": get_key}
+
+    if metadata is not None:
+        data["metadata"] = metadata
+    else:
+        data.update({"models": ["gpt-4"], "duration": "120s"})
 
     async with session.post(url, headers=headers, json=data) as response:
         status = response.status
@@ -276,20 +283,24 @@ async def chat_completion_streaming(session, key, model="gpt-4"):
     return prompt_tokens, completion_tokens
 
 
+@pytest.mark.parametrize("metadata", [{"test": "new"}, {}])
 @pytest.mark.asyncio
-async def test_key_update():
+async def test_key_update(metadata):
     """
     Create key
     Update key with new model
     Test key w/ model
     """
     async with aiohttp.ClientSession() as session:
-        key_gen = await generate_key(session=session, i=0)
+        key_gen = await generate_key(session=session, i=0, metadata={"test": "test"})
         key = key_gen["key"]
-        await update_key(
+        assert key_gen["metadata"]["test"] == "test"
+        updated_key = await update_key(
             session=session,
             get_key=key,
+            metadata=metadata,
         )
+        assert updated_key["metadata"] == metadata
         await update_proxy_budget(session=session)  # resets proxy spend
         await chat_completion(session=session, key=key)
 

tests/test_team_logging.py

@@ -62,8 +62,8 @@ async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata=
 
 
 @pytest.mark.asyncio
-@pytest.mark.flaky(retries=6, delay=1)
+@pytest.mark.flaky(retries=12, delay=2)
-async def test_team_logging():
+async def test_aaateam_logging():
     """
     -> Team 1 logs to project 1
     -> Create Key