modal_app.py

"""
Spruce backend on Modal: the model layer for an AI-native CRM a health coach
talks to instead of maintaining a Kanban board by hand.

Two genuinely small (<=4B) models, each doing the job it is good at:

  * openbmb/MiniCPM3-4B  -> the write path. It reads a raw natural-language
    update ("had a call with Max, booked his analysis call, he wants to drop
    5kg") and does three jobs: ROUTE it to the right client (or flag a new one),
    EXTRACT a structured record and a one-line timeline event, classify the
    pipeline STAGE, and HARVEST any reusable method the coach states.

  * nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16 -> the read/reason path. It ANSWERS
    natural-language questions over the client data ("what's the update on Max?",
    "who have I not contacted in a while?") and writes a grounded coaching BRIEF
    from the coach's own knowledgebase. It never messages a client and never
    introduces outside clinical claims.

Neither model does medicine. One turns language into structured records, the
other reasons over data the coach supplied. That keeps the small-model fit
honest.

The two need different runtimes (MiniCPM3 pins transformers 4.41; Nemotron's
hybrid Mamba needs vLLM's kernels), so they run as two separate Modal services
with their own images, each warm-loaded and scaling to zero.

Prize alignment (build-small-hackathon): Modal (both run here), OpenBMB
(MiniCPM3-4B), NVIDIA (Nemotron), Tiny Titan (both <=4B), Best Agent (route ->
update -> classify -> remind is agentic), Off-Brand (custom CRM UI).

Usage:
  modal deploy modal_app.py     # prints the two web endpoint URLs for the Space
"""

import json
import re

import modal

MINUTES = 60

EXTRACT_MODEL = "openbmb/MiniCPM3-4B"
COACH_MODEL = "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16"

# Elijah's client journey. The model classifies each update into one of these.
STAGES = [
    "Discovery Call",
    "Customer booked",
    "Analysis call",
    "Second call",
    "Weeks 1-4",
    "Weeks 5-8",
]

# Model-maintained fields of a client record. Dates and timeline are kept by the
# cockpit, not the model (the model has no reliable clock).
EMPTY_RECORD = {
    "stage": "",
    "goals": "",
    "current_protocol": "",
    "next_step": "",
    "flags": [],
    "follow_ups": [],
}

# MiniCPM3 remote code targets transformers 4.41; 5.x breaks it.
extract_image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "torch",
        "transformers==4.41.2",
        "accelerate",
        "sentencepiece",
        "numpy",
        "fastapi[standard]",
    )
    .env({"HF_HOME": "/cache"})
)

# Nemotron-3-Nano is a hybrid Mamba/Transformer. vLLM ships NemotronH support and
# the Mamba kernels, which plain transformers does not build cleanly.
# VLLM_USE_FLASHINFER_SAMPLER=0: we decode greedily (temperature 0), so we do not
# need FlashInfer's sampler, which would otherwise JIT-compile a CUDA kernel at
# startup and fail because the slim image has no nvcc.
coach_image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install("vllm", "fastapi[standard]")
    # Remove flashinfer so vLLM uses its native PyTorch sampler. flashinfer would
    # otherwise JIT-compile a CUDA kernel at startup and crash (no nvcc in slim).
    .run_commands("python -m pip uninstall -y flashinfer-python flashinfer || true")
    .env({"HF_HOME": "/cache", "VLLM_USE_FLASHINFER_SAMPLER": "0"})
)

app = modal.App("coach-cockpit")
hf_cache = modal.Volume.from_name("coach-cockpit-cache", create_if_missing=True)


# ---------------------------------------------------------------------------
# Service 1: MiniCPM3-4B. Route, extract, classify, harvest.
# ---------------------------------------------------------------------------
@app.cls(
    gpu="L4",
    image=extract_image,
    volumes={"/cache": hf_cache},
    timeout=20 * MINUTES,
    scaledown_window=300,
    # Keep one container warm so judges never hit a cold start during the demo
    # window. Set back to 0 after judging to stop the idle GPU burn.
    min_containers=1,
)
class Extractor:
    @modal.enter()
    def load(self):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.torch = torch
        self.tok = AutoTokenizer.from_pretrained(EXTRACT_MODEL, trust_remote_code=True)
        if self.tok.pad_token is None:
            self.tok.pad_token = self.tok.eos_token
        self.tok.padding_side = "left"
        self.model = AutoModelForCausalLM.from_pretrained(
            EXTRACT_MODEL,
            trust_remote_code=True,
            torch_dtype=torch.bfloat16,
            device_map="cuda",
        ).eval()

    def _generate(self, messages, max_new_tokens=512):
        prompt = self.tok.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        inputs = self.tok(prompt, return_tensors="pt").to("cuda")
        with self.torch.no_grad():
            out = self.model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                pad_token_id=self.tok.pad_token_id,
            )
        gen = out[:, inputs["input_ids"].shape[1]:]
        return self.tok.decode(gen[0], skip_special_tokens=True).strip()

    def _route_impl(self, text, clients):
        """Decide if this is an update or a question, and which client it is about."""
        roster = ", ".join(clients) if clients else "(no clients yet)"
        sys = (
            "You triage a health coach's inbox. Given a list of existing clients "
            "and one thing the coach typed, decide:\n"
            '  intent: "update" if the coach is reporting something that happened '
            'or a plan, "query" if the coach is asking a question.\n'
            "  client: the existing client this is about, matched from the list "
            "(match partial names, e.g. 'Max' -> 'Max Mustermann'). Empty string "
            "if it is a general question about no single client.\n"
            "  is_new: true only if intent is update and the named person is not "
            "in the list.\n"
            "Return ONLY a JSON object: "
            '{"intent": "...", "client": "...", "is_new": true/false}.'
        )
        user = f"Existing clients: {roster}\n\nCoach typed: {text}"
        out = self._generate(
            [{"role": "system", "content": sys}, {"role": "user", "content": user}],
            max_new_tokens=120,
        )
        obj = _extract_json(out) or {}
        intent = obj.get("intent", "update")
        if intent not in ("update", "query"):
            intent = "update"
        return {
            "intent": intent,
            "client": str(obj.get("client", "") or ""),
            "is_new": bool(obj.get("is_new", False)),
        }

    def _extract_impl(self, existing, new_text):
        """Merge an update into the record, classify the stage, summarize the event."""
        existing = {**EMPTY_RECORD, **(existing or {})}
        sys = (
            "You maintain a health coach's private record on one client. You are "
            "given the CURRENT record as JSON and a NEW update the coach typed. "
            "Update the record using only information present in the text. Never "
            "invent facts. Keep entries short.\n"
            f"Classify the client's current stage as exactly one of: {STAGES}. "
            "If the update does not change the stage, keep the current one.\n"
            "Also write 'event': one short past-tense line summarizing what "
            "happened in this update, for a timeline.\n"
            "Return ONLY a JSON object with these keys:\n"
            f'  "stage": one of {STAGES} or "",\n'
            '  "goals": string,\n'
            '  "current_protocol": string,\n'
            '  "next_step": string (the single next action),\n'
            '  "flags": array of short strings (things to watch),\n'
            '  "follow_ups": array of short strings (open promises),\n'
            '  "event": string (one line, past tense).'
        )
        user = (
            f"CURRENT record:\n{json.dumps(existing, ensure_ascii=False)}\n\n"
            f"NEW update:\n{new_text}"
        )
        out = self._generate(
            [{"role": "system", "content": sys}, {"role": "user", "content": user}]
        )
        obj = _extract_json(out)
        record = _coerce_record(obj, existing)
        event = ""
        if isinstance(obj, dict):
            event = str(obj.get("event", "") or "").strip()
        return {"record": record, "event": event}

    def _harvest_impl(self, new_text):
        """Pull reusable methods the coach states, to grow the knowledgebase."""
        sys = (
            "You read something a health coach wrote. Extract any general method, "
            "protocol, or recommendation that could be reused with other clients. "
            "Ignore client-specific logistics and chit-chat. Return ONLY a JSON "
            'array of objects, each with "title" (a few words) and "body" (one or '
            "two sentences). If nothing is reusable, return []."
        )
        out = self._generate(
            [{"role": "system", "content": sys}, {"role": "user", "content": new_text}],
            max_new_tokens=400,
        )
        items = _extract_json(out)
        if not isinstance(items, list):
            return []
        clean = []
        for it in items:
            if isinstance(it, dict) and it.get("title") and it.get("body"):
                clean.append({"title": str(it["title"]), "body": str(it["body"])})
        return clean

    @modal.method()
    def route(self, text, clients):
        return self._route_impl(text, clients)

    @modal.method()
    def extract(self, existing, new_text):
        return self._extract_impl(existing, new_text)

    @modal.method()
    def harvest(self, new_text):
        return self._harvest_impl(new_text)

    @modal.fastapi_endpoint(method="POST")
    def web(self, payload: dict):
        """POST one of:
        {"op":"route","text":"...","clients":[...]}
        {"op":"extract","existing":{...},"new_text":"..."}
        {"op":"harvest","new_text":"..."}"""
        op = payload.get("op", "extract")
        if op == "route":
            return self._route_impl(payload.get("text", ""), payload.get("clients", []))
        if op == "harvest":
            return {"methods": self._harvest_impl(payload.get("new_text", ""))}
        return self._extract_impl(payload.get("existing"), payload.get("new_text", ""))


# ---------------------------------------------------------------------------
# Service 2: Nemotron-3-Nano-4B on vLLM. Answer questions, write briefs.
# ---------------------------------------------------------------------------
@app.cls(
    gpu="L4",
    image=coach_image,
    volumes={"/cache": hf_cache},
    timeout=20 * MINUTES,
    scaledown_window=300,
    # vLLM cold start is ~90s — keep one container warm so the demo never eats it.
    # Set back to 0 after judging to stop the idle GPU burn.
    min_containers=1,
)
class Coach:
    @modal.enter()
    def load(self):
        from vllm import LLM, SamplingParams

        self.SamplingParams = SamplingParams
        self.llm = LLM(
            model=COACH_MODEL,
            trust_remote_code=True,
            dtype="bfloat16",
            max_model_len=8192,
            gpu_memory_utilization=0.92,
            enforce_eager=True,  # skip CUDA graph capture for a faster cold start
        )

    def _generate(self, messages, max_new_tokens=700):
        # Nemotron-3-Nano is a reasoning model: it emits <think>...</think> before
        # the answer. Give it room to think, then keep only the post-think answer.
        params = self.SamplingParams(temperature=0.0, max_tokens=max_new_tokens)
        outs = self.llm.chat(messages, params, use_tqdm=False)
        text = outs[0].outputs[0].text.strip()
        if "</think>" in text:
            text = text.split("</think>")[-1].strip()
        return text

    def _answer_impl(self, context, question):
        """Answer a natural-language question over the CRM data given as context."""
        sys = (
            "You are the assistant inside a health coach's CRM. Answer the coach's "
            "question using ONLY the data below. Be concise and direct. If the data "
            "does not contain the answer, say so plainly. Do not invent clients, "
            "dates, or facts."
        )
        user = f"CRM DATA:\n{context}\n\nQUESTION: {question}"
        return self._generate(
            [{"role": "system", "content": sys}, {"role": "user", "content": user}]
        )

    def _brief_impl(self, record, kb, question):
        """A grounded brief from the coach's own knowledgebase before they reply."""
        kb_text = (
            "\n".join(f"- {e.get('title','')}: {e.get('body','')}" for e in kb)
            if kb
            else "(no matching knowledgebase entries)"
        )
        sys = (
            "You brief a health coach before they reply to a client. You are not "
            "the coach and you do not message the client. Use ONLY the client "
            "record and the coach's own knowledgebase below. Do not introduce "
            "outside medical claims. Write short bullet sections:\n"
            "  RELEVANT METHODS: which knowledgebase entries apply and why (name them).\n"
            "  ALREADY NOTED: what the record already shows.\n"
            "  WATCH FOR: flags to respect.\n"
            "  TALKING POINTS: a few angles the coach could raise.\n"
            "End with: 'Not medical advice. Coach reviews before sending.'"
        )
        user = (
            f"CLIENT RECORD:\n{json.dumps(record or EMPTY_RECORD, ensure_ascii=False)}"
            f"\n\nKNOWLEDGEBASE:\n{kb_text}"
            f"\n\nCOACH'S QUESTION: {question or 'What should I consider before replying?'}"
        )
        return self._generate(
            [{"role": "system", "content": sys}, {"role": "user", "content": user}]
        )

    @modal.method()
    def answer(self, context, question):
        return self._answer_impl(context, question)

    @modal.method()
    def brief(self, record, kb, question=""):
        return self._brief_impl(record, kb, question)

    @modal.fastapi_endpoint(method="POST")
    def web(self, payload: dict):
        """POST one of:
        {"op":"answer","context":"...","question":"..."}
        {"op":"brief","record":{...},"kb":[{title,body}],"question":"..."}"""
        op = payload.get("op", "answer")
        if op == "brief":
            return {
                "brief": self._brief_impl(
                    payload.get("record"), payload.get("kb", []), payload.get("question", "")
                )
            }
        return {
            "answer": self._answer_impl(
                payload.get("context", ""), payload.get("question", "")
            )
        }


# ---------------------------------------------------------------------------
# Parsing helpers.
# ---------------------------------------------------------------------------
def _extract_json(text):
    """Pull the first JSON object or array out of model output, defensively."""
    if not text:
        return None
    text = re.sub(r"```(?:json)?", "", text)
    start = None
    for i, ch in enumerate(text):
        if ch in "{[":
            start = i
            break
    if start is None:
        return None
    opener = text[start]
    closer = "}" if opener == "{" else "]"
    depth = 0
    for j in range(start, len(text)):
        if text[j] == opener:
            depth += 1
        elif text[j] == closer:
            depth -= 1
            if depth == 0:
                try:
                    return json.loads(text[start : j + 1])
                except json.JSONDecodeError:
                    return None
    return None


def _coerce_record(obj, existing):
    """Force whatever the model returned into the record shape. Never raises."""
    existing = {**EMPTY_RECORD, **(existing or {})}
    if not isinstance(obj, dict):
        return {k: existing[k] for k in EMPTY_RECORD}
    out = {}
    stage = str(obj.get("stage", existing.get("stage", "")) or "")
    out["stage"] = stage if stage in STAGES else existing.get("stage", "")
    for key in ("goals", "current_protocol", "next_step"):
        out[key] = str(obj.get(key, existing.get(key, "")) or "")
    for key in ("flags", "follow_ups"):
        val = obj.get(key, existing.get(key, []))
        if isinstance(val, str):
            val = [val] if val else []
        elif not isinstance(val, list):
            val = list(existing.get(key, []))
        out[key] = [str(x) for x in val if str(x).strip()]
    return out


# ---------------------------------------------------------------------------
# Synthetic smoke test. Run with: modal run --detach modal_app.py
# (detach keeps the app alive through the slow vLLM cold start).
# ---------------------------------------------------------------------------
_SAMPLE_UPDATE = (
    "Had my discovery call with Max Mustermann today. He wants to drop 5kg before "
    "September and fix his afternoon energy crashes. Booked his analysis call for "
    "Thursday and told him to start a food log. As a rule I have clients front-load "
    "protein at breakfast when they report afternoon dips."
)


@app.local_entrypoint()
def main():
    ext = Extractor()
    coach = Coach()

    print("=== route: who is this about, update or question? ===")
    print(ext.route.remote(_SAMPLE_UPDATE, ["Dana R.", "Sam P."]))

    print("\n=== extract: structured record + stage + timeline event ===")
    res = ext.extract.remote(EMPTY_RECORD, _SAMPLE_UPDATE)
    print(json.dumps(res, indent=2, ensure_ascii=False))

    print("\n=== harvest: reusable methods for the knowledgebase ===")
    print(json.dumps(ext.harvest.remote(_SAMPLE_UPDATE), indent=2, ensure_ascii=False))

    print("\n=== answer: concierge Q&A over the record ===")
    ctx = json.dumps({"Max Mustermann": res["record"]}, ensure_ascii=False)
    print(coach.answer.remote(ctx, "What is the next step for Max?"))