add precommit

vllm_omni/entrypoints/api_server.py

@@ -3,17 +3,18 @@ API server for vLLM-omni.
 """
 
 import asyncio
-from typing import Dict, Any, List
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
+from typing import Any, Dict, List
+
 import uvicorn
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 
 from .omni_llm import AsyncOmniLLM
 
 
 class GenerateRequest(BaseModel):
     """Request model for generation."""
+
     prompts: List[str]
     max_tokens: int = 100
     temperature: float = 0.7
@@ -22,6 +23,7 @@ class GenerateRequest(BaseModel):
 
 class GenerateResponse(BaseModel):
     """Response model for generation."""
+
     outputs: List[Dict[str, Any]]
     stage_outputs: List[Dict[str, Any]] = None
 
@@ -29,7 +31,7 @@ class GenerateResponse(BaseModel):
 app = FastAPI(
     title="vLLM-omni API",
     description="Multi-modality models inference and serving",
-    version="0.1.0"
+    version="0.1.0",
 )
 
 # Global omni_llm instance
@@ -50,48 +52,51 @@ async def generate(request: GenerateRequest):
     try:
         if omni_llm is None:
             raise HTTPException(status_code=500, detail="OmniLLM not initialized")
-        
+
         # Prepare stage arguments
         if request.stage_args is None:
             # Create default stage arguments - one per stage config
             # For now, we'll process all prompts in the first stage
-            stage_args = [{
-                "prompt": " ".join(request.prompts) if request.prompts else "",
-                "max_tokens": request.max_tokens,
-                "temperature": request.temperature
-            }]
+            stage_args = [
+                {
+                    "prompt": " ".join(request.prompts) if request.prompts else "",
+                    "max_tokens": request.max_tokens,
+                    "temperature": request.temperature,
+                }
+            ]
         else:
             stage_args = request.stage_args
-        
+
         # Generate using omni_llm
         try:
             outputs = await omni_llm.generate_async(stage_args)
-        except Exception as e:
+        except Exception:
             # Fallback to synchronous generation
             outputs = omni_llm.generate(stage_args)
-        
+
         # Convert outputs to response format
         response_outputs = []
         for output in outputs:
-            if hasattr(output, 'outputs') and output.outputs:
+            if hasattr(output, "outputs") and output.outputs:
                 for out in output.outputs:
-                    response_outputs.append({
-                        "text": getattr(out, 'text', ''),
-                        "finished": getattr(out, 'finish_reason', 'length') != 'length',
-                        "tokens": getattr(out, 'token_ids', [])
-                    })
+                    response_outputs.append(
+                        {
+                            "text": getattr(out, "text", ""),
+                            "finished": getattr(out, "finish_reason", "length")
+                            != "length",
+                            "tokens": getattr(out, "token_ids", []),
+                        }
+                    )
             else:
-                response_outputs.append({
-                    "text": "",
-                    "finished": True,
-                    "tokens": []
-                })
-        
+                response_outputs.append({"text": "", "finished": True, "tokens": []})
+
         return GenerateResponse(
             outputs=response_outputs,
-            stage_outputs=[{"stage": i, "output": "processed"} for i in range(len(stage_args))]
+            stage_outputs=[
+                {"stage": i, "output": "processed"} for i in range(len(stage_args))
+            ],
         )
-    
+
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -107,36 +112,35 @@ async def get_info():
     """Get information about the service."""
     if omni_llm is None:
         return {"error": "OmniLLM not initialized"}
-    
+
     return {
         "service": "vllm-omni",
         "version": "0.1.0",
-        "num_stages": omni_llm.get_num_stages() if hasattr(omni_llm, 'get_num_stages') else 0,
+        "num_stages": omni_llm.get_num_stages()
+        if hasattr(omni_llm, "get_num_stages")
+        else 0,
         "stage_configs": [
             {
                 "stage_id": config.stage_id,
                 "engine_type": config.engine_type,
                 "model_path": config.model_path,
                 "input_modalities": config.input_modalities,
-                "output_modalities": config.output_modalities
+                "output_modalities": config.output_modalities,
             }
             for config in omni_llm.stage_configs
-        ]
+        ],
     }
 
 
-async def run_server(omni_llm_instance: AsyncOmniLLM, host: str = "0.0.0.0", port: int = 8000):
+async def run_server(
+    omni_llm_instance: AsyncOmniLLM, host: str = "0.0.0.0", port: int = 8000
+):
     """Run the API server."""
     global omni_llm
     omni_llm = omni_llm_instance
-    
-    config = uvicorn.Config(
-        app=app,
-        host=host,
-        port=port,
-        log_level="info"
-    )
-    
+
+    config = uvicorn.Config(app=app, host=host, port=port, log_level="info")
+
     server = uvicorn.Server(config)
     await server.serve()