dh-cli 0.14.0__tar.gz → 0.15.0__tar.gz

{dh_cli-0.14.0 → dh_cli-0.15.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: dh-cli
-Version: 0.14.0
+Version: 0.15.0
 Summary: Dayhoff Labs developer CLI
 Author-email: Dayhoff Labs <dev@dayhofflabs.com>
 License: # PolyForm Noncommercial License 1.0.0

{dh_cli-0.14.0 → dh_cli-0.15.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dh-cli"
-version = "0.14.0"
+version = "0.15.0"
 description = "Dayhoff Labs developer CLI"
 requires-python = ">=3.11"
 readme = "README.md"

{dh_cli-0.14.0 → dh_cli-0.15.0}/src/dh_cli/hz/mcp_server.py

@@ -6,6 +6,7 @@ land in that user's Horizyn history. Requires the ``mcp`` extra:
 """
 
 from importlib.metadata import version
+from pathlib import Path
 
 import httpx
 from mcp.server.mcpserver import MCPServer
@@ -17,6 +18,8 @@ from dh_cli.hz.auth import NotLoggedIn, get_id_token
 
 CLIENT_HEADER = f"mcp/{version('dh-cli')}"
 _READ_ONLY = ToolAnnotations(read_only_hint=True)
+_SUBMIT = ToolAnnotations(read_only_hint=False, destructive_hint=False, open_world_hint=True)
+_WRITES_FILES = ToolAnnotations(read_only_hint=False, destructive_hint=False, idempotent_hint=True)
 
 server = MCPServer(
     "horizyn",
@@ -198,3 +201,121 @@ def horizyn_protein_details(protein_id: str, screening_set: str) -> dict:
     score, full sequence, length, AlphaFold id and any extra metadata.
     """
     return _request("GET", f"/protein/{protein_id}/details", params={"screening_set": screening_set})
+
+
+@server.tool(annotations=_SUBMIT)
+def horizyn_dock_submit(
+    protein_sequence: str = "",
+    reaction_smiles: str | None = None,
+    dock_products: bool = False,
+    protein_id: str | None = None,
+    label: str | None = None,
+    predict_waters: bool = False,
+    waters_selection: str | None = None,
+    boltz_yaml: str | None = None,
+) -> dict:
+    """Submit a Boltz fold+dock job. Costs GPU time; ask the user before calling.
+
+    Folds `protein_sequence` and docks the ligands parsed from `reaction_smiles`
+    ("substrates>>products"; the substrate side by default, the product side
+    when `dock_products` is true). A job typically takes 30-120 s of GPU plus
+    about a minute for the MSA lookup, times out after 10 min, and the queue
+    holds 20 jobs. `protein_id` (UniProt) is an MSA hint; `label` names the job
+    in the user's history. `predict_waters` chains an ordered-water prediction
+    on the folded structure (`waters_selection` "confidence" or "density");
+    its job id appears as `child_job_id` in horizyn_dock_status. `boltz_yaml`
+    is advanced mode: a raw Boltz config replaces `reaction_smiles`.
+
+    Returns `job_id` and `status` ("queued"). Poll horizyn_dock_status until
+    `status` is "complete", then call horizyn_dock_result.
+    """
+    body = {
+        "protein_sequence": protein_sequence,
+        "reaction_smiles": reaction_smiles,
+        "dock_products": dock_products,
+        "protein_id": protein_id,
+        "label": label,
+        "predict_waters": predict_waters,
+        "waters_selection": waters_selection,
+        "boltz_yaml": boltz_yaml,
+    }
+    return _request("POST", "/dock/submit", json={k: v for k, v in body.items() if v is not None})
+
+
+@server.tool(annotations=_SUBMIT)
+def horizyn_dock_waters_submit(
+    structure_path: str,
+    selection: str | None = None,
+    label: str | None = None,
+) -> dict:
+    """Predict ordered water positions on an existing structure. Costs GPU time.
+
+    `structure_path` is a local `.pdb` or `.cif` file (format taken from the
+    suffix, 8 MB max); the server reads it, so never paste structure text into
+    this call. `selection` is "confidence" or "density". Returns `job_id` and
+    `status`; poll horizyn_dock_status, then horizyn_dock_result, whose
+    `waters` list and `n_waters` carry the prediction.
+    """
+    path = Path(structure_path).expanduser()
+    try:
+        structure_data = path.read_text()
+    except OSError as exc:
+        raise ToolError(f"Cannot read structure file {path}: {exc.strerror or type(exc).__name__}") from None
+
+    body = {"structure_data": structure_data, "structure_format": path.suffix, "selection": selection, "label": label}
+    return _request("POST", "/dock/waters/submit", json={k: v for k, v in body.items() if v is not None})
+
+
+@server.tool(annotations=_READ_ONLY)
+def horizyn_dock_status(job_id: str) -> dict:
+    """Poll a docking job. Wait at least 15 s between polls.
+
+    `status` is "queued", "running", "complete" or "failed"; `error` explains a
+    failure. `elapsed_seconds` and `queue_position` track progress. `msa_used`
+    false means Boltz fell back to single-sequence mode, so confidence is lower.
+    `job_type` is "fold" or "waters"; `child_job_id` names a chained waters job
+    and `parent_job_id` its fold.
+    """
+    return _request("GET", f"/dock/{job_id}/status")
+
+
+@server.tool(annotations=_WRITES_FILES)
+def horizyn_dock_result(job_id: str, output_dir: str | None = None) -> dict:
+    """Fetch a completed docking job's structure and confidence.
+
+    Returns 409 until horizyn_dock_status reports "complete", so poll first.
+    With `output_dir`, writes `{job_id}.pdb` (and `{job_id}.cif` when present)
+    there, creating the directory, and returns `pdb_path`/`cif_path` instead of
+    the raw text. Without it, `pdb_data`/`cif_data` are returned inline (large).
+    Always includes `confidence_data`, `msa_used`, `job_type`, `waters`,
+    `n_waters`, `parent_job_id` and `child_job_id`.
+    """
+    data = _request("GET", f"/dock/{job_id}/result")
+    if output_dir is None:
+        return data
+
+    out = Path(output_dir).expanduser()
+    out.mkdir(parents=True, exist_ok=True)
+    pdb_path = out / f"{job_id}.pdb"
+    pdb_path.write_text(data.pop("pdb_data", ""))
+    data["pdb_path"] = str(pdb_path)
+
+    cif_data = data.pop("cif_data", "")
+    cif_path = out / f"{job_id}.cif"
+    if cif_data:
+        cif_path.write_text(cif_data)
+    data["cif_path"] = str(cif_path) if cif_data else None
+    return data
+
+
+@server.tool(annotations=_READ_ONLY)
+def horizyn_list_docking_jobs(limit: int = 20) -> dict:
+    """List the logged-in user's recent docking jobs.
+
+    `limit` is 1-100. Items in `jobs` are ordered by `job_id`, not by time;
+    sort on `submitted_at` yourself. Each carries `job_id`, `status`,
+    `job_type`, `submitted_at`, `label`, `protein_sequence`, `reaction_smiles`
+    and `client_source`. Status here is only updated when a job is polled or
+    its result fetched; use horizyn_dock_status for the live value.
+    """
+    return _request("GET", "/me/docking-jobs", params={"limit": limit})

{dh_cli-0.14.0 → dh_cli-0.15.0}/tests/hz/live_smoke.py

@@ -21,6 +21,11 @@ EXPECTED = {
     "horizyn_query_reaction",
     "horizyn_query_protein",
     "horizyn_protein_details",
+    "horizyn_dock_submit",
+    "horizyn_dock_waters_submit",
+    "horizyn_dock_status",
+    "horizyn_dock_result",
+    "horizyn_list_docking_jobs",
 }
 
 

{dh_cli-0.14.0 → dh_cli-0.15.0}/tests/hz/test_mcp_server.py

@@ -13,13 +13,17 @@ from dh_cli.hz import mcp_server
 from dh_cli.hz.auth import NotLoggedIn
 
 TOKEN = "eyJ-test-token"
-EXPECTED_TOOLS = {
+READ_ONLY_TOOLS = {
     "horizyn_whoami",
     "horizyn_list_screening_sets",
     "horizyn_query_reaction",
     "horizyn_query_protein",
     "horizyn_protein_details",
+    "horizyn_dock_status",
+    "horizyn_list_docking_jobs",
 }
+WRITE_TOOLS = {"horizyn_dock_submit", "horizyn_dock_waters_submit", "horizyn_dock_result"}
+EXPECTED_TOOLS = READ_ONLY_TOOLS | WRITE_TOOLS
 
 
 class Recorder:
@@ -54,13 +58,35 @@ def api(monkeypatch):
 
 
 class TestRegistration:
-    def test_lists_five_read_only_tools(self):
+    def test_lists_ten_tools_with_descriptions(self):
         tools = anyio.run(mcp_server.server.list_tools)
 
         assert {t.name for t in tools} == EXPECTED_TOOLS
-        assert all(t.annotations and t.annotations.read_only_hint for t in tools)
         assert all(t.description for t in tools)
 
+    def test_read_only_annotations(self):
+        by_name = {t.name: t for t in anyio.run(mcp_server.server.list_tools)}
+
+        for name in READ_ONLY_TOOLS:
+            assert by_name[name].annotations.read_only_hint is True, name
+        for name in WRITE_TOOLS:
+            assert by_name[name].annotations.read_only_hint is False, name
+
+    def test_submit_tools_are_non_destructive_open_world(self):
+        by_name = {t.name: t for t in anyio.run(mcp_server.server.list_tools)}
+
+        for name in ("horizyn_dock_submit", "horizyn_dock_waters_submit"):
+            ann = by_name[name].annotations
+            assert ann.destructive_hint is False, name
+            assert ann.open_world_hint is True, name
+
+    def test_result_tool_is_idempotent_non_destructive(self):
+        by_name = {t.name: t for t in anyio.run(mcp_server.server.list_tools)}
+
+        ann = by_name["horizyn_dock_result"].annotations
+        assert ann.destructive_hint is False
+        assert ann.idempotent_hint is True
+
 
 class TestHeaders:
     def test_every_request_carries_auth_and_client_headers(self, api):
@@ -146,6 +172,193 @@ class TestPaths:
             mcp_server.horizyn_protein_details(protein_id="P81747")
 
 
+SEQ = "MKVLAAGIVGLLLAGQPALA"
+RXN = "CC=O.O>>CC(O)O"
+PDB = "ATOM      1  N   MET A   1      11.104   6.134  -6.504  1.00  0.00           N\nEND\n"
+CIF = "data_x\n_atom_site.id 1\n"
+
+
+class TestDockingPaths:
+    def test_dock_submit_minimal_omits_none_fields(self, api):
+        api.body = {"job_id": "j1", "status": "queued"}
+
+        result = mcp_server.horizyn_dock_submit(protein_sequence=SEQ, reaction_smiles=RXN)
+
+        assert api.request.method == "POST"
+        assert api.request.url.path == "/dock/submit"
+        assert api.sent_json == {
+            "protein_sequence": SEQ,
+            "reaction_smiles": RXN,
+            "dock_products": False,
+            "predict_waters": False,
+        }
+        assert result == {"job_id": "j1", "status": "queued"}
+
+    def test_dock_submit_optional_fields_sent_when_given(self, api):
+        mcp_server.horizyn_dock_submit(
+            protein_sequence=SEQ,
+            reaction_smiles=RXN,
+            dock_products=True,
+            protein_id="P81747",
+            label="hzmcp2-live",
+            predict_waters=True,
+            waters_selection="density",
+        )
+
+        body = api.sent_json
+        assert body["dock_products"] is True
+        assert body["protein_id"] == "P81747"
+        assert body["label"] == "hzmcp2-live"
+        assert body["predict_waters"] is True
+        assert body["waters_selection"] == "density"
+        assert "boltz_yaml" not in body
+
+    def test_dock_submit_boltz_yaml_without_sequence(self, api):
+        mcp_server.horizyn_dock_submit(boltz_yaml="version: 1\n")
+
+        body = api.sent_json
+        assert body["boltz_yaml"] == "version: 1\n"
+        assert body["protein_sequence"] == ""
+        assert "reaction_smiles" not in body
+
+    def test_waters_submit_reads_pdb_file(self, api, tmp_path):
+        path = tmp_path / "fold.pdb"
+        path.write_text(PDB)
+
+        mcp_server.horizyn_dock_waters_submit(structure_path=str(path))
+
+        assert api.request.method == "POST"
+        assert api.request.url.path == "/dock/waters/submit"
+        assert api.sent_json == {"structure_data": PDB, "structure_format": ".pdb"}
+
+    def test_waters_submit_cif_suffix_and_options(self, api, tmp_path):
+        path = tmp_path / "fold.cif"
+        path.write_text(CIF)
+
+        mcp_server.horizyn_dock_waters_submit(structure_path=str(path), selection="confidence", label="w1")
+
+        assert api.sent_json == {
+            "structure_data": CIF,
+            "structure_format": ".cif",
+            "selection": "confidence",
+            "label": "w1",
+        }
+
+    def test_waters_submit_missing_file_names_path_and_sends_nothing(self, api, tmp_path):
+        path = tmp_path / "nope.pdb"
+
+        with pytest.raises(ToolError, match=str(path)):
+            mcp_server.horizyn_dock_waters_submit(structure_path=str(path))
+        assert api.request is None
+
+    def test_dock_status_passes_response_through(self, api):
+        api.body = {
+            "job_id": "j1",
+            "status": "running",
+            "elapsed_seconds": 12.5,
+            "error": None,
+            "queue_position": None,
+            "msa_used": True,
+            "job_type": "fold",
+            "parent_job_id": None,
+            "child_job_id": None,
+        }
+
+        result = mcp_server.horizyn_dock_status(job_id="j1")
+
+        assert api.request.method == "GET"
+        assert api.request.url.path == "/dock/j1/status"
+        assert result == api.body
+
+    def test_list_docking_jobs_default_limit(self, api):
+        api.body = {"jobs": []}
+
+        result = mcp_server.horizyn_list_docking_jobs()
+
+        assert api.request.method == "GET"
+        assert api.request.url.path == "/me/docking-jobs"
+        assert api.request.url.params["limit"] == "20"
+        assert result == {"jobs": []}
+
+    def test_list_docking_jobs_custom_limit(self, api):
+        mcp_server.horizyn_list_docking_jobs(limit=3)
+
+        assert api.request.url.params["limit"] == "3"
+
+
+RESULT_BODY = {
+    "job_id": "j1",
+    "pdb_data": PDB,
+    "cif_data": CIF,
+    "confidence_data": {"confidence_score": 0.9},
+    "msa_used": True,
+    "job_type": "fold",
+    "waters": None,
+    "n_waters": None,
+    "parent_job_id": None,
+    "child_job_id": None,
+}
+
+
+class TestDockResult:
+    def test_inline_when_no_output_dir(self, api):
+        api.body = RESULT_BODY
+
+        result = mcp_server.horizyn_dock_result(job_id="j1")
+
+        assert api.request.method == "GET"
+        assert api.request.url.path == "/dock/j1/result"
+        assert result == RESULT_BODY
+
+    def test_output_dir_writes_both_files_and_returns_paths(self, api, tmp_path):
+        api.body = RESULT_BODY
+        out = tmp_path / "new" / "dir"
+
+        result = mcp_server.horizyn_dock_result(job_id="j1", output_dir=str(out))
+
+        assert (out / "j1.pdb").read_text() == PDB
+        assert (out / "j1.cif").read_text() == CIF
+        assert result["pdb_path"] == str(out / "j1.pdb")
+        assert result["cif_path"] == str(out / "j1.cif")
+        assert "pdb_data" not in result
+        assert "cif_data" not in result
+        for key in (
+            "job_id",
+            "confidence_data",
+            "msa_used",
+            "job_type",
+            "waters",
+            "n_waters",
+            "parent_job_id",
+            "child_job_id",
+        ):
+            assert result[key] == RESULT_BODY[key], key
+
+    def test_output_dir_skips_empty_cif(self, api, tmp_path):
+        api.body = {**RESULT_BODY, "cif_data": ""}
+
+        result = mcp_server.horizyn_dock_result(job_id="j1", output_dir=str(tmp_path))
+
+        assert (tmp_path / "j1.pdb").exists()
+        assert not (tmp_path / "j1.cif").exists()
+        assert result["cif_path"] is None
+        assert "cif_data" not in result
+
+    def test_409_not_complete(self, api):
+        api.status = 409
+        api.body = {"detail": "Job not complete yet"}
+
+        with pytest.raises(ToolError, match="Horizyn returned 409: Job not complete yet"):
+            mcp_server.horizyn_dock_result(job_id="j1")
+
+    def test_404_unknown_job(self, api):
+        api.status = 404
+        api.body = {"detail": "Unknown job"}
+
+        with pytest.raises(ToolError, match="Horizyn returned 404: Unknown job"):
+            mcp_server.horizyn_dock_status(job_id="nope")
+
+
 class TestTrimming:
     def test_query_response_drops_timings_and_tsne(self, api):
         api.body = {