Add type hints and use model objects in api_keys.py endpoints (#12939)

Co-authored-by: openhands <openhands@all-hands.dev>

enterprise/server/routes/api_keys.py

@@ -2,6 +2,7 @@ from datetime import UTC, datetime
 
 from fastapi import APIRouter, Depends, HTTPException, status
 from pydantic import BaseModel, field_validator
+from storage.api_key import ApiKey
 from storage.api_key_store import ApiKeyStore
 from storage.lite_llm_manager import LiteLlmManager
 from storage.org_member import OrgMember
@@ -135,9 +136,9 @@ class ApiKeyCreate(BaseModel):
 class ApiKeyResponse(BaseModel):
     id: int
     name: str | None = None
-    created_at: str
-    last_used_at: str | None = None
-    expires_at: str | None = None
+    created_at: datetime
+    last_used_at: datetime | None = None
+    expires_at: datetime | None = None
 
 
 class ApiKeyCreateResponse(ApiKeyResponse):
@@ -152,12 +153,29 @@ class ByorPermittedResponse(BaseModel):
     permitted: bool
 
 
-@api_router.get('/llm/byor/permitted', response_model=ByorPermittedResponse)
-async def check_byor_permitted(user_id: str = Depends(get_user_id)):
+class MessageResponse(BaseModel):
+    message: str
+
+
+def api_key_to_response(key: ApiKey) -> ApiKeyResponse:
+    """Convert an ApiKey model to an ApiKeyResponse."""
+    return ApiKeyResponse(
+        id=key.id,
+        name=key.name,
+        created_at=key.created_at,
+        last_used_at=key.last_used_at,
+        expires_at=key.expires_at,
+    )
+
+
+@api_router.get('/llm/byor/permitted', tags=['Keys'])
+async def check_byor_permitted(
+    user_id: str = Depends(get_user_id),
+) -> ByorPermittedResponse:
     """Check if BYOR key export is permitted for the user's current org."""
     try:
         permitted = await OrgService.check_byor_export_enabled(user_id)
-        return {'permitted': permitted}
+        return ByorPermittedResponse(permitted=permitted)
     except Exception as e:
         logger.exception(
             'Error checking BYOR export permission', extra={'error': str(e)}
@@ -168,8 +186,10 @@ async def check_byor_permitted(user_id: str = Depends(get_user_id)):
         )
 
 
-@api_router.post('', response_model=ApiKeyCreateResponse)
-async def create_api_key(key_data: ApiKeyCreate, user_id: str = Depends(get_user_id)):
+@api_router.post('', tags=['Keys'])
+async def create_api_key(
+    key_data: ApiKeyCreate, user_id: str = Depends(get_user_id)
+) -> ApiKeyCreateResponse:
     """Create a new API key for the authenticated user."""
     try:
         api_key = await api_key_store.create_api_key(
@@ -178,48 +198,29 @@ async def create_api_key(key_data: ApiKeyCreate, user_id: str = Depends(get_user
         # Get the created key details
         keys = await api_key_store.list_api_keys(user_id)
         for key in keys:
-            if key['name'] == key_data.name:
-                return {
-                    **key,
-                    'key': api_key,
-                    'created_at': (
-                        key['created_at'].isoformat() if key['created_at'] else None
-                    ),
-                    'last_used_at': (
-                        key['last_used_at'].isoformat() if key['last_used_at'] else None
-                    ),
-                    'expires_at': (
-                        key['expires_at'].isoformat() if key['expires_at'] else None
-                    ),
-                }
+            if key.name == key_data.name:
+                return ApiKeyCreateResponse(
+                    id=key.id,
+                    name=key.name,
+                    key=api_key,
+                    created_at=key.created_at,
+                    last_used_at=key.last_used_at,
+                    expires_at=key.expires_at,
+                )
     except Exception:
         logger.exception('Error creating API key')
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail='Failed to create API key',
-        )
+    raise HTTPException(
+        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+        detail='Failed to create API key',
+    )
 
 
-@api_router.get('', response_model=list[ApiKeyResponse])
-async def list_api_keys(user_id: str = Depends(get_user_id)):
+@api_router.get('', tags=['Keys'])
+async def list_api_keys(user_id: str = Depends(get_user_id)) -> list[ApiKeyResponse]:
     """List all API keys for the authenticated user."""
     try:
         keys = await api_key_store.list_api_keys(user_id)
-        return [
-            {
-                **key,
-                'created_at': (
-                    key['created_at'].isoformat() if key['created_at'] else None
-                ),
-                'last_used_at': (
-                    key['last_used_at'].isoformat() if key['last_used_at'] else None
-                ),
-                'expires_at': (
-                    key['expires_at'].isoformat() if key['expires_at'] else None
-                ),
-            }
-            for key in keys
-        ]
+        return [api_key_to_response(key) for key in keys]
     except Exception:
         logger.exception('Error listing API keys')
         raise HTTPException(
@@ -228,8 +229,10 @@ async def list_api_keys(user_id: str = Depends(get_user_id)):
         )
 
 
-@api_router.delete('/{key_id}')
-async def delete_api_key(key_id: int, user_id: str = Depends(get_user_id)):
+@api_router.delete('/{key_id}', tags=['Keys'])
+async def delete_api_key(
+    key_id: int, user_id: str = Depends(get_user_id)
+) -> MessageResponse:
     """Delete an API key."""
     try:
         # First, verify the key belongs to the user
@@ -237,7 +240,7 @@ async def delete_api_key(key_id: int, user_id: str = Depends(get_user_id)):
         key_to_delete = None
 
         for key in keys:
-            if key['id'] == key_id:
+            if key.id == key_id:
                 key_to_delete = key
                 break
 
@@ -255,7 +258,7 @@ async def delete_api_key(key_id: int, user_id: str = Depends(get_user_id)):
                 status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                 detail='Failed to delete API key',
             )
-        return {'message': 'API key deleted successfully'}
+        return MessageResponse(message='API key deleted successfully')
     except HTTPException:
         raise
     except Exception:
@@ -266,8 +269,10 @@ async def delete_api_key(key_id: int, user_id: str = Depends(get_user_id)):
         )
 
 
-@api_router.get('/llm/byor', response_model=LlmApiKeyResponse)
-async def get_llm_api_key_for_byor(user_id: str = Depends(get_user_id)):
+@api_router.get('/llm/byor', tags=['Keys'])
+async def get_llm_api_key_for_byor(
+    user_id: str = Depends(get_user_id),
+) -> LlmApiKeyResponse:
     """Get the LLM API key for BYOR (Bring Your Own Runtime) for the authenticated user.
 
     This endpoint validates that the key exists in LiteLLM before returning it.
@@ -290,7 +295,7 @@ async def get_llm_api_key_for_byor(user_id: str = Depends(get_user_id)):
             # Validate that the key is actually registered in LiteLLM
             is_valid = await LiteLlmManager.verify_key(byor_key, user_id)
             if is_valid:
-                return {'key': byor_key}
+                return LlmApiKeyResponse(key=byor_key)
             else:
                 # Key exists in DB but is invalid in LiteLLM - regenerate it
                 logger.warning(
@@ -315,7 +320,7 @@ async def get_llm_api_key_for_byor(user_id: str = Depends(get_user_id)):
                 'Successfully generated and stored new BYOR key',
                 extra={'user_id': user_id},
             )
-            return {'key': key}
+            return LlmApiKeyResponse(key=key)
         else:
             logger.error(
                 'Failed to generate new BYOR LLM API key',
@@ -337,8 +342,10 @@ async def get_llm_api_key_for_byor(user_id: str = Depends(get_user_id)):
         )
 
 
-@api_router.post('/llm/byor/refresh', response_model=LlmApiKeyResponse)
-async def refresh_llm_api_key_for_byor(user_id: str = Depends(get_user_id)):
+@api_router.post('/llm/byor/refresh', tags=['Keys'])
+async def refresh_llm_api_key_for_byor(
+    user_id: str = Depends(get_user_id),
+) -> LlmApiKeyResponse:
     """Refresh the LLM API key for BYOR (Bring Your Own Runtime) for the authenticated user.
 
     Returns 402 Payment Required if BYOR export is not enabled for the user's org.
@@ -391,7 +398,7 @@ async def refresh_llm_api_key_for_byor(user_id: str = Depends(get_user_id)):
             'BYOR LLM API key refresh completed successfully',
             extra={'user_id': user_id},
         )
-        return {'key': key}
+        return LlmApiKeyResponse(key=key)
     except HTTPException as he:
         logger.error(
             'HTTP exception during BYOR LLM API key refresh',

enterprise/storage/api_key_store.py

@@ -126,7 +126,7 @@ class ApiKeyStore:
 
             return True
 
-    async def list_api_keys(self, user_id: str) -> list[dict]:
+    async def list_api_keys(self, user_id: str) -> list[ApiKey]:
         """List all API keys for a user."""
         user = await UserStore.get_user_by_id_async(user_id)
         org_id = user.current_org_id
@@ -134,24 +134,14 @@ class ApiKeyStore:
 
     def _list_api_keys_from_db(self, user_id: str, org_id: str) -> list[ApiKey]:
         with self.session_maker() as session:
-            keys = (
+            keys: list[ApiKey] = (
                 session.query(ApiKey)
                 .filter(ApiKey.user_id == user_id)
                 .filter(ApiKey.org_id == org_id)
                 .all()
             )
 
-            return [
-                {
-                    'id': key.id,
-                    'name': key.name,
-                    'created_at': key.created_at,
-                    'last_used_at': key.last_used_at,
-                    'expires_at': key.expires_at,
-                }
-                for key in keys
-                if 'MCP_API_KEY' != key.name
-            ]
+            return [key for key in keys if key.name != 'MCP_API_KEY']
 
     async def retrieve_mcp_api_key(self, user_id: str) -> str | None:
         user = await UserStore.get_user_by_id_async(user_id)

enterprise/tests/unit/server/routes/test_api_keys.py

@@ -6,6 +6,8 @@ import httpx
 import pytest
 from fastapi import HTTPException
 from server.routes.api_keys import (
+    ByorPermittedResponse,
+    LlmApiKeyResponse,
     check_byor_permitted,
     delete_byor_key_from_litellm,
     get_llm_api_key_for_byor,
@@ -203,7 +205,7 @@ class TestGetLlmApiKeyForByor:
         result = await get_llm_api_key_for_byor(user_id=user_id)
 
         # Assert
-        assert result == {'key': new_key}
+        assert result == LlmApiKeyResponse(key=new_key)
         mock_check_enabled.assert_called_once_with(user_id)
         mock_get_key.assert_called_once_with(user_id)
         mock_generate_key.assert_called_once_with(user_id)
@@ -228,7 +230,7 @@ class TestGetLlmApiKeyForByor:
         result = await get_llm_api_key_for_byor(user_id=user_id)
 
         # Assert
-        assert result == {'key': existing_key}
+        assert result == LlmApiKeyResponse(key=existing_key)
         mock_check_enabled.assert_called_once_with(user_id)
         mock_get_key.assert_called_once_with(user_id)
         mock_verify_key.assert_called_once_with(existing_key, user_id)
@@ -265,7 +267,7 @@ class TestGetLlmApiKeyForByor:
         result = await get_llm_api_key_for_byor(user_id=user_id)
 
         # Assert
-        assert result == {'key': new_key}
+        assert result == LlmApiKeyResponse(key=new_key)
         mock_check_enabled.assert_called_once_with(user_id)
         mock_get_key.assert_called_once_with(user_id)
         mock_verify_key.assert_called_once_with(invalid_key, user_id)
@@ -305,7 +307,7 @@ class TestGetLlmApiKeyForByor:
         result = await get_llm_api_key_for_byor(user_id=user_id)
 
         # Assert
-        assert result == {'key': new_key}
+        assert result == LlmApiKeyResponse(key=new_key)
         mock_check_enabled.assert_called_once_with(user_id)
         mock_delete_key.assert_called_once_with(user_id, invalid_key)
         mock_generate_key.assert_called_once_with(user_id)
@@ -478,7 +480,7 @@ class TestCheckByorPermitted:
         result = await check_byor_permitted(user_id=user_id)
 
         # Assert
-        assert result == {'permitted': True}
+        assert result == ByorPermittedResponse(permitted=True)
         mock_check_enabled.assert_called_once_with(user_id)
 
     @pytest.mark.asyncio
@@ -493,7 +495,7 @@ class TestCheckByorPermitted:
         result = await check_byor_permitted(user_id=user_id)
 
         # Assert
-        assert result == {'permitted': False}
+        assert result == ByorPermittedResponse(permitted=False)
         mock_check_enabled.assert_called_once_with(user_id)
 
     @pytest.mark.asyncio

enterprise/tests/unit/test_api_key_store.py

@@ -265,17 +265,17 @@ async def test_list_api_keys(
     # Verify
     mock_get_user.assert_called_once_with(user_id)
     assert len(result) == 2
-    assert result[0]['id'] == 1
-    assert result[0]['name'] == 'Key 1'
-    assert result[0]['created_at'] == now
-    assert result[0]['last_used_at'] == now
-    assert result[0]['expires_at'] == now + timedelta(days=30)
+    assert result[0].id == 1
+    assert result[0].name == 'Key 1'
+    assert result[0].created_at == now
+    assert result[0].last_used_at == now
+    assert result[0].expires_at == now + timedelta(days=30)
 
-    assert result[1]['id'] == 2
-    assert result[1]['name'] == 'Key 2'
-    assert result[1]['created_at'] == now
-    assert result[1]['last_used_at'] is None
-    assert result[1]['expires_at'] is None
+    assert result[1].id == 2
+    assert result[1].name == 'Key 2'
+    assert result[1].created_at == now
+    assert result[1].last_used_at is None
+    assert result[1].expires_at is None
 
 
 @pytest.mark.asyncio