Add MLX remote server support via MLX_SERVER_URL env var

app.py

@@ -320,6 +320,45 @@ def _run_pipe(prompt: str) -> str:
     return result[0]["generated_text"]
 
 
+# ---------------------------------------------------------------------------
+# MLX Remote Server (Mac Mini via ngrok) - takes priority when MLX_SERVER_URL is set
+# ---------------------------------------------------------------------------
+_MLX_SERVER_URL = os.environ.get("MLX_SERVER_URL", "").rstrip("/")
+_mlx_available = False
+
+if _MLX_SERVER_URL:
+    try:
+        import urllib.request as _ur
+        _ur.urlopen(_MLX_SERVER_URL + "/v1/models", timeout=5)
+        _mlx_available = True
+        print("MLX remote server ready: " + _MLX_SERVER_URL)
+    except Exception as _e:
+        print("MLX server unreachable (" + str(_e) + "), falling back to local model.")
+
+
+def _run_via_mlx(messages, strict=False):
+    import urllib.request as _ur, json as _j
+    msgs = list(messages)
+    if strict:
+        msgs[-1] = dict(msgs[-1])
+        msgs[-1]["content"] += STRICT_SUFFIX
+    payload = _j.dumps({
+        "model": "mlx-community/Qwen3.5-9B-MLX-4bit",
+        "messages": msgs,
+        "max_tokens": 512,
+        "temperature": 0.0,
+    }).encode()
+    req = _ur.Request(
+        _MLX_SERVER_URL + "/v1/chat/completions",
+        data=payload,
+        headers={"Content-Type": "application/json"},
+    )
+    with _ur.urlopen(req, timeout=120) as resp:
+        data = _j.loads(resp.read())
+    return data["choices"][0]["message"]["content"]
+
+
+
 # ---------------------------------------------------------------------------
 # Sample contract content
 # ---------------------------------------------------------------------------
@@ -389,7 +428,7 @@ def analyze_contract(contract_text: str, question: str) -> tuple[str, str, str,
         return format_label_html("N/A"), "", "", "Please paste a contract above."
     if not question.strip():
         return format_label_html("N/A"), "", "", "Please enter a question."
-    if _pipe is None:
+    if _pipe is None and not _mlx_available:
         return (
             format_label_html("N/A"),
             "Model not loaded",
@@ -402,9 +441,12 @@ def analyze_contract(contract_text: str, question: str) -> tuple[str, str, str,
     messages = _build_contract_messages(contract_text, question)
 
     for attempt in range(2):
-        prompt = _apply_template(messages, strict=(attempt == 1))
         try:
-            raw = _run_pipe(prompt)
+            if _mlx_available:
+                raw = _run_via_mlx(messages, strict=(attempt == 1))
+            else:
+                prompt = _apply_template(messages, strict=(attempt == 1))
+                raw = _run_pipe(prompt)
             result = _parse_model_output(raw, question)
             label_html = format_label_html(result.label.value)
             answer = result.answer or "(none — clause is absent or not applicable)"
@@ -470,7 +512,7 @@ def analyse_bank_statement(paste_text: str, pdf_file, csv_file) -> tuple[str, st
     statement_text, error = _get_statement_text(paste_text, pdf_file, csv_file)
     if error:
         return f"**Error:** {error}", ""
-    if _pipe is None:
+    if _pipe is None and not _mlx_available:
         return (
             f"**Model not loaded.** Set `HF_MODEL_REPO` in Space secrets. Error: {model_load_error}",
             statement_text,
@@ -480,9 +522,12 @@ def analyse_bank_statement(paste_text: str, pdf_file, csv_file) -> tuple[str, st
     messages = _build_bank_messages(statement_text, "SUMMARISE")
 
     for attempt in range(2):
-        prompt = _apply_template(messages, strict=(attempt == 1))
         try:
-            raw = _run_pipe(prompt)
+            if _mlx_available:
+                raw = _run_via_mlx(messages, strict=(attempt == 1))
+            else:
+                prompt = _apply_template(messages, strict=(attempt == 1))
+                raw = _run_pipe(prompt)
             summary = _parse_summary(raw)
             lines = ["## Statement Summary", ""]
             lines.append(f"**Total Credits:** {summary.total_credits or 'N/A'}")
@@ -515,7 +560,7 @@ def bank_qa(statement_text: str, question: str) -> tuple[str, str, str, str]:
         )
     if not question.strip():
         return format_label_html("N/A"), "", "", "Please enter a question."
-    if _pipe is None:
+    if _pipe is None and not _mlx_available:
         return (
             format_label_html("N/A"), "Model not loaded", "",
             f"Model failed to load: {model_load_error}.",
@@ -525,9 +570,12 @@ def bank_qa(statement_text: str, question: str) -> tuple[str, str, str, str]:
     messages = _build_bank_messages(statement_text, question)
 
     for attempt in range(2):
-        prompt = _apply_template(messages, strict=(attempt == 1))
         try:
-            raw = _run_pipe(prompt)
+            if _mlx_available:
+                raw = _run_via_mlx(messages, strict=(attempt == 1))
+            else:
+                prompt = _apply_template(messages, strict=(attempt == 1))
+                raw = _run_pipe(prompt)
             result = _parse_model_output(raw, question)
             label_html = format_label_html(result.label.value)
             answer = result.answer or "(none — information not found in statement)"