fix: workaround socketio upstream bug when websockets are not available

This commit is contained in:
Jun Siang Cheah
2024-09-08 11:54:56 +01:00
parent 827c419251
commit 9401f6c821
2 changed files with 21 additions and 10 deletions

View File

@@ -812,6 +812,24 @@ async def update_embedding_function(request: Request, call_next):
return response
@app.middleware("http")
async def inspect_websocket(request: Request, call_next):
if (
"/ws/socket.io" in request.url.path
and request.query_params.get("transport") == "websocket"
):
upgrade = (request.headers.get("Upgrade") or "").lower()
connection = (request.headers.get("Connection") or "").lower().split(",")
# Check that there's the correct headers for an upgrade, else reject the connection
# This is to work around this upstream issue: https://github.com/miguelgrinberg/python-engineio/issues/367
if upgrade != "websocket" or "upgrade" not in connection:
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST,
content={"detail": "Invalid WebSocket upgrade request"},
)
return await call_next(request)
app.mount("/ws", socket_app)
app.mount("/ollama", ollama_app)