app.py

from fastapi import FastAPI
from fastapi.responses import Response
from pydantic import BaseModel
from prometheus_client import Counter, generate_latest, CONTENT_TYPE_LATEST
from transformers import pipeline
import gradio as gr

app = FastAPI(title="Incident Triage Bot")

triage_requests = Counter("triage_requests_total", "Total triage requests")
triage_errors = Counter("triage_errors_total", "Total triage errors")

triage_pipeline = pipeline(
    "text2text-generation",
    model="google/flan-t5-small",
    device=-1  # CPU-only
)

class Incident(BaseModel):
    description: str

@app.post("/triage")
async def triage(incident: Incident):
    triage_requests.inc()
    try:
        prompt = f"Classify this incident and suggest priority:\n\n{incident.description}"
        result = triage_pipeline(prompt, max_new_tokens=100)[0]["generated_text"]
        return {"triage": result.strip()}
    except Exception as e:
        triage_errors.inc()
        return {"error": str(e)}

@app.get("/metrics")
def metrics():
    return Response(generate_latest(), media_type=CONTENT_TYPE_LATEST)

@app.get("/")
def home():
    return {"message": "LLM Incident Triage Bot is running!"}

def gradio_triage(description):
    try:
        triage_requests.inc()
        prompt = f"Classify this incident and suggest priority:\n\n{description}"
        result = triage_pipeline(prompt, max_new_tokens=100)[0]["generated_text"]
        return result.strip()
    except Exception as e:
        triage_errors.inc()
        return f"Error: {str(e)}"

demo = gr.Interface(
    fn=gradio_triage,
    inputs=gr.Textbox(lines=5, placeholder="Describe the incident..."),
    outputs="text",
    title="Incident Triage Bot",
    description="Classify incidents and suggest priority using a lightweight Transformers model"
)

@app.get("/gradio")
def gradio_ui():
    return Response(content=demo.launch(share=False, inline=True), media_type="text/html")