feat: receive both request and body

Signed-off-by: TomeHirata <tomu.hirata@gmail.com>

examples/deployments/deployments_server/openai/config.yaml

@@ -8,7 +8,7 @@ endpoints:
         openai_api_key: $OPENAI_API_KEY
     limit:
       renewal_period: "minute"
-      calls: 1
+      calls: 10
 
   - name: completions
     endpoint_type: llm/v1/completions

mlflow/deployments/server/app.py

@@ -91,9 +91,8 @@ def _create_chat_endpoint(config: RouteConfig):
 
     # https://slowapi.readthedocs.io/en/latest/#limitations-and-known-issues
     async def _chat(
-        request: Request,
+        request: Request, payload: chat.RequestPayload
     ) -> Union[chat.ResponsePayload, chat.StreamResponsePayload]:
-        payload = await parse_request_schema(request, chat.RequestPayload)
         if payload.stream:
             return await make_streaming_response(prov.chat_stream(payload))
         else:
@@ -106,9 +105,8 @@ def _create_completions_endpoint(config: RouteConfig):
     prov = get_provider(config.model.provider)(config)
 
     async def _completions(
-        request: Request,
+        request: Request, payload: completions.RequestPayload
     ) -> Union[completions.ResponsePayload, completions.StreamResponsePayload]:
-        payload = await parse_request_schema(request, completions.RequestPayload)
         if payload.stream:
             return await make_streaming_response(prov.completions_stream(payload))
         else:
@@ -120,8 +118,9 @@ async def _completions(
 def _create_embeddings_endpoint(config: RouteConfig):
     prov = get_provider(config.model.provider)(config)
 
-    async def _embeddings(request: Request) -> embeddings.ResponsePayload:
-        payload = await parse_request_schema(request, embeddings.RequestPayload)
+    async def _embeddings(
+        request: Request, payload: embeddings.RequestPayload
+    ) -> embeddings.ResponsePayload:
         return await prov.embeddings(payload)
 
     return _embeddings

mlflow/gateway/config.py

@@ -8,7 +8,6 @@
 
 import pydantic
 import yaml
-from limits import parse
 from packaging import version
 from packaging.version import Version
 from pydantic import ConfigDict, Field, ValidationError, root_validator, validator
@@ -401,6 +400,8 @@ def validate_route_type(cls, value):
 
     @validator("limit", pre=True)
     def validate_limit(cls, value):
+        from limits import parse
+
         if value:
             limit = Limit(**value)
             try:
@@ -452,7 +453,7 @@ class Route(ConfigModel):
     route_type: str
     model: RouteModelInfo
     route_url: str
-    limit: Optional[Limit]
+    limit: Optional[Limit] = None
 
     class Config:
         if IS_PYDANTIC_V2: