app.py

import os
import sqlite3
import secrets
import requests
from fastapi import FastAPI, HTTPException, Header, Depends
from pydantic import BaseModel

app = FastAPI(title="ChatGPT OS Key Faucet")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allows all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allows all methods (GET, POST, PUT, DELETE, etc.)
    allow_headers=["*"],  # Allows all headers
)

# --- CONFIGURATION ---
TARGET_SPACE_URL = "https://coollamaceo-chatgptopensource1-0.hf.space/chat"
DB_PATH = "/data/keys.db" if os.path.exists("/data") else "keys.db"

# --- DATABASE SETUP ---
def init_db():
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    c = conn.cursor()
    c.execute('CREATE TABLE IF NOT EXISTS api_keys (key TEXT PRIMARY KEY)')
    conn.commit()
    conn.close()

init_db()

# --- DATA MODELS ---
class ChatRequest(BaseModel):
    message: str
    session_id: str = "default_user"

# --- KEY GENERATION ---
@app.get("/generate")
async def generate_public_key():
    """Anyone can call this to get a new sk- key"""
    new_key = f"sk-{secrets.token_urlsafe(24)}"
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    c = conn.cursor()
    c.execute('INSERT INTO api_keys (key) VALUES (?)', (new_key,))
    conn.commit()
    conn.close()
    return {
        "api_key": new_key,
        "instructions": "Add this to your 'X-API-Key' header to use /v1/chat"
    }

# --- KEY VERIFICATION ---
async def verify_key(x_api_key: str = Header(None)):
    if not x_api_key:
        raise HTTPException(status_code=401, detail="Missing X-API-Key header")
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    c = conn.cursor()
    c.execute('SELECT 1 FROM api_keys WHERE key = ?', (x_api_key,))
    exists = c.fetchone()
    conn.close()
    if not exists:
        raise HTTPException(status_code=403, detail="Invalid API Key")
    return x_api_key

# --- PROXY CHAT ---
@app.post("/v1/chat")
async def proxy_chat(payload: ChatRequest, api_key: str = Depends(verify_key)):
    try:
        # Forward to your 20B backend
        headers = {"Content-Type": "application/json"}
        response = requests.post(
            TARGET_SPACE_URL,
            json=payload.dict(),
            headers=headers,
            timeout=120
        )
        response.raise_for_status()  # raise exception on HTTP errors
        return response.json()
    except requests.exceptions.Timeout:
        return {"error": "The 20B engine timed out."}
    except requests.exceptions.RequestException as e:
        return {"error": f"The 20B engine is offline or returned an error: {str(e)}"}

# --- HOME ---
@app.get("/")
def home():
    return {"message": "Go to /generate to get a key, then POST to /v1/chat"}