tdpeuter/open-webui
tdpeuter/open-webui
1
Fork 0
forked from open-webui/open-webui
Code Issues Pull requests 1 Releases 2 Packages Activity Actions

chore: comment clean up

Browse source
This commit is contained in:
Timothy J. Baek 2024-03-20 19:32:25 -07:00
parent 1a240e103e
commit ed0d4b5a2a
1 changed files with 1 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
backend/apps/litellm/main.py
View file

@ -67,20 +67,15 @@ class ModifyModelsResponseMiddleware(BaseHTTPMiddleware):
response = await call_next(request) response = await call_next(request)
user = request.state.user user = request.state.user
# Check if the request is for the `/models` route
if "/models" in request.url.path: if "/models" in request.url.path:
# Ensure the response is a StreamingResponse
if isinstance(response, StreamingResponse): if isinstance(response, StreamingResponse):
# Read the content of the streaming response # Read the content of the streaming response
body = b"" body = b""
async for chunk in response.body_iterator: async for chunk in response.body_iterator:
body += chunk body += chunk
# Modify the content as needed
data = json.loads(body.decode("utf-8")) data = json.loads(body.decode("utf-8"))
print(data)
if app.state.MODEL_FILTER_ENABLED: if app.state.MODEL_FILTER_ENABLED:
if user and user.role == "user": if user and user.role == "user":
data["data"] = list( data["data"] = list(
@ -91,14 +86,11 @@ class ModifyModelsResponseMiddleware(BaseHTTPMiddleware):
) )
) )
# Example modification: Add a new key-value pair # Modified Flag
data["modified"] = True data["modified"] = True
# Return a new JSON response with the modified content
return JSONResponse(content=data) return JSONResponse(content=data)
return response return response
# Add the middleware to the app
app.add_middleware(ModifyModelsResponseMiddleware) app.add_middleware(ModifyModelsResponseMiddleware)