""" agent.py — the agentic core ---------------------------- Loads one instruct LLM with native tool-calling support and runs the bounded tool-calling loop: model -> parse 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 `{...}` 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"\s*(\{.*?\})\s*", 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".*", "", 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()