Enhance Private AI Key Creation with Metadata and Naming

- Add name and user_id parameters to LiteLLM key generation
- Update PostgresManager to support optional database name
- Improve key alias and metadata generation for better tracking
- Modify key creation process to include more contextual information

Author: Schnitzel

app/api/private_ai_keys.py

@@ -54,7 +54,9 @@ async def create_private_ai_key(
         # Create new postgres database
         postgres_manager = PostgresManager(region=region)
         key_credentials = await postgres_manager.create_database(
-            owner=current_user.email
+            owner=current_user.email,
+            name=private_ai_key.name,
+            user_id=current_user.id
         )
 
         # Store private AI key info in main application database

app/db/postgres.py

@@ -18,9 +18,14 @@ def __init__(self, region: DBRegion = None):
         else:
             raise ValueError("Region is required for PostgresManager")
 
-    async def create_database(self, owner: str) -> Dict:
+    async def create_database(self, owner: str, name: str = None, user_id: int = None) -> Dict:
+        # Generate unique database name and credentials
+        db_name = f"db_{uuid.uuid4().hex[:8]}"
+        db_user = f"user_{uuid.uuid4().hex[:8]}"
+        db_password = uuid.uuid4().hex
+        
         # Generate LiteLLM token
-        litellm_token = await self.litellm_service.create_key(email=owner)
+        litellm_token = await self.litellm_service.create_key(email=owner, name=name, user_id=user_id)
 
         # Connect to postgres and create database/user
         try:
@@ -38,11 +43,6 @@ async def create_database(self, owner: str) -> Dict:
             print(f"Unexpected error connecting to PostgreSQL: {str(e)}")
             raise
 
-        # Generate unique database name and credentials
-        db_name = f"db_{uuid.uuid4().hex[:8]}"
-        db_user = f"user_{uuid.uuid4().hex[:8]}"
-        db_password = uuid.uuid4().hex
-
         try:
             await conn.execute(f'CREATE DATABASE {db_name}')
             await conn.execute(f'CREATE USER {db_user} WITH PASSWORD \'{db_password}\'')

app/services/litellm.py

@@ -9,7 +9,7 @@ def __init__(self, api_url: str = None, api_key: str = None):
         if not self.master_key:
             raise ValueError("LiteLLM API key is required")
 
-    async def create_key(self, email: str = None) -> str:
+    async def create_key(self, email: str = None, name: str = None, user_id: int = None) -> str:
         """Create a new API key for LiteLLM"""
         try:
             request_data = {
@@ -20,10 +20,22 @@ async def create_key(self, email: str = None) -> str:
                 "spend": 0,
             }
 
-            # Add email as key_alias and metadata if provided
+            # Add email and name to key_alias and metadata if provided
             if email:
-                request_data["key_alias"] = email
-                request_data["metadata"] = {"service_account_id": email}
+                key_alias = email
+                metadata = {"service_account_id": email}
+                
+                # Add name to key_alias and metadata if provided
+                if name:
+                    key_alias = f"{email} - {name}"
+                    metadata["amazeeai_private_ai_key_name"] = name
+                
+                # Add user_id to metadata if provided
+                if user_id is not None:
+                    metadata["amazeeai_user_id"] = str(user_id)
+                
+                request_data["key_alias"] = key_alias
+                request_data["metadata"] = metadata
 
             response = requests.post(
                 f"{self.api_url}/key/generate",