app.py

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
import requests as http_requests
import json
import re
import os

app = FastAPI(title="Meta-SRE Agent", version="1.0.0")

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

MODEL_ID  = "Anvit25/meta-sre-agent"
API_URL   = f"https://api-inference.huggingface.co/models/{MODEL_ID}"
HF_TOKEN  = os.getenv("HF_TOKEN", "")
HEADERS   = {"Authorization": f"Bearer {HF_TOKEN}"}

SYSTEM_PROMPT = """You are a Senior Site Reliability Engineer (SRE) at Meta.
A developer has reported a bug in their service. Analyze the error and provide a diagnosis.

Respond ONLY with a JSON object in this exact format:
{
  "root_cause": "one sentence explaining the exact cause",
  "fix": "one sentence describing what to change",
  "severity": "P0 or P1 or P2",
  "file_hint": "filename or function where fix should be applied"
}"""


class DiagnoseRequest(BaseModel):
    error: str
    code: Optional[str] = ""
    service: Optional[str] = "unknown-service"
    context: Optional[str] = ""


class DiagnoseResponse(BaseModel):
    root_cause: str
    fix: str
    severity: str
    file_hint: str
    raw_output: str


def build_prompt(req: DiagnoseRequest) -> str:
    parts = [f"SERVICE: {req.service}"]
    parts.append(f"\nERROR:\n{req.error}")
    if req.code:
        parts.append(f"\nCODE:\n{req.code}")
    if req.context:
        parts.append(f"\nCONTEXT:\n{req.context}")
    parts.append("\nDiagnose this issue:")
    return "\n".join(parts)


def parse_response(text: str) -> dict:
    # Try to extract JSON block
    match = re.search(r'\{.*?\}', text, re.DOTALL)
    if match:
        try:
            return json.loads(match.group())
        except json.JSONDecodeError:
            pass

    # Fallback: return raw text as root_cause
    return {
        "root_cause": text.strip()[:300] if text.strip() else "Could not determine root cause",
        "fix": "Review the error and code manually",
        "severity": "P2",
        "file_hint": "unknown",
    }


@app.get("/health")
def health():
    return {"status": "ok", "model": MODEL_ID, "version": "1.0.0"}


@app.post("/diagnose", response_model=DiagnoseResponse)
def diagnose(req: DiagnoseRequest):
    if not req.error:
        raise HTTPException(status_code=400, detail="error field is required")

    user_message = build_prompt(req)

    prompt = f"<|system|>\n{SYSTEM_PROMPT}\n<|user|>\n{user_message}\n<|assistant|>\n"

    try:
        resp = http_requests.post(
            API_URL,
            headers=HEADERS,
            json={"inputs": prompt, "parameters": {"max_new_tokens": 256, "temperature": 0.1}},
            timeout=60,
        )
        resp.raise_for_status()
        data = resp.json()
        if isinstance(data, list):
            raw = data[0].get("generated_text", "")
            # Strip the prompt from the output if model echoes it
            if raw.startswith(prompt):
                raw = raw[len(prompt):]
        elif isinstance(data, dict) and "error" in data:
            raise HTTPException(status_code=500, detail=f"HF API error: {data['error']}")
        else:
            raw = str(data)
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Model inference failed: {str(e)}")

    parsed = parse_response(raw)
    return DiagnoseResponse(
        root_cause=parsed.get("root_cause", ""),
        fix=parsed.get("fix", ""),
        severity=parsed.get("severity", "P2"),
        file_hint=parsed.get("file_hint", "unknown"),
        raw_output=raw,
    )