API: Switch from request ID middleware to depends

Middleware runs on both the request and response. Therefore, streaming
responses had increased latency when processing tasks and sending
data to the client which resulted in erratic streaming behavior.

Use a depends to add request IDs since it only executes when the
request is run rather than expecting the response to be sent as well.

For the future, it would be best to think about limiting the time
between each tick of chunk data to be safe.

Signed-off-by: kingbri <bdashore3@proton.me>

common/networking.py

@@ -7,6 +7,7 @@ from fastapi import HTTPException, Request
 from loguru import logger
 from pydantic import BaseModel
 from typing import Optional
+from uuid import uuid4
 
 from common import config
 from common.utils import unwrap
@@ -100,3 +101,10 @@ def is_port_in_use(port: int) -> bool:
     test_socket.settimeout(1)
     with test_socket:
         return test_socket.connect_ex(("localhost", port)) == 0
+
+
+async def add_request_id(request: Request):
+    """FastAPI depends to add a UUID to a request's state."""
+
+    request.state.id = uuid4().hex
+    return request

endpoints/server.py

@@ -1,10 +1,10 @@
-from uuid import uuid4
 import uvicorn
-from fastapi import FastAPI, Request
+from fastapi import Depends, FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from loguru import logger
 
 from common.logger import UVICORN_LOG_CONFIG
+from common.networking import add_request_id
 from endpoints.OAI.router import router as OAIRouter
 
 app = FastAPI(
@@ -14,6 +14,7 @@ app = FastAPI(
         "This docs page is not meant to send requests! Please use a service "
         "like Postman or a frontend UI."
     ),
+    dependencies=[Depends(add_request_id)]
 )
 
 # ALlow CORS requests
@@ -26,15 +27,6 @@ app.add_middleware(
 )
 
 
-@app.middleware("http")
-async def add_request_id(request: Request, call_next):
-    """Middleware to append an ID to a request"""
-
-    request.state.id = uuid4().hex
-    response = await call_next(request)
-    return response
-
-
 def setup_app():
     """Includes the correct routers for startup"""