Update app.py

app.py

@@ -1,15 +1,27 @@
 import os
 import torch
-import gradio as gr
+import secrets
+import time
+from fastapi import FastAPI, HTTPException, Security, Depends
+from fastapi.security.api_key import APIKeyHeader
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from starlette.status import HTTP_403_FORBIDDEN
 
-# Point to the local folder created in the Dockerfile
+# --- CONFIGURATION ---
 MODEL_PATH = "/app/model"
+API_KEY_NAME = "X-API-Key"
+api_key_header = APIKeyHeader(name=API_KEY_NAME, auto_error=False)
 
-print("Loading Overflow-111.7B from Local Docker Storage...")
+# In-memory storage for keys. 
+# Note: These will reset if the Space restarts unless you use Persistent Storage.
+generated_keys = {} 
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+app = FastAPI(title="Overflow-111.7B API Manager")
 
+# --- MODEL LOADING ---
+print("Loading Overflow-111.7B Engine...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_PATH,
     trust_remote_code=True,
@@ -18,12 +30,45 @@ model = AutoModelForCausalLM.from_pretrained(
     low_cpu_mem_usage=True
 )
 
-def respond(message, history):
-    inputs = tokenizer(message, return_tensors="pt")
+class Query(BaseModel):
+    prompt: str
+    max_tokens: int = 50
+
+# --- API KEY GENERATION ---
+@app.get("/api/generate")
+async def create_new_key():
+    """Generates a new of_sk key for the user."""
+    # Generate a random string of 24 characters
+    random_hex = secrets.token_hex(12)
+    new_key = f"of_sk-{random_hex}"
+    
+    # Store with a timestamp
+    generated_keys[new_key] = {"created_at": time.time()}
+    
+    return {
+        "status": "success",
+        "api_key": new_key,
+        "instructions": f"Include this key in your request header as '{API_KEY_NAME}'"
+    }
+
+# --- SECURITY CHECK ---
+async def get_api_key(api_key_header: str = Depends(api_key_header)):
+    if api_key_header in generated_keys:
+        return api_key_header
+    raise HTTPException(
+        status_code=HTTP_403_FORBIDDEN, 
+        detail="Invalid or expired API Key. Generate one at /api/generate"
+    )
+
+@app.post("/v1/generate")
+async def generate(query: Query, api_key: str = Depends(get_api_key)):
+    inputs = tokenizer(query.prompt, return_tensors="pt")
     with torch.no_grad():
-        output_tokens = model.generate(**inputs, max_new_tokens=30)
-    return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+        output_tokens = model.generate(**inputs, max_new_tokens=query.max_tokens)
+    
+    response = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    return {"text": response}
 
-demo = gr.ChatInterface(respond)
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+@app.get("/")
+def home():
+    return {"message": "Welcome to Overflow-111.7B. Go to /api/generate to get a key."}