gateway.py

import logging
import httpx
from urllib.parse import urlencode
import json
from fastapi import FastAPI, Request, HTTPException, Depends
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session
from sqlalchemy import func
import database
from encryption import decrypt_key

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("vapi-gateway")

app = FastAPI(title="vAPI Gateway", description="Local-first virtual API gateway")

client = httpx.AsyncClient(timeout=60.0)

@app.on_event("shutdown")
async def shutdown_event():
    await client.aclose()

@app.get("/health")
async def health():
    return {"status": "ok"}

@app.api_route("/{path:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH", "OPTIONS", "HEAD"])
async def gateway_proxy(path: str, request: Request, db: Session = Depends(database.get_db)):
    virtual_key = None
    
    query_params = dict(request.query_params)
    if "key" in query_params:
        virtual_key = query_params["key"]
    elif "api_key" in query_params:
        virtual_key = query_params["api_key"]
        
    if not virtual_key:
        auth_header = request.headers.get("Authorization") or request.headers.get("authorization")
        if auth_header:
            if auth_header.startswith("Bearer "):
                virtual_key = auth_header[7:]
            else:
                virtual_key = auth_header
        else:
            virtual_key = request.headers.get("x-api-key") or request.headers.get("x-goog-api-key")
            
    if not virtual_key:
        raise HTTPException(
            status_code=401,
            detail="API Key missing. Provide via query parameter 'key' or headers ('Authorization', 'x-api-key', 'x-goog-api-key')."
        )

    db_key = db.query(database.VirtualKey).filter(database.VirtualKey.key == virtual_key).first()
    if not db_key:
        raise HTTPException(status_code=401, detail="Invalid Virtual API Key.")
    
    if not db_key.is_active:
        raise HTTPException(status_code=403, detail="Virtual API Key is disabled.")
    
    provider = db_key.provider
    if not provider:
        raise HTTPException(status_code=400, detail="Associated provider not found.")
        
    if not provider.is_active:
        raise HTTPException(status_code=403, detail="Associated provider key is disabled.")

    if db_key.daily_budget is not None or db_key.monthly_budget is not None:
        from datetime import date
        today = date.today()
        if db_key.daily_budget is not None:
            daily_usage = db.query(func.sum(database.UsageLog.cost)).filter(
                database.UsageLog.virtual_key_id == db_key.id,
                func.date(database.UsageLog.timestamp) == today.strftime('%Y-%m-%d')
            ).scalar() or 0.0
            if daily_usage >= db_key.daily_budget:
                raise HTTPException(status_code=429, detail="Daily budget exceeded for this virtual key.")
                
        if db_key.monthly_budget is not None:
            month_str = today.strftime('%Y-%m')
            monthly_usage = db.query(func.sum(database.UsageLog.cost)).filter(
                database.UsageLog.virtual_key_id == db_key.id,
                func.strftime('%Y-%m', database.UsageLog.timestamp) == month_str
            ).scalar() or 0.0
            if monthly_usage >= db_key.monthly_budget:
                raise HTTPException(status_code=429, detail="Monthly budget exceeded for this virtual key.")

    try:
        real_key = decrypt_key(provider.api_key_encrypted)
    except Exception as e:
        logger.error(f"Failed to decrypt provider key: {e}")
        raise HTTPException(status_code=500, detail="Internal server error decrypting provider API key.")

    provider_name = provider.name.lower()
    
    headers = dict(request.headers)
    headers.pop("host", None)
    headers.pop("authorization", None)
    headers.pop("x-api-key", None)
    headers.pop("x-goog-api-key", None)

    if provider_name == "gemini":
        base_url = "https://generativelanguage.googleapis.com"
        target_path = f"/{path}" if not path.startswith("/") else path
        
        query_params.pop("key", None)
        query_params.pop("api_key", None)
        
        query_params["key"] = real_key
        
        encoded_query = urlencode(query_params)
        target_url = f"{base_url}{target_path}"
        if encoded_query:
            target_url = f"{target_url}?{encoded_query}"
            
    elif provider_name == "openai":
        base_url = "https://api.openai.com"
        target_path = f"/{path}" if not path.startswith("/") else path
        
        query_params.pop("key", None)
        query_params.pop("api_key", None)
        
        headers["authorization"] = f"Bearer {real_key}"
        
        encoded_query = urlencode(query_params)
        target_url = f"{base_url}{target_path}"
        if encoded_query:
            target_url = f"{target_url}?{encoded_query}"
            
    else:
        raise HTTPException(status_code=400, detail=f"Provider type '{provider.name}' not supported by gateway.")

    body = await request.body()

    if db_key.allowed_models:
        try:
            allowed_list = json.loads(db_key.allowed_models)
            if allowed_list and isinstance(allowed_list, list):
                model_requested = None
                if provider_name == "gemini":
                    import re
                    match = re.search(r'models/([^:]+)', path)
                    if match:
                        model_requested = match.group(1)
                elif provider_name == "openai":
                    if body:
                        payload = json.loads(body)
                        model_requested = payload.get("model")
                
                if model_requested and model_requested not in allowed_list:
                    raise HTTPException(status_code=403, detail=f"Model '{model_requested}' is not allowed for this virtual key.")
        except HTTPException:
            raise
        except Exception as e:
            logger.error(f"Error checking allowed models: {e}")

    import asyncio
    max_attempts = 3
    attempt = 0
    response = None
    
    while attempt < max_attempts:
        attempt += 1
        try:
            req = client.build_request(
                method=request.method,
                url=target_url,
                headers=headers,
                content=body
            )
            response = await client.send(req, stream=True)
            
            if response.status_code not in (429, 502, 503, 504):
                break
                
            if attempt == max_attempts:
                break
                
            await response.aclose()
            await asyncio.sleep(2 ** (attempt - 1))
            
        except Exception as e:
            if attempt == max_attempts:
                logger.error(f"Error forwarding request to {provider.name}: {e}")
                try:
                    db_log = database.SessionLocal()
                    log_entry = database.UsageLog(
                        virtual_key_id=db_key.id,
                        endpoint=path,
                        method=request.method,
                        status_code=502,
                        cost=1.0
                    )
                    db_log.add(log_entry)
                    db_log.commit()
                    db_log.close()
                except Exception as log_ex:
                    logger.error(f"Failed to write usage log: {log_ex}")
                raise HTTPException(status_code=502, detail="Error communicating with upstream provider.")
            await asyncio.sleep(2 ** (attempt - 1))

    async def response_generator():
        try:
            async for chunk in response.aiter_bytes():
                yield chunk
        finally:
            await response.aclose()
            try:
                db_log = database.SessionLocal()
                log_entry = database.UsageLog(
                    virtual_key_id=db_key.id,
                    endpoint=path,
                    method=request.method,
                    status_code=response.status_code,
                    cost=1.0
                )
                db_log.add(log_entry)
                db_log.commit()
                db_log.close()
            except Exception as log_ex:
                logger.error(f"Failed to write usage log: {log_ex}")

    exclude_headers = ["content-encoding", "content-length", "transfer-encoding", "connection"]
    response_headers = {
        k: v for k, v in response.headers.items()
        if k.lower() not in exclude_headers
    }

    return StreamingResponse(
        response_generator(),
        status_code=response.status_code,
        headers=response_headers
    )