Update app.py

app.py

@@ -1,155 +1,149 @@
-"""
-Code Collab AI Backend — Fast code generation API
-Uses Qwen 2.5 Coder 0.5B GGUF for low-latency HTML/CSS/JS generation.
-"""
-
-import os
-import re
-import time
-from contextlib import asynccontextmanager
-
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
-
-# ─── Config ────────────────────────────────────────────────────────────
-MODEL_REPO = "Qwen/Qwen2.5-Coder-0.5B-Instruct-GGUF"
-MODEL_FILE = "qwen2.5-coder-0.5b-instruct-q4_k_m.gguf"
-N_CTX = 1536        # Smaller context = faster (free tier)
-N_THREADS = 2       # HF free tier has 2 vCPU
-MAX_TOKENS = 512    # Keep output short for speed
-TEMPERATURE = 0.5   # Lower = faster + more deterministic
-
-# ─── System prompt (optimized for structured output) ───────────────────
-SYSTEM_PROMPT = """You are a web code generator. Given a user request, output ONLY three fenced code blocks:
-
-```html
-(body content only, no html/head/body tags)
-```
-
-```css
-(complete styles)
-```
-
-```js
-(complete JavaScript)
-```
-
-Rules:
-- No explanations, no markdown text outside code blocks
-- If a section is not needed, output an empty code block for it
-- Write clean, modern code"""
-
-# ─── Global model reference ───────────────────────────────────────────
-llm = None
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    """Load model once at startup, keep in memory for fast inference."""
-    global llm
-    print("⬇️  Downloading model...")
-    model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
-    print(f"✅ Model downloaded: {model_path}")
-
-    print("🔄 Loading model into memory...")
-    llm = Llama(
-        model_path=model_path,
-        n_ctx=N_CTX,
-        n_threads=N_THREADS,
-        n_gpu_layers=0,  # CPU only (free HF Spaces)
-        verbose=False,
-    )
-    print("🚀 Model loaded and ready!")
-    yield
-    llm = None
-
-
-# ─── FastAPI App ───────────────────────────────────────────────────────
-app = FastAPI(title="Code Collab AI", lifespan=lifespan)
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-
-class GenerateRequest(BaseModel):
-    prompt: str
-    max_tokens: int = MAX_TOKENS
-    temperature: float = TEMPERATURE
-
-
-class GenerateResponse(BaseModel):
-    html: str
-    css: str
-    js: str
-    raw: str
-    time_ms: int
-
-
-def parse_code_blocks(text: str) -> dict:
-    """Extract HTML, CSS, JS from fenced code blocks."""
-    result = {"html": "", "css": "", "js": ""}
-
-    for match in re.finditer(r"```(\w+)\s*\n([\s\S]*?)```", text):
-        lang = match.group(1).lower()
-        code = match.group(2).strip()
-
-        if lang in ("html", "xml"):
-            result["html"] = code
-        elif lang == "css":
-            result["css"] = code
-        elif lang in ("js", "javascript"):
-            result["js"] = code
-
-    # Fallback: if no blocks found, treat as HTML
-    if not any(result.values()):
-        result["html"] = text.strip()
-
-    return result
-
-
-# ─── API Endpoints ─────────────────────────────────────────────────────
-
-@app.get("/")
-def health():
-    return {"status": "ok", "model": MODEL_REPO}
-
-
-@app.post("/generate", response_model=GenerateResponse)
-def generate(req: GenerateRequest):
-    if llm is None:
-        raise HTTPException(503, "Model not loaded yet")
-
-    if not req.prompt.strip():
-        raise HTTPException(400, "Prompt cannot be empty")
-
-    start = time.time()
-
-    output = llm.create_chat_completion(
-        messages=[
-            {"role": "system", "content": SYSTEM_PROMPT},
-            {"role": "user", "content": req.prompt},
-        ],
-        max_tokens=req.max_tokens,
-        temperature=req.temperature,
-        stop=["```\n\n", "---"],  # Stop early if model rambles
-    )
-
-    raw_text = output["choices"][0]["message"]["content"]
-    elapsed_ms = int((time.time() - start) * 1000)
-
-    parsed = parse_code_blocks(raw_text)
-
-    return GenerateResponse(
-        html=parsed["html"],
-        css=parsed["css"],
-        js=parsed["js"],
-        raw=raw_text,
-        time_ms=elapsed_ms,
-    )
+"""
+Code Collab AI Backend — Fast code generation API
+Uses Qwen 2.5 Coder 0.5B GGUF for low-latency HTML/CSS/JS generation.
+Model is loaded once at startup and stays in memory forever.
+"""
+
+import os
+import re
+import time
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from llama_cpp import Llama
+
+# ─── Config ────────────────────────────────────────────────────────────
+MODEL_PATH = os.path.join(os.path.dirname(__file__), "qwen2.5-coder-0.5b-instruct-q4_k_m.gguf")
+N_CTX = 1536        # Smaller context = faster (free tier)
+N_THREADS = 2       # HF free tier has 2 vCPU
+MAX_TOKENS = 512    # Keep output short for speed
+TEMPERATURE = 0.5   # Lower = faster + more deterministic
+
+# ─── System prompt (optimized for structured output) ───────────────────
+SYSTEM_PROMPT = """You are a web code generator. Given a user request, output ONLY three fenced code blocks:
+
+```html
+(body content only, no html/head/body tags)
+```
+
+```css
+(complete styles)
+```
+
+```js
+(complete JavaScript)
+```
+
+Rules:
+- No explanations, no markdown text outside code blocks
+- If a section is not needed, output an empty code block for it
+- Write clean, modern code"""
+
+# ─── Global model reference (loaded once, stays forever) ──────────────
+llm = None
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Load model once at startup. It stays in memory for the entire lifetime."""
+    global llm
+    print(f"🔄 Loading model from {MODEL_PATH}...")
+    llm = Llama(
+        model_path=MODEL_PATH,
+        n_ctx=N_CTX,
+        n_threads=N_THREADS,
+        n_gpu_layers=0,  # CPU only (free HF Spaces)
+        verbose=False,
+    )
+    print("🚀 Model loaded and ready! It will stay in memory forever.")
+    yield
+    # Model stays loaded — never unloaded
+
+
+# ─── FastAPI App ───────────────────────────────────────────────────────
+app = FastAPI(title="Code Collab AI", lifespan=lifespan)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+
+class GenerateRequest(BaseModel):
+    prompt: str
+    max_tokens: int = MAX_TOKENS
+    temperature: float = TEMPERATURE
+
+
+class GenerateResponse(BaseModel):
+    html: str
+    css: str
+    js: str
+    raw: str
+    time_ms: int
+
+
+def parse_code_blocks(text: str) -> dict:
+    """Extract HTML, CSS, JS from fenced code blocks."""
+    result = {"html": "", "css": "", "js": ""}
+
+    for match in re.finditer(r"```(\w+)\s*\n([\s\S]*?)```", text):
+        lang = match.group(1).lower()
+        code = match.group(2).strip()
+
+        if lang in ("html", "xml"):
+            result["html"] = code
+        elif lang == "css":
+            result["css"] = code
+        elif lang in ("js", "javascript"):
+            result["js"] = code
+
+    # Fallback: if no blocks found, treat as HTML
+    if not any(result.values()):
+        result["html"] = text.strip()
+
+    return result
+
+
+# ─── API Endpoints ─────────────────────────────────────────────────────
+
+@app.get("/")
+def health():
+    return {"status": "ok", "model": "Qwen2.5-Coder-0.5B-Instruct-GGUF", "loaded": llm is not None}
+
+
+@app.post("/generate", response_model=GenerateResponse)
+def generate(req: GenerateRequest):
+    if llm is None:
+        raise HTTPException(503, "Model not loaded yet")
+
+    if not req.prompt.strip():
+        raise HTTPException(400, "Prompt cannot be empty")
+
+    start = time.time()
+
+    output = llm.create_chat_completion(
+        messages=[
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": req.prompt},
+        ],
+        max_tokens=req.max_tokens,
+        temperature=req.temperature,
+    )
+
+    raw_text = output["choices"][0]["message"]["content"]
+    elapsed_ms = int((time.time() - start) * 1000)
+
+    parsed = parse_code_blocks(raw_text)
+
+    return GenerateResponse(
+        html=parsed["html"],
+        css=parsed["css"],
+        js=parsed["js"],
+        raw=raw_text,
+        time_ms=elapsed_ms,
+    )