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, status_code=403,
detail="Model not found", detail="Model not found",
) )
else: elif not bypass_filter:
if user.role != "admin": if user.role != "admin":
raise HTTPException( raise HTTPException(
status_code=403, status_code=403,

View File

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

View File

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