x

app/engine.py

@@ -39,7 +39,7 @@ class InferenceEngine:
         return GenerateResponse(
             text=completion,
             prompt=req.prompt,
-            model=self.settings.model_name,
+            model=self.settings.served_model_name or self.settings.model_name,
             usage_prompt_tokens=usage_prompt,
             usage_completion_tokens=usage_completion,
             usage_total_tokens=usage_prompt + usage_completion,

app/main.py

@@ -33,11 +33,16 @@ app = FastAPI(title="ROCm vLLM Inference API", version="1.0.0", lifespan=lifespa
 
 @app.get("/health", response_model=HealthResponse)
 def health(settings: Settings = Depends(get_settings)) -> HealthResponse:
-    return HealthResponse(status="ok", model=settings.model_name)
+    return HealthResponse(status="ok", model=settings.served_model_name or settings.model_name)
 
 
 @app.post("/v1/generate", response_model=GenerateResponse, dependencies=[Depends(verify_api_key)])
-def generate(req: GenerateRequest) -> GenerateResponse:
+def generate(req: GenerateRequest, settings: Settings = Depends(get_settings)) -> GenerateResponse:
     if engine is None:
         raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Engine not ready")
+    if req.max_tokens > settings.max_tokens:
+        raise HTTPException(
+            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
+            detail=f"max_tokens must be <= {settings.max_tokens}",
+        )
     return engine.generate(req)