diff --git a/common/auth.py b/common/auth.py index 0f82f3d..fa53262 100644 --- a/common/auth.py +++ b/common/auth.py @@ -10,8 +10,6 @@ from pydantic import BaseModel from loguru import logger from typing import Optional -from endpoints.OAI.types.auth import AuthPermissionResponse - class AuthKeys(BaseModel): """ @@ -82,9 +80,9 @@ async def validate_key_permission(test_key: str): test_key = test_key.split(" ")[1] if AUTH_KEYS.verify_key(test_key, "admin_key"): - return AuthPermissionResponse(permission="admin") + return "admin" elif AUTH_KEYS.verify_key(test_key, "api_key"): - return AuthPermissionResponse(permission="api") + return "api" else: raise ValueError("The provided authentication key is invalid.") diff --git a/endpoints/OAI/app.py b/endpoints/OAI/app.py index ec98ffe..28ee9d1 100644 --- a/endpoints/OAI/app.py +++ b/endpoints/OAI/app.py @@ -27,6 +27,7 @@ from common.utils import ( handle_request_error, unwrap, ) +from endpoints.OAI.types.auth import AuthPermissionResponse from endpoints.OAI.types.completion import CompletionRequest from endpoints.OAI.types.chat_completion import ChatCompletionRequest from endpoints.OAI.types.lora import ( @@ -419,8 +420,8 @@ async def get_key_permission( test_key = coalesce(x_admin_key, x_api_key, authorization) try: - response = await validate_key_permission(test_key) - return response + permission = await validate_key_permission(test_key) + return AuthPermissionResponse(permission=permission) except ValueError as exc: error_message = handle_request_error(str(exc)).error.message