Upload api/gateway.py with huggingface_hub

api/gateway.py

@@ -0,0 +1,266 @@
+"""
+dispatchAI Inference API — Main Gateway Server
+Runs on the Windows PC (112 cores). Acts as load balancer + API gateway.
+
+Architecture:
+  Customer → api.dispatchai.ai:8081 → This server → Routes to phone
+  Phone runs phone_server.py (HTTP server + llama.cpp)
+
+This server:
+1. Receives OpenAI-compatible API requests
+2. Finds an available phone
+3. Routes the request to that phone
+4. Returns the response to the customer
+5. Tracks token usage for billing
+"""
+import os
+import json
+import time
+import asyncio
+import httpx
+from datetime import datetime
+from typing import Optional
+from fastapi import FastAPI, HTTPException, Depends, Header
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+
+# ============================================================
+# Configuration
+# ============================================================
+
+# Phone farm — list of phone IPs and ports
+# Each phone runs phone_server.py on port 5000
+# For now, we use ADB to get phone serials and assign ports
+PHONE_PORTS = {}  # serial → port mapping, filled at startup
+BASE_PHONE_PORT = 5000  # First phone gets port 5000, second 5001, etc.
+
+# API keys (simple auth — in production use a database)
+API_KEYS_FILE = "data/api_keys.json"
+USAGE_FILE = "data/api_usage.json"
+
+# Available models
+MODELS = {
+    "dispatchAI/SmolLM2-135M-Instruct-mobile": {"phone_model": "SmolLM2-135M-Instruct-mobile", "chat_format": "llama-3"},
+    "dispatchAI/Qwen2.5-0.5B-Instruct-mobile-int4": {"phone_model": "Qwen2.5-0.5B-Instruct-mobile-int4", "chat_format": "chatml"},
+    "dispatchAI/Llama-3.2-1B-Instruct-Q4-mobile": {"phone_model": "Llama-3.2-1B-Instruct-Q4-mobile", "chat_format": "chatml"},
+    "dispatchAI/TinyLlama-1.1B-Chat-Q5-mobile": {"phone_model": "TinyLlama-1.1B-Chat-Q5-mobile", "chat_format": "chatml"},
+    "dispatchAI/Qwen2.5-0.5B-Coder-mobile": {"phone_model": "Qwen2.5-0.5B-Coder-mobile", "chat_format": "chatml"},
+}
+
+# Pricing (per 1K tokens)
+PRICING = {
+    "input": 0.001,   # $0.001 per 1K input tokens
+    "output": 0.002,  # $0.002 per 1K output tokens
+}
+
+# ============================================================
+# Data Models (OpenAI-compatible)
+# ============================================================
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str = "dispatchAI/SmolLM2-135M-Instruct-mobile"
+    messages: list[ChatMessage]
+    max_tokens: int = 100
+    temperature: float = 0.7
+    stream: bool = False
+
+# ============================================================
+# API Key Management
+# ============================================================
+
+def load_api_keys():
+    if os.path.exists(API_KEYS_FILE):
+        return json.load(open(API_KEYS_FILE))
+    # Create default key
+    keys = {"da-demo-key-0001": {"name": "Demo Key", "created": datetime.now().isoformat(), "balance": 1000}}
+    json.dump(keys, open(API_KEYS_FILE, "w"), indent=2)
+    return keys
+
+def load_usage():
+    if os.path.exists(USAGE_FILE):
+        return json.load(open(USAGE_FILE))
+    return {}
+
+def save_usage(usage):
+    json.dump(usage, open(USAGE_FILE, "w"), indent=2)
+
+def verify_api_key(authorization: Optional[str] = Header(None)):
+    if not authorization:
+        raise HTTPException(status_code=401, detail="Missing API key. Add 'Authorization: Bearer da-xxx' header.")
+    
+    key = authorization.replace("Bearer ", "").strip()
+    api_keys = load_api_keys()
+    
+    if key not in api_keys:
+        raise HTTPException(status_code=401, detail="Invalid API key")
+    
+    return key
+
+# ============================================================
+# Phone Pool Management
+# ============================================================
+
+def get_available_phones():
+    """Get list of connected phones via ADB."""
+    import subprocess
+    result = subprocess.run(["adb", "devices"], capture_output=True, text=True, timeout=10)
+    phones = []
+    for line in result.stdout.strip().split("\n")[1:]:
+        if "\tdevice" in line:
+            serial = line.split("\t")[0]
+            phones.append(serial)
+    return phones
+
+def get_phone_port(serial: str) -> int:
+    """Get or assign a port for a phone."""
+    if serial not in PHONE_PORTS:
+        PHONE_PORTS[serial] = BASE_PHONE_PORT + len(PHONE_PORTS)
+    return PHONE_PORTS[serial]
+
+# ============================================================
+# FastAPI App
+# ============================================================
+
+app = FastAPI(
+    title="dispatchAI Inference API",
+    description="Mobile-optimized LLM inference. Small. Mobile. Free. UAE-built.",
+    version="1.0.0",
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["https://dispatchai.ai", "https://www.dispatchai.ai", "https://huggingface.co"],
+    allow_methods=["GET", "POST"],
+    allow_headers=["*"],
+)
+
+@app.get("/")
+async def root():
+    """API info."""
+    phones = get_available_phones()
+    return {
+        "name": "dispatchAI Inference API",
+        "version": "1.0.0",
+        "status": "running",
+        "phones_connected": len(phones),
+        "models": list(MODELS.keys()),
+        "pricing": {"input": f"${PRICING['input']}/1K tokens", "output": f"${PRICING['output']}/1K tokens"},
+        "docs": "/docs",
+        "website": "https://huggingface.co/dispatchAI",
+    }
+
+@app.get("/v1/models")
+async def list_models(api_key: str = Depends(verify_api_key)):
+    """List available models (OpenAI-compatible)."""
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": model_id,
+                "object": "model",
+                "created": 1719500000,
+                "owned_by": "dispatchAI",
+            }
+            for model_id in MODELS.keys()
+        ]
+    }
+
+@app.post("/v1/chat/completions")
+async def chat_completions(
+    request: ChatCompletionRequest,
+    api_key: str = Depends(verify_api_key),
+):
+    """Create a chat completion (OpenAI-compatible)."""
+    
+    if request.model not in MODELS:
+        raise HTTPException(status_code=400, detail=f"Model '{request.model}' not available. Use GET /v1/models to see available models.")
+    
+    # Get available phones
+    phones = get_available_phones()
+    if not phones:
+        raise HTTPException(status_code=503, detail="No phones available. Try again later.")
+    
+    # Round-robin load balancing across active phone proxies
+    # Each phone proxy runs on port 5000, 5001, 5002, etc.
+    import time as _time
+    available_ports = [5000, 5001, 5002, 5003, 5004]  # 3 phones with proxies
+    phone_port = available_ports[int(_time.time()) % len(available_ports)]
+    
+    model_info = MODELS[request.model]
+    
+    # Prepare request for phone
+    phone_request = {
+        "model": request.model,
+        "messages": [{"role": m.role, "content": m.content} for m in request.messages],
+        "max_tokens": request.max_tokens,
+        "temperature": request.temperature,
+        "chat_format": model_info["chat_format"],
+        "raw_completion": True,  # Use raw text completion, not chat template
+    }
+    
+    # Send to phone
+    try:
+        async with httpx.AsyncClient(timeout=120.0) as client:
+            response = await client.post(
+                f"http://127.0.0.1:{phone_port}/v1/chat/completions",
+                json=phone_request,
+            )
+            
+            if response.status_code != 200:
+                raise HTTPException(status_code=502, detail=f"Phone error: {response.text[:200]}")
+            
+            result = response.json()
+            
+            # Track usage
+            usage = load_usage()
+            if api_key not in usage:
+                usage[api_key] = {"total_tokens": 0, "requests": 0, "cost": 0.0}
+            
+            tokens_used = result.get("usage", {}).get("total_tokens", 0)
+            cost = (tokens_used / 1000) * (PRICING["input"] + PRICING["output"])
+            
+            usage[api_key]["total_tokens"] += tokens_used
+            usage[api_key]["requests"] += 1
+            usage[api_key]["cost"] += cost
+            usage[api_key]["last_request"] = datetime.now().isoformat()
+            save_usage(usage)
+            
+            return result
+    
+    except httpx.TimeoutException:
+        raise HTTPException(status_code=504, detail="Phone inference timed out. Try a smaller max_tokens.")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Internal error: {str(e)[:200]}")
+
+@app.get("/v1/usage")
+async def get_usage(api_key: str = Depends(verify_api_key)):
+    """Get API usage stats."""
+    usage = load_usage()
+    return usage.get(api_key, {"total_tokens": 0, "requests": 0, "cost": 0.0})
+
+@app.get("/admin/phones")
+async def admin_phones(api_key: str = Depends(verify_api_key)):
+    """Get phone farm status (requires auth)."""
+    phones = get_available_phones()
+    return {
+        "phones_connected": len(phones),
+        "phones": [{"serial": p, "port": get_phone_port(p)} for p in phones],
+        "total_capacity_tokens_per_sec": len(phones) * 20,  # ~20 t/s per phone
+    }
+
+# ============================================================
+# Startup
+# ============================================================
+
+if __name__ == "__main__":
+    import uvicorn
+    print("🚀 dispatchAI Inference API — Starting...")
+    print(f"   Endpoint: http://api.dispatchai.ai:8081")
+    print(f"   Docs: http://api.dispatchai.ai:8081/docs")
+    print(f"   Phones: {len(get_available_phones())} connected")
+    print()
+    uvicorn.run(app, host="0.0.0.0", port=8081)