ALwrity/docs/llm_gateway/Extending_the_Gateway.md

Extending the LLM Gateway

This guide provides a checklist and templates for adding new providers or modalities to the ALwrity LLM Gateway.

Checklist

1. Define the Provider Interface:

   • Create a new module in backend/services/llm_providers/<modality>/.
   • Define input options dataclass (e.g., MyNewProviderOptions).
   • Implement the standard Protocol (e.g., ImageGenerationProvider).

2. Configuration:

   • Add necessary API keys to .env.example and APIKeyManager.
   • Add new provider enum to backend/models/subscription_models.py (APIProvider).

3. Pricing & Usage:

   • Add default pricing in PricingService or migration script.
   • Ensure UsageSummary table has columns for this provider (if it's a major one) or map it to a generic category.

4. Integration:

   • Register the provider in the main entry point (e.g., main_image_generation.py).
   • Update preflight.py to handle cost estimation for this provider.

5. Frontend:

   • Update billingService.ts to handle the new provider key in usage stats (if applicable).
   • Add provider icon/color in billingService.ts.

Skeleton Template (Python)

Here is a template for a new Image Generation Provider:

from __future__ import annotations
import os
from typing import Optional, Dict, Any
from dataclasses import dataclass
from .base import ImageGenerationOptions, ImageGenerationResult, ImageGenerationProvider
from utils.logger_utils import get_service_logger

logger = get_service_logger("image_generation.my_new_provider")

class MyNewProvider(ImageGenerationProvider):
    """
    My New Provider implementation.
    """

    def __init__(self, api_key: Optional[str] = None):
        self.api_key = api_key or os.getenv("MY_PROVIDER_API_KEY")
        if not self.api_key:
            raise RuntimeError("MY_PROVIDER_API_KEY is required")
        # Initialize client here

    def generate(self, options: ImageGenerationOptions) -> ImageGenerationResult:
        logger.info(f"Generating image with MyNewProvider: {options.prompt[:50]}...")
        
        try:
            # 1. Call External API
            # response = client.generate(...)
            
            # 2. Process Response (Mock)
            image_bytes = b"fake_image_data" 
            width = options.width
            height = options.height
            
            # 3. Return Standard Result
            return ImageGenerationResult(
                image_bytes=image_bytes,
                width=width,
                height=height,
                provider="my_new_provider",
                model=options.model or "default-model",
                seed=options.seed,
                metadata={"raw_response": "..."}
            )
            
        except Exception as e:
            logger.error(f"Generation failed: {e}")
            raise

Skeleton Template (Route Integration)

In main_image_generation.py:

from .image_generation.my_new_provider import MyNewProvider

def generate_image(prompt: str, provider: str, ...):
    # ... existing code ...
    
    if provider == "my_new_provider":
        service = MyNewProvider()
        result = service.generate(options)
        
    # ... existing code ...