{"repo": "chatmol/chatmol", "files": {"pymol_plugin/v2/chatmol.py": {"d376ccd803a94e2d59bd024a190e17f925e5ee84": {"t": 1785745411.3365545, "n": 59743, "h": "5a4226f171e50d6d", "c": "\"\"\"\nChatMol — Agentic PyMOL Plugin with OpenRouter and Qt5 GUI.\n\nLoad in PyMOL:\n    run /path/to/pymol_plugin/v2/chatmol.py\n\nOr from remote:\n    load https://raw.githubusercontent.com/ChatMol/ChatMol/main/pymol_plugin/v2/chatmol.py\n\nDemo prompt:\nZH: 3wzm的A链是玉米赤霉烯酮水解酶和底物的共晶复合物，作图展示酶底物(ZER)相互作用，风格参考顶级结构生物学期刊的风格，画面干净、清晰。\nEN: Chain A of 3WZM is a cocrystal complex of zearalenone hydrolase and its substrate; create a figure showing the enzyme-substrate (ZER) interaction, referencing the style of top-tier structural biology journals, with a clean and clear aesthetic.\n\"\"\"\n\nimport os\nimport json\nimport re\nimport time\nimport random\nimport base64\nimport tempfile\nimport threading\nimport requests\n\nfrom datetime import datetime\nfrom pymol import cmd, util, preset\n\n# ---------------------------------------------------------------------------\n# 1. Provider registry + LLMClient — multi-provider HTTP wrapper (no SDK)\n# ---------------------------------------------------------------------------\n\n# Each provider entry:\n#   base_url  – chat completions endpoint\n#   env_var   – environment variable for the API key\n#   models    – list of (model_id, label) for the settings UI\n#   extra_headers – provider-specific headers (optional)\nPROVIDERS = {\n    \"openrouter\": {\n        \"label\": \"OpenRouter\",\n        \"base_url\": \"https://openrouter.ai/api/v1/chat/completions\",\n        \"env_var\": \"OPENROUTER_API_KEY\",\n        \"models\": [\n            (\"openai/gpt-4o\", \"GPT-4o\"),\n            (\"openai/gpt-5.2\", \"GPT-5.2\"),\n            (\"openai/gpt-5.3-codex\", \"GPT-5.3 Codex\"),\n            (\"google/gemini-3.1-pro-preview\", \"Gemini 3.1 Pro\"),\n            (\"google/gemini-3-flash-preview\", \"Gemini 3 Flash\"),\n            (\"deepseek/deepseek-chat\", \"DeepSeek V3\"),\n        ],\n        \"vision_models\": [\n            (\"openai/gpt-4o\", \"GPT-4o\"),\n            (\"openai/gpt-5.2\", \"GPT-5.2\"),\n            (\"anthropic/claude-sonnet-4.6\", \"Claude Sonnet 4.6\"),\n            (\"google/gemini-3-flash-preview\", \"Gemini 3 Flash\"),\n            (\"google/gemini-3.1-pro-preview\", \"Gemini 3.1 Pro\"),\n        ],\n        \"extra_headers\": {\n            \"HTTP-Referer\": \"https://github.com/ChatMol/ChatMol\",\n            \"X-Title\": \"ChatMol\",\n        },\n    },\n    \"deepseek\": {\n        \"label\": \"DeepSeek\",\n        \"base_url\": \"https://api.deepseek.com/chat/completions\",\n        \"env_var\": \"DEEPSEEK_API_KEY\",\n        \"models\": [\n            (\"deepseek-chat\", \"DeepSeek V3\"),\n            (\"deepseek-reasoner\", \"DeepSeek R1\"),\n        ],\n        \"vision_models\": [],\n        \"extra_headers\": {},\n    },\n    \"kimi\": {\n        \"label\": \"Kimi (Moonshot)\",\n        \"base_url\": \"https://api.moonshot.ai/v1/chat/completions\",\n        \"env_var\": \"MOONSHOT_API_KEY\",\n        \"models\": [\n            (\"kimi-k2.5\", \"Kimi K2.5\"),\n        ],\n        \"vision_models\": [\n            (\"kimi-k2.5\", \"Kimi K2.5\"),\n        ],\n        \"extra_headers\": {},\n    },\n    \"glm\": {\n        \"label\": \"GLM (Zhipu)\",\n        \"base_url\": \"https://open.bigmodel.cn/api/paas/v4/chat/completions\",\n        \"env_var\": \"GLM_API_KEY\",\n        \"models\": [\n            (\"glm-5\", \"GLM-5\"),\n        ],\n        \"vision_models\": [],\n        \"extra_headers\": {},\n    },\n}\n\n\nclass LLMTransientError(RuntimeError):\n    \"\"\"Transient API/network errors that should be retried.\"\"\"\n\n\nclass LLMFatalError(RuntimeError):\n    \"\"\"Non-retriable API errors.\"\"\"\n\n\nclass LLMClient:\n    \"\"\"Thin HTTP client for OpenAI-compatible chat-completion APIs.\"\"\"\n\n    def __init__(self, provider_name, api_key):\n        prov = PROVIDERS.get(provider_name, PROVIDERS[\"openrouter\"])\n        self.base_url = prov[\"base_url\"]\n        self.extra_headers = prov.get(\"extra_headers\", {})\n        self.api_key = api_key\n\n    def _headers(self):\n        h = {\n            \"Content-Type\": \"application/json\",\n            \"Authorization\": f\"Bearer {self.api_key}\",\n        }\n        h.update(self.extra_headers)\n        return h\n\n    def _check_response(self, resp):\n        if resp.status_code >= 400:\n            try:\n                body = resp.json()\n                err = body.get(\"error\", {})\n                if isinstance(err, dict):\n                    detail = err.get(\"message\", \"\") or err.get(\"msg\", \"\")\n                else:\n                    detail = str(err)\n            except Exception:\n                detail = resp.text[:500] if resp.text else \"\"\n            message = (\n                f\"API error {resp.status_code} from {self.base_url}: \"\n                f\"{detail or resp.reason}\"\n            )\n            if resp.status_code in (408, 409, 425, 429) or resp.status_code >= 500:\n                raise LLMTransientError(message)\n            raise LLMFatalError(message)\n\n    def _parse_choices(self, data):\n        if \"error\" in data:\n            err = data[\"error\"]\n            msg = err.get(\"message\", str(err)) if isinstance(err, dict) else str(err)\n            raise LLMFatalError(f\"API returned error: {msg}\")\n        if \"choices\" not in data or not data[\"choices\"]:\n            raise LLMFatalError(\n                f\"Unexpected API response (no choices): \"\n                f\"{json.dumps(data, ensure_ascii=False)[:300]}\"\n            )\n        return data[\"choices\"][0]\n\n    def chat_completion(\n        self, model, messages, tools=None, temperature=0.01, max_tokens=4096\n    ):\n        payload = {\n            \"model\": model,\n            \"messages\": messages,\n            \"temperature\": temperature,\n            \"max_tokens\": max_tokens,\n        }\n        if tools:\n            payload[\"tools\"] = tools\n        try:\n            resp = requests.post(\n                self.base_url, headers=self._headers(), json=payload, timeout=120\n            )\n        except requests.ConnectionError:\n            raise LLMTransientError(f\"Connection failed: cannot reach {self.base_url}\")\n        except requests.Timeout:\n            raise LLMTransientError(f\"Request timed out after 120s to {self.base_url}\")\n        self._check_response(resp)\n        data = resp.json()\n        self._parse_choices(data)  # validate structure\n        return data\n\n    def vision_completion(self, model, text_prompt, image_base64):\n        messages = [\n            {\n                \"role\": \"user\",\n                \"content\": [\n                    {\"type\": \"text\", \"text\": text_prompt},\n                    {\n                        \"type\": \"image_url\",\n                        \"image_url\": {\"url\": f\"data:image/png;base64,{image_base64}\"},\n                    },\n                ],\n            }\n        ]\n        payload = {\n            \"model\": model,\n            \"messages\": messages,\n            \"max_tokens\": 1024,\n        }\n        try:\n            resp = requests.post(\n                self.base_url, headers=self._headers(), json=payload, timeout=120\n            )\n        except requests.ConnectionError:\n            raise LLMTransientError(f\"Connection failed: cannot reach {self.base_url}\")\n        except requests.Timeout:\n            raise LLMTransientError(f\"Request timed out after 120s to {self.base_url}\")\n        self._check_response(resp)\n        data = resp.json()\n        choice = self._parse_choices(data)\n        return choice[\"message\"][\"content\"]\n\n\n# ---------------------------------------------------------------------------\n# 2. PyMOLTools — 4 tool definitions + execution\n# ---------------------------------------------------------------------------\n\n\nBLOCKED_COMMANDS = [\n    (\"blocked destructive command\", re.compile(r\"^\\s*delete\\s+all\\s*$\", re.I)),\n    (\"blocked destructive command\", re.compile(r\"^\\s*remove\\s+all\\s*$\", re.I)),\n    (\"blocked session reset command\", re.compile(r\"^\\s*reinitialize(?:\\s|$)\", re.I)),\n    (\"blocked quit command\", re.compile(r\"^\\s*quit(?:\\s|$)\", re.I)),\n    (\"blocked script execution command\", re.compile(r\"^\\s*run(?:\\s|$)\", re.I)),\n    (\"blocked python command\", re.compile(r\"^\\s*python(?:\\s|$)\", re.I)),\n    (\n        \"blocked shell/system command\",\n        re.compile(\n            r\"^\\s*(?:!|shell(?:\\s|$)|system(?:\\s|$)|os\\.system|subprocess\\.)\", re.I\n        ),\n    ),\n]\n\nTOOL_DEFINITIONS = [\n    {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": \"inspect_session\",\n            \"description\": (\n                \"Inspect the current PyMOL session. Returns structured JSON \"\n                \"with objects, chain summaries, atom counts, and selections.\"\n            ),\n            \"parameters\": {\n                \"type\": \"object\",\n                \"properties\": {\n                    \"include_view\": {\n                        \"type\": \"boolean\",\n                        \"description\": \"Include camera view matrix.\",\n                        \"default\": False,\n                    }\n                },\n            },\n        },\n    },\n    {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": \"run_pymol_commands\",\n            \"description\": (\n                \"Execute newline-separated PyMOL commands. Use this for all \"\n                \"PyMOL operations: selections, styling, coloring, distances, \"\n                \"presets, etc. Commands are run sequentially with safety \"\n                \"guardrails (destructive/shell commands are blocked).\"\n            ),\n            \"parameters\": {\n                \"type\": \"object\",\n                \"properties\": {\n                    \"commands\": {\n                        \"type\": \"string\",\n                        \"description\": \"Newline-separated PyMOL commands.\",\n                    }\n                },\n                \"required\": [\"commands\"],\n            },\n        },\n    },\n    {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": \"render\",\n            \"description\": (\n                \"Render and export an image from the current scene. \"\n                \"Use purpose=preview for quick iterative checks, and \"\n                \"purpose=final for final publication export.\"\n            ),\n            \"parameters\": {\n                \"type\": \"object\",\n                \"properties\": {\n                    \"path\": {\n                        \"type\": \"string\",\n                        \"description\": \"Output image path.\",\n                    },\n                    \"purpose\": {\n                        \"type\": \"string\",\n                        \"enum\": [\"preview\", \"final\"],\n                        \"default\": \"preview\",\n                        \"description\": \"preview is lightweight; final is high quality.\",\n                    },\n                    \"width\": {\"type\": \"integer\", \"default\": 1400},\n                    \"height\": {\"type\": \"integer\", \"default\": 1050},\n                    \"dpi\": {\"type\": \"integer\", \"default\": 150},\n                    \"ray\": {\n                        \"type\": \"boolean\",\n                        \"description\": \"If omitted: preview=False, final=True.\",\n                    },\n                    \"transparent_bg\": {\"type\": \"boolean\", \"default\": True},\n                },\n                \"required\": [\"path\"],\n            },\n        },\n    },\n    {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": \"capture_viewport\",\n            \"description\": (\n                \"Capture current viewport and run vision model analysis. \"\n                \"Use this to visually verify your work (e.g., check that \"\n                \"styling looks correct before finalizing).\"\n            ),\n            \"parameters\": {\n                \"type\": \"object\",\n                \"properties\": {\n                    \"analysis_prompt\": {\n                        \"type\": \"string\",\n                        \"default\": \"Describe what you see in this molecular visualization.\",\n                    }\n                },\n            },\n        },\n    },\n]\n\n\ndef _to_bool(value, default=False):\n    if isinstance(value, bool):\n        return value\n    if isinstance(value, str):\n        return value.strip().lower() in (\"1\", \"true\", \"yes\", \"y\", \"on\")\n    if value is None:\n        return default\n    return bool(value)\n\n\ndef _blocked_reason(line):\n    for reason, pattern in BLOCKED_COMMANDS:\n        if pattern.search(line):\n            return reason\n    return \"\"\n\n\ndef execute_tool(tool_name, arguments, client=None, vision_model=None):\n    \"\"\"Dispatch to a tool implementation and always return JSON text.\"\"\"\n    dispatch = {\n        \"inspect_session\": _tool_inspect_session,\n        \"run_pymol_commands\": _tool_run_pymol_commands,\n        \"render\": _tool_render,\n        \"capture_viewport\": _tool_capture_viewport,\n    }\n    fn = dispatch.get(tool_name)\n    if fn is None:\n        return json.dumps(\n            {\"ok\": False, \"tool\": tool_name, \"error\": f\"Unknown tool: {tool_name}\"},\n            ensure_ascii=False,\n        )\n    try:\n        if tool_name == \"capture_viewport\":\n            result = fn(client=client, vision_model=vision_model, **(arguments or {}))\n        else:\n            result = fn(**(arguments or {}))\n    except TypeError as exc:\n        result = {\"ok\": False, \"tool\": tool_name, \"error\": f\"Invalid arguments: {exc}\"}\n    except Exception as exc:\n        result = {\"ok\": False, \"tool\": tool_name, \"error\": str(exc)}\n    return json.dumps(result, ensure_ascii=False)\n\n\ndef _tool_inspect_session(include_view=False):\n    objects = cmd.get_names(\"objects\")\n    selections = cmd.get_names(\"selections\")\n    details = []\n    for obj in objects:\n        details.append(\n            {\n                \"name\": obj,\n                \"atoms\": cmd.count_atoms(obj),\n                \"polymer_atoms\": cmd.count_atoms(f\"({obj}) and polymer\"),\n                \"hetatm_atoms\": cmd.count_atoms(f\"({obj}) and hetatm\"),\n                \"chains\": cmd.get_chains(obj),\n                \"states\": cmd.count_states(obj),\n            }\n        )\n    out = {\n        \"ok\": True,\n        \"tool\": \"inspect_session\",\n        \"object_count\": len(objects),\n        \"selection_count\": len(selections),\n        \"objects\": details,\n        \"selections\": selections,\n    }\n    if _to_bool(include_view):\n        out[\"view\"] = list(cmd.get_view())\n    return out\n\n\ndef _tool_run_pymol_commands(commands):\n    lines = [\n        ln.strip()\n        for ln in commands.splitlines()\n        if ln.strip() and not ln.strip().startswith(\"#\")\n    ]\n    executed, blocked, errors = [], [], []\n    known_names = set(cmd.get_names(\"all\"))\n    for line in lines:\n        reason = _blocked_reason(line)\n        if reason:\n            blocked.append({\"command\": line, \"reason\": reason})\n            continue\n        m_del = re.match(r\"^\\s*delete\\s+([A-Za-z0-9_]+)\\s*$\", line, re.I)\n        if m_del:\n            target = m_del.group(1)\n            if target not in known_names:\n                continue  # skip silently\n        try:\n            cmd.do(line)\n            executed.append(line)\n            if m_del:\n                known_names.discard(m_del.group(1))\n        except Exception as exc:\n            errors.append({\"command\": line, \"error\": str(exc)})\n    return {\n        \"ok\": not blocked and not errors,\n        \"tool\": \"run_pymol_commands\",\n        \"executed_count\": len(executed),\n        \"blocked_count\": len(blocked),\n        \"error_count\": len(errors),\n        \"executed\": executed,\n        \"blocked\": blocked,\n        \"errors\": errors,\n    }\n\n\ndef _tool_render(\n    path,\n    width=None,\n    height=None,\n    dpi=None,\n    ray=None,\n    transparent_bg=True,\n    purpose=\"preview\",\n):\n    out_path = os.path.abspath(os.path.expanduser(path))\n    out_dir = os.path.dirname(out_path)\n    if out_dir and not os.path.exists(out_dir):\n        os.makedirs(out_dir, exist_ok=True)\n\n    purpose = (purpose or \"preview\").strip().lower()\n    if purpose not in (\"preview\", \"final\"):\n        purpose = \"preview\"\n\n    if purpose == \"final\":\n        default_width, default_height, default_dpi = 2400, 1800, 300\n        default_ray = True\n    else:\n        default_width, default_height, default_dpi = 1400, 1050, 150\n        default_ray = False\n\n    try:\n        width = default_width if width is None else int(width)\n        height = default_height if height is None else int(height)\n        dpi = default_dpi if dpi is None else int(dpi)\n        width = max(256, width)\n        height = max(256, height)\n        dpi = max(72, dpi)\n    except (TypeError, ValueError):\n        width, height, dpi = default_width, default_height, default_dpi\n\n    if purpose == \"preview\":\n        width = min(width, 2200)\n        height = min(height, 1650)\n        dpi = min(dpi, 220)\n\n    ray_flag = 1 if _to_bool(ray, default=default_ray) else 0\n    transparent = _to_bool(transparent_bg, default=True)\n    cmd.set(\"ray_opaque_background\", 0 if transparent else 1)\n    cmd.png(out_path, width=width, height=height, dpi=dpi, ray=ray_flag, quiet=1)\n\n    exists = os.path.exists(out_path)\n    size = os.path.getsize(out_path) if exists else 0\n    return {\n        \"ok\": exists,\n        \"tool\": \"render\",\n        \"path\": out_path,\n        \"purpose\": purpose,\n        \"width\": width,\n        \"height\": height,\n        \"dpi\": dpi,\n        \"ray\": bool(ray_flag),\n        \"transparent_bg\": transparent,\n        \"bytes\": size,\n    }\n\n\ndef _tool_capture_viewport(analysis_prompt=None, client=None, vision_model=None):\n    if analysis_prompt is None:\n        analysis_prompt = \"Describe what you see in this molecular visualization.\"\n\n    tmp = tempfile.NamedTemporaryFile(suffix=\".png\", delete=False)\n    tmp_path = tmp.name\n    tmp.close()\n\n    try:\n        cmd.png(tmp_path, width=800, height=600, ray=0, quiet=1)\n        time.sleep(0.5)\n        with open(tmp_path, \"rb\") as f:\n            img_data = base64.b64encode(f.read()).decode(\"ascii\")\n    except Exception as exc:\n        return {\n            \"ok\": False,\n            \"tool\": \"capture_viewport\",\n            \"error\": f\"Failed to capture viewport: {exc}\",\n        }\n    finally:\n        try:\n            os.unlink(tmp_path)\n        except OSError:\n            pass\n\n    if not vision_model or not client:\n        return {\n            \"ok\": False,\n            \"tool\": \"capture_viewport\",\n            \"error\": \"No vision model configured. Use set_vision_model <model>.\",\n        }\n\n    try:\n        description = client.vision_completion(vision_model, analysis_prompt, img_data)\n        return {\n            \"ok\": True,\n            \"tool\": \"capture_viewport\",\n            \"analysis_prompt\": analysis_prompt,\n            \"description\": description,\n        }\n    except Exception as exc:\n        return {\n            \"ok\": False,\n            \"tool\": \"capture_viewport\",\n            \"error\": f\"Vision analysis failed: {exc}\",\n        }\n\n\n# ---------------------------------------------------------------------------\n# 3. System prompt\n# ---------------------------------------------------------------------------\n\nSYSTEM_PROMPT = \"\"\"\\\nYou are ChatMol, an expert AI assistant for PyMOL molecular visualization.\n\nYou have 4 tools:\n- `inspect_session`: See what's loaded in PyMOL (objects, chains, atoms, selections).\n- `run_pymol_commands`: Execute any PyMOL commands (selections, styling, coloring, \\\ndistances, presets, fetch, etc.). Commands are newline-separated.\n- `render`: Export an image (preview or final quality).\n- `capture_viewport`: Screenshot + vision analysis to check your work visually.\n\nWorkflow:\n1. For non-trivial requests, start with `inspect_session` to understand the current state.\n2. Use `run_pymol_commands` for all PyMOL operations. You know PyMOL well — use \\\ncmd.select, cmd.show, cmd.hide, cmd.color, cmd.set, cmd.distance, cmd.zoom, \\\ncmd.orient, util.color_chains, util.cnc, preset.ligand_sites_hq, etc.\n3. Use `capture_viewport` to verify your work visually when doing complex styling.\n4. Use `render` when the user requests an image export.\n\nStyle guidelines (publication quality):\n- White background, clean composition.\n- Cartoon as baseline representation for protein.\n- Sticks only for key residues, ligands, and interaction sites.\n- Context in gray; 1-2 accent colors (often cyan/marine + orange).\n- Surfaces used purposefully for targets/interfaces/pores, not everywhere.\n- Hydrogen bonds / polar contacts shown with dashed lines when relevant.\n- Keep decoration sparse — prioritize clarity over ornamentation.\n\nRules:\n- Never issue destructive commands (reinitialize, quit, delete all, shell commands).\n- If user input is ambiguous, ask for clarification rather than guessing.\n- Be concise in your final response.\n\"\"\"\n\n\n# ---------------------------------------------------------------------------\n# 4. ChatMolAgent — config persistence + simple agent loop\n# ---------------------------------------------------------------------------\n\n\nclass ChatMolAgent:\n    \"\"\"Simple agentic loop: call LLM, execute tools, repeat.\"\"\"\n\n    CONFIG_PATH = os.path.expanduser(\"~/.PyMOL/chatmol_config.json\")\n\n    DEFAULT_CONFIG = {\n        \"provider\": \"openrouter\",\n        \"api_keys\": {},\n        \"text_model\": \"google/gemini-3-flash-preview\",\n        \"vision_model\": \"google/gemini-3-flash-preview\",\n        \"temperature\": 0.01,\n        \"max_tokens\": 4096,\n        \"max_iterations\": 50,\n        \"max_tool_calls\": 30,\n    }\n\n    def __init__(self):\n        os.makedirs(os.path.dirname(self.CONFIG_PATH), exist_ok=True)\n        self.config = self._load_config()\n        self._reinit_client()\n        self.conversation_history = []\n\n    # -- config persistence -------------------------------------------------\n\n    def _load_config(self):\n        try:\n            with open(self.CONFIG_PATH, \"r\") as f:\n                cfg = json.load(f)\n            merged = dict(self.DEFAULT_CONFIG)\n            merged[\"api_keys\"] = dict(self.DEFAULT_CONFIG[\"api_keys\"])\n            merged.update(cfg)\n            if not isinstance(merged.get(\"api_keys\"), dict):\n                merged[\"api_keys\"] = {}\n            # Migrate old single \"api_key\" field\n            old_key = merged.pop(\"api_key\", \"\")\n            if old_key and isinstance(old_key, str):\n                prov = merged.get(\"provider\", \"openrouter\")\n                merged[\"api_keys\"].setdefault(prov, old_key)\n                self._save_config(merged)\n            # Remove legacy keys\n            for legacy in (\"human_in_the_loop\", \"multi_agent_enabled\"):\n                merged.pop(legacy, None)\n            return merged\n        except (FileNotFoundError, json.JSONDecodeError):\n            self._save_config(self.DEFAULT_CONFIG)\n            return dict(self.DEFAULT_CONFIG)\n\n    def _save_config(self, cfg=None):\n        if cfg is None:\n            cfg = self.config\n        try:\n            with open(self.CONFIG_PATH, \"w\") as f:\n                json.dump(cfg, f, indent=2)\n        except Exception as exc:\n            print(f\"Warning: could not save config: {exc}\")\n\n    def _resolve_api_key(self):\n        prov_name = self.config.get(\"provider\", \"openrouter\")\n        prov = PROVIDERS.get(prov_name, PROVIDERS[\"openrouter\"])\n        return os.getenv(prov[\"env_var\"], \"\") or self.config.get(\"api_keys\", {}).get(\n            prov_name, \"\"\n        )\n\n    def _reinit_client(self):\n        prov_name = self.config.get(\"provider\", \"openrouter\")\n        api_key = self._resolve_api_key()\n        self.client = LLMClient(prov_name, api_key)\n\n    # -- PyMOL-registered commands ------------------------------------------\n\n    def chat(self, *args):\n        \"\"\"PyMOL command: chat <message>\"\"\"\n        message = \" \".join(str(a) for a in args).strip()\n        if not message:\n            print(\"Usage: chat <message>\")\n            return\n        if not self._resolve_api_key():\n            prov = self.config.get(\"provider\", \"openrouter\")\n            env = PROVIDERS.get(prov, {}).get(\"env_var\", \"\")\n            print(f\"No API key set. Use: set_api_key <key>  (or set env var {env})\")\n            return\n        try:\n            response = self._run_agent_loop(message)\n            print(\"=\" * 50)\n            print(\"ChatMol:\", response)\n            print(\"=\" * 50)\n        except RuntimeError as exc:\n            print(f\"ChatMol error: {exc}\")\n        except requests.exceptions.RequestException as exc:\n            print(f\"ChatMol network error: {exc}\")\n        except (KeyError, IndexError, TypeError) as exc:\n            prov = self.config.get(\"provider\", \"openrouter\")\n            model = self.config.get(\"text_model\", \"\")\n            print(f\"ChatMol error: unexpected response format — {exc}\")\n            print(f\"  Provider: {prov}, Model: {model}\")\n            print(\n                \"  The model may not support tool calling, \"\n                \"or the API returned an unexpected structure.\"\n            )\n        except Exception as exc:\n            print(f\"ChatMol error ({type(exc).__name__}): {exc}\")\n\n    def set_api_key(self, api_key=\"\"):\n        \"\"\"PyMOL command: set_api_key <key>\"\"\"\n        api_key = api_key.strip()\n        if not api_key:\n            prov_name = self.config.get(\"provider\", \"openrouter\")\n            label = PROVIDERS.get(prov_name, {}).get(\"label\", prov_name)\n            print(f\"Usage: set_api_key <your_{label}_api_key>\")\n            return\n        prov_name = self.config.get(\"provider\", \"openrouter\")\n        if not isinstance(self.config.get(\"api_keys\"), dict):\n            self.config[\"api_keys\"] = {}\n        self.config[\"api_keys\"][prov_name] = api_key\n        self._save_config()\n        self._reinit_client()\n        label = PROVIDERS.get(prov_name, {}).get(\"label\", prov_name)\n        print(f\"{label} API key saved.\")\n\n    def set_provider(self, provider_name=\"\"):\n        \"\"\"PyMOL command: set_provider <name>\"\"\"\n        provider_name = provider_name.strip().lower()\n        if not provider_name or provider_name not in PROVIDERS:\n            print(f\"Available providers: {', '.join(PROVIDERS.keys())}\")\n            if provider_name:\n                print(f\"Unknown provider: {provider_name}\")\n            return\n        self.config[\"provider\"] = provider_name\n        prov = PROVIDERS[provider_name]\n        if prov[\"models\"]:\n            self.config[\"text_model\"] = prov[\"models\"][0][0]\n        if prov[\"vision_models\"]:\n            self.config[\"vision_model\"] = prov[\"vision_models\"][0][0]\n        else:\n            self.config[\"vision_model\"] = \"\"\n        self._save_config()\n        self._reinit_client()\n        print(f\"Provider set to: {prov['label']}\")\n        print(f\"  Text model: {self.config['text_model']}\")\n        if self.config[\"vision_model\"]:\n            print(f\"  Vision model: {self.config['vision_model']}\")\n        print(f\"  Set API key via: set_api_key <key>  (or env {prov['env_var']})\")\n\n    def set_model(self, model_name=\"\"):\n        \"\"\"PyMOL command: set_model <model>\"\"\"\n        model_name = model_name.strip()\n        if not model_name:\n            print(f\"Current model: {self.config['text_model']}\")\n            print(\"Usage: set_model <provider/model>\")\n            return\n        self.config[\"text_model\"] = model_name\n        self._save_config()\n        print(f\"Text model set to: {model_name}\")\n\n    def set_vision_model(self, model_name=\"\"):\n        \"\"\"PyMOL command: set_vision_model <model>\"\"\"\n        model_name = model_name.strip()\n        if not model_name:\n            print(f\"Current vision model: {self.config['vision_model']}\")\n            print(\"Usage: set_vision_model <provider/model>\")\n            return\n        self.config[\"vision_model\"] = model_name\n        self._save_config()\n        print(f\"Vision model set to: {model_name}\")\n\n    def reset_conversation(self):\n        \"\"\"PyMOL command: reset_conversation\"\"\"\n        self.conversation_history = []\n        print(\"Conversation history cleared.\")\n\n    def save_conversation(self, filename=\"\"):\n        \"\"\"PyMOL command: save_conversation [filename]\"\"\"\n        if not filename:\n            filename = (\n                f\"chatmol_conversation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json\"\n            )\n        with open(filename, \"w\", encoding=\"utf-8\") as f:\n            json.dump(self.conversation_history, f, indent=2)\n        print(f\"Conversation saved to {filename}\")\n\n    def load_conversation(self, filename=\"\"):\n        \"\"\"PyMOL command: load_conversation <filename>\"\"\"\n        if not filename:\n            print(\"Usage: load_conversation <filename.json>\")\n            return\n        try:\n            with open(filename, \"r\", encoding=\"utf-8\") as f:\n                self.conversation_history = json.load(f)\n            print(\n                f\"Loaded conversation from {filename} \"\n                f\"({len(self.conversation_history)} messages)\"\n            )\n        except Exception as exc:\n            print(f\"Failed to load conversation: {exc}\")\n\n    def show_config(self):\n        \"\"\"PyMOL command: chatmol_config\"\"\"\n        prov_name = self.config.get(\"provider\", \"openrouter\")\n        prov = PROVIDERS.get(prov_name, {})\n        key = self._resolve_api_key()\n        masked = (key[:8] + \"...\" + key[-4:]) if key and len(key) > 12 else \"(not set)\"\n        has_env = bool(os.getenv(prov.get(\"env_var\", \"\"), \"\"))\n        print(\"ChatMol configuration:\")\n        print(f\"  provider: {prov.get('label', prov_name)} ({prov_name})\")\n        print(f\"  api_key: {masked} ({'env var' if has_env else 'config file'})\")\n        print(f\"  text_model: {self.config.get('text_model', '')}\")\n        print(f\"  vision_model: {self.config.get('vision_model', '') or '(none)'}\")\n        print(f\"  temperature: {self.config.get('temperature', 0.01)}\")\n        print(f\"  max_tokens: {self.config.get('max_tokens', 4096)}\")\n        print(f\"  max_iterations: {self.config.get('max_iterations', 50)}\")\n        print(f\"  max_tool_calls: {self.config.get('max_tool_calls', 30)}\")\n        stored = [k for k, v in self.config.get(\"api_keys\", {}).items() if v]\n        if stored:\n            print(f\"  keys stored for: {', '.join(stored)}\")\n\n    # -- agentic loop -------------------------------------------------------\n\n    def _chat_completion_with_retry(\n        self, model, messages, tool_defs, temperature, max_tokens, phase_callback=None\n    ):\n        max_attempts = 4\n        for attempt in range(1, max_attempts + 1):\n            try:\n                return self.client.chat_completion(\n                    model,\n                    messages,\n                    tools=tool_defs,\n                    temperature=temperature,\n                    max_tokens=max_tokens,\n                )\n            except LLMTransientError as exc:\n                if attempt >= max_attempts:\n                    raise RuntimeError(\n                        f\"Transient API errors persisted after {max_attempts} attempts: {exc}\"\n                    )\n                delay = min(8.0, float(2 ** (attempt - 1))) + random.uniform(0.0, 0.4)\n                print(\n                    f\"  [Retry {attempt}/{max_attempts}] transient error: \"\n                    f\"{exc}. Sleeping {delay:.1f}s.\"\n                )\n                if phase_callback:\n                    phase_callback(\"Retrying API request\")\n                time.sleep(delay)\n\n    @staticmethod\n    def _parse_assistant_message(data):\n        choices = data.get(\"choices\")\n        if not isinstance(choices, list) or not choices:\n            raise RuntimeError(\"Invalid API response: missing choices.\")\n        message = choices[0].get(\"message\")\n        if not isinstance(message, dict):\n            raise RuntimeError(\"Invalid API response: missing assistant message.\")\n        return message\n\n    @staticmethod\n    def _normalize_content(content):\n        if isinstance(content, (list, dict)):\n            return json.dumps(content, ensure_ascii=False)\n        if content is None:\n            return \"\"\n        if not isinstance(content, str):\n            return str(content)\n        return content\n\n    @staticmethod\n    def _parse_tool_calls(assistant_msg):\n        raw_calls = assistant_msg.get(\"tool_calls\") or []\n        parsed = []\n        for tc in raw_calls:\n            if not isinstance(tc, dict):\n                raise RuntimeError(\"Invalid tool call entry from model.\")\n            tc_id = tc.get(\"id\")\n            fn = tc.get(\"function\") or {}\n            fn_name = fn.get(\"name\")\n            raw_args = fn.get(\"arguments\", \"{}\")\n            if not tc_id or not fn_name:\n                raise RuntimeError(\n                    \"Invalid tool call payload (missing id or function name).\"\n                )\n            if isinstance(raw_args, dict):\n                fn_args = raw_args\n            else:\n                raw_args = \"{}\" if raw_args is None else str(raw_args).strip()\n                if not raw_args:\n                    raw_args = \"{}\"\n                try:\n                    fn_args = json.loads(raw_args)\n                except json.JSONDecodeError as exc:\n                    raise RuntimeError(\n                        f\"Tool arguments not valid JSON for {fn_name}: {exc}\"\n                    )\n            if not isinstance(fn_args, dict):\n                raise RuntimeError(f\"Tool arguments for {fn_name} must be an object.\")\n            parsed.append(\n                {\"id\": tc_id, \"name\": fn_name, \"arguments\": fn_args, \"raw\": tc}\n            )\n        return parsed\n\n    def _make_tool_executor(self, tool_executor_override=None):\n        \"\"\"Return a tool executor function that passes client/vision_model.\"\"\"\n        if tool_executor_override:\n            return tool_executor_override\n        client = self.client\n        vision_model = self.config.get(\"vision_model\", \"\")\n\n        def _executor(tool_name, arguments):\n            return execute_tool(\n                tool_name, arguments, client=client, vision_model=vision_model\n            )\n\n        return _executor\n\n    def _run_agent_loop_internal(\n        self, message, tool_executor=None, phase_callback=None\n    ):\n        self.conversation_history.append({\"role\": \"user\", \"content\": message})\n\n        model = self.config[\"text_model\"]\n        temperature = self.config.get(\"temperature\", 0.01)\n        max_tokens = self.config.get(\"max_tokens\", 4096)\n        max_iterations = max(1, int(self.config.get(\"max_iterations\", 50) or 50))\n        max_tool_calls = max(8, int(self.config.get(\"max_tool_calls\", 30) or 30))\n        executor = self._make_tool_executor(tool_executor)\n        tool_calls_used = 0\n\n        for iteration in range(max_iterations):\n            if phase_callback:\n                phase_callback(\"Thinking\" if iteration == 0 else \"Thinking more\")\n\n            messages = [{\"role\": \"system\", \"content\": SYSTEM_PROMPT}]\n            messages += self.conversation_history\n\n            data = self._chat_completion_with_retry(\n                model,\n                messages,\n                TOOL_DEFINITIONS,\n                temperature,\n                max_tokens,\n                phase_callback,\n            )\n            assistant_msg = self._parse_assistant_message(data)\n            parsed_calls = self._parse_tool_calls(assistant_msg)\n            content = self._normalize_content(assistant_msg.get(\"content\", None))\n\n            history_entry = {\"role\": \"assistant\", \"content\": content}\n            if parsed_calls:\n                history_entry[\"tool_calls\"] = [tc[\"raw\"] for tc in parsed_calls]\n\n            if parsed_calls:\n                tool_entries = []\n                for tc in parsed_calls:\n                    fn_name = tc[\"name\"]\n                    fn_args = tc[\"arguments\"]\n\n                    if phase_callback:\n                        phase_callback(f\"Running {fn_name}\")\n                    print(\n                        f\"  [Tool: {fn_name}] \"\n                        f\"{json.dumps(fn_args, ensure_ascii=False)[:120]}\"\n                    )\n\n                    if tool_calls_used >= max_tool_calls:\n                        result = json.dumps(\n                            {\n                                \"ok\": False,\n                                \"tool\": fn_name,\n                                \"error\": (\n                                    \"Tool-call budget reached. Stop calling tools \"\n                                    \"and provide your final response.\"\n                                ),\n                            },\n                            ensure_ascii=False,\n                        )\n                    else:\n                        try:\n                            result = executor(fn_name, fn_args)\n                        except Exception as exc:\n                            result = json.dumps(\n                                {\"ok\": False, \"tool\": fn_name, \"error\": str(exc)},\n                                ensure_ascii=False,\n                            )\n\n                    tool_entries.append(\n                        {\n                            \"role\": \"tool\",\n                            \"tool_call_id\": tc[\"id\"],\n                            \"content\": result,\n                        }\n                    )\n                    tool_calls_used += 1\n\n                self.conversation_history.append(history_entry)\n                self.conversation_history.extend(tool_entries)\n                continue\n\n            # No tool calls — this is the final response\n            self.conversation_history.append(history_entry)\n            return (content or \"\").strip()\n\n        return (\n            \"Agent loop reached maximum iterations \"\n            f\"({max_iterations}) without a final response.\"\n        )\n\n    def _run_agent_loop(self, message):\n        return self._run_agent_loop_internal(message, tool_executor=None)\n\n    def run_agent_loop_with_callback(self, message, tool_executor, phase_callback=None):\n        return self._run_agent_loop_internal(\n            message, tool_executor=tool_executor, phase_callback=phase_callback\n        )\n\n\n# ---------------------------------------------------------------------------\n# 5. Qt5 GUI (loaded conditionally)\n# ---------------------------------------------------------------------------\n\n_HAS_QT = False\ntry:\n    from pymol.Qt import QtWidgets, QtCore, QtGui  # noqa: E402\n\n    QDockWidget = QtWidgets.QDockWidget\n    QWidget = QtWidgets.QWidget\n    QVBoxLayout = QtWidgets.QVBoxLayout\n    QHBoxLayout = QtWidgets.QHBoxLayout\n    QTextEdit = QtWidgets.QTextEdit\n    QLineEdit = QtWidgets.QLineEdit\n    QPushButton = QtWidgets.QPushButton\n    QLabel = QtWidgets.QLabel\n    QDialog = QtWidgets.QDialog\n    QFormLayout = QtWidgets.QFormLayout\n    QDoubleSpinBox = QtWidgets.QDoubleSpinBox\n    QSpinBox = QtWidgets.QSpinBox\n    QComboBox = QtWidgets.QComboBox\n    QApplication = QtWidgets.QApplication\n\n    Qt = QtCore.Qt\n    QThread = QtCore.QThread\n    QTimer = QtCore.QTimer\n    QPropertyAnimation = QtCore.QPropertyAnimation\n    QEasingCurve = QtCore.QEasingCurve\n    Signal = getattr(QtCore, \"Signal\", None) or getattr(QtCore, \"pyqtSignal\")\n    Property = getattr(QtCore, \"Property\", None) or getattr(QtCore, \"pyqtProperty\")\n\n    QColor = QtGui.QColor\n    QPainter = QtGui.QPainter\n\n    _HAS_QT = True\nexcept ImportError:\n    pass\n\n\nif _HAS_QT:\n\n    # -- ThinkingIndicator --------------------------------------------------\n\n    class ThinkingIndicator(QWidget):\n        \"\"\"Three dots that pulse with a breathing animation.\"\"\"\n\n        _DOT_COUNT = 3\n        _DOT_RADIUS = 4\n        _DOT_SPACING = 14\n\n        def __init__(self, parent=None):\n            super().__init__(parent)\n            self._opacity = 0.3\n            self._phase_text = \"Thinking\"\n            self.setFixedHeight(28)\n            self._anim = QPropertyAnimation(self, b\"dot_opacity\")\n            self._anim.setDuration(900)\n            self._anim.setStartValue(0.3)\n            self._anim.setEndValue(1.0)\n            self._anim.setEasingCurve(QEasingCurve.InOutSine)\n            self._anim.setLoopCount(-1)\n            self._anim.finished.connect(lambda: None)\n            self._forward = True\n            self._cycle_timer = QTimer(self)\n            self._cycle_timer.setInterval(900)\n            self._cycle_timer.timeout.connect(self._toggle_direction)\n            self.hide()\n\n        def _toggle_direction(self):\n            self._forward = not self._forward\n            if self._forward:\n                self._anim.setStartValue(0.3)\n                self._anim.setEndValue(1.0)\n            else:\n                self._anim.setStartValue(1.0)\n                self._anim.setEndValue(0.3)\n            self._anim.start()\n\n        @Property(float)\n        def dot_opacity(self):\n            return self._opacity\n\n        @dot_opacity.setter\n        def dot_opacity(self, value):\n            self._opacity = value\n            self.update()\n\n        def start(self, phase_text=\"Thinking\"):\n            self._phase_text = phase_text\n            self._anim.setStartValue(0.3)\n            self._anim.setEndValue(1.0)\n            self._forward = True\n            self._anim.start()\n            self._cycle_timer.start()\n            self.show()\n\n        def stop(self):\n            self._anim.stop()\n            self._cycle_timer.stop()\n            self.hide()\n\n        def set_phase(self, text):\n            self._phase_text = text\n            self.update()\n\n        def paintEvent(self, _event):\n            p = QPainter(self)\n            p.setRenderHint(QPainter.Antialiasing)\n            p.setPen(QColor(120, 120, 120))\n            p.drawText(8, 18, self._phase_text)\n            label_width = p.fontMetrics().horizontalAdvance(self._phase_text) + 16\n            color = QColor(70, 130, 230)\n            for i in range(self._DOT_COUNT):\n                offset = i * 0.15\n                alpha = max(0.0, min(1.0, self._opacity - offset))\n                color.setAlphaF(alpha)\n                p.setBrush(color)\n                p.setPen(Qt.NoPen)\n                x = label_width + i * self._DOT_SPACING\n                p.drawEllipse(x, 10, self._DOT_RADIUS * 2, self._DOT_RADIUS * 2)\n            p.end()\n\n    # -- AgentWorker --------------------------------------------------------\n\n    class AgentWorker(QThread):\n        \"\"\"Runs the agentic loop off the main thread.\"\"\"\n\n        finished = Signal(str)\n        error = Signal(str)\n        tool_executed = Signal(str, str)\n        request_tool_execution = Signal(str, str, str)\n        phase_changed = Signal(str)\n\n        def __init__(self, agent, message, parent=None):\n            super().__init__(parent)\n            self.agent = agent\n            self.message = message\n            self._tool_result = None\n            self._tool_event = threading.Event()\n\n        def run(self):\n            try:\n                response = self.agent.run_agent_loop_with_callback(\n                    self.message, self._tool_executor, phase_callback=self._on_phase\n                )\n                self.finished.emit(response)\n            except Exception as exc:\n                self.error.emit(str(exc))\n\n        def _on_phase(self, phase_text):\n            self.phase_changed.emit(phase_text)\n\n        def _tool_executor(self, tool_name, arguments):\n            self._tool_event.clear()\n            self._tool_result = None\n            args_json = json.dumps(arguments, ensure_ascii=False)\n            self.request_tool_execution.emit(\"\", tool_name, args_json)\n            self._tool_event.wait(timeout=60)\n            result = self._tool_result or \"Tool execution timed out.\"\n            self.tool_executed.emit(tool_name, result[:200])\n            return result\n\n        def deliver_tool_result(self, result):\n            self._tool_result = result\n            self._tool_event.set()\n\n    # -- ChatMolSettingsDialog ----------------------------------------------\n\n    class ChatMolSettingsDialog(QDialog):\n\n        def __init__(self, agent, parent=None):\n            super().__init__(parent)\n            self.agent = agent\n            self.setWindowTitle(\"ChatMol Settings\")\n            self.setMinimumWidth(460)\n            self._build_ui()\n\n        def _build_ui(self):\n            layout = QFormLayout(self)\n\n            # Provider selector\n            self.provider_combo = QComboBox()\n            for key, prov in PROVIDERS.items():\n                self.provider_combo.addItem(prov[\"label\"], key)\n            cur_prov = self.agent.config.get(\"provider\", \"openrouter\")\n            idx = self.provider_combo.findData(cur_prov)\n            if idx >= 0:\n                self.provider_combo.setCurrentIndex(idx)\n            self.provider_combo.currentIndexChanged.connect(self._on_provider_changed)\n            layout.addRow(\"Provider:\", self.provider_combo)\n\n            # API key\n            self.api_key_edit = QLineEdit()\n            self.api_key_edit.setEchoMode(QLineEdit.Password)\n            self.api_key_edit.setText(\n                self.agent.config.get(\"api_keys\", {}).get(cur_prov, \"\")\n            )\n            self.env_hint = QLabel()\n            self.env_hint.setStyleSheet(\"color: grey; font-size: 10px;\")\n            key_layout = QVBoxLayout()\n            key_layout.setSpacing(2)\n            key_layout.addWidget(self.api_key_edit)\n            key_layout.addWidget(self.env_hint)\n            layout.addRow(\"API Key:\", key_layout)\n\n            # Text model\n            self.text_model_combo = QComboBox()\n            self.text_model_combo.setEditable(True)\n            layout.addRow(\"Text Model:\", self.text_model_combo)\n\n            # Vision model\n            self.vision_model_combo = QComboBox()\n            self.vision_model_combo.setEditable(True)\n            layout.addRow(\"Vision Model:\", self.vision_model_combo)\n\n            # Temperature\n            self.temp_spin = QDoubleSpinBox()\n            self.temp_spin.setRange(0.0, 2.0)\n            self.temp_spin.setSingleStep(0.05)\n            self.temp_spin.setValue(self.agent.config.get(\"temperature\", 0.01))\n            layout.addRow(\"Temperature:\", self.temp_spin)\n\n            # Max tokens\n            self.max_tokens_spin = QSpinBox()\n            self.max_tokens_spin.setRange(256, 16384)\n            self.max_tokens_spin.setSingleStep(256)\n            self.max_tokens_spin.setValue(self.agent.config.get(\"max_tokens\", 4096))\n            layout.addRow(\"Max Tokens:\", self.max_tokens_spin)\n\n            # Max iterations\n            self.max_iterations_spin = QSpinBox()\n            self.max_iterations_spin.setRange(1, 500)\n            self.max_iterations_spin.setSingleStep(1)\n            self.max_iterations_spin.setValue(\n                self.agent.config.get(\"max_iterations\", 50)\n            )\n            layout.addRow(\"Max Iterations:\", self.max_iterations_spin)\n\n            # Max tool calls per request\n            self.max_tool_calls_spin = QSpinBox()\n            self.max_tool_calls_spin.setRange(8, 200)\n            self.max_tool_calls_spin.setSingleStep(1)\n            self.max_tool_calls_spin.setValue(\n                self.agent.config.get(\"max_tool_calls\", 30)\n            )\n            layout.addRow(\"Max Tool Calls:\", self.max_tool_calls_spin)\n\n            # Buttons\n            btn_layout = QHBoxLayout()\n            save_btn = QPushButton(\"Save\")\n            save_btn.clicked.connect(self._on_save)\n            cancel_btn = QPushButton(\"Cancel\")\n            cancel_btn.clicked.connect(self.reject)\n            btn_layout.addStretch()\n            btn_layout.addWidget(save_btn)\n            btn_layout.addWidget(cancel_btn)\n            layout.addRow(btn_layout)\n\n            self._populate_models(cur_prov)\n\n        def _on_provider_changed(self, _index):\n            prov_key = self.provider_combo.currentData()\n            stored_key = self.agent.config.get(\"api_keys\", {}).get(prov_key, \"\")\n            self.api_key_edit.setText(stored_key)\n            self._populate_models(prov_key)\n\n        def _populate_models(self, prov_key):\n            prov = PROVIDERS.get(prov_key, PROVIDERS[\"openrouter\"])\n\n            env_var = prov.get(\"env_var\", \"\")\n            has_env = bool(os.getenv(env_var, \"\"))\n            if has_env:\n                self.env_hint.setText(f\"Using env var {env_var} (overrides this field)\")\n            else:\n                self.env_hint.setText(f\"Or set env var: {env_var}\")\n\n            # Text models\n            self.text_model_combo.clear()\n            for model_id, label in prov[\"models\"]:\n                self.text_model_combo.addItem(f\"{label}  ({model_id})\", model_id)\n            saved = self.agent.config.get(\"text_model\", \"\")\n            restored = False\n            for i in range(self.text_model_combo.count()):\n                if self.text_model_combo.itemData(i) == saved:\n                    self.text_model_combo.setCurrentIndex(i)\n                    restored = True\n                    break\n            if not restored and saved:\n                self.text_model_combo.setEditText(saved)\n\n            # Vision models\n            self.vision_model_combo.clear()\n            self.vision_model_combo.addItem(\"(none)\", \"\")\n            for model_id, label in prov.get(\"vision_models\", []):\n                self.vision_model_combo.addItem(f\"{label}  ({model_id})\", model_id)\n            saved_v = self.agent.config.get(\"vision_model\", \"\")\n            restored_v = False\n            for i in range(self.vision_model_combo.count()):\n                if self.vision_model_combo.itemData(i) == saved_v:\n                    self.vision_model_combo.setCurrentIndex(i)\n                    restored_v = True\n                    break\n            if not restored_v and saved_v:\n                self.vision_model_combo.setEditText(saved_v)\n\n        def _on_save(self):\n            prov_key = self.provider_combo.currentData()\n            self.agent.config[\"provider\"] = prov_key\n\n            if not isinstance(self.agent.config.get(\"api_keys\"), dict):\n                self.agent.config[\"api_keys\"] = {}\n            self.agent.config[\"api_keys\"][prov_key] = self.api_key_edit.text().strip()\n\n            text_idx = self.text_model_combo.currentIndex()\n            text_data = self.text_model_combo.itemData(text_idx)\n            self.agent.config[\"text_model\"] = (\n                text_data if text_data else self.text_model_combo.currentText().strip()\n            )\n\n            vis_idx = self.vision_model_combo.currentIndex()\n            vis_data = self.vision_model_combo.itemData(vis_idx)\n            self.agent.config[\"vision_model\"] = (\n                vis_data\n                if vis_data is not None\n                else self.vision_model_combo.currentText().strip()\n            )\n\n            self.agent.config[\"temperature\"] = self.temp_spin.value()\n            self.agent.config[\"max_tokens\"] = self.max_tokens_spin.value()\n            self.agent.config[\"max_iterations\"] = self.max_iterations_spin.value()\n            self.agent.config[\"max_tool_calls\"] = self.max_tool_calls_spin.value()\n            self.agent._save_config()\n            self.agent._reinit_client()\n            self.accept()\n\n    # -- ChatMolChatBar -----------------------------------------------------\n\n    class ChatMolChatBar(QDockWidget):\n\n        def __init__(self, agent, parent=None):\n            super().__init__(\"ChatMol\", parent)\n            self.agent = agent\n            self._worker = None\n            self._build_ui()\n\n        def _build_ui(self):\n            container = QWidget()\n            vbox = QVBoxLayout(container)\n            vbox.setContentsMargins(4, 4, 4, 4)\n            vbox.setSpacing(4)\n\n            # Top bar\n            top_bar = QHBoxLayout()\n            top_bar.addWidget(QLabel(\"<b>ChatMol</b>\"))\n            top_bar.addStretch()\n            settings_btn = QPushButton(\"Settings\")\n            settings_btn.setFixedWidth(70)\n            settings_btn.clicked.connect(self._open_settings)\n            top_bar.addWidget(settings_btn)\n            clear_btn = QPushButton(\"Clear\")\n            clear_btn.setFixedWidth(50)\n            clear_btn.clicked.connect(self._clear_chat)\n            top_bar.addWidget(clear_btn)\n            vbox.addLayout(top_bar)\n\n            # Execution trace panel\n            self.trace_display = QTextEdit()\n            self.trace_display.setReadOnly(True)\n            self.trace_display.setMaximumHeight(92)\n            self.trace_display.setStyleSheet(\n                \"font-family: Menlo, Consolas, monospace; font-size:10px; \"\n                \"color:#263238; background:#FAFAFA;\"\n            )\n            vbox.addWidget(self.trace_display)\n\n            # Chat history\n            self.chat_display = QTextEdit()\n            self.chat_display.setReadOnly(True)\n            self.chat_display.setMinimumHeight(120)\n            vbox.addWidget(self.chat_display)\n\n            # Thinking indicator\n            self.thinking = ThinkingIndicator()\n            vbox.addWidget(self.thinking)\n\n            # Input bar\n            input_bar = QHBoxLayout()\n            self.input_edit = QLineEdit()\n            self.input_edit.setPlaceholderText(\"Type a message...\")\n            self.input_edit.returnPressed.connect(self._on_send)\n            input_bar.addWidget(self.input_edit)\n            self.send_btn = QPushButton(\"Send\")\n            self.send_btn.setFixedWidth(50)\n            self.send_btn.clicked.connect(self._on_send)\n            input_bar.addWidget(self.send_btn)\n            self.model_label = QLabel()\n            self.model_label.setStyleSheet(\"color: grey; font-size: 10px;\")\n            self._update_model_label()\n            input_bar.addWidget(self.model_label)\n            vbox.addLayout(input_bar)\n\n            self.setWidget(container)\n\n        def _update_model_label(self):\n            prov_name = self.agent.config.get(\"provider\", \"openrouter\")\n            prov_label = PROVIDERS.get(prov_name, {}).get(\"label\", prov_name)\n            model = self.agent.config.get(\"text_model\", \"\")\n            short = model.split(\"/\")[-1] if \"/\" in model else model\n            self.model_label.setText(f\"{short} ({prov_label})\")\n\n        def _open_settings(self):\n            dlg = ChatMolSettingsDialog(self.agent, self)\n            if dlg.exec_():\n                self._update_model_label()\n\n        def _clear_chat(self):\n            self.chat_display.clear()\n            self.trace_display.clear()\n            self.agent.reset_conversation()\n\n        def _append_html(self, html):\n            self.chat_display.append(html)\n            sb = self.chat_display.verticalScrollBar()\n            sb.setValue(sb.maximum())\n\n        def _append_trace(self, text):\n            self.trace_display.append(_escape_html(text))\n            sb = self.trace_display.verticalScrollBar()\n            sb.setValue(sb.maximum())\n\n        def _on_send(self):\n            text = self.input_edit.text().strip()\n            if not text:\n                return\n\n            if not self.agent._resolve_api_key():\n                self._append_html(\n                    '<span style=\"color:red;\">No API key set. '\n                    \"Open Settings to configure your API key.</span>\"\n                )\n                return\n\n            self.input_edit.clear()\n            self.input_edit.setEnabled(False)\n            self.send_btn.setEnabled(False)\n            self.trace_display.clear()\n\n            self._append_html(\n                f'<b style=\"color:#2962FF;\">You:</b> {_escape_html(text)}'\n            )\n\n            self.thinking.start(\"Thinking\")\n\n            self._worker = AgentWorker(self.agent, text)\n            self._worker.request_tool_execution.connect(\n                self._execute_tool_on_main_thread\n            )\n            self._worker.phase_changed.connect(self._on_phase_changed)\n            self._worker.tool_executed.connect(self._on_tool_executed)\n            self._worker.finished.connect(self._on_agent_finished)\n            self._worker.error.connect(self._on_agent_error)\n            self._worker.start()\n\n        def _on_phase_changed(self, phase_text):\n            self.thinking.set_phase(phase_text)\n            if \"Retry\" in phase_text:\n                self._append_trace(f\"[phase] {phase_text}\")\n\n        def _on_tool_executed(self, tool_name, result_preview):\n            self._append_trace(f\"[tool:{tool_name}] {result_preview}\")\n\n        def _execute_tool_on_main_thread(self, _call_id, tool_name, args_json):\n            try:\n                arguments = json.loads(args_json)\n            except (json.JSONDecodeError, TypeError):\n                arguments = {}\n            self._append_trace(f\"[call:{tool_name}] {args_json[:180]}\")\n            try:\n                result = execute_tool(\n                    tool_name,\n                    arguments,\n                    client=self.agent.client,\n                    vision_model=self.agent.config.get(\"vision_model\", \"\"),\n                )\n            except Exception as exc:\n                result = f\"Tool error: {exc}\"\n            if self._worker is not None:\n                self._worker.deliver_tool_result(result)\n\n        def _on_agent_finished(self, response):\n            self.thinking.stop()\n            self._append_html(f\"<b>ChatMol:</b> {_escape_html(response)}\")\n            self.input_edit.setEnabled(True)\n            self.send_btn.setEnabled(True)\n            self.input_edit.setFocus()\n\n        def _on_agent_error(self, error_msg):\n            self.thinking.stop()\n            prov = self.agent.config.get(\"provider\", \"?\")\n            model = self.agent.config.get(\"text_model\", \"?\")\n            self._append_trace(f\"[error] {error_msg}\")\n            self._append_html(\n                f'<span style=\"color:red;\"><b>Error</b> '\n                f\"[{_escape_html(prov)}/{_escape_html(model)}]: \"\n                f\"{_escape_html(error_msg)}</span>\"\n            )\n            self.input_edit.setEnabled(True)\n            self.send_btn.setEnabled(True)\n            self.input_edit.setFocus()\n\n\ndef _escape_html(text):\n    return (\n        text.replace(\"&\", \"&amp;\")\n        .replace(\"<\", \"&lt;\")\n        .replace(\">\", \"&gt;\")\n        .replace(\"\\n\", \"<br>\")\n    )\n\n\n# ---------------------------------------------------------------------------\n# 6. Module-level initialisation\n# ---------------------------------------------------------------------------\n\n_agent = ChatMolAgent()\n\ncmd.extend(\"chat\", _agent.chat)\ncmd.extend(\"set_provider\", _agent.set_provider)\ncmd.extend(\"set_api_key\", _agent.set_api_key)\ncmd.extend(\"set_model\", _agent.set_model)\ncmd.extend(\"set_vision_model\", _agent.set_vision_model)\ncmd.extend(\"reset_conversation\", _agent.reset_conversation)\ncmd.extend(\"save_conversation\", _agent.save_conversation)\ncmd.extend(\"load_conversation\", _agent.load_conversation)\ncmd.extend(\"chatmol_config\", _agent.show_config)\n\n\ndef chatmol_settings():\n    \"\"\"Open the ChatMol settings dialog (requires Qt).\"\"\"\n    if not _HAS_QT:\n        print(\n            \"Qt not available. Configure via commands: \"\n            \"set_api_key, set_model, set_vision_model\"\n        )\n        return\n    app = QApplication.instance()\n    if app is None:\n        print(\"No Qt application running.\")\n        return\n    dlg = ChatMolSettingsDialog(_agent)\n    dlg.exec_()\n\n\ndef chatmol_gui():\n    \"\"\"Manually open the ChatMol chat bar (requires Qt).\"\"\"\n    _init_gui()\n\n\ncmd.extend(\"chatmol_settings\", chatmol_settings)\ncmd.extend(\"chatmol_gui\", chatmol_gui)\n\n\n_chatbar = None\n\n\ndef _init_gui():\n    \"\"\"Find PyMOL's main window and dock the ChatMol chat bar.\"\"\"\n    global _chatbar\n    if not _HAS_QT:\n        print(\n            \"Qt not available — ChatMol GUI disabled. \"\n            \"Use command-line: chat <message>\"\n        )\n        return\n    if _chatbar is not None:\n        _chatbar.show()\n        _chatbar.raise_()\n        return\n\n    app = QApplication.instance()\n    if app is None:\n        return\n\n    main_win = None\n    for w in app.topLevelWidgets():\n        if w.inherits(\"QMainWindow\") and w.isVisible():\n            main_win = w\n            break\n\n    if main_win is None:\n        print(\"ChatMol: could not find PyMOL main window for docking.\")\n        return\n\n    _chatbar = ChatMolChatBar(_agent, main_win)\n    main_win.addDockWidget(Qt.BottomDockWidgetArea, _chatbar)\n    print(\"ChatMol chat bar loaded.\")\n\n\nif _HAS_QT:\n    try:\n        QTimer.singleShot(1500, _init_gui)\n    except Exception:\n        pass\n\nprint(\n    \"ChatMol plugin loaded. Commands: chat, set_provider, set_api_key, \"\n    \"set_model, set_vision_model, reset_conversation, \"\n    \"chatmol_config, chatmol_settings, chatmol_gui\"\n)\n"}, "d36af966b330c7283cfd209f4100cd7d12a37763": {"t": 1785752931.6582787, "n": 59743, "h": "5a4226f171e50d6d", "c": "\"\"\"\nChatMol — Agentic PyMOL Plugin with OpenRouter and Qt5 GUI.\n\nLoad in PyMOL:\n    run /path/to/pymol_plugin/v2/chatmol.py\n\nOr from remote:\n    load https://raw.githubusercontent.com/ChatMol/ChatMol/main/pymol_plugin/v2/chatmol.py\n\nDemo prompt:\nZH: 3wzm的A链是玉米赤霉烯酮水解酶和底物的共晶复合物，作图展示酶底物(ZER)相互作用，风格参考顶级结构生物学期刊的风格，画面干净、清晰。\nEN: Chain A of 3WZM is a cocrystal complex of zearalenone hydrolase and its substrate; create a figure showing the enzyme-substrate (ZER) interaction, referencing the style of top-tier structural biology journals, with a clean and clear aesthetic.\n\"\"\"\n\nimport os\nimport json\nimport re\nimport time\nimport random\nimport base64\nimport tempfile\nimport threading\nimport requests\n\nfrom datetime import datetime\nfrom pymol import cmd, util, preset\n\n# ---------------------------------------------------------------------------\n# 1. Provider registry + LLMClient — multi-provider HTTP wrapper (no SDK)\n# ---------------------------------------------------------------------------\n\n# Each provider entry:\n#   base_url  – chat completions endpoint\n#   env_var   – environment variable for the API key\n#   models    – list of (model_id, label) for the settings UI\n#   extra_headers – provider-specific headers (optional)\nPROVIDERS = {\n    \"openrouter\": {\n        \"label\": \"OpenRouter\",\n        \"base_url\": \"https://openrouter.ai/api/v1/chat/completions\",\n        \"env_var\": \"OPENROUTER_API_KEY\",\n        \"models\": [\n            (\"openai/gpt-4o\", \"GPT-4o\"),\n            (\"openai/gpt-5.2\", \"GPT-5.2\"),\n            (\"openai/gpt-5.3-codex\", \"GPT-5.3 Codex\"),\n            (\"google/gemini-3.1-pro-preview\", \"Gemini 3.1 Pro\"),\n            (\"google/gemini-3-flash-preview\", \"Gemini 3 Flash\"),\n            (\"deepseek/deepseek-chat\", \"DeepSeek V3\"),\n        ],\n        \"vision_models\": [\n            (\"openai/gpt-4o\", \"GPT-4o\"),\n            (\"openai/gpt-5.2\", \"GPT-5.2\"),\n            (\"anthropic/claude-sonnet-4.6\", \"Claude Sonnet 4.6\"),\n            (\"google/gemini-3-flash-preview\", \"Gemini 3 Flash\"),\n            (\"google/gemini-3.1-pro-preview\", \"Gemini 3.1 Pro\"),\n        ],\n        \"extra_headers\": {\n            \"HTTP-Referer\": \"https://github.com/ChatMol/ChatMol\",\n            \"X-Title\": \"ChatMol\",\n        },\n    },\n    \"deepseek\": {\n        \"label\": \"DeepSeek\",\n        \"base_url\": \"https://api.deepseek.com/chat/completions\",\n        \"env_var\": \"DEEPSEEK_API_KEY\",\n        \"models\": [\n            (\"deepseek-chat\", \"DeepSeek V3\"),\n            (\"deepseek-reasoner\", \"DeepSeek R1\"),\n        ],\n        \"vision_models\": [],\n        \"extra_headers\": {},\n    },\n    \"kimi\": {\n        \"label\": \"Kimi (Moonshot)\",\n        \"base_url\": \"https://api.moonshot.ai/v1/chat/completions\",\n        \"env_var\": \"MOONSHOT_API_KEY\",\n        \"models\": [\n            (\"kimi-k2.5\", \"Kimi K2.5\"),\n        ],\n        \"vision_models\": [\n            (\"kimi-k2.5\", \"Kimi K2.5\"),\n        ],\n        \"extra_headers\": {},\n    },\n    \"glm\": {\n        \"label\": \"GLM (Zhipu)\",\n        \"base_url\": \"https://open.bigmodel.cn/api/paas/v4/chat/completions\",\n        \"env_var\": \"GLM_API_KEY\",\n        \"models\": [\n            (\"glm-5\", \"GLM-5\"),\n        ],\n        \"vision_models\": [],\n        \"extra_headers\": {},\n    },\n}\n\n\nclass LLMTransientError(RuntimeError):\n    \"\"\"Transient API/network errors that should be retried.\"\"\"\n\n\nclass LLMFatalError(RuntimeError):\n    \"\"\"Non-retriable API errors.\"\"\"\n\n\nclass LLMClient:\n    \"\"\"Thin HTTP client for OpenAI-compatible chat-completion APIs.\"\"\"\n\n    def __init__(self, provider_name, api_key):\n        prov = PROVIDERS.get(provider_name, PROVIDERS[\"openrouter\"])\n        self.base_url = prov[\"base_url\"]\n        self.extra_headers = prov.get(\"extra_headers\", {})\n        self.api_key = api_key\n\n    def _headers(self):\n        h = {\n            \"Content-Type\": \"application/json\",\n            \"Authorization\": f\"Bearer {self.api_key}\",\n        }\n        h.update(self.extra_headers)\n        return h\n\n    def _check_response(self, resp):\n        if resp.status_code >= 400:\n            try:\n                body = resp.json()\n                err = body.get(\"error\", {})\n                if isinstance(err, dict):\n                    detail = err.get(\"message\", \"\") or err.get(\"msg\", \"\")\n                else:\n                    detail = str(err)\n            except Exception:\n                detail = resp.text[:500] if resp.text else \"\"\n            message = (\n                f\"API error {resp.status_code} from {self.base_url}: \"\n                f\"{detail or resp.reason}\"\n            )\n            if resp.status_code in (408, 409, 425, 429) or resp.status_code >= 500:\n                raise LLMTransientError(message)\n            raise LLMFatalError(message)\n\n    def _parse_choices(self, data):\n        if \"error\" in data:\n            err = data[\"error\"]\n            msg = err.get(\"message\", str(err)) if isinstance(err, dict) else str(err)\n            raise LLMFatalError(f\"API returned error: {msg}\")\n        if \"choices\" not in data or not data[\"choices\"]:\n            raise LLMFatalError(\n                f\"Unexpected API response (no choices): \"\n                f\"{json.dumps(data, ensure_ascii=False)[:300]}\"\n            )\n        return data[\"choices\"][0]\n\n    def chat_completion(\n        self, model, messages, tools=None, temperature=0.01, max_tokens=4096\n    ):\n        payload = {\n            \"model\": model,\n            \"messages\": messages,\n            \"temperature\": temperature,\n            \"max_tokens\": max_tokens,\n        }\n        if tools:\n            payload[\"tools\"] = tools\n        try:\n            resp = requests.post(\n                self.base_url, headers=self._headers(), json=payload, timeout=120\n            )\n        except requests.ConnectionError:\n            raise LLMTransientError(f\"Connection failed: cannot reach {self.base_url}\")\n        except requests.Timeout:\n            raise LLMTransientError(f\"Request timed out after 120s to {self.base_url}\")\n        self._check_response(resp)\n        data = resp.json()\n        self._parse_choices(data)  # validate structure\n        return data\n\n    def vision_completion(self, model, text_prompt, image_base64):\n        messages = [\n            {\n                \"role\": \"user\",\n                \"content\": [\n                    {\"type\": \"text\", \"text\": text_prompt},\n                    {\n                        \"type\": \"image_url\",\n                        \"image_url\": {\"url\": f\"data:image/png;base64,{image_base64}\"},\n                    },\n                ],\n            }\n        ]\n        payload = {\n            \"model\": model,\n            \"messages\": messages,\n            \"max_tokens\": 1024,\n        }\n        try:\n            resp = requests.post(\n                self.base_url, headers=self._headers(), json=payload, timeout=120\n            )\n        except requests.ConnectionError:\n            raise LLMTransientError(f\"Connection failed: cannot reach {self.base_url}\")\n        except requests.Timeout:\n            raise LLMTransientError(f\"Request timed out after 120s to {self.base_url}\")\n        self._check_response(resp)\n        data = resp.json()\n        choice = self._parse_choices(data)\n        return choice[\"message\"][\"content\"]\n\n\n# ---------------------------------------------------------------------------\n# 2. PyMOLTools — 4 tool definitions + execution\n# ---------------------------------------------------------------------------\n\n\nBLOCKED_COMMANDS = [\n    (\"blocked destructive command\", re.compile(r\"^\\s*delete\\s+all\\s*$\", re.I)),\n    (\"blocked destructive command\", re.compile(r\"^\\s*remove\\s+all\\s*$\", re.I)),\n    (\"blocked session reset command\", re.compile(r\"^\\s*reinitialize(?:\\s|$)\", re.I)),\n    (\"blocked quit command\", re.compile(r\"^\\s*quit(?:\\s|$)\", re.I)),\n    (\"blocked script execution command\", re.compile(r\"^\\s*run(?:\\s|$)\", re.I)),\n    (\"blocked python command\", re.compile(r\"^\\s*python(?:\\s|$)\", re.I)),\n    (\n        \"blocked shell/system command\",\n        re.compile(\n            r\"^\\s*(?:!|shell(?:\\s|$)|system(?:\\s|$)|os\\.system|subprocess\\.)\", re.I\n        ),\n    ),\n]\n\nTOOL_DEFINITIONS = [\n    {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": \"inspect_session\",\n            \"description\": (\n                \"Inspect the current PyMOL session. Returns structured JSON \"\n                \"with objects, chain summaries, atom counts, and selections.\"\n            ),\n            \"parameters\": {\n                \"type\": \"object\",\n                \"properties\": {\n                    \"include_view\": {\n                        \"type\": \"boolean\",\n                        \"description\": \"Include camera view matrix.\",\n                        \"default\": False,\n                    }\n                },\n            },\n        },\n    },\n    {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": \"run_pymol_commands\",\n            \"description\": (\n                \"Execute newline-separated PyMOL commands. Use this for all \"\n                \"PyMOL operations: selections, styling, coloring, distances, \"\n                \"presets, etc. Commands are run sequentially with safety \"\n                \"guardrails (destructive/shell commands are blocked).\"\n            ),\n            \"parameters\": {\n                \"type\": \"object\",\n                \"properties\": {\n                    \"commands\": {\n                        \"type\": \"string\",\n                        \"description\": \"Newline-separated PyMOL commands.\",\n                    }\n                },\n                \"required\": [\"commands\"],\n            },\n        },\n    },\n    {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": \"render\",\n            \"description\": (\n                \"Render and export an image from the current scene. \"\n                \"Use purpose=preview for quick iterative checks, and \"\n                \"purpose=final for final publication export.\"\n            ),\n            \"parameters\": {\n                \"type\": \"object\",\n                \"properties\": {\n                    \"path\": {\n                        \"type\": \"string\",\n                        \"description\": \"Output image path.\",\n                    },\n                    \"purpose\": {\n                        \"type\": \"string\",\n                        \"enum\": [\"preview\", \"final\"],\n                        \"default\": \"preview\",\n                        \"description\": \"preview is lightweight; final is high quality.\",\n                    },\n                    \"width\": {\"type\": \"integer\", \"default\": 1400},\n                    \"height\": {\"type\": \"integer\", \"default\": 1050},\n                    \"dpi\": {\"type\": \"integer\", \"default\": 150},\n                    \"ray\": {\n                        \"type\": \"boolean\",\n                        \"description\": \"If omitted: preview=False, final=True.\",\n                    },\n                    \"transparent_bg\": {\"type\": \"boolean\", \"default\": True},\n                },\n                \"required\": [\"path\"],\n            },\n        },\n    },\n    {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": \"capture_viewport\",\n            \"description\": (\n                \"Capture current viewport and run vision model analysis. \"\n                \"Use this to visually verify your work (e.g., check that \"\n                \"styling looks correct before finalizing).\"\n            ),\n            \"parameters\": {\n                \"type\": \"object\",\n                \"properties\": {\n                    \"analysis_prompt\": {\n                        \"type\": \"string\",\n                        \"default\": \"Describe what you see in this molecular visualization.\",\n                    }\n                },\n            },\n        },\n    },\n]\n\n\ndef _to_bool(value, default=False):\n    if isinstance(value, bool):\n        return value\n    if isinstance(value, str):\n        return value.strip().lower() in (\"1\", \"true\", \"yes\", \"y\", \"on\")\n    if value is None:\n        return default\n    return bool(value)\n\n\ndef _blocked_reason(line):\n    for reason, pattern in BLOCKED_COMMANDS:\n        if pattern.search(line):\n            return reason\n    return \"\"\n\n\ndef execute_tool(tool_name, arguments, client=None, vision_model=None):\n    \"\"\"Dispatch to a tool implementation and always return JSON text.\"\"\"\n    dispatch = {\n        \"inspect_session\": _tool_inspect_session,\n        \"run_pymol_commands\": _tool_run_pymol_commands,\n        \"render\": _tool_render,\n        \"capture_viewport\": _tool_capture_viewport,\n    }\n    fn = dispatch.get(tool_name)\n    if fn is None:\n        return json.dumps(\n            {\"ok\": False, \"tool\": tool_name, \"error\": f\"Unknown tool: {tool_name}\"},\n            ensure_ascii=False,\n        )\n    try:\n        if tool_name == \"capture_viewport\":\n            result = fn(client=client, vision_model=vision_model, **(arguments or {}))\n        else:\n            result = fn(**(arguments or {}))\n    except TypeError as exc:\n        result = {\"ok\": False, \"tool\": tool_name, \"error\": f\"Invalid arguments: {exc}\"}\n    except Exception as exc:\n        result = {\"ok\": False, \"tool\": tool_name, \"error\": str(exc)}\n    return json.dumps(result, ensure_ascii=False)\n\n\ndef _tool_inspect_session(include_view=False):\n    objects = cmd.get_names(\"objects\")\n    selections = cmd.get_names(\"selections\")\n    details = []\n    for obj in objects:\n        details.append(\n            {\n                \"name\": obj,\n                \"atoms\": cmd.count_atoms(obj),\n                \"polymer_atoms\": cmd.count_atoms(f\"({obj}) and polymer\"),\n                \"hetatm_atoms\": cmd.count_atoms(f\"({obj}) and hetatm\"),\n                \"chains\": cmd.get_chains(obj),\n                \"states\": cmd.count_states(obj),\n            }\n        )\n    out = {\n        \"ok\": True,\n        \"tool\": \"inspect_session\",\n        \"object_count\": len(objects),\n        \"selection_count\": len(selections),\n        \"objects\": details,\n        \"selections\": selections,\n    }\n    if _to_bool(include_view):\n        out[\"view\"] = list(cmd.get_view())\n    return out\n\n\ndef _tool_run_pymol_commands(commands):\n    lines = [\n        ln.strip()\n        for ln in commands.splitlines()\n        if ln.strip() and not ln.strip().startswith(\"#\")\n    ]\n    executed, blocked, errors = [], [], []\n    known_names = set(cmd.get_names(\"all\"))\n    for line in lines:\n        reason = _blocked_reason(line)\n        if reason:\n            blocked.append({\"command\": line, \"reason\": reason})\n            continue\n        m_del = re.match(r\"^\\s*delete\\s+([A-Za-z0-9_]+)\\s*$\", line, re.I)\n        if m_del:\n            target = m_del.group(1)\n            if target not in known_names:\n                continue  # skip silently\n        try:\n            cmd.do(line)\n            executed.append(line)\n            if m_del:\n                known_names.discard(m_del.group(1))\n        except Exception as exc:\n            errors.append({\"command\": line, \"error\": str(exc)})\n    return {\n        \"ok\": not blocked and not errors,\n        \"tool\": \"run_pymol_commands\",\n        \"executed_count\": len(executed),\n        \"blocked_count\": len(blocked),\n        \"error_count\": len(errors),\n        \"executed\": executed,\n        \"blocked\": blocked,\n        \"errors\": errors,\n    }\n\n\ndef _tool_render(\n    path,\n    width=None,\n    height=None,\n    dpi=None,\n    ray=None,\n    transparent_bg=True,\n    purpose=\"preview\",\n):\n    out_path = os.path.abspath(os.path.expanduser(path))\n    out_dir = os.path.dirname(out_path)\n    if out_dir and not os.path.exists(out_dir):\n        os.makedirs(out_dir, exist_ok=True)\n\n    purpose = (purpose or \"preview\").strip().lower()\n    if purpose not in (\"preview\", \"final\"):\n        purpose = \"preview\"\n\n    if purpose == \"final\":\n        default_width, default_height, default_dpi = 2400, 1800, 300\n        default_ray = True\n    else:\n        default_width, default_height, default_dpi = 1400, 1050, 150\n        default_ray = False\n\n    try:\n        width = default_width if width is None else int(width)\n        height = default_height if height is None else int(height)\n        dpi = default_dpi if dpi is None else int(dpi)\n        width = max(256, width)\n        height = max(256, height)\n        dpi = max(72, dpi)\n    except (TypeError, ValueError):\n        width, height, dpi = default_width, default_height, default_dpi\n\n    if purpose == \"preview\":\n        width = min(width, 2200)\n        height = min(height, 1650)\n        dpi = min(dpi, 220)\n\n    ray_flag = 1 if _to_bool(ray, default=default_ray) else 0\n    transparent = _to_bool(transparent_bg, default=True)\n    cmd.set(\"ray_opaque_background\", 0 if transparent else 1)\n    cmd.png(out_path, width=width, height=height, dpi=dpi, ray=ray_flag, quiet=1)\n\n    exists = os.path.exists(out_path)\n    size = os.path.getsize(out_path) if exists else 0\n    return {\n        \"ok\": exists,\n        \"tool\": \"render\",\n        \"path\": out_path,\n        \"purpose\": purpose,\n        \"width\": width,\n        \"height\": height,\n        \"dpi\": dpi,\n        \"ray\": bool(ray_flag),\n        \"transparent_bg\": transparent,\n        \"bytes\": size,\n    }\n\n\ndef _tool_capture_viewport(analysis_prompt=None, client=None, vision_model=None):\n    if analysis_prompt is None:\n        analysis_prompt = \"Describe what you see in this molecular visualization.\"\n\n    tmp = tempfile.NamedTemporaryFile(suffix=\".png\", delete=False)\n    tmp_path = tmp.name\n    tmp.close()\n\n    try:\n        cmd.png(tmp_path, width=800, height=600, ray=0, quiet=1)\n        time.sleep(0.5)\n        with open(tmp_path, \"rb\") as f:\n            img_data = base64.b64encode(f.read()).decode(\"ascii\")\n    except Exception as exc:\n        return {\n            \"ok\": False,\n            \"tool\": \"capture_viewport\",\n            \"error\": f\"Failed to capture viewport: {exc}\",\n        }\n    finally:\n        try:\n            os.unlink(tmp_path)\n        except OSError:\n            pass\n\n    if not vision_model or not client:\n        return {\n            \"ok\": False,\n            \"tool\": \"capture_viewport\",\n            \"error\": \"No vision model configured. Use set_vision_model <model>.\",\n        }\n\n    try:\n        description = client.vision_completion(vision_model, analysis_prompt, img_data)\n        return {\n            \"ok\": True,\n            \"tool\": \"capture_viewport\",\n            \"analysis_prompt\": analysis_prompt,\n            \"description\": description,\n        }\n    except Exception as exc:\n        return {\n            \"ok\": False,\n            \"tool\": \"capture_viewport\",\n            \"error\": f\"Vision analysis failed: {exc}\",\n        }\n\n\n# ---------------------------------------------------------------------------\n# 3. System prompt\n# ---------------------------------------------------------------------------\n\nSYSTEM_PROMPT = \"\"\"\\\nYou are ChatMol, an expert AI assistant for PyMOL molecular visualization.\n\nYou have 4 tools:\n- `inspect_session`: See what's loaded in PyMOL (objects, chains, atoms, selections).\n- `run_pymol_commands`: Execute any PyMOL commands (selections, styling, coloring, \\\ndistances, presets, fetch, etc.). Commands are newline-separated.\n- `render`: Export an image (preview or final quality).\n- `capture_viewport`: Screenshot + vision analysis to check your work visually.\n\nWorkflow:\n1. For non-trivial requests, start with `inspect_session` to understand the current state.\n2. Use `run_pymol_commands` for all PyMOL operations. You know PyMOL well — use \\\ncmd.select, cmd.show, cmd.hide, cmd.color, cmd.set, cmd.distance, cmd.zoom, \\\ncmd.orient, util.color_chains, util.cnc, preset.ligand_sites_hq, etc.\n3. Use `capture_viewport` to verify your work visually when doing complex styling.\n4. Use `render` when the user requests an image export.\n\nStyle guidelines (publication quality):\n- White background, clean composition.\n- Cartoon as baseline representation for protein.\n- Sticks only for key residues, ligands, and interaction sites.\n- Context in gray; 1-2 accent colors (often cyan/marine + orange).\n- Surfaces used purposefully for targets/interfaces/pores, not everywhere.\n- Hydrogen bonds / polar contacts shown with dashed lines when relevant.\n- Keep decoration sparse — prioritize clarity over ornamentation.\n\nRules:\n- Never issue destructive commands (reinitialize, quit, delete all, shell commands).\n- If user input is ambiguous, ask for clarification rather than guessing.\n- Be concise in your final response.\n\"\"\"\n\n\n# ---------------------------------------------------------------------------\n# 4. ChatMolAgent — config persistence + simple agent loop\n# ---------------------------------------------------------------------------\n\n\nclass ChatMolAgent:\n    \"\"\"Simple agentic loop: call LLM, execute tools, repeat.\"\"\"\n\n    CONFIG_PATH = os.path.expanduser(\"~/.PyMOL/chatmol_config.json\")\n\n    DEFAULT_CONFIG = {\n        \"provider\": \"openrouter\",\n        \"api_keys\": {},\n        \"text_model\": \"google/gemini-3-flash-preview\",\n        \"vision_model\": \"google/gemini-3-flash-preview\",\n        \"temperature\": 0.01,\n        \"max_tokens\": 4096,\n        \"max_iterations\": 50,\n        \"max_tool_calls\": 30,\n    }\n\n    def __init__(self):\n        os.makedirs(os.path.dirname(self.CONFIG_PATH), exist_ok=True)\n        self.config = self._load_config()\n        self._reinit_client()\n        self.conversation_history = []\n\n    # -- config persistence -------------------------------------------------\n\n    def _load_config(self):\n        try:\n            with open(self.CONFIG_PATH, \"r\") as f:\n                cfg = json.load(f)\n            merged = dict(self.DEFAULT_CONFIG)\n            merged[\"api_keys\"] = dict(self.DEFAULT_CONFIG[\"api_keys\"])\n            merged.update(cfg)\n            if not isinstance(merged.get(\"api_keys\"), dict):\n                merged[\"api_keys\"] = {}\n            # Migrate old single \"api_key\" field\n            old_key = merged.pop(\"api_key\", \"\")\n            if old_key and isinstance(old_key, str):\n                prov = merged.get(\"provider\", \"openrouter\")\n                merged[\"api_keys\"].setdefault(prov, old_key)\n                self._save_config(merged)\n            # Remove legacy keys\n            for legacy in (\"human_in_the_loop\", \"multi_agent_enabled\"):\n                merged.pop(legacy, None)\n            return merged\n        except (FileNotFoundError, json.JSONDecodeError):\n            self._save_config(self.DEFAULT_CONFIG)\n            return dict(self.DEFAULT_CONFIG)\n\n    def _save_config(self, cfg=None):\n        if cfg is None:\n            cfg = self.config\n        try:\n            with open(self.CONFIG_PATH, \"w\") as f:\n                json.dump(cfg, f, indent=2)\n        except Exception as exc:\n            print(f\"Warning: could not save config: {exc}\")\n\n    def _resolve_api_key(self):\n        prov_name = self.config.get(\"provider\", \"openrouter\")\n        prov = PROVIDERS.get(prov_name, PROVIDERS[\"openrouter\"])\n        return os.getenv(prov[\"env_var\"], \"\") or self.config.get(\"api_keys\", {}).get(\n            prov_name, \"\"\n        )\n\n    def _reinit_client(self):\n        prov_name = self.config.get(\"provider\", \"openrouter\")\n        api_key = self._resolve_api_key()\n        self.client = LLMClient(prov_name, api_key)\n\n    # -- PyMOL-registered commands ------------------------------------------\n\n    def chat(self, *args):\n        \"\"\"PyMOL command: chat <message>\"\"\"\n        message = \" \".join(str(a) for a in args).strip()\n        if not message:\n            print(\"Usage: chat <message>\")\n            return\n        if not self._resolve_api_key():\n            prov = self.config.get(\"provider\", \"openrouter\")\n            env = PROVIDERS.get(prov, {}).get(\"env_var\", \"\")\n            print(f\"No API key set. Use: set_api_key <key>  (or set env var {env})\")\n            return\n        try:\n            response = self._run_agent_loop(message)\n            print(\"=\" * 50)\n            print(\"ChatMol:\", response)\n            print(\"=\" * 50)\n        except RuntimeError as exc:\n            print(f\"ChatMol error: {exc}\")\n        except requests.exceptions.RequestException as exc:\n            print(f\"ChatMol network error: {exc}\")\n        except (KeyError, IndexError, TypeError) as exc:\n            prov = self.config.get(\"provider\", \"openrouter\")\n            model = self.config.get(\"text_model\", \"\")\n            print(f\"ChatMol error: unexpected response format — {exc}\")\n            print(f\"  Provider: {prov}, Model: {model}\")\n            print(\n                \"  The model may not support tool calling, \"\n                \"or the API returned an unexpected structure.\"\n            )\n        except Exception as exc:\n            print(f\"ChatMol error ({type(exc).__name__}): {exc}\")\n\n    def set_api_key(self, api_key=\"\"):\n        \"\"\"PyMOL command: set_api_key <key>\"\"\"\n        api_key = api_key.strip()\n        if not api_key:\n            prov_name = self.config.get(\"provider\", \"openrouter\")\n            label = PROVIDERS.get(prov_name, {}).get(\"label\", prov_name)\n            print(f\"Usage: set_api_key <your_{label}_api_key>\")\n            return\n        prov_name = self.config.get(\"provider\", \"openrouter\")\n        if not isinstance(self.config.get(\"api_keys\"), dict):\n            self.config[\"api_keys\"] = {}\n        self.config[\"api_keys\"][prov_name] = api_key\n        self._save_config()\n        self._reinit_client()\n        label = PROVIDERS.get(prov_name, {}).get(\"label\", prov_name)\n        print(f\"{label} API key saved.\")\n\n    def set_provider(self, provider_name=\"\"):\n        \"\"\"PyMOL command: set_provider <name>\"\"\"\n        provider_name = provider_name.strip().lower()\n        if not provider_name or provider_name not in PROVIDERS:\n            print(f\"Available providers: {', '.join(PROVIDERS.keys())}\")\n            if provider_name:\n                print(f\"Unknown provider: {provider_name}\")\n            return\n        self.config[\"provider\"] = provider_name\n        prov = PROVIDERS[provider_name]\n        if prov[\"models\"]:\n            self.config[\"text_model\"] = prov[\"models\"][0][0]\n        if prov[\"vision_models\"]:\n            self.config[\"vision_model\"] = prov[\"vision_models\"][0][0]\n        else:\n            self.config[\"vision_model\"] = \"\"\n        self._save_config()\n        self._reinit_client()\n        print(f\"Provider set to: {prov['label']}\")\n        print(f\"  Text model: {self.config['text_model']}\")\n        if self.config[\"vision_model\"]:\n            print(f\"  Vision model: {self.config['vision_model']}\")\n        print(f\"  Set API key via: set_api_key <key>  (or env {prov['env_var']})\")\n\n    def set_model(self, model_name=\"\"):\n        \"\"\"PyMOL command: set_model <model>\"\"\"\n        model_name = model_name.strip()\n        if not model_name:\n            print(f\"Current model: {self.config['text_model']}\")\n            print(\"Usage: set_model <provider/model>\")\n            return\n        self.config[\"text_model\"] = model_name\n        self._save_config()\n        print(f\"Text model set to: {model_name}\")\n\n    def set_vision_model(self, model_name=\"\"):\n        \"\"\"PyMOL command: set_vision_model <model>\"\"\"\n        model_name = model_name.strip()\n        if not model_name:\n            print(f\"Current vision model: {self.config['vision_model']}\")\n            print(\"Usage: set_vision_model <provider/model>\")\n            return\n        self.config[\"vision_model\"] = model_name\n        self._save_config()\n        print(f\"Vision model set to: {model_name}\")\n\n    def reset_conversation(self):\n        \"\"\"PyMOL command: reset_conversation\"\"\"\n        self.conversation_history = []\n        print(\"Conversation history cleared.\")\n\n    def save_conversation(self, filename=\"\"):\n        \"\"\"PyMOL command: save_conversation [filename]\"\"\"\n        if not filename:\n            filename = (\n                f\"chatmol_conversation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json\"\n            )\n        with open(filename, \"w\", encoding=\"utf-8\") as f:\n            json.dump(self.conversation_history, f, indent=2)\n        print(f\"Conversation saved to {filename}\")\n\n    def load_conversation(self, filename=\"\"):\n        \"\"\"PyMOL command: load_conversation <filename>\"\"\"\n        if not filename:\n            print(\"Usage: load_conversation <filename.json>\")\n            return\n        try:\n            with open(filename, \"r\", encoding=\"utf-8\") as f:\n                self.conversation_history = json.load(f)\n            print(\n                f\"Loaded conversation from {filename} \"\n                f\"({len(self.conversation_history)} messages)\"\n            )\n        except Exception as exc:\n            print(f\"Failed to load conversation: {exc}\")\n\n    def show_config(self):\n        \"\"\"PyMOL command: chatmol_config\"\"\"\n        prov_name = self.config.get(\"provider\", \"openrouter\")\n        prov = PROVIDERS.get(prov_name, {})\n        key = self._resolve_api_key()\n        masked = (key[:8] + \"...\" + key[-4:]) if key and len(key) > 12 else \"(not set)\"\n        has_env = bool(os.getenv(prov.get(\"env_var\", \"\"), \"\"))\n        print(\"ChatMol configuration:\")\n        print(f\"  provider: {prov.get('label', prov_name)} ({prov_name})\")\n        print(f\"  api_key: {masked} ({'env var' if has_env else 'config file'})\")\n        print(f\"  text_model: {self.config.get('text_model', '')}\")\n        print(f\"  vision_model: {self.config.get('vision_model', '') or '(none)'}\")\n        print(f\"  temperature: {self.config.get('temperature', 0.01)}\")\n        print(f\"  max_tokens: {self.config.get('max_tokens', 4096)}\")\n        print(f\"  max_iterations: {self.config.get('max_iterations', 50)}\")\n        print(f\"  max_tool_calls: {self.config.get('max_tool_calls', 30)}\")\n        stored = [k for k, v in self.config.get(\"api_keys\", {}).items() if v]\n        if stored:\n            print(f\"  keys stored for: {', '.join(stored)}\")\n\n    # -- agentic loop -------------------------------------------------------\n\n    def _chat_completion_with_retry(\n        self, model, messages, tool_defs, temperature, max_tokens, phase_callback=None\n    ):\n        max_attempts = 4\n        for attempt in range(1, max_attempts + 1):\n            try:\n                return self.client.chat_completion(\n                    model,\n                    messages,\n                    tools=tool_defs,\n                    temperature=temperature,\n                    max_tokens=max_tokens,\n                )\n            except LLMTransientError as exc:\n                if attempt >= max_attempts:\n                    raise RuntimeError(\n                        f\"Transient API errors persisted after {max_attempts} attempts: {exc}\"\n                    )\n                delay = min(8.0, float(2 ** (attempt - 1))) + random.uniform(0.0, 0.4)\n                print(\n                    f\"  [Retry {attempt}/{max_attempts}] transient error: \"\n                    f\"{exc}. Sleeping {delay:.1f}s.\"\n                )\n                if phase_callback:\n                    phase_callback(\"Retrying API request\")\n                time.sleep(delay)\n\n    @staticmethod\n    def _parse_assistant_message(data):\n        choices = data.get(\"choices\")\n        if not isinstance(choices, list) or not choices:\n            raise RuntimeError(\"Invalid API response: missing choices.\")\n        message = choices[0].get(\"message\")\n        if not isinstance(message, dict):\n            raise RuntimeError(\"Invalid API response: missing assistant message.\")\n        return message\n\n    @staticmethod\n    def _normalize_content(content):\n        if isinstance(content, (list, dict)):\n            return json.dumps(content, ensure_ascii=False)\n        if content is None:\n            return \"\"\n        if not isinstance(content, str):\n            return str(content)\n        return content\n\n    @staticmethod\n    def _parse_tool_calls(assistant_msg):\n        raw_calls = assistant_msg.get(\"tool_calls\") or []\n        parsed = []\n        for tc in raw_calls:\n            if not isinstance(tc, dict):\n                raise RuntimeError(\"Invalid tool call entry from model.\")\n            tc_id = tc.get(\"id\")\n            fn = tc.get(\"function\") or {}\n            fn_name = fn.get(\"name\")\n            raw_args = fn.get(\"arguments\", \"{}\")\n            if not tc_id or not fn_name:\n                raise RuntimeError(\n                    \"Invalid tool call payload (missing id or function name).\"\n                )\n            if isinstance(raw_args, dict):\n                fn_args = raw_args\n            else:\n                raw_args = \"{}\" if raw_args is None else str(raw_args).strip()\n                if not raw_args:\n                    raw_args = \"{}\"\n                try:\n                    fn_args = json.loads(raw_args)\n                except json.JSONDecodeError as exc:\n                    raise RuntimeError(\n                        f\"Tool arguments not valid JSON for {fn_name}: {exc}\"\n                    )\n            if not isinstance(fn_args, dict):\n                raise RuntimeError(f\"Tool arguments for {fn_name} must be an object.\")\n            parsed.append(\n                {\"id\": tc_id, \"name\": fn_name, \"arguments\": fn_args, \"raw\": tc}\n            )\n        return parsed\n\n    def _make_tool_executor(self, tool_executor_override=None):\n        \"\"\"Return a tool executor function that passes client/vision_model.\"\"\"\n        if tool_executor_override:\n            return tool_executor_override\n        client = self.client\n        vision_model = self.config.get(\"vision_model\", \"\")\n\n        def _executor(tool_name, arguments):\n            return execute_tool(\n                tool_name, arguments, client=client, vision_model=vision_model\n            )\n\n        return _executor\n\n    def _run_agent_loop_internal(\n        self, message, tool_executor=None, phase_callback=None\n    ):\n        self.conversation_history.append({\"role\": \"user\", \"content\": message})\n\n        model = self.config[\"text_model\"]\n        temperature = self.config.get(\"temperature\", 0.01)\n        max_tokens = self.config.get(\"max_tokens\", 4096)\n        max_iterations = max(1, int(self.config.get(\"max_iterations\", 50) or 50))\n        max_tool_calls = max(8, int(self.config.get(\"max_tool_calls\", 30) or 30))\n        executor = self._make_tool_executor(tool_executor)\n        tool_calls_used = 0\n\n        for iteration in range(max_iterations):\n            if phase_callback:\n                phase_callback(\"Thinking\" if iteration == 0 else \"Thinking more\")\n\n            messages = [{\"role\": \"system\", \"content\": SYSTEM_PROMPT}]\n            messages += self.conversation_history\n\n            data = self._chat_completion_with_retry(\n                model,\n                messages,\n                TOOL_DEFINITIONS,\n                temperature,\n                max_tokens,\n                phase_callback,\n            )\n            assistant_msg = self._parse_assistant_message(data)\n            parsed_calls = self._parse_tool_calls(assistant_msg)\n            content = self._normalize_content(assistant_msg.get(\"content\", None))\n\n            history_entry = {\"role\": \"assistant\", \"content\": content}\n            if parsed_calls:\n                history_entry[\"tool_calls\"] = [tc[\"raw\"] for tc in parsed_calls]\n\n            if parsed_calls:\n                tool_entries = []\n                for tc in parsed_calls:\n                    fn_name = tc[\"name\"]\n                    fn_args = tc[\"arguments\"]\n\n                    if phase_callback:\n                        phase_callback(f\"Running {fn_name}\")\n                    print(\n                        f\"  [Tool: {fn_name}] \"\n                        f\"{json.dumps(fn_args, ensure_ascii=False)[:120]}\"\n                    )\n\n                    if tool_calls_used >= max_tool_calls:\n                        result = json.dumps(\n                            {\n                                \"ok\": False,\n                                \"tool\": fn_name,\n                                \"error\": (\n                                    \"Tool-call budget reached. Stop calling tools \"\n                                    \"and provide your final response.\"\n                                ),\n                            },\n                            ensure_ascii=False,\n                        )\n                    else:\n                        try:\n                            result = executor(fn_name, fn_args)\n                        except Exception as exc:\n                            result = json.dumps(\n                                {\"ok\": False, \"tool\": fn_name, \"error\": str(exc)},\n                                ensure_ascii=False,\n                            )\n\n                    tool_entries.append(\n                        {\n                            \"role\": \"tool\",\n                            \"tool_call_id\": tc[\"id\"],\n                            \"content\": result,\n                        }\n                    )\n                    tool_calls_used += 1\n\n                self.conversation_history.append(history_entry)\n                self.conversation_history.extend(tool_entries)\n                continue\n\n            # No tool calls — this is the final response\n            self.conversation_history.append(history_entry)\n            return (content or \"\").strip()\n\n        return (\n            \"Agent loop reached maximum iterations \"\n            f\"({max_iterations}) without a final response.\"\n        )\n\n    def _run_agent_loop(self, message):\n        return self._run_agent_loop_internal(message, tool_executor=None)\n\n    def run_agent_loop_with_callback(self, message, tool_executor, phase_callback=None):\n        return self._run_agent_loop_internal(\n            message, tool_executor=tool_executor, phase_callback=phase_callback\n        )\n\n\n# ---------------------------------------------------------------------------\n# 5. Qt5 GUI (loaded conditionally)\n# ---------------------------------------------------------------------------\n\n_HAS_QT = False\ntry:\n    from pymol.Qt import QtWidgets, QtCore, QtGui  # noqa: E402\n\n    QDockWidget = QtWidgets.QDockWidget\n    QWidget = QtWidgets.QWidget\n    QVBoxLayout = QtWidgets.QVBoxLayout\n    QHBoxLayout = QtWidgets.QHBoxLayout\n    QTextEdit = QtWidgets.QTextEdit\n    QLineEdit = QtWidgets.QLineEdit\n    QPushButton = QtWidgets.QPushButton\n    QLabel = QtWidgets.QLabel\n    QDialog = QtWidgets.QDialog\n    QFormLayout = QtWidgets.QFormLayout\n    QDoubleSpinBox = QtWidgets.QDoubleSpinBox\n    QSpinBox = QtWidgets.QSpinBox\n    QComboBox = QtWidgets.QComboBox\n    QApplication = QtWidgets.QApplication\n\n    Qt = QtCore.Qt\n    QThread = QtCore.QThread\n    QTimer = QtCore.QTimer\n    QPropertyAnimation = QtCore.QPropertyAnimation\n    QEasingCurve = QtCore.QEasingCurve\n    Signal = getattr(QtCore, \"Signal\", None) or getattr(QtCore, \"pyqtSignal\")\n    Property = getattr(QtCore, \"Property\", None) or getattr(QtCore, \"pyqtProperty\")\n\n    QColor = QtGui.QColor\n    QPainter = QtGui.QPainter\n\n    _HAS_QT = True\nexcept ImportError:\n    pass\n\n\nif _HAS_QT:\n\n    # -- ThinkingIndicator --------------------------------------------------\n\n    class ThinkingIndicator(QWidget):\n        \"\"\"Three dots that pulse with a breathing animation.\"\"\"\n\n        _DOT_COUNT = 3\n        _DOT_RADIUS = 4\n        _DOT_SPACING = 14\n\n        def __init__(self, parent=None):\n            super().__init__(parent)\n            self._opacity = 0.3\n            self._phase_text = \"Thinking\"\n            self.setFixedHeight(28)\n            self._anim = QPropertyAnimation(self, b\"dot_opacity\")\n            self._anim.setDuration(900)\n            self._anim.setStartValue(0.3)\n            self._anim.setEndValue(1.0)\n            self._anim.setEasingCurve(QEasingCurve.InOutSine)\n            self._anim.setLoopCount(-1)\n            self._anim.finished.connect(lambda: None)\n            self._forward = True\n            self._cycle_timer = QTimer(self)\n            self._cycle_timer.setInterval(900)\n            self._cycle_timer.timeout.connect(self._toggle_direction)\n            self.hide()\n\n        def _toggle_direction(self):\n            self._forward = not self._forward\n            if self._forward:\n                self._anim.setStartValue(0.3)\n                self._anim.setEndValue(1.0)\n            else:\n                self._anim.setStartValue(1.0)\n                self._anim.setEndValue(0.3)\n            self._anim.start()\n\n        @Property(float)\n        def dot_opacity(self):\n            return self._opacity\n\n        @dot_opacity.setter\n        def dot_opacity(self, value):\n            self._opacity = value\n            self.update()\n\n        def start(self, phase_text=\"Thinking\"):\n            self._phase_text = phase_text\n            self._anim.setStartValue(0.3)\n            self._anim.setEndValue(1.0)\n            self._forward = True\n            self._anim.start()\n            self._cycle_timer.start()\n            self.show()\n\n        def stop(self):\n            self._anim.stop()\n            self._cycle_timer.stop()\n            self.hide()\n\n        def set_phase(self, text):\n            self._phase_text = text\n            self.update()\n\n        def paintEvent(self, _event):\n            p = QPainter(self)\n            p.setRenderHint(QPainter.Antialiasing)\n            p.setPen(QColor(120, 120, 120))\n            p.drawText(8, 18, self._phase_text)\n            label_width = p.fontMetrics().horizontalAdvance(self._phase_text) + 16\n            color = QColor(70, 130, 230)\n            for i in range(self._DOT_COUNT):\n                offset = i * 0.15\n                alpha = max(0.0, min(1.0, self._opacity - offset))\n                color.setAlphaF(alpha)\n                p.setBrush(color)\n                p.setPen(Qt.NoPen)\n                x = label_width + i * self._DOT_SPACING\n                p.drawEllipse(x, 10, self._DOT_RADIUS * 2, self._DOT_RADIUS * 2)\n            p.end()\n\n    # -- AgentWorker --------------------------------------------------------\n\n    class AgentWorker(QThread):\n        \"\"\"Runs the agentic loop off the main thread.\"\"\"\n\n        finished = Signal(str)\n        error = Signal(str)\n        tool_executed = Signal(str, str)\n        request_tool_execution = Signal(str, str, str)\n        phase_changed = Signal(str)\n\n        def __init__(self, agent, message, parent=None):\n            super().__init__(parent)\n            self.agent = agent\n            self.message = message\n            self._tool_result = None\n            self._tool_event = threading.Event()\n\n        def run(self):\n            try:\n                response = self.agent.run_agent_loop_with_callback(\n                    self.message, self._tool_executor, phase_callback=self._on_phase\n                )\n                self.finished.emit(response)\n            except Exception as exc:\n                self.error.emit(str(exc))\n\n        def _on_phase(self, phase_text):\n            self.phase_changed.emit(phase_text)\n\n        def _tool_executor(self, tool_name, arguments):\n            self._tool_event.clear()\n            self._tool_result = None\n            args_json = json.dumps(arguments, ensure_ascii=False)\n            self.request_tool_execution.emit(\"\", tool_name, args_json)\n            self._tool_event.wait(timeout=60)\n            result = self._tool_result or \"Tool execution timed out.\"\n            self.tool_executed.emit(tool_name, result[:200])\n            return result\n\n        def deliver_tool_result(self, result):\n            self._tool_result = result\n            self._tool_event.set()\n\n    # -- ChatMolSettingsDialog ----------------------------------------------\n\n    class ChatMolSettingsDialog(QDialog):\n\n        def __init__(self, agent, parent=None):\n            super().__init__(parent)\n            self.agent = agent\n            self.setWindowTitle(\"ChatMol Settings\")\n            self.setMinimumWidth(460)\n            self._build_ui()\n\n        def _build_ui(self):\n            layout = QFormLayout(self)\n\n            # Provider selector\n            self.provider_combo = QComboBox()\n            for key, prov in PROVIDERS.items():\n                self.provider_combo.addItem(prov[\"label\"], key)\n            cur_prov = self.agent.config.get(\"provider\", \"openrouter\")\n            idx = self.provider_combo.findData(cur_prov)\n            if idx >= 0:\n                self.provider_combo.setCurrentIndex(idx)\n            self.provider_combo.currentIndexChanged.connect(self._on_provider_changed)\n            layout.addRow(\"Provider:\", self.provider_combo)\n\n            # API key\n            self.api_key_edit = QLineEdit()\n            self.api_key_edit.setEchoMode(QLineEdit.Password)\n            self.api_key_edit.setText(\n                self.agent.config.get(\"api_keys\", {}).get(cur_prov, \"\")\n            )\n            self.env_hint = QLabel()\n            self.env_hint.setStyleSheet(\"color: grey; font-size: 10px;\")\n            key_layout = QVBoxLayout()\n            key_layout.setSpacing(2)\n            key_layout.addWidget(self.api_key_edit)\n            key_layout.addWidget(self.env_hint)\n            layout.addRow(\"API Key:\", key_layout)\n\n            # Text model\n            self.text_model_combo = QComboBox()\n            self.text_model_combo.setEditable(True)\n            layout.addRow(\"Text Model:\", self.text_model_combo)\n\n            # Vision model\n            self.vision_model_combo = QComboBox()\n            self.vision_model_combo.setEditable(True)\n            layout.addRow(\"Vision Model:\", self.vision_model_combo)\n\n            # Temperature\n            self.temp_spin = QDoubleSpinBox()\n            self.temp_spin.setRange(0.0, 2.0)\n            self.temp_spin.setSingleStep(0.05)\n            self.temp_spin.setValue(self.agent.config.get(\"temperature\", 0.01))\n            layout.addRow(\"Temperature:\", self.temp_spin)\n\n            # Max tokens\n            self.max_tokens_spin = QSpinBox()\n            self.max_tokens_spin.setRange(256, 16384)\n            self.max_tokens_spin.setSingleStep(256)\n            self.max_tokens_spin.setValue(self.agent.config.get(\"max_tokens\", 4096))\n            layout.addRow(\"Max Tokens:\", self.max_tokens_spin)\n\n            # Max iterations\n            self.max_iterations_spin = QSpinBox()\n            self.max_iterations_spin.setRange(1, 500)\n            self.max_iterations_spin.setSingleStep(1)\n            self.max_iterations_spin.setValue(\n                self.agent.config.get(\"max_iterations\", 50)\n            )\n            layout.addRow(\"Max Iterations:\", self.max_iterations_spin)\n\n            # Max tool calls per request\n            self.max_tool_calls_spin = QSpinBox()\n            self.max_tool_calls_spin.setRange(8, 200)\n            self.max_tool_calls_spin.setSingleStep(1)\n            self.max_tool_calls_spin.setValue(\n                self.agent.config.get(\"max_tool_calls\", 30)\n            )\n            layout.addRow(\"Max Tool Calls:\", self.max_tool_calls_spin)\n\n            # Buttons\n            btn_layout = QHBoxLayout()\n            save_btn = QPushButton(\"Save\")\n            save_btn.clicked.connect(self._on_save)\n            cancel_btn = QPushButton(\"Cancel\")\n            cancel_btn.clicked.connect(self.reject)\n            btn_layout.addStretch()\n            btn_layout.addWidget(save_btn)\n            btn_layout.addWidget(cancel_btn)\n            layout.addRow(btn_layout)\n\n            self._populate_models(cur_prov)\n\n        def _on_provider_changed(self, _index):\n            prov_key = self.provider_combo.currentData()\n            stored_key = self.agent.config.get(\"api_keys\", {}).get(prov_key, \"\")\n            self.api_key_edit.setText(stored_key)\n            self._populate_models(prov_key)\n\n        def _populate_models(self, prov_key):\n            prov = PROVIDERS.get(prov_key, PROVIDERS[\"openrouter\"])\n\n            env_var = prov.get(\"env_var\", \"\")\n            has_env = bool(os.getenv(env_var, \"\"))\n            if has_env:\n                self.env_hint.setText(f\"Using env var {env_var} (overrides this field)\")\n            else:\n                self.env_hint.setText(f\"Or set env var: {env_var}\")\n\n            # Text models\n            self.text_model_combo.clear()\n            for model_id, label in prov[\"models\"]:\n                self.text_model_combo.addItem(f\"{label}  ({model_id})\", model_id)\n            saved = self.agent.config.get(\"text_model\", \"\")\n            restored = False\n            for i in range(self.text_model_combo.count()):\n                if self.text_model_combo.itemData(i) == saved:\n                    self.text_model_combo.setCurrentIndex(i)\n                    restored = True\n                    break\n            if not restored and saved:\n                self.text_model_combo.setEditText(saved)\n\n            # Vision models\n            self.vision_model_combo.clear()\n            self.vision_model_combo.addItem(\"(none)\", \"\")\n            for model_id, label in prov.get(\"vision_models\", []):\n                self.vision_model_combo.addItem(f\"{label}  ({model_id})\", model_id)\n            saved_v = self.agent.config.get(\"vision_model\", \"\")\n            restored_v = False\n            for i in range(self.vision_model_combo.count()):\n                if self.vision_model_combo.itemData(i) == saved_v:\n                    self.vision_model_combo.setCurrentIndex(i)\n                    restored_v = True\n                    break\n            if not restored_v and saved_v:\n                self.vision_model_combo.setEditText(saved_v)\n\n        def _on_save(self):\n            prov_key = self.provider_combo.currentData()\n            self.agent.config[\"provider\"] = prov_key\n\n            if not isinstance(self.agent.config.get(\"api_keys\"), dict):\n                self.agent.config[\"api_keys\"] = {}\n            self.agent.config[\"api_keys\"][prov_key] = self.api_key_edit.text().strip()\n\n            text_idx = self.text_model_combo.currentIndex()\n            text_data = self.text_model_combo.itemData(text_idx)\n            self.agent.config[\"text_model\"] = (\n                text_data if text_data else self.text_model_combo.currentText().strip()\n            )\n\n            vis_idx = self.vision_model_combo.currentIndex()\n            vis_data = self.vision_model_combo.itemData(vis_idx)\n            self.agent.config[\"vision_model\"] = (\n                vis_data\n                if vis_data is not None\n                else self.vision_model_combo.currentText().strip()\n            )\n\n            self.agent.config[\"temperature\"] = self.temp_spin.value()\n            self.agent.config[\"max_tokens\"] = self.max_tokens_spin.value()\n            self.agent.config[\"max_iterations\"] = self.max_iterations_spin.value()\n            self.agent.config[\"max_tool_calls\"] = self.max_tool_calls_spin.value()\n            self.agent._save_config()\n            self.agent._reinit_client()\n            self.accept()\n\n    # -- ChatMolChatBar -----------------------------------------------------\n\n    class ChatMolChatBar(QDockWidget):\n\n        def __init__(self, agent, parent=None):\n            super().__init__(\"ChatMol\", parent)\n            self.agent = agent\n            self._worker = None\n            self._build_ui()\n\n        def _build_ui(self):\n            container = QWidget()\n            vbox = QVBoxLayout(container)\n            vbox.setContentsMargins(4, 4, 4, 4)\n            vbox.setSpacing(4)\n\n            # Top bar\n            top_bar = QHBoxLayout()\n            top_bar.addWidget(QLabel(\"<b>ChatMol</b>\"))\n            top_bar.addStretch()\n            settings_btn = QPushButton(\"Settings\")\n            settings_btn.setFixedWidth(70)\n            settings_btn.clicked.connect(self._open_settings)\n            top_bar.addWidget(settings_btn)\n            clear_btn = QPushButton(\"Clear\")\n            clear_btn.setFixedWidth(50)\n            clear_btn.clicked.connect(self._clear_chat)\n            top_bar.addWidget(clear_btn)\n            vbox.addLayout(top_bar)\n\n            # Execution trace panel\n            self.trace_display = QTextEdit()\n            self.trace_display.setReadOnly(True)\n            self.trace_display.setMaximumHeight(92)\n            self.trace_display.setStyleSheet(\n                \"font-family: Menlo, Consolas, monospace; font-size:10px; \"\n                \"color:#263238; background:#FAFAFA;\"\n            )\n            vbox.addWidget(self.trace_display)\n\n            # Chat history\n            self.chat_display = QTextEdit()\n            self.chat_display.setReadOnly(True)\n            self.chat_display.setMinimumHeight(120)\n            vbox.addWidget(self.chat_display)\n\n            # Thinking indicator\n            self.thinking = ThinkingIndicator()\n            vbox.addWidget(self.thinking)\n\n            # Input bar\n            input_bar = QHBoxLayout()\n            self.input_edit = QLineEdit()\n            self.input_edit.setPlaceholderText(\"Type a message...\")\n            self.input_edit.returnPressed.connect(self._on_send)\n            input_bar.addWidget(self.input_edit)\n            self.send_btn = QPushButton(\"Send\")\n            self.send_btn.setFixedWidth(50)\n            self.send_btn.clicked.connect(self._on_send)\n            input_bar.addWidget(self.send_btn)\n            self.model_label = QLabel()\n            self.model_label.setStyleSheet(\"color: grey; font-size: 10px;\")\n            self._update_model_label()\n            input_bar.addWidget(self.model_label)\n            vbox.addLayout(input_bar)\n\n            self.setWidget(container)\n\n        def _update_model_label(self):\n            prov_name = self.agent.config.get(\"provider\", \"openrouter\")\n            prov_label = PROVIDERS.get(prov_name, {}).get(\"label\", prov_name)\n            model = self.agent.config.get(\"text_model\", \"\")\n            short = model.split(\"/\")[-1] if \"/\" in model else model\n            self.model_label.setText(f\"{short} ({prov_label})\")\n\n        def _open_settings(self):\n            dlg = ChatMolSettingsDialog(self.agent, self)\n            if dlg.exec_():\n                self._update_model_label()\n\n        def _clear_chat(self):\n            self.chat_display.clear()\n            self.trace_display.clear()\n            self.agent.reset_conversation()\n\n        def _append_html(self, html):\n            self.chat_display.append(html)\n            sb = self.chat_display.verticalScrollBar()\n            sb.setValue(sb.maximum())\n\n        def _append_trace(self, text):\n            self.trace_display.append(_escape_html(text))\n            sb = self.trace_display.verticalScrollBar()\n            sb.setValue(sb.maximum())\n\n        def _on_send(self):\n            text = self.input_edit.text().strip()\n            if not text:\n                return\n\n            if not self.agent._resolve_api_key():\n                self._append_html(\n                    '<span style=\"color:red;\">No API key set. '\n                    \"Open Settings to configure your API key.</span>\"\n                )\n                return\n\n            self.input_edit.clear()\n            self.input_edit.setEnabled(False)\n            self.send_btn.setEnabled(False)\n            self.trace_display.clear()\n\n            self._append_html(\n                f'<b style=\"color:#2962FF;\">You:</b> {_escape_html(text)}'\n            )\n\n            self.thinking.start(\"Thinking\")\n\n            self._worker = AgentWorker(self.agent, text)\n            self._worker.request_tool_execution.connect(\n                self._execute_tool_on_main_thread\n            )\n            self._worker.phase_changed.connect(self._on_phase_changed)\n            self._worker.tool_executed.connect(self._on_tool_executed)\n            self._worker.finished.connect(self._on_agent_finished)\n            self._worker.error.connect(self._on_agent_error)\n            self._worker.start()\n\n        def _on_phase_changed(self, phase_text):\n            self.thinking.set_phase(phase_text)\n            if \"Retry\" in phase_text:\n                self._append_trace(f\"[phase] {phase_text}\")\n\n        def _on_tool_executed(self, tool_name, result_preview):\n            self._append_trace(f\"[tool:{tool_name}] {result_preview}\")\n\n        def _execute_tool_on_main_thread(self, _call_id, tool_name, args_json):\n            try:\n                arguments = json.loads(args_json)\n            except (json.JSONDecodeError, TypeError):\n                arguments = {}\n            self._append_trace(f\"[call:{tool_name}] {args_json[:180]}\")\n            try:\n                result = execute_tool(\n                    tool_name,\n                    arguments,\n                    client=self.agent.client,\n                    vision_model=self.agent.config.get(\"vision_model\", \"\"),\n                )\n            except Exception as exc:\n                result = f\"Tool error: {exc}\"\n            if self._worker is not None:\n                self._worker.deliver_tool_result(result)\n\n        def _on_agent_finished(self, response):\n            self.thinking.stop()\n            self._append_html(f\"<b>ChatMol:</b> {_escape_html(response)}\")\n            self.input_edit.setEnabled(True)\n            self.send_btn.setEnabled(True)\n            self.input_edit.setFocus()\n\n        def _on_agent_error(self, error_msg):\n            self.thinking.stop()\n            prov = self.agent.config.get(\"provider\", \"?\")\n            model = self.agent.config.get(\"text_model\", \"?\")\n            self._append_trace(f\"[error] {error_msg}\")\n            self._append_html(\n                f'<span style=\"color:red;\"><b>Error</b> '\n                f\"[{_escape_html(prov)}/{_escape_html(model)}]: \"\n                f\"{_escape_html(error_msg)}</span>\"\n            )\n            self.input_edit.setEnabled(True)\n            self.send_btn.setEnabled(True)\n            self.input_edit.setFocus()\n\n\ndef _escape_html(text):\n    return (\n        text.replace(\"&\", \"&amp;\")\n        .replace(\"<\", \"&lt;\")\n        .replace(\">\", \"&gt;\")\n        .replace(\"\\n\", \"<br>\")\n    )\n\n\n# ---------------------------------------------------------------------------\n# 6. Module-level initialisation\n# ---------------------------------------------------------------------------\n\n_agent = ChatMolAgent()\n\ncmd.extend(\"chat\", _agent.chat)\ncmd.extend(\"set_provider\", _agent.set_provider)\ncmd.extend(\"set_api_key\", _agent.set_api_key)\ncmd.extend(\"set_model\", _agent.set_model)\ncmd.extend(\"set_vision_model\", _agent.set_vision_model)\ncmd.extend(\"reset_conversation\", _agent.reset_conversation)\ncmd.extend(\"save_conversation\", _agent.save_conversation)\ncmd.extend(\"load_conversation\", _agent.load_conversation)\ncmd.extend(\"chatmol_config\", _agent.show_config)\n\n\ndef chatmol_settings():\n    \"\"\"Open the ChatMol settings dialog (requires Qt).\"\"\"\n    if not _HAS_QT:\n        print(\n            \"Qt not available. Configure via commands: \"\n            \"set_api_key, set_model, set_vision_model\"\n        )\n        return\n    app = QApplication.instance()\n    if app is None:\n        print(\"No Qt application running.\")\n        return\n    dlg = ChatMolSettingsDialog(_agent)\n    dlg.exec_()\n\n\ndef chatmol_gui():\n    \"\"\"Manually open the ChatMol chat bar (requires Qt).\"\"\"\n    _init_gui()\n\n\ncmd.extend(\"chatmol_settings\", chatmol_settings)\ncmd.extend(\"chatmol_gui\", chatmol_gui)\n\n\n_chatbar = None\n\n\ndef _init_gui():\n    \"\"\"Find PyMOL's main window and dock the ChatMol chat bar.\"\"\"\n    global _chatbar\n    if not _HAS_QT:\n        print(\n            \"Qt not available — ChatMol GUI disabled. \"\n            \"Use command-line: chat <message>\"\n        )\n        return\n    if _chatbar is not None:\n        _chatbar.show()\n        _chatbar.raise_()\n        return\n\n    app = QApplication.instance()\n    if app is None:\n        return\n\n    main_win = None\n    for w in app.topLevelWidgets():\n        if w.inherits(\"QMainWindow\") and w.isVisible():\n            main_win = w\n            break\n\n    if main_win is None:\n        print(\"ChatMol: could not find PyMOL main window for docking.\")\n        return\n\n    _chatbar = ChatMolChatBar(_agent, main_win)\n    main_win.addDockWidget(Qt.BottomDockWidgetArea, _chatbar)\n    print(\"ChatMol chat bar loaded.\")\n\n\nif _HAS_QT:\n    try:\n        QTimer.singleShot(1500, _init_gui)\n    except Exception:\n        pass\n\nprint(\n    \"ChatMol plugin loaded. Commands: chat, set_provider, set_api_key, \"\n    \"set_model, set_vision_model, reset_conversation, \"\n    \"chatmol_config, chatmol_settings, chatmol_gui\"\n)\n"}}}}