agentpay_sdk/client.py

# Simple client to interact with AgentHub API 

import requests
from typing import Optional, Dict, Any
import json # For decoding JSON response

from .types import ConsumptionResult
from .exceptions import AgentPayAPIError, AgentPayConnectionError # Import custom exceptions

API_BASE_URL = "https://api.agentpay.me"
API_VERSION_PREFIX = "/v1"

class AgentPayClient:
    """Client for interacting with the AgentPay API.

    Handles authentication using a Service Token and provides methods
    for validating user API keys and consuming balance.
    """

    def __init__(
        self,
        service_token: str,
        api_url: Optional[str] = None,
    ):
        """Initializes the AgentPay Client.

        Args:
            service_token: The unique token issued to your service by AgentHub.
            api_url: Optional. Manually override the base API URL.
        """
        if not service_token:
            raise ValueError("service_token is required.")
        
        self.service_token = service_token
        
        if api_url:
            # Use manual override if provided
            self.base_api_url = api_url.rstrip('/')
        else:
            self.base_api_url = API_BASE_URL
            
        # Add the /api/v1 prefix
        self.base_api_url += API_VERSION_PREFIX
        
        # Prepare headers for requests
        self.headers = {
            "Authorization": f"Bearer {self.service_token}",
            "Content-Type": "application/json",
            "Accept": "application/json",
        } 

    def _request(self, method: str, endpoint: str, data: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Internal helper to make requests to the AgentPay API.

        Args:
            method: HTTP method (e.g., "POST", "GET").
            endpoint: API endpoint path (e.g., "/balances/consume").
            data: Optional dictionary payload for POST/PUT/PATCH requests.

        Returns:
            The JSON response body as a dictionary.

        Raises:
            AgentPayConnectionError: If a network error occurs.
            AgentPayAPIError: If the API returns an error status code.
        """
        full_url = self.base_api_url + endpoint
        
        try:
            response = requests.request(
                method=method,
                url=full_url,
                headers=self.headers,
                json=data, # requests library handles JSON serialization
                timeout=10 # Add a reasonable timeout (e.g., 10 seconds)
            )
            
            # Raise exception for bad status codes (4xx or 5xx)
            response.raise_for_status()

            # Attempt to parse successful response as JSON
            # Handle cases where response might be empty (e.g., 204 No Content)
            if response.status_code == 204:
                return {} # Return empty dict for No Content
            
            return response.json()
        
        except requests.exceptions.Timeout as e:
            raise AgentPayConnectionError(f"Request timed out connecting to {full_url}: {e}") from e
        except requests.exceptions.ConnectionError as e:
             raise AgentPayConnectionError(f"Could not connect to {full_url}: {e}") from e
        except requests.exceptions.HTTPError as e:
            # Handle API errors (4xx, 5xx)
            status_code = e.response.status_code
            error_code = None
            error_message = None
            try:
                # Attempt to parse error details from response body
                error_data = e.response.json()
                error_code = error_data.get("error_code")
                error_message = error_data.get("error_message") or error_data.get("detail") # FastAPI uses detail sometimes
            except json.JSONDecodeError:
                # If response body is not JSON or empty
                error_message = e.response.text or f"HTTP error {status_code}"
                
            raise AgentPayAPIError(
                status_code=status_code, 
                error_code=error_code, 
                error_message=error_message
            ) from e
        except json.JSONDecodeError as e:
             # Handle cases where a 2xx response is not valid JSON
             raise AgentPayAPIError(
                 status_code=response.status_code, # Use the status code from the original response
                 error_message=f"Failed to decode successful JSON response: {e}"
            ) from e
        except requests.exceptions.RequestException as e:
            # Catch any other requests-related errors
            raise AgentPayConnectionError(f"An unexpected network error occurred: {e}") from e 

    def consume(
        self,
        api_key: str,
        amount_cents: int,
        usage_event_id: str,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> ConsumptionResult:
        """Consumes balance for a user associated with an API key.

        Args:
            api_key: The User's API key to consume balance against.
            amount_cents: The amount to consume in cents (must be >= 0).
            usage_event_id: A unique identifier (e.g., UUID string) for this specific 
                            consumption attempt, used for idempotency.
            metadata: Optional dictionary containing additional data to associate
                      with the usage event.

        Returns:
            A ConsumptionResult object indicating success or failure.
        """
        if amount_cents < 0:
            # Prevent making an API call with an invalid amount
            # The API would reject this anyway, but better to catch early.
            return ConsumptionResult(
                success=False,
                error_code="INVALID_INPUT",
                error_message="Consumption amount cannot be negative."
            )
            
        endpoint = "/balances/consume"
        payload = {
            "api_key": api_key,
            "amount_cents": amount_cents,
            "usage_event_id": usage_event_id,
        }
        if metadata:
            payload["metadata"] = metadata
            
        try:
            # _request raises AgentPayAPIError for 4xx/5xx responses
            # and AgentPayConnectionError for network issues.
            response_data = self._request("POST", endpoint, data=payload)
            
            # If _request completes without raising an exception, it implies a 2xx response.
            # The API's 200 OK response for consume just confirms success.
            return ConsumptionResult(success=True)
        
        except AgentPayAPIError as e:
            # API returned a non-2xx status code (e.g., 402, 403, 409)
            return ConsumptionResult(
                success=False,
                error_code=e.error_code, 
                error_message=e.error_message
            )
        except AgentPayConnectionError as e:
            # Network error occurred during the request
            return ConsumptionResult(
                success=False,
                error_code="NETWORK_ERROR", 
                error_message=str(e)
            )
        except Exception as e:
            # Catch any other unexpected errors during the process
            return ConsumptionResult(
                success=False,
                error_code="SDK_ERROR", 
                error_message=f"An unexpected error occurred in the SDK: {e}"
            ) 

    def validate_api_key(self, api_key: str) -> bool:
        """Validates if a User API Key is active and valid for this service.

        Args:
            api_key: The User's API key to validate.

        Returns:
            True if the API key is valid for the service, False otherwise.
            Returns False if any API or network error occurs during validation.
        """
        endpoint = "/api_keys/validate"
        payload = {"api_key": api_key}
        
        try:
            response_data = self._request("POST", endpoint, data=payload)
            
            # Expecting a response like {"is_valid": true | false}
            is_valid = response_data.get("is_valid")
            
            if isinstance(is_valid, bool):
                return is_valid
            else:
                # Unexpected response format from the API
                # Log this? For now, treat unexpected format as invalid.
                print(f"Warning: Unexpected response format from /validate endpoint: {response_data}")
                return False
        
        except (AgentPayAPIError, AgentPayConnectionError) as e:
            # If any API or network error occurs, treat the key as invalid for safety.
            # Log the error for debugging if needed.
            print(f"API or Connection Error during validation: {e}")
            return False
        except Exception as e:
            # Catch any other unexpected errors
            print(f"Unexpected SDK error during validation: {e}")
            return False