use FireworkS SDK for Secret automation

Author: dphuang2

eval_protocol/platform_api.py

@@ -1,17 +1,17 @@
 # eval_protocol/platform_api.py
 import logging
 import sys
-from typing import Any, Dict, Optional
+from typing import Optional
 
-import requests
 from dotenv import find_dotenv, load_dotenv
 
 from eval_protocol.auth import (
     get_fireworks_account_id,
     get_fireworks_api_base,
     get_fireworks_api_key,
 )
-from eval_protocol.common_utils import get_user_agent
+from fireworks.types import Secret
+from fireworks import Fireworks, FireworksError, NotFoundError, InternalServerError
 
 logger = logging.getLogger(__name__)
 
@@ -88,47 +88,31 @@ def create_or_update_fireworks_secret(
     resolved_api_key = api_key or get_fireworks_api_key()
     resolved_api_base = api_base or get_fireworks_api_base()
     resolved_account_id = account_id  # Must be provided
+    client = Fireworks(api_key=resolved_api_key, account_id=resolved_account_id, base_url=resolved_api_base)
 
     if not all([resolved_api_key, resolved_api_base, resolved_account_id]):
         logger.error("Missing Fireworks API key, base URL, or account ID for creating/updating secret.")
         return False
 
-    headers = {
-        "Authorization": f"Bearer {resolved_api_key}",
-        "Content-Type": "application/json",
-        "User-Agent": get_user_agent(),
-    }
-
-    # The secret_id for GET/PATCH/DELETE operations is the key_name.
-    # The 'name' field in the gatewaySecret model for POST/PATCH is a bit ambiguous.
-    # For POST (create), the body is gatewaySecret, which has 'name', 'keyName', 'value'.
-    # 'name' in POST body is likely just the 'keyName' or 'secret_id' for creation context,
-    # as the full resource name 'accounts/.../secrets/...' is server-generated.
-    # Let's assume for POST, we send 'keyName' and 'value'.
-    # For PATCH, the path contains {secret_id} which is the key_name. The body is also gatewaySecret.
-
     # Check if secret exists using GET (path uses normalized resource id)
     resource_id = _normalize_secret_resource_id(key_name)
     secret_exists = False
     try:
-        url = f"{resolved_api_base}/v1/accounts/{resolved_account_id}/secrets/{resource_id}"
-        response = requests.get(url, headers=headers, timeout=10)
-        if response.status_code == 200:
+        secret = client.secrets.get(resource_id)
+        if secret:
             secret_exists = True
             logger.info(f"Secret '{key_name}' already exists. Will attempt to update.")
-        elif response.status_code == 404:
-            logger.info(f"Secret '{key_name}' does not exist. Will attempt to create.")
-            secret_exists = False
-        elif response.status_code == 500:  # As per user feedback, 500 on GET might mean not found
-            logger.warning(
-                f"Received 500 error when checking for secret '{key_name}'. Assuming it does not exist and will attempt to create. Response: {response.text}"
-            )
-            secret_exists = False
-        else:
-            logger.error(f"Error checking for secret '{key_name}': {response.status_code} - {response.text}")
-            return False
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Request exception while checking for secret '{key_name}': {e}")
+    except NotFoundError:
+        # Secret doesn't exist, proceed with creation
+        secret_exists = False
+    except InternalServerError as e:
+        # As per user feedback, 500 on GET might mean not found, treat as not found
+        logger.warning(
+            f"Received 500 error when checking for secret '{key_name}'. Assuming it doesn't exist. Response: {e}"
+        )
+        secret_exists = False
+    except FireworksError as e:
+        logger.error(f"Error checking for secret '{key_name}': {e}")
         return False
 
     if secret_exists:
@@ -144,31 +128,15 @@ def create_or_update_fireworks_secret(
             )
             payload_key_name = "EP_SECRET"  # Fallback, though unlikely needed with .upper()
 
-        payload = {"keyName": payload_key_name, "value": secret_value}
         try:
-            logger.debug(f"PATCH payload for '{key_name}': {payload}")
-            url = f"{resolved_api_base}/v1/accounts/{resolved_account_id}/secrets/{resource_id}"
-            response = requests.patch(url, json=payload, headers=headers, timeout=30)
-            response.raise_for_status()
+            logger.debug(f"PATCH payload for '{key_name}': key_name={payload_key_name}")
+            client.secrets.update(resource_id, key_name=payload_key_name, value=secret_value)
             logger.info(f"Successfully updated secret '{key_name}' on Fireworks platform.")
             return True
-        except requests.exceptions.HTTPError as e:
-            logger.error(f"HTTP error updating secret '{key_name}': {e.response.status_code} - {e.response.text}")
-            return False
-        except requests.exceptions.RequestException as e:
-            logger.error(f"Request exception updating secret '{key_name}': {e}")
+        except FireworksError as e:
+            logger.error(f"Error updating secret '{key_name}': {e}")
             return False
     else:
-        # Create new secret (POST)
-        # Body for POST is gatewaySecret. 'name' field in payload is the resource path.
-        # Let's assume for POST, the 'name' in payload can be omitted or is the key_name.
-        # The API should ideally use 'keyName' from URL or a specific 'secretId' in payload for creation if 'name' is server-assigned.
-        # Given the Swagger, 'name' is required in gatewaySecret.
-        # Let's try with 'name' being the 'key_name' for the payload, as the full path is not known yet.
-        # This might need adjustment based on actual API behavior.
-        # Construct the full 'name' path for the POST payload as per Swagger's title for 'name'
-        full_resource_name_for_payload = f"accounts/{resolved_account_id}/secrets/{resource_id}"
-
         # Transform key_name for payload "keyName" field: uppercase and underscores
         payload_key_name = key_name.upper().replace("-", "_")
         if not payload_key_name or not payload_key_name[0].isupper():
@@ -177,26 +145,12 @@ def create_or_update_fireworks_secret(
             )
             payload_key_name = "EP_SECRET"
 
-        payload = {
-            "name": full_resource_name_for_payload,  # This 'name' is the resource path
-            "keyName": payload_key_name,  # This 'keyName' is the specific field with new rules
-            "value": secret_value,
-        }
         try:
-            logger.debug(f"POST payload for '{key_name}': {payload}")
-            url = f"{resolved_api_base}/v1/accounts/{resolved_account_id}/secrets"
-            response = requests.post(url, json=payload, headers=headers, timeout=30)
-            response.raise_for_status()
-            logger.info(
-                f"Successfully created secret '{key_name}' on Fireworks platform. Full name: {response.json().get('name')}"
-            )
+            logger.debug(f"POST payload for '{key_name}': {payload_key_name}")
+            client.secrets.create(key_name=payload_key_name, value=secret_value, name=resource_id)
             return True
-        except requests.exceptions.HTTPError as e:
-            logger.error(f"HTTP error creating secret '{key_name}': {e.response.status_code} - {e.response.text}")
-            # If error is due to 'name' field, this log will show it.
-            return False
-        except requests.exceptions.RequestException as e:
-            logger.error(f"Request exception creating secret '{key_name}': {e}")
+        except FireworksError as e:
+            logger.error(f"Error creating secret '{key_name}': {e}")
             return False
 
 
@@ -205,7 +159,7 @@ def get_fireworks_secret(
     key_name: str,  # This is the identifier for the secret
     api_key: Optional[str] = None,
     api_base: Optional[str] = None,
-) -> Optional[Dict[str, Any]]:
+) -> Optional[Secret]:
     """
     Retrieves a secret from the Fireworks AI platform by its keyName.
     Note: This typically does not return the secret's actual value for security reasons,
@@ -215,30 +169,25 @@ def get_fireworks_secret(
     resolved_api_base = api_base or get_fireworks_api_base()
     resolved_account_id = account_id
 
-    if not all([resolved_api_key, resolved_api_base, resolved_account_id]):
-        logger.error("Missing Fireworks API key, base URL, or account ID for getting secret.")
-        return None
-
-    headers = {
-        "Authorization": f"Bearer {resolved_api_key}",
-        "User-Agent": get_user_agent(),
-    }
+    client = Fireworks(api_key=resolved_api_key, account_id=resolved_account_id, base_url=resolved_api_base)
     resource_id = _normalize_secret_resource_id(key_name)
 
     try:
-        url = f"{resolved_api_base}/v1/accounts/{resolved_account_id}/secrets/{resource_id}"
-        response = requests.get(url, headers=headers, timeout=10)
-        if response.status_code == 200:
+        secret = client.secrets.get(resource_id)
+        if secret:
             logger.info(f"Successfully retrieved secret '{key_name}'.")
-            return response.json()
-        elif response.status_code == 404:
-            logger.info(f"Secret '{key_name}' not found.")
-            return None
-        else:
-            logger.error(f"Error getting secret '{key_name}': {response.status_code} - {response.text}")
-            return None
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Request exception while getting secret '{key_name}': {e}")
+            return secret
+    except NotFoundError:
+        logger.info(f"Secret '{key_name}' not found.")
+        return None
+    except InternalServerError as e:
+        # As per user feedback, 500 on GET might mean not found
+        logger.warning(
+            f"Received 500 error when getting secret '{key_name}'. Assuming it doesn't exist. Response: {e}"
+        )
+        return None
+    except FireworksError as e:
+        logger.error(f"Error getting secret '{key_name}': {e}")
         return None
 
 
@@ -259,33 +208,24 @@ def delete_fireworks_secret(
         logger.error("Missing Fireworks API key, base URL, or account ID for deleting secret.")
         return False
 
-    headers = {
-        "Authorization": f"Bearer {resolved_api_key}",
-        "User-Agent": get_user_agent(),
-    }
+    client = Fireworks(api_key=resolved_api_key, account_id=resolved_account_id, base_url=resolved_api_base)
     resource_id = _normalize_secret_resource_id(key_name)
 
     try:
-        url = f"{resolved_api_base}/v1/accounts/{resolved_account_id}/secrets/{resource_id}"
-        response = requests.delete(url, headers=headers, timeout=30)
-        if response.status_code == 200 or response.status_code == 204:  # 204 No Content is also success for DELETE
-            logger.info(f"Successfully deleted secret '{key_name}'.")
-            return True
-        elif response.status_code == 404:
-            logger.info(f"Secret '{key_name}' not found, nothing to delete.")
-            return True
-        elif (
-            response.status_code == 500
-        ):  # As per user feedback, 500 on GET might mean not found, apply same logic for DELETE
-            logger.warning(
-                f"Received 500 error when deleting secret '{key_name}'. Assuming it might not have existed. Response: {response.text}"
-            )
-            return True  # Consider deletion successful if it results in non-existence
-        else:
-            logger.error(f"Error deleting secret '{key_name}': {response.status_code} - {response.text}")
-            return False
-    except requests.exceptions.RequestException as e:
-        logger.error(f"Request exception while deleting secret '{key_name}': {e}")
+        client.secrets.delete(resource_id, account_id=resolved_account_id)
+        logger.info(f"Successfully deleted secret '{key_name}'.")
+        return True
+    except NotFoundError:
+        logger.info(f"Secret '{key_name}' not found, nothing to delete.")
+        return True
+    except InternalServerError as e:
+        # As per user feedback, 500 on GET might mean not found, apply same logic for DELETE
+        logger.warning(
+            f"Received 500 error when deleting secret '{key_name}'. Assuming it might not have existed. Response: {e}"
+        )
+        return True  # Consider deletion successful if it results in non-existence
+    except FireworksError as e:
+        logger.error(f"Error deleting secret '{key_name}': {e}")
         return False
 
 
@@ -319,8 +259,6 @@ def delete_fireworks_secret(
         logger.error(
             "CRITICAL: FIREWORKS_API_KEY and FIREWORKS_API_BASE must be correctly set in environment or .env file to run this test."
         )
-        import sys  # Make sure sys is imported if using sys.exit
-
         sys.exit(1)
 
     test_secret_key_name = "rewardkit-test-secret-delete-me"  # Changed to be valid
@@ -331,7 +269,7 @@ def delete_fireworks_secret(
 
     # 1. Ensure it doesn't exist initially (or delete if it does from a previous failed run)
     logger.info(f"\n[Test Step 0] Attempting to delete '{test_secret_key_name}' if it exists (cleanup)...")
-    delete_fireworks_secret(test_account_id, test_secret_key_name)
+    delete_fireworks_secret(account_id=test_account_id, key_name=test_secret_key_name)
     retrieved = get_fireworks_secret(test_account_id, test_secret_key_name)
     if retrieved is None:
         logger.info(f"Confirmed secret '{test_secret_key_name}' does not exist before creation test.")
@@ -351,8 +289,11 @@ def delete_fireworks_secret(
         logger.info(f"Retrieved secret metadata: {retrieved_after_create}")
         # Assert against the transformed keyName that's expected in the payload/response body
         expected_payload_key_name = test_secret_key_name.upper().replace("-", "_")
-        assert retrieved_after_create.get("keyName") == expected_payload_key_name
-        assert retrieved_after_create.get("value") == test_secret_value  # Also check value if returned
+        assert retrieved_after_create.key_name == expected_payload_key_name
+        # Note: value is typically not returned in GET responses for security reasons
+        # The value field will be None or empty string, so we don't assert on it
+        if retrieved_after_create.value:
+            logger.info(f"Note: Secret value was returned (unusual): {retrieved_after_create.value[:10]}...")
     else:
         logger.error(f"Failed to retrieve secret '{test_secret_key_name}' after creation.")