fix: arena access control

This commit is contained in:
Timothy Jaeryang Baek 2024-11-18 07:40:37 -08:00
parent 269151cd2c
commit f37d847521
3 changed files with 62 additions and 38 deletions

View File

@ -958,7 +958,7 @@ async def generate_chat_completion(
status_code=403,
detail="Model not found",
)
else:
elif not bypass_filter:
if user.role != "admin":
raise HTTPException(
status_code=403,

View File

@ -510,7 +510,7 @@ async def generate_chat_completion(
status_code=403,
detail="Model not found",
)
else:
elif not bypass_filter:
if user.role != "admin":
raise HTTPException(
status_code=403,

View File

@ -557,6 +557,19 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
model_info = Models.get_model_by_id(model["id"])
if user.role == "user":
if model.get("arena"):
if not has_access(
user.id,
type="read",
access_control=model.get("info", {})
.get("meta", {})
.get("access_control", {}),
):
raise HTTPException(
status_code=403,
detail="Model not found",
)
else:
if not model_info:
return JSONResponse(
status_code=status.HTTP_404_NOT_FOUND,
@ -1160,8 +1173,22 @@ async def generate_chat_completions(
)
model = models[model_id]
# Check if user has access to the model
if user.role == "user":
if not bypass_filter and user.role == "user":
if model.get("arena"):
if not has_access(
user.id,
type="read",
access_control=model.get("info", {})
.get("meta", {})
.get("access_control", {}),
):
raise HTTPException(
status_code=403,
detail="Model not found",
)
else:
model_info = Models.get_model_by_id(model_id)
if not model_info:
raise HTTPException(
@ -1186,9 +1213,7 @@ async def generate_chat_completions(
model_ids = [
model["id"]
for model in await get_all_models()
if model.get("owned_by") != "arena"
and not model.get("info", {}).get("meta", {}).get("hidden", False)
and model["id"] not in model_ids
if model.get("owned_by") != "arena" and model["id"] not in model_ids
]
selected_model_id = None
@ -1199,7 +1224,6 @@ async def generate_chat_completions(
model["id"]
for model in await get_all_models()
if model.get("owned_by") != "arena"
and not model.get("info", {}).get("meta", {}).get("hidden", False)
]
selected_model_id = random.choice(model_ids)