Route Nemotron through Modal

modal_refine.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import os
+import requests
 import modal
 from fastapi import Request
 
@@ -8,50 +10,54 @@ app = modal.App("paper2lab-nemotron")
 image = (
     modal.Image.debian_slim(python_version="3.11")
     .pip_install("requests", "fastapi")
-    .add_local_dir("src/paper2lab", remote_path="/root/paper2lab")
 )
 
 secret = modal.Secret.from_name("nvidia-api-key")
 
+NVIDIA_CHAT_URL = "https://integrate.api.nvidia.com/v1/chat/completions"
+DEFAULT_MODEL = "nvidia/nemotron-3-nano-30b-a3b"
 
-@app.function(
-    image=image,
-    secrets=[secret],
-    timeout=300,
-)
+
+@app.function(image=image, secrets=[secret], timeout=300)
 @modal.fastapi_endpoint(method="POST")
 async def refine_remote(request: Request):
-    from paper2lab.inference.nemotron_refiner import refine_with_nemotron
-
     body = await request.json()
 
-    llm_evidence_pack = body.get("llm_evidence_pack")
-    model = body.get("model", "nvidia/nemotron-3-nano-30b-a3b")
-    return_comparison = body.get("return_comparison", True)
-
-    if not llm_evidence_pack:
-        return {
-            "status": "error",
-            "error": "Missing llm_evidence_pack",
-        }
-
-    try:
-        result = refine_with_nemotron(
-            llm_evidence_pack=llm_evidence_pack,
-            model=model,
-            return_comparison=return_comparison,
-        )
-        return {
-            "status": "ok",
-            "result": result,
-        }
-    except Exception as exc:
-        return {
-            "status": "error",
-            "error": str(exc),
-        }
-
-
-@app.local_entrypoint()
-def main():
-    print("Deploy this app with: modal deploy modal_refine.py")
+    prompt = body.get("prompt")
+    model = body.get("model", DEFAULT_MODEL)
+
+    if not prompt:
+        return {"status": "error", "error": "Missing prompt"}
+
+    api_key = os.environ["NVIDIA_API_KEY"]
+
+    response = requests.post(
+        NVIDIA_CHAT_URL,
+        headers={
+            "Authorization": f"Bearer {api_key}",
+            "Content-Type": "application/json",
+        },
+        json={
+            "model": model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": "You are a precise scientific JSON refiner. Return only valid JSON. No markdown.",
+                },
+                {"role": "user", "content": prompt},
+            ],
+            "temperature": 0.1,
+            "top_p": 0.7,
+            "max_tokens": 8192,
+        },
+        timeout=180,
+    )
+
+    if not response.ok:
+        return {"status": "error", "error": response.text[:1000]}
+
+    data = response.json()
+    return {
+        "status": "ok",
+        "content": data["choices"][0]["message"]["content"],
+    }

src/paper2lab/inference/nemotron_refiner.py

@@ -445,6 +445,39 @@ def diff_cards(before: Dict[str, Any], after: Dict[str, Any]) -> Dict[str, Any]:
 
 
 def _call_nvidia(prompt: str, model: str, timeout: int = 180) -> str:
+    modal_url = os.getenv("MODAL_REFINE_URL")
+    print("MODAL_REFINE_URL:", modal_url)
+
+    # =========================
+    # Use Modal if configured
+    # =========================
+    if modal_url:
+        response = requests.post(
+            modal_url,
+            json={
+                "prompt": prompt,
+                "model": model,
+            },
+            timeout=timeout,
+        )
+
+        if not response.ok:
+            raise RuntimeError(
+                f"Modal error {response.status_code}: {response.text[:1000]}"
+            )
+
+        data = response.json()
+
+        if data.get("status") != "ok":
+            raise RuntimeError(
+                f"Modal refinement failed: {data}"
+            )
+
+        return data["content"]
+
+    # =========================
+    # Fallback: Direct NVIDIA
+    # =========================
     api_key = (
         os.getenv("NVIDIA_API_KEY")
         or os.getenv("NVIDIA_API_KEY".lower())
@@ -452,14 +485,19 @@ def _call_nvidia(prompt: str, model: str, timeout: int = 180) -> str:
     )
 
     if not api_key:
-        raise RuntimeError("Missing NVIDIA_API_KEY. Set it in your environment before using refinement_mode='nemotron'.")
+        raise RuntimeError(
+            "Missing NVIDIA_API_KEY. Set it in your environment before using refinement_mode='nemotron'."
+        )
 
     payload = {
         "model": model,
         "messages": [
             {
                 "role": "system",
-                "content": "You are a precise scientific JSON refiner. Return only valid JSON. No markdown.",
+                "content": (
+                    "You are a precise scientific JSON refiner. "
+                    "Return only valid JSON. No markdown."
+                ),
             },
             {
                 "role": "user",
@@ -491,7 +529,9 @@ def _call_nvidia(prompt: str, model: str, timeout: int = 180) -> str:
     try:
         return data["choices"][0]["message"]["content"]
     except Exception as exc:
-        raise RuntimeError(f"Unexpected NVIDIA response: {data}") from exc
+        raise RuntimeError(
+            f"Unexpected NVIDIA response: {data}"
+        ) from exc
 
 
 def refine_with_nemotron(