Wire app to configurable model backend

Co-authored-by: Codex <codex@openai.com>

app.py

@@ -1,5 +1,10 @@
+import os
+from functools import lru_cache
+from pathlib import Path
+
 import gradio as gr
 
+from jawbreaker.analyzers import build_llama_cpp_analyzer, heuristic_analyzer, prediction_to_analysis
 from jawbreaker.render import render_analysis_html, render_memory_html
 from jawbreaker.schema import ScamAnalysis
 
@@ -11,17 +16,115 @@ EXAMPLES = [
 ]
 
 
+def app_theme() -> gr.Theme:
+    return gr.themes.Soft(
+        primary_hue="red",
+        secondary_hue="slate",
+        neutral_hue="zinc",
+        radius_size="sm",
+    )
+
+
+def app_css() -> str:
+    return Path("style.css").read_text(encoding="utf-8")
+
+
+def _env_int(name: str, default: int | None = None) -> int | None:
+    value = os.getenv(name)
+    if value is None or value == "":
+        return default
+    return int(value)
+
+
+def _env_bool(name: str, default: bool | None = None) -> bool | None:
+    value = os.getenv(name)
+    if value is None or value == "":
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+@lru_cache(maxsize=1)
+def get_analyzer():
+    backend = os.getenv("JAWBREAKER_BACKEND", "heuristic").strip().lower()
+    if backend == "heuristic":
+        return heuristic_analyzer
+
+    if backend == "llama-cpp":
+        model_path = resolve_model_path()
+        return build_llama_cpp_analyzer(
+            model_path,
+            chat_format=os.getenv("JAWBREAKER_CHAT_FORMAT") or None,
+            n_ctx=_env_int("JAWBREAKER_N_CTX", 2048) or 2048,
+            n_threads=_env_int("JAWBREAKER_N_THREADS"),
+            n_gpu_layers=_env_int("JAWBREAKER_N_GPU_LAYERS", 0) or 0,
+            n_batch=_env_int("JAWBREAKER_N_BATCH", 512) or 512,
+            n_ubatch=_env_int("JAWBREAKER_N_UBATCH", 512) or 512,
+            offload_kqv=_env_bool("JAWBREAKER_OFFLOAD_KQV", True),
+            op_offload=_env_bool("JAWBREAKER_OP_OFFLOAD"),
+            max_tokens=_env_int("JAWBREAKER_MAX_TOKENS", 512) or 512,
+            temperature=float(os.getenv("JAWBREAKER_TEMPERATURE", "0")),
+        )
+
+    raise ValueError(f"Unsupported JAWBREAKER_BACKEND: {backend}")
+
+
+def resolve_model_path() -> Path:
+    model_path = Path(os.getenv("JAWBREAKER_MODEL_PATH", "models/qwen3-4b-gguf/Qwen3-4B-Q4_K_M.gguf"))
+    if model_path.exists():
+        return model_path
+
+    repo_id = os.getenv("JAWBREAKER_MODEL_REPO", "Qwen/Qwen3-4B-GGUF")
+    filename = os.getenv("JAWBREAKER_MODEL_FILE", "Qwen3-4B-Q4_K_M.gguf")
+    cache_dir = os.getenv("JAWBREAKER_MODEL_CACHE")
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError as exc:
+        raise RuntimeError("huggingface_hub is required to download the configured model file.") from exc
+
+    return Path(hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir))
+
+
+def run_analysis(message: str, memory: list[dict] | None) -> ScamAnalysis:
+    memory = memory or []
+    prediction = get_analyzer()(message)
+    similar_memory = ScamAnalysis.from_heuristics(message, memory).similar_memory
+    return prediction_to_analysis(prediction, similar_memory=similar_memory)
+
+
 def analyze_message(message: str, memory: list[dict] | None) -> tuple[str, str, list[dict]]:
     memory = memory or []
-    analysis = ScamAnalysis.from_heuristics(message, memory)
+    try:
+        analysis = run_analysis(message, memory)
+    except Exception as exc:
+        analysis = analysis_error(exc)
     return render_analysis_html(message, analysis), render_memory_html(analysis, memory), memory
 
 
+def analysis_error(exc: Exception) -> ScamAnalysis:
+    return ScamAnalysis(
+        risk_level="needs_check",
+        scam_type="analysis_error",
+        summary="Jawbreaker could not finish the model scan in this environment.",
+        tactics=["runtime error"],
+        safest_action="Do not click links or reply yet. Verify through an official app, website, or known phone number.",
+        trusted_person_message=f"Can you check this for me? Jawbreaker hit an analysis error: {exc}",
+        scam_dna={
+            "Impersonates": "Unknown",
+            "Pressure": "Unknown",
+            "Ask": "Unknown",
+            "Risk": "Could not analyze",
+        },
+    )
+
+
 def remember_current(message: str, memory: list[dict] | None) -> tuple[str, list[dict]]:
     memory = memory or []
-    analysis = ScamAnalysis.from_heuristics(message, memory)
     if not message.strip():
         return "Paste a message first.", memory
+    try:
+        analysis = run_analysis(message, memory)
+    except Exception as exc:
+        analysis = analysis_error(exc)
 
     memory.append(
         {
@@ -36,15 +139,7 @@ def remember_current(message: str, memory: list[dict] | None) -> tuple[str, list
 
 
 def build_app() -> gr.Blocks:
-    css = open("style.css", "r", encoding="utf-8").read()
-    theme = gr.themes.Soft(
-        primary_hue="red",
-        secondary_hue="slate",
-        neutral_hue="zinc",
-        radius_size="sm",
-    )
-
-    with gr.Blocks(title="Jawbreaker", theme=theme, css=css) as demo:
+    with gr.Blocks(title="Jawbreaker") as demo:
         memory_state = gr.State([])
 
         gr.HTML(
@@ -100,5 +195,4 @@ def build_app() -> gr.Blocks:
 
 
 if __name__ == "__main__":
-    build_app().launch()
-
+    build_app().launch(theme=app_theme(), css=app_css())

eval/run_eval.py

@@ -51,6 +51,10 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--n-ctx", type=int, default=4096)
     parser.add_argument("--n-threads", type=int)
     parser.add_argument("--n-gpu-layers", type=int, default=0)
+    parser.add_argument("--n-batch", type=int, default=512)
+    parser.add_argument("--n-ubatch", type=int, default=512)
+    parser.add_argument("--offload-kqv", action=argparse.BooleanOptionalAction, default=True)
+    parser.add_argument("--op-offload", action=argparse.BooleanOptionalAction)
     parser.add_argument("--max-tokens", type=int, default=512)
     parser.add_argument("--temperature", type=float, default=0.0)
     return parser.parse_args()
@@ -131,6 +135,10 @@ def build_analyzer(args: argparse.Namespace):
             n_ctx=args.n_ctx,
             n_threads=args.n_threads,
             n_gpu_layers=args.n_gpu_layers,
+            n_batch=args.n_batch,
+            n_ubatch=args.n_ubatch,
+            offload_kqv=args.offload_kqv,
+            op_offload=args.op_offload,
             max_tokens=args.max_tokens,
             temperature=args.temperature,
         )

jawbreaker/analyzers.py

@@ -30,6 +30,46 @@ def analysis_to_prediction(analysis: ScamAnalysis) -> Prediction:
     }
 
 
+def prediction_to_analysis(prediction: Prediction, *, similar_memory: str = "") -> ScamAnalysis:
+    risk_level = prediction.get("risk_level")
+    if risk_level not in {"dangerous", "suspicious", "needs_check", "safe"}:
+        risk_level = "needs_check"
+
+    tactics = prediction.get("tactics", [])
+    if not isinstance(tactics, list):
+        tactics = []
+
+    scam_dna = prediction.get("scam_dna", {})
+    if not isinstance(scam_dna, dict):
+        scam_dna = {}
+
+    return ScamAnalysis(
+        risk_level=str(risk_level),
+        scam_type=str(prediction.get("scam_type", "unknown")),
+        summary=str(prediction.get("summary", "This message should be checked before anyone acts.")),
+        tactics=[str(tactic) for tactic in tactics],
+        safest_action=str(
+            prediction.get(
+                "safest_action",
+                "Do not click links or reply. Verify through an official app, website, or known phone number.",
+            )
+        ),
+        trusted_person_message=str(
+            prediction.get(
+                "trusted_person_message",
+                "Can you check this for me before I respond or click anything?",
+            )
+        ),
+        scam_dna={
+            "Impersonates": str(scam_dna.get("impersonates", "")),
+            "Pressure": str(scam_dna.get("pressure", "")),
+            "Ask": str(scam_dna.get("ask", "")),
+            "Risk": str(scam_dna.get("risk", "")),
+        },
+        similar_memory=similar_memory,
+    )
+
+
 def heuristic_analyzer(message: str) -> Prediction:
     return analysis_to_prediction(ScamAnalysis.from_heuristics(message))
 
@@ -61,6 +101,10 @@ def build_llama_cpp_analyzer(
     n_ctx: int = 4096,
     n_threads: int | None = None,
     n_gpu_layers: int = 0,
+    n_batch: int = 512,
+    n_ubatch: int = 512,
+    offload_kqv: bool = True,
+    op_offload: bool | None = None,
     max_tokens: int = 512,
     temperature: float = 0.0,
 ) -> Analyzer:
@@ -75,8 +119,13 @@ def build_llama_cpp_analyzer(
         "model_path": str(model_path),
         "n_ctx": n_ctx,
         "n_gpu_layers": n_gpu_layers,
+        "n_batch": n_batch,
+        "n_ubatch": n_ubatch,
+        "offload_kqv": offload_kqv,
         "verbose": False,
     }
+    if op_offload is not None:
+        kwargs["op_offload"] = op_offload
     if chat_format:
         kwargs["chat_format"] = chat_format
     if n_threads is not None:
@@ -132,4 +181,3 @@ def write_predictions(path: Path, rows: Iterable[dict], predictions: dict[str, P
         case_id = row["id"]
         lines.append(json.dumps({"id": case_id, "prediction": predictions[case_id]}, ensure_ascii=True))
     path.write_text("\n".join(lines) + "\n", encoding="utf-8")
-

requirements.txt

@@ -1,3 +1,4 @@
 gradio==6.16.0
+huggingface-hub==1.18.0
+llama-cpp-python==0.3.26
 pytest==8.4.2
-

tests/test_schema.py

@@ -1,4 +1,5 @@
 from jawbreaker.schema import ScamAnalysis
+from jawbreaker.analyzers import prediction_to_analysis
 
 
 def test_family_impersonation_is_dangerous() -> None:
@@ -19,3 +20,26 @@ def test_legitimate_fraud_alert_needs_check_not_dangerous() -> None:
     assert analysis.risk_level == "needs_check"
     assert analysis.scam_type == "possible_legitimate_alert"
 
+
+def test_prediction_to_analysis_normalizes_model_json() -> None:
+    analysis = prediction_to_analysis(
+        {
+            "risk_level": "not_valid",
+            "scam_type": "package_phishing",
+            "summary": "Fake delivery notice.",
+            "tactics": ["fake authority"],
+            "safest_action": "Open the official carrier website yourself.",
+            "trusted_person_message": "Can you check this?",
+            "scam_dna": {
+                "impersonates": "USPS",
+                "pressure": "Held package",
+                "ask": "Open link",
+                "risk": "credential theft",
+            },
+        },
+        similar_memory="This resembles a saved pattern.",
+    )
+
+    assert analysis.risk_level == "needs_check"
+    assert analysis.scam_dna["Impersonates"] == "USPS"
+    assert analysis.similar_memory == "This resembles a saved pattern."