indoxrouter 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

indoxRouter/client.py

@@ -1,286 +1,632 @@
-from typing import Dict, Optional, Any, List
+"""
+IndoxRouter Client Module
+
+This module provides a client for interacting with the IndoxRouter API, which serves as a unified
+interface to multiple AI providers and models. The client handles authentication, rate limiting,
+error handling, and provides a standardized response format across different AI services.
+
+The Client class offers methods for:
+- Authentication and session management
+- Making API requests with automatic token refresh
+- Accessing AI capabilities: chat completions, text completions, embeddings, and image generation
+- Retrieving information about available providers and models
+- Monitoring usage statistics
+
+Usage example:
+    ```python
+    from indoxRouter import Client
+    
+    # Initialize client with API key
+    client = Client(api_key="your_api_key")
+    
+    # Get available models
+    models = client.models()
+    
+    # Generate a chat completion
+    response = client.chat([
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Tell me a joke."}
+    ], provider="openai", model="gpt-3.5-turbo")
+    
+    # Generate text embeddings
+    embeddings = client.embeddings("This is a sample text")
+    
+    # Clean up resources when done
+    client.close()
+    ```
+
+The client can also be used as a context manager:
+    ```python
+    with Client(api_key="your_api_key") as client:
+        response = client.chat([{"role": "user", "content": "Hello!"}])
+    ```
+"""
+
 import os
-import sys
-import json
+import logging
+from datetime import datetime, timedelta
+from typing import Dict, List, Any, Optional, Union
 import requests
-
-# Add the parent directory to the path to make imports work correctly
-current_dir = os.path.dirname(os.path.abspath(__file__))
-parent_dir = os.path.dirname(current_dir)
-sys.path.append(parent_dir)
-
-# Use absolute imports
-from indoxRouter.utils.auth import AuthManager
-from indoxRouter.providers.base_provider import BaseProvider
-
+import jwt
+from uuid import uuid4
+
+from .exceptions import (
+    AuthenticationError,
+    NetworkError,
+    RateLimitError,
+    ProviderError,
+    ModelNotFoundError,
+    ProviderNotFoundError,
+    InvalidParametersError
+)
+from .models import (
+    ChatMessage,
+    ChatResponse,
+    CompletionResponse,
+    EmbeddingResponse,
+    ImageResponse,
+    ModelInfo,
+)
+from .config import get_config
+from .client_resourses import (
+    Chat,
+    Completions,
+    Embeddings,
+    Images,
+    Models,
+)
+from .constants import (
+    DEFAULT_API_VERSION,
+    DEFAULT_TIMEOUT,
+    DEFAULT_BASE_URL,
+    ERROR_INVALID_API_KEY,
+    ERROR_MODEL_NOT_FOUND,
+    ERROR_PROVIDER_NOT_FOUND,
+    ERROR_INVALID_PARAMETERS,
+    ERROR_RATE_LIMIT,
+    ERROR_REQUEST_FAILED,
+)
+
+logger = logging.getLogger(__name__)
 
 class Client:
     """
-    Client for making API requests to the IndoxRouter API.
+    Client for the IndoxRouter API that provides a unified interface to multiple AI providers.
+    
+    The Client class handles:
+    - Authentication and token management with automatic refresh
+    - Rate limiting and quota tracking
+    - Standardized error handling across providers
+    - Consistent response formatting
+    
+    This client provides access to various AI capabilities including:
+    - Chat completions (chat)
+    - Text completions (completion)
+    - Text embeddings (embeddings)
+    - Image generation (image)
+    
+    It also offers methods to retrieve information about available providers and models,
+    as well as usage statistics for the authenticated user.
+    
+    The client can be used directly or as a context manager with the 'with' statement.
     """
 
-    def __init__(self, api_key: str, base_url: str = None):
-        """
-        Initialize the client.
-
-        Args:
-            api_key: API key for authentication
-            base_url: Base URL for the API (default: http://localhost:8000)
-        """
-        self.api_key = api_key
-        self.base_url = base_url or "http://localhost:8000"
-        self.auth_manager = AuthManager()
-
-        # Verify the API key
-        self.user_data = self.auth_manager.verify_api_key(api_key)
-        if not self.user_data:
-            raise ValueError("Invalid API key")
-
-    def generate(
+    def __init__(
         self,
-        provider: str,
-        model: str,
-        prompt: str,
-        temperature: float = 0.7,
-        max_tokens: int = 1000,
-        **kwargs,
-    ) -> str:
+        api_key: Optional[str] = None,
+        base_url: str = f"{DEFAULT_BASE_URL}/{DEFAULT_API_VERSION}",
+        timeout: int = DEFAULT_TIMEOUT,
+        auto_refresh: bool = True,
+    ):
         """
-        Generate a response from a model.
+        Initialize the client with authentication and session management
 
         Args:
-            provider: Provider name
-            model: Model name
-            prompt: Prompt
-            temperature: Temperature
-            max_tokens: Maximum tokens
-            **kwargs: Additional parameters
-
-        Returns:
-            Model response
+            api_key: User's API key (default: INDOX_ROUTER_API_KEY env var)
+            base_url: Base URL for the API server
+            timeout: Request timeout in seconds
+            auto_refresh: Enable automatic token refresh
         """
-        url = f"{self.base_url}/v1/completions"
+        # Authentication setup
+        self.api_key = api_key or os.getenv("INDOX_ROUTER_API_KEY")
+        if not self.api_key:
+            raise AuthenticationError(ERROR_INVALID_API_KEY)
+
+        self.base_url = base_url
+        self.timeout = timeout
+        self.auto_refresh = auto_refresh
+        self.config = get_config()
+
+        # Session state management
+        self.session = requests.Session()
+        self._auth_token = None
+        self._token_expiry = None
+        self.user_info = None
+        self.rate_limits = {}
+
+        # Initialize resources
+        self._init_resources()
+        self._authenticate()
+
+    def _init_resources(self):
+        """Initialize resource controllers"""
+        self._chat = Chat(self)
+        self._completions = Completions(self)
+        self._embeddings = Embeddings(self)
+        self._images = Images(self)
+        self._models = Models(self)
+
+        # Backward compatibility
+        self.chat = self._chat
+        self.completions = self._completions
+        self.embeddings = self._embeddings
+        self.images = self._images
+        self.models = self._models
+
+    def _authenticate(self):
+        """Full authentication flow with the API server"""
+        try:
+            # Get authentication token
+            auth_response = self.session.post(
+                f"{self.base_url}/auth/token",
+                json={"api_key": self.api_key},
+                timeout=self.timeout
+            )
+            auth_response.raise_for_status()
+
+            auth_data = auth_response.json()
+            self._process_auth_data(auth_data)
+
+            logger.info("Authenticated as user: %s", self.user_info['email'])
+
+        except requests.exceptions.RequestException as e:
+            raise NetworkError(f"Authentication failed: {str(e)}") from e
+        except jwt.PyJWTError as e:
+            raise AuthenticationError(f"Token validation failed: {str(e)}") from e
+        except KeyError as e:
+            raise AuthenticationError(f"Invalid auth response: {str(e)}") from e
+
+    def _process_auth_data(self, auth_data: dict):
+        """Process authentication response data"""
+        # Validate and decode JWT
+        decoded_token = jwt.decode(
+            auth_data['access_token'],
+            self.config.JWT_PUBLIC_KEY,
+            algorithms=["RS256"],
+            audience="indox-router-api"
+        )
 
-        headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {self.api_key}",
-        }
+        # Update session state
+        self._auth_token = auth_data['access_token']
+        self._token_expiry = datetime.fromtimestamp(decoded_token['exp'])
+        self.user_info = decoded_token['user']
+        self.rate_limits = decoded_token.get('rate_limits', {})
+
+        # Set default headers
+        self.session.headers.update({
+            "Authorization": f"Bearer {self._auth_token}",
+            "Content-Type": "application/json"
+        })
+
+    def _check_auth(self):
+        """Validate authentication state"""
+        if datetime.now() >= self._token_expiry:
+            if self.auto_refresh:
+                self._refresh_token()
+            else:
+                raise AuthenticationError("Session expired")
+
+    def _refresh_token(self):
+        """Refresh access token using refresh token"""
+        try:
+            refresh_response = self.session.post(
+                f"{self.base_url}/auth/refresh",
+                timeout=self.timeout
+            )
+            refresh_response.raise_for_status()
 
-        data = {
-            "provider": provider,
-            "model": model,
-            "prompt": prompt,
-            "temperature": temperature,
-            "max_tokens": max_tokens,
-            **kwargs,
-        }
+            refresh_data = refresh_response.json()
+            self._process_auth_data(refresh_data)
 
-        response = requests.post(url, headers=headers, json=data)
+            logger.debug("Successfully refreshed authentication token")
 
-        if response.status_code != 200:
-            error_message = (
-                response.json().get("error", {}).get("message", "Unknown error")
-            )
-            raise Exception(f"Error: {error_message}")
+        except requests.exceptions.RequestException as e:
+            raise AuthenticationError(f"Token refresh failed: {str(e)}") from e
 
-        return response.json().get("choices", [{}])[0].get("text", "")
+    def _check_rate_limit(self, endpoint: str):
+        """Check rate limits for specific endpoint"""
+        limits = self.rate_limits.get(endpoint, {})
+        if limits.get('remaining', 1) <= 0:
+            reset_time = datetime.fromtimestamp(limits.get('reset', datetime.now().timestamp()))
+            raise RateLimitError(
+                f"{ERROR_RATE_LIMIT} for {endpoint}. Resets at {reset_time}",
+                reset_time=reset_time
+            )
 
-    def list_models(self, provider: Optional[str] = None) -> List[Dict[str, Any]]:
+    def _make_request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
         """
-        List available models.
-
-        Args:
-            provider: Provider name (optional)
-
-        Returns:
-            List of models
+        Unified request handler with:
+        - Authentication checks
+        - Rate limiting
+        - Error handling
+        - Response standardization
         """
-        url = f"{self.base_url}/v1/models"
+        self._check_auth()
+        self._check_rate_limit(endpoint)
 
-        if provider:
-            url += f"?provider={provider}"
+        request_id = uuid4().hex
+        url = f"{self.base_url}/{endpoint}"
+        start_time = datetime.now()
 
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
+        try:
+            response = self.session.request(
+                method=method,
+                url=url,
+                timeout=self.timeout,
+                **kwargs
+            )
+            duration = (datetime.now() - start_time).total_seconds()
+
+            # Update rate limits from headers
+            self._update_rate_limits(response.headers)
+
+            response.raise_for_status()
+            return self._format_success(response.json(), request_id, duration, endpoint)
+
+        except requests.exceptions.RequestException as e:
+            error_response = self._format_error(e, request_id, duration, endpoint)
+            
+            # Map HTTP errors to appropriate exception types
+            if hasattr(e, 'response') and e.response is not None:
+                status_code = e.response.status_code
+                error_data = {}
+                
+                try:
+                    error_data = e.response.json().get('error', {})
+                except ValueError:
+                    pass
+                
+                error_type = error_data.get('type', '')
+                
+                if status_code == 404:
+                    if 'provider' in error_type.lower():
+                        raise ProviderNotFoundError(f"{ERROR_PROVIDER_NOT_FOUND}: {error_data.get('message', str(e))}")
+                    elif 'model' in error_type.lower():
+                        raise ModelNotFoundError(f"{ERROR_MODEL_NOT_FOUND}: {error_data.get('message', str(e))}")
+                elif status_code == 400:
+                    raise InvalidParametersError(f"{ERROR_INVALID_PARAMETERS}: {error_data.get('message', str(e))}")
+                elif status_code == 429:
+                    reset_time = datetime.fromtimestamp(error_data.get('reset', datetime.now().timestamp() + 60))
+                    raise RateLimitError(f"{ERROR_RATE_LIMIT}: {error_data.get('message', str(e))}", reset_time=reset_time)
+                elif status_code >= 500:
+                    raise ProviderError(f"{ERROR_REQUEST_FAILED}: {error_data.get('message', str(e))}")
+                
+            raise NetworkError(f"Request failed: {str(e)}")
+
+    def _update_rate_limits(self, headers: Dict[str, str]):
+        """Update rate limits from response headers"""
+        for key, value in headers.items():
+            if key.startswith('X-Ratelimit-'):
+                endpoint = key.split('-')[-1]
+                self.rate_limits[endpoint] = {
+                    'limit': int(headers[f'X-Ratelimit-Limit-{endpoint}']),
+                    'remaining': int(headers[f'X-Ratelimit-Remaining-{endpoint}']),
+                    'reset': int(headers[f'X-Ratelimit-Reset-{endpoint}'])
+                }
+
+    def _format_success(self, data: dict, request_id: str, duration: float, endpoint: str) -> Dict[str, Any]:
+        """Standard success response format"""
+        return {
+            "request_id": request_id,
+            "data": data.get('result'),
+            "metadata": {
+                "endpoint": endpoint,
+                "duration": duration,
+                "timestamp": datetime.now().isoformat(),
+                **data.get('metadata', {})
+            },
+            "error": None
         }
 
-        response = requests.get(url, headers=headers)
+    def _format_error(self, error: Exception, request_id: str, duration: float, endpoint: str) -> Dict[str, Any]:
+        """Standard error response format"""
+        error_info = {
+            "request_id": request_id,
+            "data": None,
+            "metadata": {
+                "endpoint": endpoint,
+                "duration": duration,
+                "timestamp": datetime.now().isoformat()
+            },
+            "error": {
+                "type": error.__class__.__name__,
+                "message": str(error),
+                "details": {}
+            }
+        }
 
-        if response.status_code != 200:
-            error_message = (
-                response.json().get("error", {}).get("message", "Unknown error")
-            )
-            raise Exception(f"Error: {error_message}")
+        if isinstance(error, requests.exceptions.HTTPError):
+            error_info["error"]["code"] = error.response.status_code
+            error_info["error"]["details"] = error.response.json().get('error', {})
 
-        return response.json().get("data", [])
+        logger.error("Request %s failed: %s", request_id, error)
+        return error_info
 
-    def list_providers(self) -> List[str]:
+    # Resource proxy methods
+    def providers(self) -> List[Dict[str, Any]]:
         """
-        List available providers.
-
+        Get a list of all available AI providers.
+        
         Returns:
-            List of providers
+            List[Dict[str, Any]]: A list of provider information dictionaries, each containing
+            details such as provider ID, name, description, and supported features.
+        
+        Example:
+            ```python
+            providers = client.providers()
+            for provider in providers:
+                print(f"{provider['id']}: {provider['name']}")
+            ```
         """
-        url = f"{self.base_url}/v1/providers"
+        return self._make_request('GET', 'providers')['data']
 
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
-        }
-
-        response = requests.get(url, headers=headers)
-
-        if response.status_code != 200:
-            error_message = (
-                response.json().get("error", {}).get("message", "Unknown error")
-            )
-            raise Exception(f"Error: {error_message}")
-
-        return response.json().get("data", [])
-
-    def _parse_model_name(self, model_name: str) -> tuple:
+    def models(self, provider: Optional[str] = None) -> Dict[str, List[Dict[str, Any]]]:
         """
-        Parse model name into provider and model parts
-
+        Get available models, optionally filtered by provider.
+        
         Args:
-            model_name: Full model name (e.g., 'openai/gpt-4')
-
+            provider (Optional[str]): Provider ID to filter models by. If None, returns models from all providers.
+        
         Returns:
-            Tuple of (provider_name, model_part)
+            Dict[str, List[Dict[str, Any]]]: A dictionary mapping provider IDs to lists of model information.
+            Each model information dictionary contains details such as model ID, name, capabilities, and pricing.
+        
+        Example:
+            ```python
+            # Get all models
+            all_models = client.models()
+            
+            # Get models from a specific provider
+            openai_models = client.models(provider="openai")
+            ```
         """
-        if "/" not in model_name:
-            raise ValueError(
-                f"Invalid model name format: {model_name}. Expected format: 'provider/model'"
-            )
-
-        provider_name, model_part = model_name.split("/", 1)
-        return provider_name, model_part
+        params = {"provider": provider} if provider else None
+        return self._make_request('GET', 'models', params=params)['data']
 
-    def _load_provider_class(self, provider_name: str):
+    def model_info(self, provider: str, model: str) -> ModelInfo:
         """
-        Dynamically load provider class
-
+        Get detailed information about a specific model.
+        
         Args:
-            provider_name: Name of the provider
-
+            provider (str): Provider ID (e.g., "openai", "anthropic")
+            model (str): Model ID (e.g., "gpt-4", "claude-2")
+        
         Returns:
-            Provider class
+            ModelInfo: Detailed information about the model, including capabilities,
+            pricing, rate limits, and other provider-specific details.
+        
+        Raises:
+            ModelNotFoundError: If the specified model doesn't exist or isn't available
+            ProviderNotFoundError: If the specified provider doesn't exist
+        
+        Example:
+            ```python
+            model_details = client.model_info(provider="openai", model="gpt-4")
+            print(f"Context window: {model_details.context_window} tokens")
+            ```
         """
         try:
-            # Import the provider module dynamically
-            module_path = f".providers.{provider_name}"
-            provider_module = __import__(
-                module_path, fromlist=["Provider"], globals=globals()
-            )
-            return provider_module.Provider
-        except (ImportError, AttributeError) as e:
-            raise ValueError(f"Provider not supported: {provider_name}") from e
+            return self._make_request('GET', f'models/{provider}/{model}')['data']
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 404:
+                raise ModelNotFoundError(f"{ERROR_MODEL_NOT_FOUND}: '{model}' for provider '{provider}'")
+            raise
 
-    def _get_provider(self, model_name: str) -> BaseProvider:
+    def chat(self, messages: List[Union[Dict[str, str], ChatMessage]], **kwargs) -> ChatResponse:
         """
-        Get provider instance with cached credentials
-
+        Generate a chat completion from a series of messages.
+        
         Args:
-            model_name: Full model name (e.g., 'openai/gpt-4')
-
+            messages (List[Union[Dict[str, str], ChatMessage]]): A list of messages, where each message
+                is either a dictionary with 'role' and 'content' keys, or a ChatMessage object.
+            **kwargs: Additional parameters to pass to the provider, such as:
+                - provider (str): Provider ID (e.g., "openai", "anthropic")
+                - model (str): Model ID (e.g., "gpt-4", "claude-2")
+                - temperature (float): Controls randomness (0.0-1.0)
+                - max_tokens (int): Maximum number of tokens to generate
+                - stream (bool): Whether to stream the response
+        
         Returns:
-            Provider instance
+            ChatResponse: The generated chat completion response, containing the assistant's message,
+            usage statistics, and other metadata.
+        
+        Example:
+            ```python
+            response = client.chat([
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Tell me about AI."}
+            ], provider="openai", model="gpt-4")
+            
+            print(response.message.content)
+            ```
         """
-        if model_name in self.provider_cache:
-            return self.provider_cache[model_name]
-
-        provider_name, model_part = self._parse_model_name(model_name)
-        provider_class = self._load_provider_class(provider_name)
+        return self._chat(messages, **kwargs)
 
-        # Get provider API key from secure storage
-        provider_api_key = self._get_provider_credentials(provider_name)
-
-        instance = provider_class(api_key=provider_api_key, model_name=model_part)
-        self.provider_cache[model_name] = instance
-        return instance
-
-    def _get_provider_credentials(self, provider_name: str) -> str:
+    def completion(self, prompt: str, **kwargs) -> CompletionResponse:
         """
-        Retrieve provider API key from secure storage
-
+        Generate a text completion from a prompt.
+        
         Args:
-            provider_name: Name of the provider
-
+            prompt (str): The text prompt to complete
+            **kwargs: Additional parameters to pass to the provider, such as:
+                - provider (str): Provider ID (e.g., "openai", "cohere")
+                - model (str): Model ID (e.g., "text-davinci-003", "command")
+                - temperature (float): Controls randomness (0.0-1.0)
+                - max_tokens (int): Maximum number of tokens to generate
+                - stream (bool): Whether to stream the response
+        
         Returns:
-            Provider API key
-        """
-        # Implement your secure credential storage (e.g., AWS Secrets Manager)
-        # Example using environment variables:
-        env_var = f"{provider_name.upper()}_API_KEY"
-        if env_var not in os.environ:
-            raise ValueError(
-                f"Missing API key for provider: {provider_name}. Set {env_var} environment variable."
+            CompletionResponse: The generated completion response, containing the completed text,
+            usage statistics, and other metadata.
+        
+        Example:
+            ```python
+            response = client.completion(
+                "Once upon a time,", 
+                provider="openai", 
+                model="text-davinci-003"
             )
-
-        return os.environ[env_var]
-
-    def generate(self, model_name: str, prompt: str, **kwargs) -> Dict[str, Any]:
+            
+            print(response.text)
+            ```
         """
-        Generate completion with credit handling
+        return self._completions(prompt, **kwargs)
 
+    def embeddings(self, text: Union[str, List[str]], **kwargs) -> EmbeddingResponse:
+        """
+        Generate embeddings for text.
+        
         Args:
-            model_name: Provider/model name (e.g., 'openai/gpt-4')
-            prompt: User input prompt
-            **kwargs: Generation parameters
-
+            text (Union[str, List[str]]): Text or list of texts to generate embeddings for
+            **kwargs: Additional parameters to pass to the provider, such as:
+                - provider (str): Provider ID (e.g., "openai", "cohere")
+                - model (str): Model ID (e.g., "text-embedding-ada-002")
+                - dimensions (int): Desired dimensionality of the embeddings (if supported)
+        
         Returns:
-            Dictionary with response and credit information
+            EmbeddingResponse: The generated embeddings response, containing the vector representations,
+            usage statistics, and other metadata.
+        
+        Example:
+            ```python
+            response = client.embeddings(
+                ["Hello world", "AI is amazing"],
+                provider="openai",
+                model="text-embedding-ada-002"
+            )
+            
+            for i, embedding in enumerate(response.embeddings):
+                print(f"Embedding {i} has {len(embedding)} dimensions")
+            ```
         """
-        provider = self._get_provider(model_name)
+        return self._embeddings(text, **kwargs)
 
-        # Estimate max possible cost
-        max_tokens = kwargs.get("max_tokens", 2048)
-        estimated_cost = provider.estimate_cost(prompt, max_tokens)
-
-        # Check balance
-        if self.user_data["balance"] < estimated_cost:
-            raise ValueError(
-                f"Insufficient credits. Required: {estimated_cost:.6f}, Available: {self.user_data['balance']:.6f}"
+    def image(self, prompt: str, **kwargs) -> ImageResponse:
+        """
+        Generate an image from a text prompt.
+        
+        Args:
+            prompt (str): The text description of the image to generate
+            **kwargs: Additional parameters to pass to the provider, such as:
+                - provider (str): Provider ID (e.g., "openai", "stability")
+                - model (str): Model ID (e.g., "dall-e-3", "stable-diffusion-xl")
+                - size (str): Image size (e.g., "1024x1024")
+                - quality (str): Image quality (e.g., "standard", "hd")
+                - style (str): Image style (e.g., "vivid", "natural")
+                - response_format (str): Format of the response (e.g., "url", "b64_json")
+        
+        Returns:
+            ImageResponse: The generated image response, containing the image data or URLs,
+            usage statistics, and other metadata.
+        
+        Example:
+            ```python
+            response = client.image(
+                "A serene landscape with mountains and a lake",
+                provider="openai",
+                model="dall-e-3",
+                size="1024x1024"
             )
-
-        # Make API call
-        response = provider.generate(prompt, **kwargs)
-
-        # Deduct actual cost
-        success = self.auth_manager.deduct_credits(
-            self.user_data["id"], response["cost"]
-        )
-
-        if not success:
-            raise RuntimeError("Credit deduction failed")
-
-        # Get updated user data
-        self.user_data = self.auth_manager.get_user_by_id(self.user_data["id"])
-
-        return {
-            "text": response["text"],
-            "cost": response["cost"],
-            "remaining_credits": self.user_data["balance"],
-            "model": model_name,
-        }
-
-    def get_balance(self) -> float:
+            
+            print(f"Image URL: {response.url}")
+            ```
         """
-        Get current user balance
+        return self._images(prompt, **kwargs)
 
+    # Additional features
+    def get_usage(self) -> Dict[str, Any]:
+        """
+        Get current usage statistics for the authenticated user.
+        
         Returns:
-            Current credit balance
+            Dict[str, Any]: Usage statistics including token counts, request counts,
+            billing information, and quota limits.
+        
+        Example:
+            ```python
+            usage = client.get_usage()
+            print(f"Total tokens used: {usage['total_tokens']}")
+            print(f"Remaining quota: {usage['remaining_quota']}")
+            ```
         """
-        # Refresh user data to get the latest balance
-        self.user_data = self.auth_manager.get_user_by_id(self.user_data["id"])
-        return self.user_data["balance"]
+        return self._make_request('GET', 'usage')['data']
 
     def get_user_info(self) -> Dict[str, Any]:
         """
-        Get current user information
+        Get information about the authenticated user.
+        
+        Returns:
+            Dict[str, Any]: User information including ID, name, email, account type,
+            subscription details, and other user-specific data.
+        
+        Example:
+            ```python
+            user = client.get_user_info()
+            print(f"User ID: {user['id']}")
+            print(f"Account type: {user['account_type']}")
+            ```
+        """
+        return self.user_info.copy()
 
+    def close(self):
+        """
+        Clean up client resources and close the session.
+        
+        This method should be called when the client is no longer needed to ensure
+        proper cleanup of resources, particularly the HTTP session.
+        
+        Example:
+            ```python
+            client = Client(api_key="your_api_key")
+            # Use the client...
+            client.close()  # Clean up when done
+            ```
+        """
+        self.session.close()
+        logger.info("Client session closed")
+
+    def __enter__(self):
+        """
+        Enter the context manager.
+        
+        This method enables the client to be used as a context manager with the 'with' statement.
+        
         Returns:
-            User information dictionary
+            Client: The client instance.
+        
+        Example:
+            ```python
+            with Client(api_key="your_api_key") as client:
+                # Use the client within this block
+                response = client.chat([{"role": "user", "content": "Hello!"}])
+            # Client is automatically closed when exiting the block
+            ```
         """
-        # Refresh user data
-        self.user_data = self.auth_manager.get_user_by_id(self.user_data["id"])
-        return self.user_data
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """
+        Exit the context manager.
+        
+        This method is called when exiting a 'with' block. It ensures the client
+        is properly closed, even if an exception occurs within the block.
+        
+        Args:
+            exc_type: The exception type, if an exception was raised in the with block, otherwise None
+            exc_val: The exception value, if an exception was raised in the with block, otherwise None
+            exc_tb: The traceback, if an exception was raised in the with block, otherwise None
+        """
+        self.close()
+
+# Backward compatibility
+IndoxRouter = Client