app.py

"""
Code Collab AI Backend — Fast code generation API (Groq Edition)
Now with Conversation History Support!
"""

import os
import re
import time
from typing import List

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from groq import Groq

# ─── Config ────────────────────────────────────────────────────────────
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
MODEL_NAME = "qwen/qwen3-32b" 
MAX_TOKENS = 2048   # Increased slightly for iterative changes 
TEMPERATURE = 0.5    

# ─── System prompt ─────────────────────────────────────────────────────
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
- When the user asks for edits, rewrite the FULL code blocks with the changes applied."""

# ─── Groq Client Setup ─────────────────────────────────────────────────
if not GROQ_API_KEY:
    print("⚠️ WARNING: GROQ_API_KEY environment variable is not set!")

client = Groq(api_key=GROQ_API_KEY)

# ─── FastAPI App ───────────────────────────────────────────────────────
app = FastAPI(title="Code Collab AI")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# New Message model to handle history
class Message(BaseModel):
    role: str
    content: str

class GenerateRequest(BaseModel):
    messages: List[Message]
    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", 
        "provider": "Groq", 
        "model": MODEL_NAME,
        "api_key_configured": GROQ_API_KEY is not None
    }

@app.post("/generate", response_model=GenerateResponse)
def generate(req: GenerateRequest):
    if not GROQ_API_KEY:
        raise HTTPException(status_code=500, detail="Groq API key is missing on the server")

    if not req.messages:
        raise HTTPException(status_code=400, detail="Message history cannot be empty")

    start = time.time()

    # Prepend the system prompt to the conversation history
    api_messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for msg in req.messages:
        api_messages.append({"role": msg.role, "content": msg.content})

    try:
        output = client.chat.completions.create(
            model=MODEL_NAME,
            messages=api_messages,
            max_tokens=req.max_tokens,
            temperature=req.temperature,
        )
    except Exception as e:
        raise HTTPException(status_code=502, detail=f"Groq API error: {str(e)}")

    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,
    )