refactor: enhance private AI key creation with error handling and cleanup

- Implemented error handling during the creation of private AI keys to ensure proper cleanup of resources (LiteLLM token and vector database) in case of failures.
- Added logging for cleanup operations and preserved original exceptions for better debugging.
- Updated tests to cover various failure scenarios, ensuring that resources are cleaned up correctly and appropriate error messages are returned to the user.

Author: PhiRho

app/api/private_ai_keys.py

@@ -191,38 +191,89 @@ async def create_private_ai_key(
     Note: You must be authenticated to use this endpoint.
     Only admins can create keys for other users or teams.
     """
-    # First create the LiteLLM token
-    llm_token = await create_llm_token(private_ai_key, current_user, user_role, db, store_result=False)
-
-    # Then create the vector database
-    vector_db = VectorDBCreate(
-        region_id=private_ai_key.region_id,
-        name=private_ai_key.name,
-        owner_id=private_ai_key.owner_id,
-        team_id=private_ai_key.team_id
-    )
-    db_info = await create_vector_db(vector_db, current_user, user_role, db, store_result=False)
-
-    # Store private AI key info in main application database
-    new_key = DBPrivateAIKey(
-        database_name=db_info.database_name,
-        name=db_info.name,
-        database_host=db_info.database_host,
-        database_username=db_info.database_username,
-        database_password=db_info.database_password,
-        litellm_token=llm_token.litellm_token,
-        litellm_api_url=llm_token.litellm_api_url,
-        owner_id=db_info.owner_id,
-        team_id=db_info.team_id,
-        region_id=private_ai_key.region_id
-    )
-    db.add(new_key)
-    db.commit()
-    db.refresh(new_key)
+    llm_token = None
+    db_info = None
+    region = None
 
-    key_data = new_key.to_dict()
+    try:
+        # Get the region first for cleanup purposes
+        region = db.query(DBRegion).filter(
+            DBRegion.id == private_ai_key.region_id,
+            DBRegion.is_active == True
+        ).first()
+        if not region:
+            raise HTTPException(
+                status_code=status.HTTP_404_NOT_FOUND,
+                detail="Region not found or inactive"
+            )
+
+        # First create the LiteLLM token
+        llm_token = await create_llm_token(private_ai_key, current_user, user_role, db, store_result=False)
 
-    return PrivateAIKey.model_validate(key_data)
+        # Then create the vector database
+        vector_db = VectorDBCreate(
+            region_id=private_ai_key.region_id,
+            name=private_ai_key.name,
+            owner_id=private_ai_key.owner_id,
+            team_id=private_ai_key.team_id
+        )
+        db_info = await create_vector_db(vector_db, current_user, user_role, db, store_result=False)
+
+        # Store private AI key info in main application database
+        new_key = DBPrivateAIKey(
+            database_name=db_info.database_name,
+            name=db_info.name,
+            database_host=db_info.database_host,
+            database_username=db_info.database_username,
+            database_password=db_info.database_password,
+            litellm_token=llm_token.litellm_token,
+            litellm_api_url=llm_token.litellm_api_url,
+            owner_id=db_info.owner_id,
+            team_id=db_info.team_id,
+            region_id=private_ai_key.region_id
+        )
+        db.add(new_key)
+        db.commit()
+        db.refresh(new_key)
+
+        key_data = new_key.to_dict()
+
+        return PrivateAIKey.model_validate(key_data)
+
+    except Exception as e:
+        logger.error(f"Failed to create private AI key: {str(e)}", exc_info=True)
+        db.rollback()
+
+        # Cleanup resources on failure
+        try:
+            if llm_token and region:
+                # Delete LiteLLM token
+                litellm_service = LiteLLMService(
+                    api_url=region.litellm_api_url,
+                    api_key=region.litellm_api_key
+                )
+                await litellm_service.delete_key(llm_token.litellm_token)
+                logger.info(f"Cleaned up LiteLLM token after failure")
+        except Exception as cleanup_error:
+            logger.error(f"Failed to cleanup LiteLLM token: {str(cleanup_error)}", exc_info=True)
+
+        try:
+            if db_info and region:
+                # Delete vector database
+                postgres_manager = PostgresManager(region=region)
+                await postgres_manager.delete_database(db_info.database_name)
+                logger.info(f"Cleaned up vector database after failure")
+        except Exception as cleanup_error:
+            logger.error(f"Failed to cleanup vector database: {str(cleanup_error)}", exc_info=True)
+
+        # Re-raise the original exception
+        if isinstance(e, HTTPException):
+            raise e
+        else:
+            raise HTTPException(
+                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
+                detail=f"Failed to create private AI key: {str(e)}"
+            )
 
 @router.post("/token", response_model=LiteLLMToken)
 async def create_llm_token(

tests/test_private_ai.py

@@ -1,9 +1,10 @@
 import pytest
-from unittest.mock import patch
+from unittest.mock import patch, Mock
 from app.db.models import DBPrivateAIKey, DBTeam, DBUser, DBProduct, DBTeamProduct
 from datetime import datetime, UTC
 from app.core.security import get_password_hash
 from requests.exceptions import HTTPError
+from fastapi import status, HTTPException
 from app.core.resource_limits import (
     DEFAULT_MAX_SPEND,
     DEFAULT_RPM_PER_KEY,
@@ -1905,3 +1906,190 @@ def test_delete_private_ai_key_litellm_service_unavailable(mock_post, client, ad
         json={"keys": [test_key.litellm_token]}
     )
     db.commit()
+
+@patch("app.services.litellm.requests.post")
+@patch("app.db.postgres.PostgresManager.create_database")
+def test_create_private_ai_key_cleanup_on_vector_db_failure(mock_create_db, mock_post, client, test_token, test_region, test_user):
+    """
+    Given a user creates a private AI key
+    When the vector database creation fails after LiteLLM token is created
+    Then the LiteLLM token should be cleaned up and an error returned
+    """
+    # Mock successful LiteLLM token creation
+    mock_post.return_value.status_code = 200
+    mock_post.return_value.json.return_value = {"key": "test-private-key-123"}
+    mock_post.return_value.raise_for_status.return_value = None
+
+    # Mock vector database creation failure
+    mock_create_db.side_effect = Exception("Database creation failed")
+
+    response = client.post(
+        "/private-ai-keys/",
+        headers={"Authorization": f"Bearer {test_token}"},
+        json={
+            "region_id": test_region.id,
+            "name": "Test AI Key"
+        }
+    )
+
+    # Verify the response indicates failure
+    assert response.status_code == 500
+    assert "Failed to create vector database" in response.json()["detail"]
+
+    # Verify LiteLLM token was cleaned up (delete API was called)
+    # First call is for token creation, second call is for cleanup
+    assert mock_post.call_count == 2
+
+    # Verify the cleanup call
+    cleanup_call = mock_post.call_args_list[1]
+    assert cleanup_call[0][0] == f"{test_region.litellm_api_url}/key/delete"
+    assert cleanup_call[1]["headers"]["Authorization"] == f"Bearer {test_region.litellm_api_key}"
+    assert cleanup_call[1]["json"]["keys"] == ["test-private-key-123"]
+
+    # Verify no key was stored in the database
+    stored_keys = client.get(
+        "/private-ai-keys/",
+        headers={"Authorization": f"Bearer {test_token}"}
+    ).json()
+    assert len([k for k in stored_keys if k["name"] == "Test AI Key"]) == 0
+
+@patch("app.services.litellm.requests.post")
+@patch("app.db.postgres.PostgresManager.create_database")
+@patch("app.db.postgres.PostgresManager.delete_database")
+@patch("sqlalchemy.orm.Session.commit")
+def test_create_private_ai_key_cleanup_on_db_storage_failure(mock_commit, mock_delete_db, mock_create_db, mock_post, client, test_token, test_region, test_user):
+    """
+    Given a user creates a private AI key
+    When the database storage fails after both LiteLLM token and vector DB are created
+    Then both resources should be cleaned up and an error returned
+    """
+    # Mock successful LiteLLM token creation
+    mock_post.return_value.status_code = 200
+    mock_post.return_value.json.return_value = {"key": "test-private-key-123"}
+    mock_post.return_value.raise_for_status.return_value = None
+
+    # Mock successful vector database creation
+    mock_create_db.return_value = {
+        "database_name": "test_db_123",
+        "database_host": "test-host",
+        "database_username": "test_user",
+        "database_password": "test_pass"
+    }
+
+    # Mock successful vector database deletion
+    mock_delete_db.return_value = None
+
+    # Mock database storage failure
+    mock_commit.side_effect = Exception("Database storage failed")
+
+    response = client.post(
+        "/private-ai-keys/",
+        headers={"Authorization": f"Bearer {test_token}"},
+        json={
+            "region_id": test_region.id,
+            "name": "Test AI Key"
+        }
+    )
+
+    # Verify the response indicates failure
+    assert response.status_code == 500
+    assert "Failed to create private AI key" in response.json()["detail"]
+
+    # Verify LiteLLM token was cleaned up
+    assert mock_post.call_count == 2
+    cleanup_call = mock_post.call_args_list[1]
+    assert cleanup_call[0][0] == f"{test_region.litellm_api_url}/key/delete"
+    assert cleanup_call[1]["json"]["keys"] == ["test-private-key-123"]
+
+    # Verify vector database was cleaned up
+    mock_delete_db.assert_called_once_with("test_db_123")
+
+    # Verify no key was stored in the database
+    stored_keys = client.get(
+        "/private-ai-keys/",
+        headers={"Authorization": f"Bearer {test_token}"}
+    ).json()
+    assert len([k for k in stored_keys if k["name"] == "Test AI Key"]) == 0
+
+@patch("app.services.litellm.requests.post")
+@patch("app.db.postgres.PostgresManager.create_database")
+def test_create_private_ai_key_cleanup_failure_handling(mock_create_db, mock_post, client, test_token, test_region, test_user):
+    """
+    Given a user creates a private AI key
+    When the cleanup process itself fails
+    Then the original error should still be returned to the user
+    """
+    # Mock successful LiteLLM token creation
+    mock_post.return_value.status_code = 200
+    mock_post.return_value.json.return_value = {"key": "test-private-key-123"}
+    mock_post.return_value.raise_for_status.return_value = None
+
+    # Mock vector database creation failure
+    mock_create_db.side_effect = Exception("Database creation failed")
+
+    # Mock cleanup failure - the second call to requests.post will be for cleanup
+    # First call succeeds, second call fails
+    mock_post.side_effect = [
+        Mock(
+            status_code=200,
+            json=Mock(return_value={"key": "test-private-key-123"}),
+            raise_for_status=Mock(return_value=None)
+        ),
+        Mock(
+            status_code=500,
+            raise_for_status=Mock(side_effect=HTTPError("Cleanup failed"))
+        )
+    ]
+
+    response = client.post(
+        "/private-ai-keys/",
+        headers={"Authorization": f"Bearer {test_token}"},
+        json={
+            "region_id": test_region.id,
+            "name": "Test AI Key"
+        }
+    )
+
+    # Verify the response indicates the original failure, not cleanup failure
+    assert response.status_code == 500
+    assert "Failed to create vector database" in response.json()["detail"]
+    assert "Database creation failed" in response.json()["detail"]
+
+    # Verify cleanup was attempted
+    assert mock_post.call_count == 2
+
+@patch("app.services.litellm.requests.post")
+@patch("app.db.postgres.PostgresManager.create_database")
+def test_create_private_ai_key_http_exception_preservation(mock_create_db, mock_post, client, test_token, test_region, test_user):
+    """
+    Given a user creates a private AI key
+    When an HTTPException is raised during creation
+    Then the original HTTPException should be preserved and returned
+    """
+    # Mock successful LiteLLM token creation
+    mock_post.return_value.status_code = 200
+    mock_post.return_value.json.return_value = {"key": "test-private-key-123"}
+    mock_post.return_value.raise_for_status.return_value = None
+
+    # Mock vector database creation failure with HTTPException
+    mock_create_db.side_effect = HTTPException(
+        status_code=status.HTTP_400_BAD_REQUEST,
+        detail="Invalid database configuration"
+    )
+
+    response = client.post(
+        "/private-ai-keys/",
+        headers={"Authorization": f"Bearer {test_token}"},
+        json={
+            "region_id": test_region.id,
+            "name": "Test AI Key"
+        }
+    )
+
+    # Verify the original HTTPException is preserved
+    assert response.status_code == 500
+    assert "Failed to create vector database" in response.json()["detail"]
+    assert "Invalid database configuration" in response.json()["detail"]
+
+    # Verify cleanup was attempted
+    assert mock_post.call_count == 2