agent.py

"""
agent.py — the agentic core
----------------------------
Loads one instruct LLM with native tool-calling support and runs the
bounded tool-calling loop: model -> parse <tool_call> blocks -> execute
the matching sandboxed function from tools.py -> feed the result back as
a "tool" turn -> repeat until the model answers in plain text.

Works for both the text chat tab and the realtime voice tab (voice just
transcribes to text first, then calls `run_agent_turn`, then speaks the
returned text).

Model choice
------------
Any open instruct model whose chat template supports `tools=` and emits
Hermes/Qwen-style `<tool_call>{...}</tool_call>` blocks will work here
out of the box (Qwen2.5/Qwen3-Instruct, Mistral-Nemo-Instruct, etc.).
Default is Qwen2.5-7B-Instruct: solid multilingual coverage (including
Malayalam) and a well-tested tool-calling template. Swap MODEL_ID via
the MEDGUIDE_MODEL_ID env var without touching any other code.
"""

from __future__ import annotations

import json
import os
import re
import threading

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

from tools import TOOLS, TOOL_REGISTRY

try:
    import spaces
    GPU_DECORATOR = spaces.GPU
except ImportError:  # local / CPU dev fallback
    def GPU_DECORATOR(fn=None, **kwargs):
        return fn if fn is not None else (lambda f: f)

MODEL_ID = os.environ.get("MEDGUIDE_MODEL_ID", "Qwen/Qwen2.5-7B-Instruct")
MAX_TOOL_ROUNDS = 4
MAX_NEW_TOKENS = 512

SYSTEM_PROMPT = """You are MedGuide, a general health-information assistant \
available by text and by realtime voice, in English and in Malayalam \
(മലയാളം), for a "pre-visit" conversation before someone sees a clinician.

You are NOT a doctor, nurse, or pharmacist. You cannot diagnose conditions, \
prescribe medication, or replace professional medical care.

Language rule: always reply in the SAME language the user just used \
(English or Malayalam). If they mix languages, mirror the mix naturally. \
Keep sentences short and simple -- your text may be read aloud by a \
text-to-speech engine.

Rules you must always follow:
1. If the user describes symptoms that could indicate a medical emergency, \
call the symptom_urgency_triage tool (translate symptom phrases to English \
for the tool call itself). If it returns "emergency", clearly and \
immediately tell them, in their own language, to call emergency services \
(108 in India, 911 in the US, or their local equivalent) or go to the \
nearest emergency room -- before anything else.
2. Use your tools to produce concrete numbers (BMI, unit conversions, \
vital-sign reference ranges, interaction awareness notes, triage level) \
instead of guessing or calculating them yourself. Never state a number a \
tool could have computed without calling the tool.
3. Never provide specific dosing instructions for prescription or \
controlled substances. You may share general, widely-published OTC dosing \
information but always tell the user to confirm against the package label \
or a pharmacist.
4. Always make clear you provide general information only and encourage \
the user to consult a licensed healthcare professional for diagnosis, \
treatment, or prescriptions.
5. Be warm, clear, and plain-spoken. No alarming language unless the \
situation is genuinely urgent per rule 1.
6. This is a pre-visit intake conversation: where useful, ask brief \
follow-up questions (one at a time) to gather chief complaint, duration, \
severity, and relevant vitals, so a clinician can be handed a clean summary \
later.
"""

_TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)

_tokenizer = None
_model = None
_lock = threading.Lock()


def _load():
    global _tokenizer, _model
    if _model is not None:
        return
    with _lock:
        if _model is not None:
            return
        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        _model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            dtype=torch.bfloat16,
            device_map="auto",
        )
        _model.eval()


def _extract_tool_calls(text: str):
    calls = []
    for blob in _TOOL_CALL_RE.findall(text):
        try:
            obj = json.loads(blob)
        except json.JSONDecodeError:
            continue
        name = obj.get("name")
        args = obj.get("arguments", {})
        if isinstance(args, str):
            try:
                args = json.loads(args)
            except json.JSONDecodeError:
                args = {}
        calls.append({"name": name, "arguments": args})
    return calls


def _generate(messages, tools=True, stream_to: TextIteratorStreamer | None = None):
    _load()
    text = _tokenizer.apply_chat_template(
        messages,
        tools=TOOLS if tools else None,
        tokenize=False,
        add_generation_prompt=True,
    )
    inputs = _tokenizer(text, return_tensors="pt").to(_model.device)
    gen_kwargs = dict(
        **inputs,
        max_new_tokens=MAX_NEW_TOKENS,
        do_sample=True,
        temperature=0.6,
        top_p=0.9,
        pad_token_id=_tokenizer.eos_token_id,
    )
    if stream_to is not None:
        gen_kwargs["streamer"] = stream_to
        thread = threading.Thread(target=_model.generate, kwargs=gen_kwargs)
        thread.start()
        return thread
    with torch.no_grad():
        out = _model.generate(**gen_kwargs)
    gen_tokens = out[0][inputs["input_ids"].shape[-1]:]
    return _tokenizer.decode(gen_tokens, skip_special_tokens=True)


def _history_to_messages(history: list) -> list:
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in history or []:
        role, content = turn.get("role"), turn.get("content")
        if role in ("user", "assistant") and content:
            messages.append({"role": role, "content": content})
    return messages


@GPU_DECORATOR(duration=60)
def run_agent_turn(user_message: str, history: list):
    """Runs the bounded tool-calling loop for one user turn.

    Returns (answer_text, trace) where trace is a list of human-readable
    "tool used" strings for display in the UI.
    """
    messages = _history_to_messages(history)
    messages.append({"role": "user", "content": user_message})

    trace = []
    answer = ""
    for _round in range(MAX_TOOL_ROUNDS):
        answer = _generate(messages, tools=True)
        calls = _extract_tool_calls(answer)
        if not calls:
            # strip any stray tool-call tags the model didn't close properly
            return re.sub(r"</?tool_call>.*", "", answer, flags=re.DOTALL).strip(), trace

        messages.append({"role": "assistant", "content": answer})
        for call in calls:
            name, args = call["name"], call["arguments"]
            if name not in TOOL_REGISTRY:
                response = {"error": f"'{name}' is not an available tool."}
            else:
                try:
                    response = TOOL_REGISTRY[name](**args)
                except TypeError as e:
                    response = {"error": f"Bad arguments for {name}: {e}"}
            trace.append(f"🔧 {name}({json.dumps(args, ensure_ascii=False)}) → "
                         f"{json.dumps(response, ensure_ascii=False)}")
            messages.append({
                "role": "tool",
                "name": name,
                "content": json.dumps(response, ensure_ascii=False),
            })

    return (answer.strip() or
            "I wasn't able to finish reasoning about that — could you "
            "rephrase or simplify your question?"), trace


@GPU_DECORATOR(duration=90)
def generate_previsit_summary(history: list, language: str = "English") -> str:
    """Turns the conversation so far into a short, structured note a
    clinician can skim before the visit. No new medical claims -- purely
    a summarization of what was already said/computed in this chat."""
    convo = "\n".join(
        f"{t.get('role', '?').upper()}: {t.get('content', '')}"
        for t in (history or []) if t.get("content")
    )
    lang_instruction = (
        "Write the summary in Malayalam." if language == "Malayalam"
        else "Write the summary in English."
    )
    prompt = (
        "Summarize the pre-visit conversation below into a short clinical "
        "handoff note with these headings: Chief Complaint, Duration/Onset, "
        "Reported Symptoms, Vitals/Numbers Discussed (only ones explicitly "
        "mentioned), Triage Level (if computed), Notes for Clinician. "
        "Do not invent any information that wasn't in the conversation. "
        f"{lang_instruction}\n\n--- CONVERSATION ---\n{convo}"
    )
    messages = [
        {"role": "system", "content": "You write concise, accurate clinical "
                                       "handoff summaries. You never invent "
                                       "facts not present in the source text."},
        {"role": "user", "content": prompt},
    ]
    return _generate(messages, tools=False).strip()