1 line
123 KiB
JSON
1 line
123 KiB
JSON
{"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(\"&\", \"&\")\n .replace(\"<\", \"<\")\n .replace(\">\", \">\")\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(\"&\", \"&\")\n .replace(\"<\", \"<\")\n .replace(\">\", \">\")\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"}}}} |