feat: OpenAI-compatible API with token auth and admin management

admin_router.py

@@ -0,0 +1,93 @@
+from fastapi import APIRouter, HTTPException, BackgroundTasks
+from pydantic import BaseModel
+from typing import List, Optional
+import secrets
+import uuid
+
+from db import get_supabase
+
+router = APIRouter(prefix="/admin", tags=["Admin"])
+
+# --- Models ---
+
+class APIKey(BaseModel):
+    id: str
+    name: str
+    token: str
+    usage_tokens: int
+    limit_tokens: int
+    created_at: str
+    is_active: bool
+
+class CreateKeyRequest(BaseModel):
+    name: str
+    limit_tokens: Optional[int] = 1000000
+
+# --- Endpoints ---
+
+@router.get("/keys", response_model=List[APIKey])
+async def list_keys():
+    """List all API keys."""
+    supabase = get_supabase()
+    if not supabase:
+        raise HTTPException(status_code=503, detail="Database unavailable")
+        
+    try:
+        res = supabase.table("api_keys").select("*").order("created_at", desc=True).execute()
+        return res.data
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.post("/keys", response_model=APIKey)
+async def create_key(req: CreateKeyRequest):
+    """Create a new API key."""
+    supabase = get_supabase()
+    if not supabase:
+        raise HTTPException(status_code=503, detail="Database unavailable")
+    
+    # Generate a secure token
+    token = f"sk-kai-{secrets.token_urlsafe(16)}"
+    
+    new_key = {
+        "name": req.name,
+        "token": token,
+        "limit_tokens": req.limit_tokens,
+        "usage_tokens": 0,
+        "is_active": True
+    }
+    
+    try:
+        res = supabase.table("api_keys").insert(new_key).execute()
+        if res.data:
+            return res.data[0]
+        raise HTTPException(status_code=500, detail="Failed to create key")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.delete("/keys/{key_id}")
+async def revoke_key(key_id: str):
+    """Revoke (delete) an API key."""
+    supabase = get_supabase()
+    if not supabase:
+        raise HTTPException(status_code=503, detail="Database unavailable")
+        
+    try:
+        # Check if exists first? Or just delete.
+        # Hard delete for now, or soft delete if we had is_active column logic in router update, but delete is cleaner for management
+        res = supabase.table("api_keys").delete().eq("id", key_id).execute()
+        return {"status": "success", "deleted": key_id}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@router.post("/keys/{key_id}/reset")
+async def reset_usage(key_id: str):
+    """Reset usage for a key."""
+    supabase = get_supabase()
+    if not supabase:
+        raise HTTPException(status_code=503, detail="Database unavailable")
+        
+    try:
+        supabase.table("api_keys").update({"usage_tokens": 0}).eq("id", key_id).execute()
+        return {"status": "reset"}
+    except Exception as e:
+         raise HTTPException(status_code=500, detail=str(e))

config.py

@@ -63,6 +63,9 @@ POLLINATIONS_MODEL_NAMES = {
     "midijourney": "midijourney",
 }
 
+# API Keys
+DEMO_API_KEY = "sk-kai-demo-public"
+
 # Models per provider (for /models endpoint)
 PROVIDER_MODELS = {
     "g4f": [

db.py

@@ -0,0 +1,19 @@
+import logging
+from supabase import create_client, Client
+from config import SUPABASE_URL, SUPABASE_KEY
+
+logger = logging.getLogger("kai_api.db")
+
+try:
+    if SUPABASE_URL and SUPABASE_KEY:
+        supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
+        logger.info("✅ Supabase client initialized")
+    else:
+        supabase = None
+        logger.warning("⚠️ Supabase credentials missing (check config.py)")
+except Exception as e:
+    supabase = None
+    logger.error(f"❌ Failed to initialize Supabase: {e}")
+
+def get_supabase() -> Client:
+    return supabase

main.py

@@ -38,8 +38,17 @@ from models import (
     HealthResponse,
     ProviderHealth,
 )
-from engine import AIEngine
-from search_engine import SearchEngine
+from models import (
+    ChatRequest,
+    ChatResponse,
+    ErrorResponse,
+    ModelsResponse,
+    HealthResponse,
+    ProviderHealth,
+)
+from services import engine, search_engine
+from v1_router import router as v1_router
+from admin_router import router as admin_router
 
 # ---------- Logging ----------
 logging.basicConfig(
@@ -69,9 +78,13 @@ app.add_middleware(
     allow_headers=CORS_HEADERS,
 )
 
-# AI Engine (initialized once, but each request is stateless internally)
-engine = AIEngine()
-search_engine = SearchEngine()
+# AI Engine (initialized via services.py)
+# engine = AIEngine() -> Moved to services.py
+# search_engine = SearchEngine() -> Moved to services.py
+
+# Include OpenAI Router
+app.include_router(v1_router)
+app.include_router(admin_router)
 
 
 # ---------- Admin Routes ----------

services.py

@@ -0,0 +1,6 @@
+from engine import AIEngine
+from search_engine import SearchEngine
+
+# Singleton instances to be shared across modules
+engine = AIEngine()
+search_engine = SearchEngine()

static/admin.html

@@ -206,7 +206,33 @@
         </div>
     </div>
 
-    <!-- Removed Charts & Tables (Moved to Main Page) -->
+    <!-- API Key Management Section -->
+    <div class="card" style="margin-bottom: 30px;">
+        <div style="display:flex; justify-content:space-between; align-items:center; margin-bottom:15px;">
+            <h2>API Key Management</h2>
+            <button onclick="createKey()"
+                style="background: var(--success); color: white; border: none; padding: 8px 16px; border-radius: 6px; font-weight: 600; cursor: pointer;">
+                + Create New Key
+            </button>
+        </div>
+
+        <table id="keys-table">
+            <thead>
+                <tr>
+                    <th>Name</th>
+                    <th>Token Prefix</th>
+                    <th>Usage / Limit</th>
+                    <th>Created</th>
+                    <th>Action</th>
+                </tr>
+            </thead>
+            <tbody id="keys-list">
+                <tr>
+                    <td colspan="5" style="text-align:center; padding:20px;">Loading keys...</td>
+                </tr>
+            </tbody>
+        </table>
+    </div>
 
     <div id="error-console"
         style="display:none; background:#ef4444; color:white; padding:10px; margin-bottom:20px; border-radius:8px; font-family:monospace;">
@@ -280,6 +306,78 @@
                 btn.disabled = false;
             }
         }
+        async function loadKeys() {
+            try {
+                const res = await fetch('/admin/keys');
+                const keys = await res.json();
+
+                const tbody = document.getElementById('keys-list');
+                tbody.innerHTML = '';
+
+                keys.forEach(key => {
+                    const usagePercent = key.limit_tokens > 0 ? Math.round((key.usage_tokens / key.limit_tokens) * 100) : 0;
+                    const color = usagePercent > 90 ? 'red' : (usagePercent > 50 ? 'orange' : '#22c55e');
+
+                    const row = `
+                        <tr>
+                            <td style="font-weight:bold; color:white;">${key.name}</td>
+                            <td style="font-family:monospace; color:var(--text-muted);">${key.token.substring(0, 10)}...</td>
+                            <td>
+                                <div style="display:flex; align-items:center; gap:10px;">
+                                    <span style="min-width:100px;">${key.usage_tokens.toLocaleString()} / ${key.limit_tokens.toLocaleString()}</span>
+                                    <div style="width:100px; height:6px; background:#333; border-radius:3px; overflow:hidden;">
+                                        <div style="height:100%; width:${Math.min(usagePercent, 100)}%; background:${color};"></div>
+                                    </div>
+                                </div>
+                            </td>
+                            <td>${new Date(key.created_at).toLocaleDateString()}</td>
+                            <td>
+                                <button onclick="revokeKey('${key.id}')" style="background:var(--error); color:white; border:none; padding:4px 8px; border-radius:4px; cursor:pointer;">Revoke</button>
+                            </td>
+                        </tr>
+                    `;
+                    tbody.innerHTML += row;
+                });
+            } catch (e) {
+                console.error("Failed to load keys", e);
+            }
+        }
+
+        async function createKey() {
+            const name = prompt("Enter Name for new API Key (e.g. 'John Doe'):");
+            if (!name) return;
+
+            try {
+                const res = await fetch('/admin/keys', {
+                    method: 'POST',
+                    headers: { 'Content-Type': 'application/json' },
+                    body: JSON.stringify({ name: name, limit_tokens: 1000000 })
+                });
+
+                if (res.ok) {
+                    const key = await res.json();
+                    alert(`✅ Key Created!\n\nToken: ${key.token}\n\nSAVE THIS NOW. IT IS SHOWN ONLY ONCE.`);
+                    loadKeys();
+                } else {
+                    alert("Failed to create key");
+                }
+            } catch (e) {
+                alert("Error: " + e.message);
+            }
+        }
+
+        async function revokeKey(id) {
+            if (!confirm("Are you sure you want to delete this key? Access will be immediately revoked.")) return;
+            try {
+                await fetch(`/admin/keys/${id}`, { method: 'DELETE' });
+                loadKeys();
+            } catch (e) {
+                alert("Error: " + e.message);
+            }
+        }
+
+        // Initial Load
+        loadKeys();
     </script>
 </body>
 

supabase_schema.sql

@@ -0,0 +1,13 @@
+-- Create API Keys Table
+create table public.api_keys (
+  id uuid default gen_random_uuid() primary key,
+  created_at timestamp with time zone default timezone('utc'::text, now()) not null,
+  name text not null,
+  token text not null unique,
+  usage_tokens bigint default 0,
+  limit_tokens bigint default 1000000, -- Default 1M tokens
+  is_active boolean default true
+);
+
+-- Indexes for performance
+create index idx_api_keys_token on public.api_keys(token);

utils.py

@@ -0,0 +1,37 @@
+import logging
+
+logger = logging.getLogger("kai_api.utils")
+
+def estimate_tokens(text: str) -> int:
+    """
+    Estimate token count using a simple rule of thumb:
+    1 word = ~1.33 tokens (English).
+    Or roughly 4 chars = 1 token.
+    
+    We'll use (len(text) / 4) as a fast approximation.
+    Minimum 1 token if text exists.
+    """
+    if not text:
+        return 0
+    
+    count = int(len(text) / 4)
+    return max(1, count)
+
+def calculate_usage(messages: list[dict], response_text: str) -> dict:
+    """
+    Calculate prompt_tokens and completion_tokens.
+    """
+    prompt_text = ""
+    for msg in messages:
+        content = msg.get("content", "")
+        if isinstance(content, str):
+            prompt_text += content + "\n"
+            
+    prompt_tokens = estimate_tokens(prompt_text)
+    completion_tokens = estimate_tokens(response_text)
+    
+    return {
+        "prompt_tokens": prompt_tokens,
+        "completion_tokens": completion_tokens,
+        "total_tokens": prompt_tokens + completion_tokens
+    }

v1_router.py

@@ -0,0 +1,212 @@
+from fastapi import APIRouter, Depends, HTTPException, Header, BackgroundTasks
+from pydantic import BaseModel, Field
+from typing import List, Optional, Union, Dict, Any
+import time
+import uuid
+
+from config import DEMO_API_KEY
+from db import get_supabase
+from services import engine
+from utils import calculate_usage
+
+# Initialize Router
+router = APIRouter()
+# engine is imported from services
+
+# --- Pydantic Models (OpenAI Spec) ---
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+    name: Optional[str] = None
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    n: Optional[int] = 1
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, List[str]]] = None
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
+    logit_bias: Optional[Dict[str, float]] = None
+    user: Optional[str] = None
+    
+    # Custom fields for our API (optional)
+    provider: Optional[str] = None 
+
+class ChatCompletionChoice(BaseModel):
+    index: int
+    message: ChatMessage
+    finish_reason: Optional[str] = "stop"
+
+class UsageInfo(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: UsageInfo
+
+
+# --- Auth Dependency ---
+
+async def verify_api_key(
+    authorization: Optional[str] = Header(None),
+    x_api_key: Optional[str] = Header(None)
+):
+    """
+    Verify Bearer Token or X-API-KEY.
+    Returns: key_data (dict) or None (if demo key)
+    Raises: HTTPException if invalid
+    """
+    token = None
+    if authorization:
+        parts = authorization.split()
+        if len(parts) == 2 and parts[0].lower() == "bearer":
+            token = parts[1]
+    
+    if not token and x_api_key:
+        token = x_api_key
+        
+    if not token:
+        raise HTTPException(status_code=401, detail="Missing API Key")
+
+    # 1. Check Demo Key
+    if token == DEMO_API_KEY:
+        return {"id": "demo", "name": "Demo User", "limit_tokens": -1}
+
+    # 2. Check Database
+    supabase = get_supabase()
+    if not supabase:
+        # Fallback if DB is down but key matches verified format? No, safer to reject.
+        raise HTTPException(status_code=503, detail="Auth service unavailable")
+
+    try:
+        # Check if key exists and is active
+        res = supabase.table("api_keys").select("*").eq("token", token).execute()
+        
+        if not res.data:
+            raise HTTPException(status_code=401, detail="Invalid API Key")
+            
+        key_data = res.data[0]
+        
+        if not key_data.get("is_active", True):
+            raise HTTPException(status_code=403, detail="API Key is inactive")
+            
+        # Check limits
+        # Note: We check limit BEFORE processing, but update usage AFTER (bg task)
+        current_usage = key_data.get("usage_tokens", 0)
+        limit = key_data.get("limit_tokens", 0)
+        
+        if limit > 0 and current_usage >= limit:
+             raise HTTPException(status_code=429, detail="Quota exceeded")
+             
+        return key_data
+
+    except Exception as e:
+        print(f"Auth Error: {e}")
+        raise HTTPException(status_code=500, detail="Auth Error")
+
+# --- Background Task for Usage Update ---
+
+def update_usage_stats(key_id: str, tokens: int):
+    """Increment token usage in DB."""
+    if key_id == "demo":
+        return # Don't track demo usage in DB (or maybe track in a separate table later)
+        
+    supabase = get_supabase()
+    if supabase and tokens > 0:
+        try:
+            # Atomic increment? Supabase (Postgres) supports it via RPC or simple update if inaccurate is okay.
+            # Best practice: use RPC. For now simple update read-modify-write (concurrency risk but okay for low volume)
+            # Actually, let's just do a simple increment if possible, or fetch-add
+            
+            # Since we can't easily do RPC without creating it in SQL first, let's just do Python-side increment
+            # (Valid since we have the key_data from verification, but it might be stale)
+            
+            # Better: create an RPC function later. For now, just logging it.
+            # The implementation plan implied we track it.
+            
+            # Let's try to get fresh usage and update.
+            current = supabase.table("api_keys").select("usage_tokens").eq("id", key_id).execute()
+            if current.data:
+                new_total = (current.data[0]['usage_tokens'] or 0) + tokens
+                supabase.table("api_keys").update({"usage_tokens": new_total}).eq("id", key_id).execute()
+                
+        except Exception as e:
+            print(f"Failed to update usage for {key_id}: {e}")
+
+# --- Endpoint ---
+
+@router.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+async def chat_completions(
+    request: ChatCompletionRequest, 
+    background_tasks: BackgroundTasks,
+    key_data: dict = Depends(verify_api_key)
+):
+    """
+    OpenAI-compatible Chat Completion Endpoint.
+    """
+    # Convert messages list to simple prompt (or keep as list if engine supports it)
+    # Our engine currently takes a single prompt string + optional system prompt.
+    
+    system_prompt = None
+    user_prompt = ""
+    
+    # Simple conversion logic
+    for m in request.messages:
+        if m.role == "system":
+            system_prompt = m.content
+        elif m.role == "user":
+            if user_prompt:
+                user_prompt += f"\n\n[User]: {m.content}"
+            else:
+                user_prompt = m.content
+        elif m.role == "assistant":
+            user_prompt += f"\n\n[Assistant]: {m.content}"
+            
+    # Call Engine
+    provider = request.provider or "auto"
+    
+    try:
+        result = await engine.chat(
+            prompt=user_prompt,
+            model=request.model,
+            provider=provider,
+            system_prompt=system_prompt
+        )
+        
+        response_text = result["response"]
+        actual_model = result["model"]
+        
+        # Calculate Usage
+        usage = calculate_usage([m.dict() for m in request.messages], response_text)
+        
+        # Background: Update DB
+        background_tasks.add_task(update_usage_stats, key_data["id"], usage["total_tokens"])
+        
+        # Construct Response
+        return ChatCompletionResponse(
+            id=f"chatcmpl-{uuid.uuid4().hex[:8]}",
+            created=int(time.time()),
+            model=actual_model,
+            choices=[
+                ChatCompletionChoice(
+                    index=0,
+                    message=ChatMessage(role="assistant", content=response_text),
+                    finish_reason="stop"
+                )
+            ],
+            usage=UsageInfo(**usage)
+        )
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))