|
@@ -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)
|