diff --git a/.github/workflows/codex-selfhost-review.yml b/.github/workflows/codex-selfhost-review.yml new file mode 100644 index 0000000..17b2274 --- /dev/null +++ b/.github/workflows/codex-selfhost-review.yml @@ -0,0 +1,90 @@ +name: Codex Self-Hosted Review + +on: + pull_request: + types: [opened, synchronize, reopened, ready_for_review] + +permissions: + contents: read + pull-requests: write + +concurrency: + group: codex-review-pr-${{ github.event.pull_request.number }} + cancel-in-progress: false + +jobs: + codex-review: + if: >- + github.event.pull_request.draft == false && + github.event.pull_request.head.repo.full_name == github.repository + runs-on: [self-hosted, chemflow] + timeout-minutes: 20 + + steps: + - name: Checkout + uses: actions/checkout@v4 + with: + fetch-depth: 0 + + - name: Ensure Codex CLI is available + run: | + source /home/tengcc/soft/miniconda3/etc/profile.d/conda.sh + conda activate chemflow + export PATH="/home/tengcc/.npm-global/bin:${PATH}" + echo "PATH=${PATH}" + echo "HOME=${HOME}" + which codex + codex --version + codex review --help + + - name: Ensure PR base branch is fetched + if: github.event_name == 'pull_request' + run: | + git fetch --no-tags origin "${{ github.base_ref }}:${{ github.base_ref }}" || true + + - name: Run Codex review + shell: bash + run: | + set -euo pipefail + source /home/tengcc/soft/miniconda3/etc/profile.d/conda.sh + conda activate chemflow + export PATH="/home/tengcc/.npm-global/bin:${PATH}" + + codex review --base "${{ github.base_ref }}" > review.stdout 2> review.stderr + + if [ -s review.stdout ]; then + cp review.stdout review.md + else + awk ' + BEGIN { capture = 0; block = "" } + /^codex$/ { capture = 1; block = ""; next } + capture { block = block $0 ORS } + END { printf "%s", block } + ' review.stderr > review.md + fi + + if [ ! -s review.md ]; then + { + echo "Codex review produced no parseable output." + echo + echo "Stderr tail:" + tail -n 120 review.stderr + } > review.md + fi + + - name: Add job summary + run: | + cat review.md >> "$GITHUB_STEP_SUMMARY" + + - name: Comment on PR + uses: actions/github-script@v7 + with: + script: | + const fs = require('fs'); + const body = fs.readFileSync('review.md', 'utf8').slice(0, 65000); + await github.rest.issues.createComment({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: context.issue.number, + body + }); diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 074e448..f69fbe2 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -4,10 +4,9 @@ on: push: tags: - "v*" - workflow_dispatch: jobs: - publish: + build: runs-on: ubuntu-latest permissions: contents: read @@ -24,8 +23,24 @@ jobs: run: python -m build - name: Check distributions run: python -m twine check dist/* + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: python-package-distributions + path: dist/ + + publish-pypi: + needs: build + runs-on: ubuntu-latest + environment: pypi + permissions: + contents: read + id-token: write + steps: + - name: Download distributions + uses: actions/download-artifact@v4 + with: + name: python-package-distributions + path: dist/ - name: Publish to PyPI - env: - TWINE_USERNAME: __token__ - TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }} - run: python -m twine upload dist/* + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index d1995a5..d1b959b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -11,7 +11,7 @@ jobs: strategy: fail-fast: false matrix: - python-version: ["3.10", "3.11", "3.12"] + python-version: ["3.9", "3.10", "3.11", "3.12"] steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 @@ -20,7 +20,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install -e .[dev] + python -m pip install -e .[dev,notebook] - name: Run tests run: pytest - name: Build package diff --git a/.gitignore b/.gitignore index 87f139f..ff9687f 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,8 @@ __pycache__/ *.py[cod] *.so .pytest_cache/ +.ipynb_checkpoints/ +.playwright-cli/ .venv/ .venv*/ dist/ @@ -9,5 +11,7 @@ build/ *.egg-info/ .coverage htmlcov/ +output/ +Untitled*.ipynb .DS_Store AGENTS.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 3a33b72..b98f81b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,12 +4,20 @@ All notable changes to `chemflow-client` will be documented in this file. ## Unreleased -- Added `DEFAULT_BASE_URL` and made `https://chemcloud.info` the default SDK endpoint for public client entry points. -- Added click-to-toggle atom selection to `Chat3DWidget`, plus `get_selected_atom_indices()` and `clear_selection()`. - ## 0.1.0 - Initial public release. - Added one-shot `chat3d(...)` API. - Added stateful `ChemFlow3DClient` with one-step local undo. - Added optional `Chat3DWidget` notebook cell widget. +- Added `DEFAULT_BASE_URL` and made `https://chemcloud.info` the default SDK endpoint for public client entry points. +- Added `CHEMFLOW_API_KEY` / `CHEMFLOW_BASE_URL` environment variable support across the public client entry points. +- Added click-to-toggle atom selection to `Chat3DWidget`, plus `get_selected_atom_indices()` and `clear_selection()`. +- Added an explicit notebook waiting state after `Send`, and changed widget request handling to surface failures in widget UI instead of raising by default. +- Changed widget `Send` handling to run in the background, and added `Chat3DWidget.chat_async(...)` so notebook cells can continue running while the request is in flight. +- Changed the notebook waiting UX from a direct status banner to a chat-inline assistant thinking bubble. +- Fixed notebook async chat completion so background results are scheduled back onto the Jupyter kernel loop instead of updating widget state from the worker thread. +- Allowed `ChemFlow3DClient.start()`, `Chat3DWidget()`, and one-shot `chat3d(atoms=None, ...)` to begin from an empty workspace and generate structure through chat. +- Hardened `Chat3DWidget` message rendering to avoid executing HTML from user prompts or backend responses inside notebooks. +- Lowered the minimum supported Python version from 3.10 to 3.9. +- Updated contributor validation guidance to run the full test suite with notebook extras, and added Python 3.9 to the CI test matrix. diff --git a/README.md b/README.md index 5a4c76b..b57343a 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ Public Python client for ChemFlow 3D chat editing. -Default public service URL: `https://chemcloud.info` +Supported Python versions: 3.9+ ## Install @@ -16,7 +16,15 @@ Notebook widget support is optional: pip install "chemflow-client[notebook]" ``` -## One-shot Usage +For contributors who want to run the full test suite: + +```bash +pip install -e ".[dev,notebook]" +``` + +## Usage + +Blocking Python API: ```python from ase.build import molecule @@ -26,44 +34,52 @@ atoms = molecule("H2O") updated_atoms, text = chat3d( atoms, "change the H-O-H angle to 110 degrees", - api_key="cfsk_xxx", +) + +generated_atoms, text = chat3d( + atoms=None, + prompt="generate methane", ) ``` -Override `base_url` only when targeting a self-hosted ChemFlow deployment. +## JupyterLab -## Stateful Usage +Async notebook widget: ```python -from ase.build import molecule -from chemflow_client import ChemFlow3DClient +from chemflow_client import Chat3DWidget -client = ChemFlow3DClient( - api_key="cfsk_xxx", -) -client.start(molecule("NH3")) +widget = Chat3DWidget() +widget +``` -atoms, text = client.chat("rotate one hydrogen slightly outward") -atoms = client.undo() -client.close() +```python +widget.get_atoms() ``` -## Notebook Widget +![JupyterLab widget demo](docs/assets/chemflow-widget-demo.gif) + + +## Configure + +Create an API key at . + +You can configure the client with environment variables: + +```bash +export CHEMFLOW_API_KEY="cfsk_xxx" +``` + +You can also pass configuration as arguments to `chat3d(...)` or `Chat3DWidget(...)`: ```python -from IPython.display import display -from ase.build import molecule -from chemflow_client import Chat3DWidget +from chemflow_client import Chat3DWidget, chat3d -widget = Chat3DWidget( - molecule("CH4"), +updated_atoms, text = chat3d( + atoms=None, + prompt="generate methane", api_key="cfsk_xxx", ) -display(widget) -latest_atoms = widget.get_atoms() -selected_atoms = widget.get_selected_atom_indices() +widget = Chat3DWidget(api_key="cfsk_xxx") ``` - -The widget is a cell output widget, not a full JupyterLab sidebar extension. -Clicking atoms in the widget toggles selection with a light-yellow highlight similar to the web viewer. diff --git a/RELEASE.md b/RELEASE.md index e4de2e3..28e0070 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -1,36 +1,69 @@ # Release Process +This repository publishes the `chemflow-client` package directly to PyPI. + ## Prerequisites -- Clean git working tree -- PyPI token configured as `TWINE_PASSWORD` -- Username set to `__token__` +- A clean git checkout based on the release commit you want to publish. +- PyPI project `chemflow-client` created on PyPI, or a pending Trusted Publisher configured for the first release. +- A GitHub environment named `pypi`. +- Trusted Publishing configured on PyPI for this GitHub repository. +- Optional but recommended: require manual approval on the `pypi` environment. + +## Trusted Publishing Setup + +Configure the following publisher in PyPI before the first release: + +- Owner: `SingletC` +- Repository: `chemflow-client` +- Workflow filename: `publish.yml` (the workflow file is `.github/workflows/publish.yml`) +- Environment name: `pypi` + +The workflow uses `pypa/gh-action-pypi-publish@release/v1` with GitHub OIDC. No long-lived API token is required after Trusted Publishing is configured. ## Local Validation +Run validation from a clean checkout. Do not publish from a working tree that contains local feature changes that are not part of the release. + +If a local `build/` directory already exists in the repository root, it can shadow the installed `build` package and break `python -m build`. Remove local build artifacts first, or run the commands in a fresh clone. + ```bash python -m venv .venv . .venv/bin/activate -pip install -U pip build twine pytest -pip install -e .[dev] -pytest +python -m pip install --upgrade pip +python -m pip install -e .[dev,notebook] +rm -rf build dist *.egg-info +pytest -q python -m build python -m twine check dist/* ``` -## Version Bump +## Versioning + +For a new release: 1. Update `version` in `pyproject.toml`. -2. Add a new section to `CHANGELOG.md`. -3. Commit and tag: `git tag vX.Y.Z`. +2. Promote the relevant notes from `CHANGELOG.md` into a new version section. +3. Commit the release changes. + +For the first public release, keep the version at `0.1.0` unless release content changes. -## Publish +## PyPI Release + +After local validation succeeds: ```bash -TWINE_USERNAME=__token__ TWINE_PASSWORD= python -m twine upload dist/* +git push origin dev +git tag v0.1.0 +git push origin v0.1.0 ``` -## GitHub Release +Pushing the version tag triggers the `publish` workflow. The `build` and `publish-pypi` jobs will publish to PyPI. + +## Post-release Checks -- Push the version tag. -- The `publish.yml` workflow will build and publish from the tag if repository secrets are configured. +- Confirm the PyPI project page renders correctly. +- Install `chemflow-client==0.1.0` from PyPI in a clean virtual environment. +- Install `"chemflow-client[notebook]==0.1.0"` in a clean virtual environment. +- Verify `from chemflow_client import chat3d, ChemFlow3DClient, Chat3DWidget, DEFAULT_BASE_URL` succeeds. +- Verify the GitHub tag matches the released version. diff --git a/docs/assets/chemflow-widget-demo.gif b/docs/assets/chemflow-widget-demo.gif new file mode 100644 index 0000000..f7d2c9c Binary files /dev/null and b/docs/assets/chemflow-widget-demo.gif differ diff --git a/pyproject.toml b/pyproject.toml index 3c45d1c..689b081 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ name = "chemflow-client" version = "0.1.0" description = "Public Python client for ChemFlow 3D chat editing" readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.9" license = "MIT" authors = [ { name = "ChemFlow" } @@ -17,6 +17,7 @@ classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Science/Research", "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", diff --git a/src/chemflow_client/__init__.py b/src/chemflow_client/__init__.py index 4cdeaa1..2d135d0 100644 --- a/src/chemflow_client/__init__.py +++ b/src/chemflow_client/__init__.py @@ -1,7 +1,7 @@ """Public ChemFlow Python client.""" from .client import ChemFlow3DClient, chat3d -from .constants import DEFAULT_BASE_URL +from .constants import CHEMFLOW_API_KEY_ENV_VAR, CHEMFLOW_BASE_URL_ENV_VAR, DEFAULT_BASE_URL from .exceptions import ( ChemFlowConfigurationError, ChemFlowError, @@ -16,6 +16,8 @@ "ChemFlow3DClient", "Chat3DWidget", "DEFAULT_BASE_URL", + "CHEMFLOW_BASE_URL_ENV_VAR", + "CHEMFLOW_API_KEY_ENV_VAR", "ChemFlowError", "ChemFlowConfigurationError", "ChemFlowHttpError", diff --git a/src/chemflow_client/api.py b/src/chemflow_client/api.py index 9dcf164..8647c7a 100644 --- a/src/chemflow_client/api.py +++ b/src/chemflow_client/api.py @@ -2,12 +2,12 @@ from __future__ import annotations -from typing import Any +from typing import Any, Optional import httpx -from .constants import DEFAULT_BASE_URL -from .exceptions import ChemFlowConfigurationError, ChemFlowHttpError +from .config import resolve_api_key, resolve_base_url +from .exceptions import ChemFlowHttpError from .types import Chat3DRequest, Chat3DResponse @@ -17,22 +17,19 @@ class ChemFlowApi: def __init__( self, *, - base_url: str = DEFAULT_BASE_URL, - api_key: str, + base_url: Optional[str] = None, + api_key: Optional[str] = None, timeout: float = 300.0, - transport: httpx.BaseTransport | None = None, + transport: Optional[httpx.BaseTransport] = None, ) -> None: - normalized_base_url = (base_url or "").strip().rstrip("/") - if not normalized_base_url: - raise ChemFlowConfigurationError("base_url is required") - if not api_key: - raise ChemFlowConfigurationError("api_key is required") + normalized_base_url = resolve_base_url(base_url) + resolved_api_key = resolve_api_key(api_key) self._client = httpx.Client( base_url=normalized_base_url, timeout=timeout, transport=transport, headers={ - "X-ChemFlow-Api-Key": api_key, + "X-ChemFlow-Api-Key": resolved_api_key, "User-Agent": "chemflow-client/0.1.0", "Accept": "application/json", }, diff --git a/src/chemflow_client/ase_adapter.py b/src/chemflow_client/ase_adapter.py index d7ed73f..c9af797 100644 --- a/src/chemflow_client/ase_adapter.py +++ b/src/chemflow_client/ase_adapter.py @@ -26,6 +26,10 @@ def to_payload(atoms: Atoms) -> AtomsPayload: masses = atoms.get_masses().tolist() except Exception: masses = None + if tags == []: + tags = None + if masses == []: + masses = None cell = atoms.get_cell().array.tolist() if atoms.get_cell() is not None else None pbc = atoms.get_pbc().tolist() if atoms.get_pbc() is not None else None return AtomsPayload( @@ -53,6 +57,8 @@ def from_payload(payload: AtomsPayload) -> Atoms: @staticmethod def to_xyz_text(atoms: Atoms) -> str: + if len(atoms) == 0: + return "" lines = [str(len(atoms)), "ChemFlow Client"] for symbol, position in zip(atoms.get_chemical_symbols(), atoms.get_positions()): lines.append( diff --git a/src/chemflow_client/client.py b/src/chemflow_client/client.py index 8023f20..b7f0c6f 100644 --- a/src/chemflow_client/client.py +++ b/src/chemflow_client/client.py @@ -2,15 +2,21 @@ from __future__ import annotations -from typing import Optional +from dataclasses import dataclass +from typing import Optional, Tuple from ase import Atoms from .api import ChemFlowApi from .ase_adapter import AseAtomsAdapter -from .constants import DEFAULT_BASE_URL from .exceptions import ChemFlowResponseError, ChemFlowStateError -from .types import Chat3DRequest +from .types import Chat3DRequest, Chat3DResponse + + +@dataclass(frozen=True) +class _PreparedChat: + request: Chat3DRequest + previous_atoms: Atoms class ChemFlow3DClient: @@ -19,9 +25,9 @@ class ChemFlow3DClient: def __init__( self, *, - base_url: str = DEFAULT_BASE_URL, - api_key: str, - model: str | None = None, + base_url: Optional[str] = None, + api_key: Optional[str] = None, + model: Optional[str] = None, timeout: float = 300.0, ) -> None: self._api = ChemFlowApi(base_url=base_url, api_key=api_key, timeout=timeout) @@ -35,8 +41,9 @@ def __init__( def session_id(self) -> Optional[str]: return self._session_id - def start(self, atoms: Atoms) -> None: - self._atoms = AseAtomsAdapter.copy_atoms(atoms) + def start(self, atoms: Optional[Atoms] = None) -> None: + """Start a session from an existing structure or an empty workspace.""" + self._atoms = AseAtomsAdapter.copy_atoms(atoms) if atoms is not None else Atoms() self._previous_atoms = None self._session_id = None self._needs_sync = False @@ -50,7 +57,7 @@ def set_atoms(self, atoms: Atoms) -> None: self._atoms = AseAtomsAdapter.copy_atoms(atoms) self._needs_sync = True - def chat(self, prompt: str) -> tuple[Atoms, str]: + def _prepare_chat(self, prompt: str) -> _PreparedChat: if self._atoms is None: raise ChemFlowStateError("No structure has been loaded. Call start(atoms) first.") normalized_prompt = (prompt or "").strip() @@ -62,25 +69,37 @@ def chat(self, prompt: str) -> tuple[Atoms, str]: outgoing_atoms = AseAtomsAdapter.to_payload(self._atoms) previous_atoms = AseAtomsAdapter.copy_atoms(self._atoms) - response = self._api.chat3d( - Chat3DRequest( + return _PreparedChat( + request=Chat3DRequest( prompt=normalized_prompt, session_id=self._session_id, atoms=outgoing_atoms, model=self._model, - ) + ), + previous_atoms=previous_atoms, ) + + def _apply_chat_response( + self, + prepared_chat: _PreparedChat, + response: Chat3DResponse, + ) -> Tuple[Atoms, str]: if response.error: raise ChemFlowResponseError(response.error, response=response) if response.atoms is None: raise ChemFlowResponseError("ChemFlow response did not include atoms.", response=response) - self._previous_atoms = previous_atoms + self._previous_atoms = prepared_chat.previous_atoms self._atoms = AseAtomsAdapter.from_payload(response.atoms) self._session_id = response.session_id self._needs_sync = False return self.get_atoms(), response.assistant_message + def chat(self, prompt: str) -> Tuple[Atoms, str]: + prepared_chat = self._prepare_chat(prompt) + response = self._api.chat3d(prepared_chat.request) + return self._apply_chat_response(prepared_chat, response) + def undo(self) -> Atoms: if self._previous_atoms is None: raise ChemFlowStateError("No previous committed structure is available for undo.") @@ -103,15 +122,15 @@ def __exit__(self, exc_type, exc, tb) -> None: def chat3d( - atoms: Atoms, + atoms: Optional[Atoms], prompt: str, *, - base_url: str = DEFAULT_BASE_URL, - api_key: str, - model: str | None = None, + base_url: Optional[str] = None, + api_key: Optional[str] = None, + model: Optional[str] = None, timeout: float = 300.0, -) -> tuple[Atoms, str]: - """Run a one-shot 3D chat edit and return updated atoms plus assistant text.""" +) -> Tuple[Atoms, str]: + """Run a one-shot 3D chat edit from existing atoms or an empty workspace.""" with ChemFlow3DClient( base_url=base_url, api_key=api_key, diff --git a/src/chemflow_client/config.py b/src/chemflow_client/config.py new file mode 100644 index 0000000..8153764 --- /dev/null +++ b/src/chemflow_client/config.py @@ -0,0 +1,30 @@ +"""Configuration helpers for chemflow-client.""" + +from __future__ import annotations + +import os +from typing import Optional + +from .constants import ( + CHEMFLOW_API_KEY_ENV_VAR, + CHEMFLOW_BASE_URL_ENV_VAR, + DEFAULT_BASE_URL, +) +from .exceptions import ChemFlowConfigurationError + + +def resolve_base_url(base_url: Optional[str]) -> str: + normalized = (base_url or os.getenv(CHEMFLOW_BASE_URL_ENV_VAR) or DEFAULT_BASE_URL).strip() + normalized = normalized.rstrip("/") + if not normalized: + raise ChemFlowConfigurationError("base_url is required") + return normalized + + +def resolve_api_key(api_key: Optional[str]) -> str: + normalized = (api_key or os.getenv(CHEMFLOW_API_KEY_ENV_VAR) or "").strip() + if not normalized: + raise ChemFlowConfigurationError( + "api_key is required. Pass api_key=... or set CHEMFLOW_API_KEY." + ) + return normalized diff --git a/src/chemflow_client/constants.py b/src/chemflow_client/constants.py index 5ffa59f..40494a5 100644 --- a/src/chemflow_client/constants.py +++ b/src/chemflow_client/constants.py @@ -1,4 +1,5 @@ """Public constants for chemflow-client.""" DEFAULT_BASE_URL = "https://chemcloud.info" - +CHEMFLOW_BASE_URL_ENV_VAR = "CHEMFLOW_BASE_URL" +CHEMFLOW_API_KEY_ENV_VAR = "CHEMFLOW_API_KEY" diff --git a/src/chemflow_client/widget.py b/src/chemflow_client/widget.py index 4f200fd..c304ebe 100644 --- a/src/chemflow_client/widget.py +++ b/src/chemflow_client/widget.py @@ -2,13 +2,16 @@ from __future__ import annotations +import asyncio import json +import threading +from typing import Any, Callable, Optional, Tuple, Union from ase import Atoms from .ase_adapter import AseAtomsAdapter from .client import ChemFlow3DClient -from .constants import DEFAULT_BASE_URL +from .exceptions import ChemFlowError try: import anywidget @@ -43,8 +46,25 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i return [*normalized, atom_index] +def _format_widget_error_message(error: Union[Exception, str]) -> str: + if isinstance(error, str): + message = error + elif isinstance(error, ChemFlowError): + message = getattr(error, "message", "") or str(error) + else: + message = str(error) + + normalized = (message or "").strip() + if normalized: + return normalized + if isinstance(error, Exception): + return error.__class__.__name__ + return "Request failed" + + _WIDGET_ESM = r''' const THREE_DMOL_URL = "https://cdn.jsdelivr.net/npm/3dmol@2.4.2/build/3Dmol-min.js"; +const THINKING_PLACEHOLDER = "Thinking harder..."; async function ensure3Dmol() { if (window.$3Dmol) { @@ -81,12 +101,27 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i } } -function buildMessageHtml(messages) { - return messages.map((entry) => { - const role = entry.role || "assistant"; - const text = entry.text || ""; - return `
${role}
${text}
`; - }).join(""); +function normalizeMessageRole(role) { + if (role === "user" || role === "system" || role === "assistant") { + return role; + } + return "assistant"; +} + +function buildDisplayMessages(model, optimisticBusy) { + const messages = parseMessages(model.get("messages_json")); + const busy = optimisticBusy || Boolean(model.get("busy")); + if (!busy) { + return messages; + } + return [ + ...messages, + { + role: "assistant", + text: THINKING_PLACEHOLDER, + thinking: true, + }, + ]; } function readAtomIntegerField(atom, field) { @@ -154,28 +189,31 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i el.innerHTML = `
@@ -189,8 +227,8 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i
- - + +
@@ -205,10 +243,13 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i const messagesEl = el.querySelector(".cf-messages"); const formEl = el.querySelector(".cf-form"); const inputEl = el.querySelector(".cf-input"); + const sendButton = el.querySelector(".cf-button-send"); const undoButton = el.querySelector(".cf-button-secondary"); let viewer = null; let currentModel = null; let currentXyzText = ""; + let optimisticBusy = false; + let optimisticPrompt = ""; let disposed = false; async function syncViewer() { @@ -251,13 +292,35 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i } function renderMessages() { - const messages = parseMessages(model.get("messages_json")); - messagesEl.innerHTML = buildMessageHtml(messages); + const messages = buildDisplayMessages(model, optimisticBusy); + const fragment = document.createDocumentFragment(); + messages.forEach((entry) => { + const role = normalizeMessageRole(entry.role); + const text = String(entry.text || ""); + const messageEl = document.createElement("div"); + messageEl.className = `cf-msg cf-msg-${role}`; + if (entry.thinking) { + messageEl.classList.add("cf-msg-thinking"); + } + + const roleEl = document.createElement("div"); + roleEl.className = "cf-role"; + roleEl.textContent = role; + + const textEl = document.createElement("div"); + textEl.className = "cf-text"; + textEl.textContent = text; + + messageEl.append(roleEl, textEl); + fragment.appendChild(messageEl); + }); + messagesEl.replaceChildren(fragment); messagesEl.scrollTop = messagesEl.scrollHeight; } function renderSelection() { const selected = parseSelectedAtomIndices(model.get("selected_atom_indices_json")); + const busy = optimisticBusy || Boolean(model.get("busy")); if (selected.length === 0) { selectionEl.textContent = "Selected atoms: none"; clearSelectionButton.disabled = true; @@ -265,7 +328,7 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i } selectionEl.textContent = `Selected atoms (${selected.length}): ${selected.map((index) => index + 1).join(", ")}`; - clearSelectionButton.disabled = false; + clearSelectionButton.disabled = busy; } function renderStatus() { @@ -275,6 +338,18 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i statusEl.textContent = errorText || statusText; } + function renderBusyState() { + const busy = optimisticBusy || Boolean(model.get("busy")); + + inputEl.disabled = busy; + sendButton.disabled = busy; + sendButton.textContent = "Send"; + undoButton.disabled = busy; + + renderSelection(); + renderMessages(); + } + function handleResize() { if (!viewer) { return; @@ -289,6 +364,9 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i if (!prompt) { return; } + optimisticBusy = true; + optimisticPrompt = prompt; + renderBusyState(); inputEl.value = ""; model.send({ type: "chat", prompt }); }); @@ -303,6 +381,15 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i window.addEventListener("resize", handleResize); + model.on("change:busy", () => { + const busy = Boolean(model.get("busy")); + optimisticBusy = busy; + if (!busy) { + optimisticPrompt = ""; + } + renderBusyState(); + }); + model.on("change:pending_prompt", renderBusyState); model.on("change:xyz_text", syncViewer); model.on("change:selected_atom_indices_json", syncViewer); model.on("change:selected_atom_indices_json", renderSelection); @@ -311,6 +398,7 @@ def _toggle_selected_atom_index(values: list[int], atom_index: int, max_atoms: i model.on("change:error_text", renderStatus); syncViewer(); + renderBusyState(); renderSelection(); renderMessages(); renderStatus(); @@ -343,28 +431,40 @@ class Chat3DWidget(anywidget.AnyWidget): messages_json = traitlets.Unicode("[]").tag(sync=True) status_text = traitlets.Unicode("Ready").tag(sync=True) error_text = traitlets.Unicode("").tag(sync=True) + busy = traitlets.Bool(False).tag(sync=True) + pending_prompt = traitlets.Unicode("").tag(sync=True) selected_atom_indices_json = traitlets.Unicode("[]").tag(sync=True) def __init__( self, - atoms: Atoms, + atoms: Optional[Atoms] = None, *, - base_url: str = DEFAULT_BASE_URL, - api_key: str, - model: str | None = None, + base_url: Optional[str] = None, + api_key: Optional[str] = None, + model: Optional[str] = None, timeout: float = 300.0, ) -> None: super().__init__() + self._client_lock = threading.RLock() + self._main_thread = threading.current_thread() + self._worker_thread: Optional[threading.Thread] = None + self._closed = False + try: + self._main_loop = asyncio.get_running_loop() + except RuntimeError: + self._main_loop = None + self._kernel_io_loop = self._resolve_kernel_io_loop() self._client = ChemFlow3DClient( base_url=base_url, api_key=api_key, model=model, timeout=timeout, ) - self._client.start(atoms) + with self._client_lock: + self._client.start(atoms) self._messages: list[dict[str, str]] = [] self._selected_atom_indices: list[int] = [] - self._sync_xyz(self._client.get_atoms()) + self._sync_xyz(self.get_atoms()) self._sync_selection() self.on_msg(self._handle_frontend_message) @@ -377,6 +477,44 @@ def _sync_messages(self) -> None: def _sync_selection(self) -> None: self.selected_atom_indices_json = json.dumps(self._selected_atom_indices) + def _set_busy_state(self, is_busy: bool, pending_prompt: str = "") -> None: + self.busy = bool(is_busy) + self.pending_prompt = pending_prompt if is_busy else "" + + @staticmethod + def _resolve_kernel_io_loop() -> Optional[Any]: + try: + from IPython import get_ipython + + ip = get_ipython() + except Exception: + return None + kernel = getattr(ip, "kernel", None) + return getattr(kernel, "io_loop", None) + + def _run_on_main_thread(self, callback: Callable[[], None]) -> None: + if self._closed: + return + if threading.current_thread() is self._main_thread: + callback() + return + + if self._main_loop is not None: + try: + self._main_loop.call_soon_threadsafe(callback) + return + except RuntimeError: + pass + + if self._kernel_io_loop is not None: + try: + self._kernel_io_loop.add_callback(callback) + return + except Exception: + pass + + callback() + def _append_message(self, role: str, text: str) -> None: normalized_text = (text or "").strip() if not normalized_text: @@ -384,12 +522,29 @@ def _append_message(self, role: str, text: str) -> None: self._messages.append({"role": role, "text": normalized_text}) self._sync_messages() + def _set_error_state( + self, + error: Union[Exception, str], + *, + append_message: bool = True, + clear_busy: bool = True, + status: str = "Request failed", + ) -> str: + message = _format_widget_error_message(error) + if clear_busy: + self._set_busy_state(False) + self.status_text = status + self.error_text = message + if append_message: + self._append_message("system", f"Request failed: {message}") + return message + def _clear_selection_state(self) -> None: self._selected_atom_indices = [] self._sync_selection() def _toggle_selection_state(self, atom_index: int) -> list[int]: - atom_count = len(self._client.get_atoms()) + atom_count = len(self.get_atoms()) self._selected_atom_indices = _toggle_selected_atom_index( self._selected_atom_indices, atom_index, @@ -398,50 +553,165 @@ def _toggle_selection_state(self, atom_index: int) -> list[int]: self._sync_selection() return self.get_selected_atom_indices() + def _finalize_background_chat_success(self, atoms: Atoms, text: str) -> None: + self._worker_thread = None + if self._closed: + return + self._append_message("assistant", text) + self._sync_xyz(atoms) + self._clear_selection_state() + self._set_busy_state(False) + self.error_text = "" + self.status_text = "Ready" + + def _finalize_background_chat_error(self, exc: Exception) -> None: + self._worker_thread = None + if self._closed: + return + self._set_error_state(exc) + + def _background_chat_worker(self, prompt: str) -> None: + try: + with self._client_lock: + prepared_chat = self._client._prepare_chat(prompt) + response = self._client._api.chat3d(prepared_chat.request) + with self._client_lock: + if self._closed: + self._worker_thread = None + return + atoms, text = self._client._apply_chat_response(prepared_chat, response) + except Exception as exc: + self._run_on_main_thread(lambda exc=exc: self._finalize_background_chat_error(exc)) + return + + self._run_on_main_thread( + lambda atoms=atoms, text=text: self._finalize_background_chat_success(atoms, text) + ) + + def chat_async(self, prompt: str, *, raise_errors: bool = False) -> bool: + normalized_prompt = (prompt or "").strip() + if not normalized_prompt: + message = self._set_error_state("prompt is required", append_message=False, clear_busy=False) + if raise_errors: + raise ValueError(message) + return False + + if self.busy: + message = self._set_error_state( + "Another request is already in progress.", + append_message=False, + clear_busy=False, + status="Busy", + ) + if raise_errors: + raise RuntimeError(message) + return False + + self.error_text = "" + self.status_text = "Working..." + self._set_busy_state(True, normalized_prompt) + self._append_message("user", normalized_prompt) + + try: + worker = threading.Thread( + target=self._background_chat_worker, + args=(normalized_prompt,), + name="chemflow-widget-chat", + daemon=True, + ) + self._worker_thread = worker + worker.start() + except Exception as exc: + self._worker_thread = None + self._set_error_state(exc) + if raise_errors: + raise + return False + + return True + def _handle_frontend_message(self, _, content, buffers) -> None: del buffers message_type = (content or {}).get("type") try: if message_type == "chat": - self.chat(str((content or {}).get("prompt") or "")) + self.chat_async(str((content or {}).get("prompt") or "")) return if message_type == "undo": self.undo() return if message_type == "toggle_selection": + if self.busy: + return atom_index = int((content or {}).get("atom_index")) self._toggle_selection_state(atom_index) return if message_type == "clear_selection": + if self.busy: + return self.clear_selection() return except Exception as exc: - self.status_text = "Error" - self.error_text = str(exc) + self._set_error_state(exc) - def chat(self, prompt: str) -> tuple[Atoms, str]: + def chat(self, prompt: str, *, raise_errors: bool = False) -> Tuple[Atoms, str]: normalized_prompt = (prompt or "").strip() if not normalized_prompt: - raise ValueError("prompt is required") + message = self._set_error_state("prompt is required", append_message=False, clear_busy=False) + if raise_errors: + raise ValueError(message) + return self.get_atoms(), "" + if self.busy: + message = self._set_error_state( + "Another request is already in progress.", + append_message=False, + clear_busy=False, + status="Busy", + ) + if raise_errors: + raise RuntimeError(message) + return self.get_atoms(), "" self.error_text = "" - self.status_text = "Running..." + self.status_text = "Working..." + self._set_busy_state(True, normalized_prompt) self._append_message("user", normalized_prompt) - atoms, text = self._client.chat(normalized_prompt) + + try: + with self._client_lock: + atoms, text = self._client.chat(normalized_prompt) + except Exception as exc: + self._set_error_state(exc) + if raise_errors: + raise + return self.get_atoms(), "" + self._append_message("assistant", text) self._sync_xyz(atoms) self._clear_selection_state() + self._set_busy_state(False) self.status_text = "Ready" return self.get_atoms(), text def set_atoms(self, atoms: Atoms) -> None: - self._client.set_atoms(atoms) - self._sync_xyz(self._client.get_atoms()) + if self.busy: + self._set_error_state( + "Cannot replace atoms while a request is in progress.", + append_message=False, + clear_busy=False, + status="Busy", + ) + return + with self._client_lock: + self._client.set_atoms(atoms) + current_atoms = self._client.get_atoms() + self._sync_xyz(current_atoms) self._clear_selection_state() self.error_text = "" self.status_text = "Ready" def get_atoms(self) -> Atoms: - return self._client.get_atoms() + with self._client_lock: + return self._client.get_atoms() def get_selected_atom_indices(self) -> list[int]: return list(self._selected_atom_indices) @@ -452,8 +722,27 @@ def clear_selection(self) -> list[int]: self.status_text = "Ready" return self.get_selected_atom_indices() - def undo(self) -> Atoms: - atoms = self._client.undo() + def undo(self, *, raise_errors: bool = False) -> Atoms: + if self.busy: + message = self._set_error_state( + "Another request is already in progress.", + append_message=False, + clear_busy=False, + status="Busy", + ) + if raise_errors: + raise RuntimeError(message) + return self.get_atoms() + + try: + with self._client_lock: + atoms = self._client.undo() + except Exception as exc: + self._set_error_state(exc) + if raise_errors: + raise + return self.get_atoms() + self._append_message("system", "Reverted to the previous committed structure.") self._sync_xyz(atoms) self._clear_selection_state() @@ -462,4 +751,6 @@ def undo(self) -> Atoms: return self.get_atoms() def close(self) -> None: - self._client.close() + self._closed = True + with self._client_lock: + self._client.close() diff --git a/tests/test_ase_adapter.py b/tests/test_ase_adapter.py index baac0a0..776658d 100644 --- a/tests/test_ase_adapter.py +++ b/tests/test_ase_adapter.py @@ -21,3 +21,20 @@ def test_ase_adapter_round_trip_preserves_core_fields(): assert restored.get_pbc().tolist() == [True, False, False] assert restored.get_tags().tolist() == [1, 2] assert restored.get_masses().tolist() == [12.0, 16.0] + + +def test_ase_adapter_preserves_cell_and_pbc_for_empty_workspace(): + atoms = Atoms( + cell=[[8.0, 0.0, 0.0], [0.0, 9.0, 0.0], [0.0, 0.0, 10.0]], + pbc=[True, True, False], + ) + + payload = AseAtomsAdapter.to_payload(atoms) + restored = AseAtomsAdapter.from_payload(payload) + + assert payload.symbols == [] + assert payload.positions == [] + assert payload.cell == [[8.0, 0.0, 0.0], [0.0, 9.0, 0.0], [0.0, 0.0, 10.0]] + assert payload.pbc == [True, True, False] + assert restored.get_cell().array.tolist() == [[8.0, 0.0, 0.0], [0.0, 9.0, 0.0], [0.0, 0.0, 10.0]] + assert restored.get_pbc().tolist() == [True, True, False] diff --git a/tests/test_client.py b/tests/test_client.py index 9e45e8b..858866e 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,6 +1,11 @@ from ase import Atoms -from chemflow_client import DEFAULT_BASE_URL +from chemflow_client import ( + CHEMFLOW_API_KEY_ENV_VAR, + CHEMFLOW_BASE_URL_ENV_VAR, + DEFAULT_BASE_URL, +) +from chemflow_client.api import ChemFlowApi from chemflow_client.client import ChemFlow3DClient, chat3d from chemflow_client.exceptions import ChemFlowStateError from chemflow_client.types import AtomsPayload, Chat3DResponse @@ -40,6 +45,38 @@ def fake_chat3d(self, request): assert text == "Applied edit." +def test_chat3d_can_start_from_empty_workspace(monkeypatch): + def fake_chat3d(self, request): + payload = request.atoms + assert payload is not None + assert payload.symbols == [] + assert payload.positions == [] + return Chat3DResponse( + session_id="session-empty", + assistant_message="Generated methane.", + atoms=AtomsPayload( + symbols=["C"], + positions=[[0.0, 0.0, 0.0]], + cell=payload.cell, + pbc=payload.pbc, + tags=payload.tags, + masses=payload.masses, + ), + changed=True, + ) + + monkeypatch.setattr("chemflow_client.api.ChemFlowApi.chat3d", fake_chat3d) + + updated_atoms, text = chat3d( + None, + "generate methane", + api_key="cfsk_test_key", + ) + + assert updated_atoms.get_chemical_symbols() == ["C"] + assert text == "Generated methane." + + def test_stateful_client_set_atoms_and_undo(monkeypatch): calls = [] @@ -86,6 +123,40 @@ def fake_chat3d(self, request): client.close() +def test_stateful_client_can_start_empty_and_generate(monkeypatch): + def fake_chat3d(self, request): + payload = request.atoms + assert payload is not None + assert payload.symbols == [] + assert payload.positions == [] + return Chat3DResponse( + session_id="session-empty-client", + assistant_message="Generated helium.", + atoms=AtomsPayload( + symbols=["He"], + positions=[[0.0, 0.0, 0.0]], + cell=payload.cell, + pbc=payload.pbc, + tags=payload.tags, + masses=payload.masses, + ), + changed=True, + ) + + monkeypatch.setattr("chemflow_client.api.ChemFlowApi.chat3d", fake_chat3d) + + client = ChemFlow3DClient(api_key="cfsk_test_key") + client.start() + + try: + assert len(client.get_atoms()) == 0 + updated_atoms, text = client.chat("generate helium") + assert updated_atoms.get_chemical_symbols() == ["He"] + assert text == "Generated helium." + finally: + client.close() + + def test_undo_requires_previous_committed_structure(): client = ChemFlow3DClient(base_url="http://localhost:8000", api_key="cfsk_test_key") client.start(Atoms(symbols=["Ne"], positions=[[0.0, 0.0, 0.0]])) @@ -105,3 +176,50 @@ def test_client_uses_public_default_base_url(): assert str(client._api._client.base_url).rstrip("/") == DEFAULT_BASE_URL finally: client.close() + + +def test_api_and_client_can_resolve_configuration_from_environment(monkeypatch): + monkeypatch.setenv(CHEMFLOW_API_KEY_ENV_VAR, "cfsk_env_key") + monkeypatch.setenv(CHEMFLOW_BASE_URL_ENV_VAR, "http://env.example:8000/") + + api = ChemFlowApi() + client = ChemFlow3DClient() + + try: + assert str(api._client.base_url).rstrip("/") == "http://env.example:8000" + assert api._client.headers["X-ChemFlow-Api-Key"] == "cfsk_env_key" + assert str(client._api._client.base_url).rstrip("/") == "http://env.example:8000" + assert client._api._client.headers["X-ChemFlow-Api-Key"] == "cfsk_env_key" + finally: + api.close() + client.close() + + +def test_chat3d_can_use_environment_api_key(monkeypatch): + def fake_chat3d(self, request): + payload = request.atoms + assert payload is not None + return Chat3DResponse( + session_id="session-env", + assistant_message="Applied edit from env.", + atoms=AtomsPayload( + symbols=list(payload.symbols), + positions=[list(row) for row in payload.positions], + cell=payload.cell, + pbc=payload.pbc, + tags=payload.tags, + masses=payload.masses, + ), + changed=False, + ) + + monkeypatch.setenv(CHEMFLOW_API_KEY_ENV_VAR, "cfsk_env_key") + monkeypatch.setattr("chemflow_client.api.ChemFlowApi.chat3d", fake_chat3d) + + atoms, text = chat3d( + Atoms(symbols=["He"], positions=[[0.0, 0.0, 0.0]]), + "leave it as is", + ) + + assert atoms.get_chemical_symbols() == ["He"] + assert text == "Applied edit from env." diff --git a/tests/test_widget_helpers.py b/tests/test_widget_helpers.py index b9ddad2..f8670bd 100644 --- a/tests/test_widget_helpers.py +++ b/tests/test_widget_helpers.py @@ -1,4 +1,21 @@ +import importlib +import json +import sys +import threading +import time +import types + +import pytest +from ase import Atoms + +traitlets = pytest.importorskip("traitlets", reason="widget tests require notebook extras") + +from chemflow_client import CHEMFLOW_API_KEY_ENV_VAR, CHEMFLOW_BASE_URL_ENV_VAR +from chemflow_client.exceptions import ChemFlowHttpError +from chemflow_client.types import AtomsPayload, Chat3DResponse from chemflow_client.widget import ( + _WIDGET_ESM, + _format_widget_error_message, _normalize_selected_atom_indices, _toggle_selected_atom_index, ) @@ -16,3 +33,255 @@ def test_toggle_selected_atom_index_matches_click_to_toggle_behavior(): selected = _toggle_selected_atom_index(selected, atom_index=4, max_atoms=5) assert selected == [3, 4] + + +def test_format_widget_error_message_prefers_structured_message(): + error = ChemFlowHttpError(502, "upstream timeout") + + assert _format_widget_error_message(error) == "upstream timeout" + + +def test_widget_esm_uses_chat_thinking_placeholder_instead_of_waiting_banner(): + assert "Thinking harder..." in _WIDGET_ESM + assert "Waiting for ChemFlow response:" not in _WIDGET_ESM + + +def test_widget_esm_keeps_fixed_shell_height_and_scrollable_message_area(): + assert 'height: min(720px, 78vh);' in _WIDGET_ESM + assert 'overflow-y: auto;' in _WIDGET_ESM + assert 'min-height: 0;' in _WIDGET_ESM + + +def test_widget_esm_renders_messages_via_text_nodes_instead_of_inner_html(): + assert "textEl.textContent = text;" in _WIDGET_ESM + assert "roleEl.textContent = role;" in _WIDGET_ESM + assert "messagesEl.replaceChildren(fragment);" in _WIDGET_ESM + assert "messagesEl.innerHTML = buildMessageHtml(messages);" not in _WIDGET_ESM + + +def _reload_widget_module_with_fake_anywidget(monkeypatch): + fake_module = types.ModuleType("anywidget") + + class FakeAnyWidget(traitlets.HasTraits): + def on_msg(self, handler): + self._message_handler = handler + + fake_module.AnyWidget = FakeAnyWidget + monkeypatch.setitem(sys.modules, "anywidget", fake_module) + + import chemflow_client.widget as widget_module + + return importlib.reload(widget_module) + + +def test_widget_chat_handles_backend_errors_without_raising(monkeypatch): + widget_module = _reload_widget_module_with_fake_anywidget(monkeypatch) + widget = widget_module.Chat3DWidget( + Atoms(symbols=["He"], positions=[[0.0, 0.0, 0.0]]), + api_key="cfsk_test_key", + ) + + def fake_chat(_prompt): + raise ChemFlowHttpError(503, "service unavailable") + + monkeypatch.setattr(widget._client, "chat", fake_chat) + + atoms, text = widget.chat("move it") + + assert text == "" + assert atoms.get_chemical_symbols() == ["He"] + assert widget.status_text == "Request failed" + assert widget.error_text == "service unavailable" + assert widget.busy is False + assert widget.pending_prompt == "" + + messages = json.loads(widget.messages_json) + assert messages[0] == {"role": "user", "text": "move it"} + assert messages[1] == {"role": "system", "text": "Request failed: service unavailable"} + + +def test_widget_chat_can_still_raise_when_requested(monkeypatch): + widget_module = _reload_widget_module_with_fake_anywidget(monkeypatch) + widget = widget_module.Chat3DWidget( + Atoms(symbols=["He"], positions=[[0.0, 0.0, 0.0]]), + api_key="cfsk_test_key", + ) + + def fake_chat(_prompt): + raise ChemFlowHttpError(504, "gateway timeout") + + monkeypatch.setattr(widget._client, "chat", fake_chat) + + with pytest.raises(ChemFlowHttpError): + widget.chat("move it", raise_errors=True) + + +def test_widget_chat_async_returns_immediately_and_applies_result(monkeypatch): + widget_module = _reload_widget_module_with_fake_anywidget(monkeypatch) + widget = widget_module.Chat3DWidget( + Atoms(symbols=["He"], positions=[[0.0, 0.0, 0.0]]), + api_key="cfsk_test_key", + ) + + started = threading.Event() + release = threading.Event() + + def fake_chat3d(_request): + started.set() + release.wait(1.0) + return Chat3DResponse( + session_id="sess-1", + assistant_message="Applied asynchronously.", + atoms=AtomsPayload( + symbols=["He"], + positions=[[1.5, 0.0, 0.0]], + ), + changed=True, + ) + + monkeypatch.setattr(widget._client._api, "chat3d", fake_chat3d) + + scheduled = widget.chat_async("move it async") + + assert scheduled is True + assert widget.busy is True + assert widget.pending_prompt == "move it async" + assert started.wait(0.2) + + release.set() + deadline = time.time() + 1.0 + while widget.busy and time.time() < deadline: + time.sleep(0.01) + + assert widget.busy is False + assert widget.pending_prompt == "" + assert widget.status_text == "Ready" + assert "Applied asynchronously." in widget.messages_json + + +def test_widget_chat_async_keeps_state_reads_and_close_responsive(monkeypatch): + widget_module = _reload_widget_module_with_fake_anywidget(monkeypatch) + widget = widget_module.Chat3DWidget( + Atoms(symbols=["He"], positions=[[0.0, 0.0, 0.0]]), + api_key="cfsk_test_key", + ) + + started = threading.Event() + release = threading.Event() + + def fake_chat3d(_request): + started.set() + release.wait(1.0) + return Chat3DResponse( + session_id="sess-2", + assistant_message="Applied asynchronously.", + atoms=AtomsPayload( + symbols=["He"], + positions=[[2.0, 0.0, 0.0]], + ), + changed=True, + ) + + monkeypatch.setattr(widget._client._api, "chat3d", fake_chat3d) + + assert widget.chat_async("move it async") is True + assert started.wait(0.2) + + read_started = time.perf_counter() + atoms = widget.get_atoms() + read_elapsed = time.perf_counter() - read_started + + close_started = time.perf_counter() + widget.close() + close_elapsed = time.perf_counter() - close_started + + assert atoms.get_chemical_symbols() == ["He"] + assert read_elapsed < 0.1 + assert close_elapsed < 0.1 + + release.set() + deadline = time.time() + 1.0 + while widget._worker_thread is not None and time.time() < deadline: + time.sleep(0.01) + + assert widget._worker_thread is None + + +def test_widget_can_initialize_empty_workspace(monkeypatch): + widget_module = _reload_widget_module_with_fake_anywidget(monkeypatch) + widget = widget_module.Chat3DWidget(api_key="cfsk_test_key") + + try: + assert len(widget.get_atoms()) == 0 + assert widget.xyz_text == "" + assert widget.status_text == "Ready" + finally: + widget.close() + + +def test_widget_can_generate_from_empty_workspace(monkeypatch): + widget_module = _reload_widget_module_with_fake_anywidget(monkeypatch) + widget = widget_module.Chat3DWidget(api_key="cfsk_test_key") + + def fake_chat(_prompt): + updated = Atoms(symbols=["He"], positions=[[0.0, 0.0, 0.0]]) + widget._client._atoms = updated.copy() + return updated, "Generated helium." + + monkeypatch.setattr(widget._client, "chat", fake_chat) + + atoms, text = widget.chat("generate helium") + + try: + assert atoms.get_chemical_symbols() == ["He"] + assert text == "Generated helium." + assert widget.xyz_text.startswith("1\nChemFlow Client\nHe ") + finally: + widget.close() + + +def test_widget_can_resolve_configuration_from_environment(monkeypatch): + widget_module = _reload_widget_module_with_fake_anywidget(monkeypatch) + monkeypatch.setenv(CHEMFLOW_API_KEY_ENV_VAR, "cfsk_env_key") + monkeypatch.setenv(CHEMFLOW_BASE_URL_ENV_VAR, "http://env.example:8000/") + + widget = widget_module.Chat3DWidget(Atoms(symbols=["He"], positions=[[0.0, 0.0, 0.0]])) + + try: + assert str(widget._client._api._client.base_url).rstrip("/") == "http://env.example:8000" + assert widget._client._api._client.headers["X-ChemFlow-Api-Key"] == "cfsk_env_key" + finally: + widget.close() + + +def test_widget_run_on_main_thread_uses_stored_kernel_io_loop(monkeypatch): + widget_module = _reload_widget_module_with_fake_anywidget(monkeypatch) + widget = widget_module.Chat3DWidget( + Atoms(symbols=["He"], positions=[[0.0, 0.0, 0.0]]), + api_key="cfsk_test_key", + ) + + callbacks = [] + invoked = [] + + class FakeIOLoop: + def add_callback(self, callback): + callbacks.append(callback) + + widget._main_loop = None + widget._kernel_io_loop = FakeIOLoop() + + worker = threading.Thread( + target=lambda: widget._run_on_main_thread(lambda: invoked.append("done")), + daemon=True, + ) + worker.start() + worker.join(timeout=1.0) + + assert invoked == [] + assert len(callbacks) == 1 + + callbacks[0]() + assert invoked == ["done"] + + widget.close()