From 0eedc8ca146dea5d532a8022f4bef8a8c4b899b7 Mon Sep 17 00:00:00 2001 From: kingbri Date: Mon, 22 Jul 2024 12:19:46 -0400 Subject: [PATCH] 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 --- common/networking.py | 8 ++++++++ endpoints/server.py | 14 +++----------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/common/networking.py b/common/networking.py index 5540105..706cb60 100644 --- a/common/networking.py +++ b/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 diff --git a/endpoints/server.py b/endpoints/server.py index ec59455..0a29e15 100644 --- a/endpoints/server.py +++ b/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"""