Files
ALwrity/docs/llm_gateway/Extending_the_Gateway.md

97 lines
3.2 KiB
Markdown

# 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**:
```python
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`:
```python
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 ...
```