How to Cook Up a Jev Model at Home

@app.post("/mask")
def mask_incident(payload: dict):
    text = payload.get("text", "")
    
    # Stage 1: Zero-decode logit check (~30ms)
    has_leak, confidence = score_password_presence(text)
    if not has_leak:
        return {"sanitized": text, "masked": False}

    # Stage 2: Targeted extraction (only runs on confirmed leaks)
    sanitized, secrets = extract_and_mask_secrets(text)
    return {"sanitized": sanitized, "masked": True, "redacted_count": len(secrets)}

# Pre-resolve token IDs for binary scoring
YES_ID = tokenizer.encode("Yes", add_special_tokens=False)[0]
NO_ID  = tokenizer.encode("No", add_special_tokens=False)[0]

def score_password_presence(text: str) -> tuple[bool, float]:
    prompt = build_chat_prompt(f"Does this ticket expose a plaintext password?\nTicket: {text}")
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        logits = model(**inputs).logits[0, -1, [YES_ID, NO_ID]]
        probs = torch.softmax(logits, dim=-1)

    has_password = probs[0] > probs[1]
    return has_password.item(), probs.max().item()

def extract_and_mask_secrets(text: str) -> tuple[str, list[str]]:
    prompt = build_chat_prompt(f"Extract all plaintext passwords as comma-separated values:\n{text}")
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=30, do_sample=False)
    
    raw = tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    secrets = [s.strip(" '\"") for s in raw.split(",") if len(s.strip()) >= 4]

    sanitized = text
    for secret in set(secrets):
        if secret in sanitized:
            sanitized = sanitized.replace(secret, "[REDACTED_PASSWORD]")

    return sanitized, secrets

"""
ServiceNow Privacy Gateway: Zero-Decode Credential Masker
Combines Jev fast-logit classification (~30ms) with targeted multi-span LLM redaction.
"""

import re
import time
from typing import List, Optional

import torch
import uvicorn
from fastapi import FastAPI
from pydantic import BaseModel, Field
from transformers import AutoModelForCausalLM, AutoTokenizer

# =====================================================================
# 1. Model Initialization
# =====================================================================
MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"
device = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Loading {MODEL_ID} on {device}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto" if torch.cuda.is_available() else None,
)
model.eval()

# Resolve single-token IDs for binary logit scoring (zero generation decode)
YES_TOKEN_ID = tokenizer.encode("Yes", add_special_tokens=False)[0]
NO_TOKEN_ID = tokenizer.encode("No", add_special_tokens=False)[0]
CANDIDATE_IDS = [YES_TOKEN_ID, NO_TOKEN_ID]

# Guardrails: Common keywords and short stopwords to protect against over-redaction
STOP_WORDS = {"user", "admin", "pass", "password", "help", "login", "root", "error", "desk", "none"}


# =====================================================================
# 2. Stage 1: Jev Fast Logit Filter (~30-40ms)
# =====================================================================
def build_jev_prompt(text: str) -> str:
    """Builds a ChatML prompt with few-shot ServiceNow incident anchors."""
    messages = [
        {
            "role": "system",
            "content": (
                "You are an IT security gatekeeper. Answer 'Yes' if the ticket contains an "
                "actual plaintext password or secret credential informally shared by a user. "
                "Answer 'No' if it merely discusses passwords (e.g., reset links, locked accounts, hashes). "
                "Answer strictly 'Yes' or 'No'."
            ),
        },
        {
            "role": "user",
            "content": "Does this expose a plaintext password?\nTicket: VPN locked. Current password is 'Summer2024!' please reset.",
        },
        {"role": "assistant", "content": "Yes"},
        {
            "role": "user",
            "content": "Does this expose a plaintext password?\nTicket: Password reset link expired for user bwayne.",
        },
        {"role": "assistant", "content": "No"},
        {"role": "user", "content": f"Does this expose a plaintext password?\nTicket: {text}"},
    ]
    return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)


def score_password_presence(text: str) -> tuple[bool, float, float]:
    """Runs a single forward pass without autoregressive decode to score candidate logits."""
    prompt = build_jev_prompt(text)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    start = time.perf_counter()
    with torch.no_grad():
        outputs = model(**inputs)
        candidate_logits = outputs.logits[0, -1, CANDIDATE_IDS]
        probs = torch.softmax(candidate_logits, dim=-1).tolist()
    elapsed_ms = (time.perf_counter() - start) * 1000

    prob_yes, prob_no = probs[0], probs[1]
    has_password = prob_yes > prob_no
    confidence = prob_yes if has_password else prob_no
    return has_password, confidence, elapsed_ms


# =====================================================================
# 3. Stage 2: Targeted Multi-Span Redaction (Runs ONLY on Leaks)
# =====================================================================
def extract_and_mask_secrets(text: str) -> tuple[str, List[str], float]:
    """Isolates active credentials and replaces each with [REDACTED_PASSWORD]."""
    # Remove existing redaction placeholders so the LLM is not anchored by past tags
    prompt_text = re.sub(r"\[REDACTED_PASSWORD\]", " ", text)

    messages = [
        {
            "role": "system",
            "content": (
                "You are an IT redaction engine. Extract all plaintext passwords, PINs, "
                "or secrets from the ticket. Do NOT include usernames, labels, or stopwords. "
                "Separate multiple secrets with commas. If none exist, output 'None'."
            ),
        },
        {
            "role": "user",
            "content": "Extract secrets:\nTried old pass 'Pass123' and temp pass 'Temp#99' neither worked.",
        },
        {"role": "assistant", "content": "Pass123, Temp#99"},
        {"role": "user", "content": f"Extract secrets:\n{prompt_text}"},
    ]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    start = time.perf_counter()
    with torch.no_grad():
        gen_ids = model.generate(
            **inputs,
            max_new_tokens=35,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=False,
        )
    elapsed_ms = (time.perf_counter() - start) * 1000

    new_tokens = gen_ids[0][inputs.input_ids.shape[1] :]
    extracted = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    if "\n" in extracted:
        extracted = extracted.split("\n")[0]

    # Parse and validate candidates against the source text
    raw_candidates = [s.strip("'\"` \n\t.,") for s in re.split(r"[,;]+", extracted) if s.strip()]
    valid_secrets = []

    for cand in raw_candidates:
        if not cand or cand.lower() in STOP_WORDS or len(cand) < 4:
            continue
        if cand in text:
            valid_secrets.append(cand)

    # Deduplicate while preserving order
    unique_secrets = list(dict.fromkeys(valid_secrets))

    # Apply redactions
    sanitized = text
    for secret in unique_secrets:
        sanitized = sanitized.replace(secret, "[REDACTED_PASSWORD]")

    return sanitized, unique_secrets, elapsed_ms


# =====================================================================
# 4. FastAPI Application & REST Endpoints
# =====================================================================
app = FastAPI(
    title="ServiceNow Security Gateway",
    description="Zero-decode pre-ingestion credential masking for ServiceNow incidents",
)


class MaskRequest(BaseModel):
    ticket_id: Optional[str] = "INC_TEMP"
    text: str = Field(..., description="Combined short description and incident body")


class MaskResponse(BaseModel):
    ticket_id: str
    has_credential: bool
    confidence: float
    sanitized_text: str
    redacted_secrets: List[str]
    stage_breakdown: dict


@app.post("/api/v1/mask", response_model=MaskResponse)
def mask_incident(req: MaskRequest):
    if not req.text.strip():
        return MaskResponse(
            ticket_id=req.ticket_id,
            has_credential=False,
            confidence=1.0,
            sanitized_text="",
            redacted_secrets=[],
            stage_breakdown={"stage1_jev_ms": 0.0, "stage2_decode_ms": 0.0},
        )

    # Stage 1: Ultra-fast Jev forward pass (~30ms)
    has_cred, confidence, jev_ms = score_password_presence(req.text)

    # 95%+ of tickets are safe: fast pass-through with 0 tokens generated
    if not has_cred:
        return MaskResponse(
            ticket_id=req.ticket_id,
            has_credential=False,
            confidence=round(confidence, 4),
            sanitized_text=req.text,
            redacted_secrets=[],
            stage_breakdown={"stage1_jev_ms": round(jev_ms, 2), "stage2_decode_ms": 0.0},
        )

    # Stage 2: Targeted multi-span redaction (runs ONLY on confirmed leaks)
    sanitized, secrets, decode_ms = extract_and_mask_secrets(req.text)

    return MaskResponse(
        ticket_id=req.ticket_id,
        has_credential=True,
        confidence=round(confidence, 4),
        sanitized_text=sanitized,
        redacted_secrets=secrets,
        stage_breakdown={
            "stage1_jev_ms": round(jev_ms, 2),
            "stage2_decode_ms": round(decode_ms, 2),
        },
    )


@app.get("/health")
def health():
    return {"status": "healthy", "model": MODEL_ID, "device": device}


# =====================================================================
# 5. Standalone Execution & Test
# =====================================================================
if __name__ == "__main__":
    # Run the server on port 8000
    # Command to test:
    # curl -X POST "http://localhost:8000/api/v1/mask" \
    #      -H "Content-Type: application/json" \
    #      -d '{"ticket_id": "INC0010011", "text": "VPN locked, tried Pass123! and Temp#9988 for user admin"}'
    uvicorn.run(app, host="0.0.0.0", port=8000)