Merge pull request #4043 from BerriAI/litellm_enforce_params

[Feat] Enterprise - Enforce Params in request to LiteLLM Proxy

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

@@ -14,6 +14,7 @@ Features:
 - ✅ [SSO for Admin UI](./ui.md#✨-enterprise-features)
 - ✅ [Audit Logs](#audit-logs)
 - ✅ [Tracking Spend for Custom Tags](#tracking-spend-for-custom-tags)
+- ✅ [Enforce Required Params for LLM Requests (ex. Reject requests missing ["metadata"]["generation_name"])](#enforce-required-params-for-llm-requests)
 - ✅ [Content Moderation with LLM Guard, LlamaGuard, Google Text Moderations](#content-moderation)
 - ✅ [Prompt Injection Detection (with LakeraAI API)](#prompt-injection-detection---lakeraai)
 - ✅ [Custom Branding + Routes on Swagger Docs](#swagger-docs---custom-routes--branding)
@@ -204,6 +205,109 @@ curl -X GET "http://0.0.0.0:4000/spend/tags" \
 ```
 
 
+## Enforce Required Params for LLM Requests
+Use this when you want to enforce all requests to include certain params. Example you need all requests to include the `user` and `["metadata]["generation_name"]` params.
+
+**Step 1** Define all Params you want to enforce on config.yaml
+
+This means `["user"]` and `["metadata]["generation_name"]` are required in all LLM Requests to LiteLLM
+
+```yaml
+general_settings:
+  master_key: sk-1234
+  enforced_params:  
+    - user
+    - metadata.generation_name
+```
+
+Start LiteLLM Proxy
+
+**Step 2 Verify if this works**
+
+<Tabs>
+
+<TabItem value="bad" label="Invalid Request (No `user` passed)">
+
+```shell
+curl --location 'http://localhost:4000/chat/completions' \
+    --header 'Authorization: Bearer sk-5fmYeaUEbAMpwBNT-QpxyA' \
+    --header 'Content-Type: application/json' \
+    --data '{
+    "model": "gpt-3.5-turbo",
+    "messages": [
+        {
+        "role": "user",
+        "content": "hi"
+        }
+    ]
+}'
+```
+
+Expected Response 
+
+```shell
+{"error":{"message":"Authentication Error, BadRequest please pass param=user in request body. This is a required param","type":"auth_error","param":"None","code":401}}% 
+```
+
+</TabItem>
+
+<TabItem value="bad2" label="Invalid Request (No `metadata` passed)">
+
+```shell
+curl --location 'http://localhost:4000/chat/completions' \
+    --header 'Authorization: Bearer sk-5fmYeaUEbAMpwBNT-QpxyA' \
+    --header 'Content-Type: application/json' \
+    --data '{
+    "model": "gpt-3.5-turbo",
+    "user": "gm",
+    "messages": [
+        {
+        "role": "user",
+        "content": "hi"
+        }
+    ],
+   "metadata": {}
+}'
+```
+
+Expected Response 
+
+```shell
+{"error":{"message":"Authentication Error, BadRequest please pass param=[metadata][generation_name] in request body. This is a required param","type":"auth_error","param":"None","code":401}}% 
+```
+
+
+</TabItem>
+<TabItem value="good" label="Valid Request">
+
+```shell
+curl --location 'http://localhost:4000/chat/completions' \
+    --header 'Authorization: Bearer sk-5fmYeaUEbAMpwBNT-QpxyA' \
+    --header 'Content-Type: application/json' \
+    --data '{
+    "model": "gpt-3.5-turbo",
+    "user": "gm",
+    "messages": [
+        {
+        "role": "user",
+        "content": "hi"
+        }
+    ],
+   "metadata": {"generation_name": "prod-app"}
+}'
+```
+
+Expected Response
+
+```shell
+{"id":"chatcmpl-9XALnHqkCBMBKrOx7Abg0hURHqYtY","choices":[{"finish_reason":"stop","index":0,"message":{"content":"Hello! How can I assist you today?","role":"assistant"}}],"created":1717691639,"model":"gpt-3.5-turbo-0125","object":"chat.completion","system_fingerprint":null,"usage":{"completion_tokens":9,"prompt_tokens":8,"total_tokens":17}}%  
+```
+
+</TabItem>
+</Tabs>
+
+
+
 
 
 

litellm/proxy/auth/auth_checks.py

@@ -97,6 +97,40 @@ def common_checks(
             raise Exception(
                 f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}"
             )
+    if general_settings.get("enforced_params", None) is not None:
+        # Enterprise ONLY Feature
+        # we already validate if user is premium_user when reading the config
+        # Add an extra premium_usercheck here too, just incase
+        from litellm.proxy.proxy_server import premium_user, CommonProxyErrors
+
+        if premium_user is not True:
+            raise ValueError(
+                "Trying to use `enforced_params`"
+                + CommonProxyErrors.not_premium_user.value
+            )
+
+        if route in LiteLLMRoutes.openai_routes.value:
+            # loop through each enforced param
+            # example enforced_params ['user', 'metadata', 'metadata.generation_name']
+            for enforced_param in general_settings["enforced_params"]:
+                _enforced_params = enforced_param.split(".")
+                if len(_enforced_params) == 1:
+                    if _enforced_params[0] not in request_body:
+                        raise ValueError(
+                            f"BadRequest please pass param={_enforced_params[0]} in request body. This is a required param"
+                        )
+                elif len(_enforced_params) == 2:
+                    # this is a scenario where user requires request['metadata']['generation_name'] to exist
+                    if _enforced_params[0] not in request_body:
+                        raise ValueError(
+                            f"BadRequest please pass param={_enforced_params[0]} in request body. This is a required param"
+                        )
+                    if _enforced_params[1] not in request_body[_enforced_params[0]]:
+                        raise ValueError(
+                            f"BadRequest please pass param=[{_enforced_params[0]}][{_enforced_params[1]}] in request body. This is a required param"
+                        )
+
+        pass
     # 7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
     if (
         litellm.max_budget > 0

litellm/proxy/proxy_config.yaml

@@ -21,7 +21,10 @@ model_list:
 
 general_settings:
   master_key: sk-1234
+  enforced_params:  
+    - user
+    - metadata
+    - metadata.generation_name
 
 litellm_settings:
-  callbacks: ["otel"]
   store_audit_logs: true

litellm/proxy/proxy_server.py

@@ -2878,6 +2878,16 @@ class ProxyConfig:
             )
             health_check_interval = general_settings.get("health_check_interval", 300)
 
+            ## check if user has set a premium feature in general_settings
+            if (
+                general_settings.get("enforced_params") is not None
+                and premium_user is not True
+            ):
+                raise ValueError(
+                    "Trying to use `enforced_params`"
+                    + CommonProxyErrors.not_premium_user.value
+                )
+
         router_params: dict = {
             "cache_responses": litellm.cache
             != None,  # cache if user passed in cache values

litellm/tests/test_key_generate_prisma.py

@@ -2248,3 +2248,55 @@ async def test_create_update_team(prisma_client):
     assert _team_info["budget_reset_at"] is not None and isinstance(
         _team_info["budget_reset_at"], datetime.datetime
     )
+
+
+@pytest.mark.asyncio()
+async def test_enforced_params(prisma_client):
+    setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
+    setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
+    from litellm.proxy.proxy_server import general_settings
+
+    general_settings["enforced_params"] = [
+        "user",
+        "metadata",
+        "metadata.generation_name",
+    ]
+
+    await litellm.proxy.proxy_server.prisma_client.connect()
+    request = NewUserRequest()
+    key = await new_user(request)
+    print(key)
+
+    generated_key = key.key
+    bearer_token = "Bearer " + generated_key
+
+    request = Request(scope={"type": "http"})
+    request._url = URL(url="/chat/completions")
+
+    # Case 1: Missing user
+    async def return_body():
+        return b'{"model": "gemini-pro-vision"}'
+
+    request.body = return_body
+    try:
+        await user_api_key_auth(request=request, api_key=bearer_token)
+        pytest.fail(f"This should have failed!. IT's an invalid request")
+    except Exception as e:
+        assert (
+            "BadRequest please pass param=user in request body. This is a required param"
+            in e.message
+        )
+
+    # Case 2: Missing metadata["generation_name"]
+    async def return_body_2():
+        return b'{"model": "gemini-pro-vision", "user": "1234", "metadata": {}}'
+
+    request.body = return_body_2
+    try:
+        await user_api_key_auth(request=request, api_key=bearer_token)
+        pytest.fail(f"This should have failed!. IT's an invalid request")
+    except Exception as e:
+        assert (
+            "Authentication Error, BadRequest please pass param=[metadata][generation_name] in request body"
+            in e.message
+        )