dh-cli 0.17.3__tar.gz → 0.17.4__tar.gz

{dh_cli-0.17.3 → dh_cli-0.17.4}/PKG-INFO

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

{dh_cli-0.17.3 → dh_cli-0.17.4}/pyproject.toml

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

dh_cli-0.17.4/src/dh_cli/hz/mcp_catalogue.py

@@ -0,0 +1,117 @@
+"""One Horizyn MCP catalogue for stdio and authenticated HTTP adapters."""
+
+import json
+from typing import Any
+
+from mcp.server.auth.provider import TokenVerifier
+from mcp.server.auth.settings import AuthSettings
+from mcp.server.mcpserver import MCPServer
+from mcp.types import ToolAnnotations
+
+from dh_cli.hz.mcp_tools import HorizynTools
+
+READ_ONLY = ToolAnnotations(read_only_hint=True, destructive_hint=False, idempotent_hint=True, open_world_hint=True)
+SUBMIT = ToolAnnotations(read_only_hint=False, destructive_hint=False, idempotent_hint=False, open_world_hint=True)
+LOCAL_FILES = ToolAnnotations(read_only_hint=False, destructive_hint=False, idempotent_hint=True)
+
+
+def create_server(
+    tools: HorizynTools,
+    *,
+    local_files: bool = False,
+    token_verifier: TokenVerifier | None = None,
+    auth: AuthSettings | None = None,
+) -> MCPServer:
+    """Bind a caller's request implementation; never load credentials here.
+
+    Remote adapters omit filesystem tools/arguments. Only the local CLI opts in.
+    AuthSettings and TokenVerifier belong to the HTTP adapter, not the tool layer.
+    """
+    server = MCPServer(
+        "horizyn",
+        instructions=(
+            "Horizyn: Dayhoff Labs' enzyme and reaction similarity search. Tools query "
+            "production as the authenticated user. Similarity scores are cosine "
+            "similarities in [0, 1]; higher means more similar. Ask before GPU submissions."
+        ),
+        token_verifier=token_verifier,
+        auth=auth,
+    )
+    for name in (
+        "horizyn_whoami",
+        "horizyn_list_screening_sets",
+        "horizyn_query_reaction",
+        "horizyn_query_protein",
+        "horizyn_protein_details",
+        "horizyn_dock_status",
+        "horizyn_list_docking_jobs",
+        "horizyn_health",
+        "horizyn_detailed_health",
+        "horizyn_validate_reaction",
+        "horizyn_validate_compound",
+        "horizyn_list_cofactors",
+        "horizyn_get_query_history",
+        "horizyn_get_protein_references",
+    ):
+        server.add_tool(getattr(tools, name), annotations=READ_ONLY, structured_output=True)
+    server.add_tool(tools.horizyn_dock_submit, annotations=SUBMIT, structured_output=True)
+    server.add_tool(tools.horizyn_submit_ordered_waters, annotations=SUBMIT, structured_output=True)
+
+    if local_files:
+        server.add_tool(tools.horizyn_dock_waters_submit, annotations=SUBMIT, structured_output=True)
+        server.add_tool(tools.horizyn_dock_result, annotations=LOCAL_FILES, structured_output=True)
+    else:
+
+        @server.tool(annotations=READ_ONLY, structured_output=True)
+        async def horizyn_dock_result(job_id: str) -> dict[str, Any]:
+            """Fetch a completed job's full result without writing server files.
+
+            Prefer horizyn_get_job_result for a compact summary with resource URIs.
+            """
+            return await tools.horizyn_dock_result(job_id)
+
+    @server.tool(annotations=READ_ONLY, structured_output=True)
+    async def horizyn_get_job_result(job_id: str) -> dict[str, Any]:
+        """Summarize a completed job and return resource URIs for large artifacts."""
+        result = await tools.horizyn_dock_result(job_id)
+        root = f"horizyn://prod/jobs/{job_id}"
+        return {
+            "job_id": job_id,
+            "job_type": result.get("job_type", "fold"),
+            "msa_used": result.get("msa_used"),
+            "n_waters": result.get("n_waters"),
+            "has_pdb": bool(result.get("pdb_data")),
+            "pdb_bytes": len(result.get("pdb_data", "").encode()),
+            "has_cif": bool(result.get("cif_data")),
+            "cif_bytes": len(result.get("cif_data", "").encode()),
+            "has_confidence": bool(result.get("confidence_data")),
+            "resources": {
+                "result": f"{root}/result",
+                "pdb": f"{root}/structure.pdb",
+                "cif": f"{root}/structure.cif",
+                "confidence": f"{root}/confidence.json",
+            },
+        }
+
+    @server.resource("horizyn://prod/jobs/{job_id}/result")
+    async def job_result(job_id: str) -> str:
+        """Full JSON result for a completed Horizyn job."""
+        return json.dumps(await tools.horizyn_dock_result(job_id), indent=2, sort_keys=True)
+
+    @server.resource("horizyn://prod/jobs/{job_id}/structure.pdb")
+    async def job_pdb(job_id: str) -> str:
+        """PDB structure produced by a completed Horizyn job."""
+        return str((await tools.horizyn_dock_result(job_id)).get("pdb_data", ""))
+
+    @server.resource("horizyn://prod/jobs/{job_id}/structure.cif")
+    async def job_cif(job_id: str) -> str:
+        """mmCIF structure produced by a completed Horizyn job."""
+        return str((await tools.horizyn_dock_result(job_id)).get("cif_data", ""))
+
+    @server.resource("horizyn://prod/jobs/{job_id}/confidence.json")
+    async def job_confidence(job_id: str) -> str:
+        """Confidence metrics produced by a completed Horizyn job."""
+        result = await tools.horizyn_dock_result(job_id)
+        return json.dumps(result.get("confidence_data", {}), indent=2, sort_keys=True)
+
+    return server

dh_cli-0.17.4/src/dh_cli/hz/mcp_server.py

@@ -0,0 +1,70 @@
+"""Horizyn MCP server over stdio (`dh hz mcp serve`). Prod only.
+
+Every request carries the logged-in user's own Cognito ID token, so queries
+land in that user's Horizyn history. Requires the ``mcp`` extra:
+``pip install 'dh-cli[mcp]'``.
+"""
+
+from functools import partial
+from importlib.metadata import version
+
+import anyio
+import httpx
+from mcp.server.mcpserver.exceptions import ToolError
+
+from dh_cli.hz import API_URLS
+from dh_cli.hz.auth import NotLoggedIn, get_id_token
+from dh_cli.hz.mcp_catalogue import create_server
+from dh_cli.hz.mcp_tools import HorizynTools
+
+CLIENT_HEADER = f"mcp/{version('dh-cli')}"
+_client = httpx.Client(base_url=API_URLS["prod"], timeout=120)
+
+
+def _request(method: str, path: str, **kwargs) -> dict:
+    try:
+        token = get_id_token()
+    except NotLoggedIn as exc:
+        raise ToolError(str(exc)) from None
+
+    headers = {"Authorization": f"Bearer {token}", "X-Horizyn-Client": CLIENT_HEADER}
+    try:
+        resp = _client.request(method, path, headers=headers, **kwargs)
+    except httpx.HTTPError as exc:
+        raise ToolError(f"Horizyn request failed: {type(exc).__name__}: {exc}") from None
+
+    if resp.status_code == 503:
+        raise ToolError("Horizyn is starting or unavailable, retry shortly")
+    if resp.status_code >= 400:
+        raise ToolError(_error_text(resp))
+    return resp.json()
+
+
+def _error_text(resp: httpx.Response) -> str:
+    try:
+        body = resp.json()
+    except ValueError:
+        body = {}
+    if not isinstance(body, dict):
+        body = {}
+    parts = [p for p in (body.get("error"), body.get("message")) if p]
+    detail = body.get("detail")
+    if isinstance(detail, str):
+        parts.append(detail)
+    elif isinstance(detail, list):
+        # FastAPI validation errors; take loc + msg only, never the echoed `input`.
+        for err in detail:
+            if isinstance(err, dict) and err.get("msg"):
+                loc = ".".join(str(x) for x in err.get("loc", []) if x != "body")
+                parts.append(f"{loc}: {err['msg']}" if loc else err["msg"])
+    text = "; ".join(parts) if parts else resp.reason_phrase
+    return f"Horizyn returned {resp.status_code}: {text}"
+
+
+async def _async_request(method: str, path: str, **kwargs) -> dict:
+    # Credential refresh and local HTTP are blocking; keep the MCP event loop free.
+    return await anyio.to_thread.run_sync(partial(_request, method, path, **kwargs))
+
+
+tools = HorizynTools(_async_request)
+server = create_server(tools, local_files=True)

dh_cli-0.17.4/src/dh_cli/hz/mcp_tools.py

@@ -0,0 +1,327 @@
+"""Transport-neutral Horizyn operations. Requests are supplied by each caller."""
+
+from collections.abc import Awaitable, Callable
+from pathlib import Path
+from typing import Any, Literal
+from urllib.parse import quote
+
+from mcp.server.mcpserver.exceptions import ToolError
+
+Request = Callable[..., Awaitable[dict[str, Any]]]
+
+
+def _path_identifier(value: str) -> str:
+    if value in ("", ".", ".."):
+        raise ToolError("Horizyn identifier must not be empty, '.' or '..'")
+    return quote(value, safe="")
+
+
+class HorizynTools:
+    def __init__(self, request: Request):
+        self.request = request
+
+    async def _query(self, path: str, body: dict, filters: dict) -> dict[str, Any]:
+        body.update({k: v for k, v in filters.items() if v is not None})
+        data = await self.request("POST", path, json=body)
+        data.pop("timings_ms", None)
+        data.pop("tsne", None)
+        return data
+
+    async def horizyn_whoami(
+        self,
+    ) -> dict[str, Any]:
+        """Identify the logged-in Horizyn user.
+
+        Returns email, tier (admin, alpha, or beta) and the feature flags that tier
+        grants. Use this first to confirm the server is authenticated.
+        """
+        return await self.request("GET", "/me")
+
+    async def horizyn_list_screening_sets(self, screening_set_type: str | None = None) -> dict[str, Any]:
+        """List the screening sets available for similarity queries.
+
+        Each entry gives the set's id, type (protein or reaction), size, model
+        version and whether it is the current default (`is_default`).
+        `screening_set_type` optionally filters to "protein" or "reaction".
+        """
+        params = {"screening_set_type": screening_set_type} if screening_set_type else None
+        return await self.request("GET", "/screening_sets", params=params)
+
+    async def horizyn_query_reaction(
+        self,
+        smiles: str,
+        screening_set: str | None = None,
+        top_k: int = 50,
+        page: int = 1,
+        page_size: int = 20,
+        ec_include: list[str] | None = None,
+        ec_exclude: list[str] | None = None,
+        tc_include: list[str] | None = None,
+        tc_exclude: list[str] | None = None,
+        cofactor_include: list[str] | None = None,
+        cofactor_exclude: list[str] | None = None,
+        min_length: int | None = None,
+        max_length: int | None = None,
+        cluster_by: Literal["ec3", "ec4"] | None = None,
+        cluster_size: int | None = None,
+        fetch_metadata: bool | None = None,
+    ) -> dict[str, Any]:
+        """Find proteins (or reactions) similar to a reaction SMILES.
+
+        `smiles` is a reaction SMILES such as "CC=O.O>>CC(O)O". When
+        `screening_set` is omitted Horizyn searches its current default protein set
+        (the response's `screening_set` names it; see horizyn_list_screening_sets).
+        Pass a reaction set such as horizyn_v1_1_rxns to find similar reactions
+        instead. `score` on each result is the cosine similarity between the query
+        and the hit in Horizyn's embedding space (0-1, higher is more similar).
+        `top_k` (max 1000) bounds the similarity search; `page`/`page_size`
+        paginate the result list.
+
+        All filters are optional. `ec_include`/`ec_exclude` and
+        `tc_include`/`tc_exclude` take hierarchical prefixes ("1", "1.1", "1.1.1",
+        "1.1.1.1"; "__NONE__" matches unannotated). `cofactor_include`/
+        `cofactor_exclude` take cofactor names ("NAD", "Zn"; "__ANY__"/"__NONE__").
+        `min_length`/`max_length` bound sequence length. `filter_stats` in the
+        response reports how many hits the filters kept.
+        """
+        body = {"smiles": smiles, "top_k": top_k, "page": page, "page_size": page_size}
+        filters = {
+            "screening_set": screening_set,
+            "ec_include": ec_include,
+            "ec_exclude": ec_exclude,
+            "tc_include": tc_include,
+            "tc_exclude": tc_exclude,
+            "cofactor_include": cofactor_include,
+            "cofactor_exclude": cofactor_exclude,
+            "min_length": min_length,
+            "max_length": max_length,
+            "cluster_by": cluster_by,
+            "cluster_size": cluster_size,
+            "fetch_metadata": fetch_metadata,
+        }
+        return await self._query("/query/reaction", body, filters)
+
+    async def horizyn_query_protein(
+        self,
+        sequence: str,
+        screening_set: str | None = None,
+        top_k: int = 50,
+        page: int = 1,
+        page_size: int = 20,
+        ec_include: list[str] | None = None,
+        ec_exclude: list[str] | None = None,
+        tc_include: list[str] | None = None,
+        tc_exclude: list[str] | None = None,
+        cofactor_include: list[str] | None = None,
+        cofactor_exclude: list[str] | None = None,
+        min_length: int | None = None,
+        max_length: int | None = None,
+        cluster_by: Literal["ec3", "ec4"] | None = None,
+        cluster_size: int | None = None,
+        fetch_metadata: bool | None = None,
+    ) -> dict[str, Any]:
+        """Find proteins (or reactions) similar to a protein sequence.
+
+        `sequence` is a one-letter amino-acid string. Same screening sets, scoring,
+        pagination and optional filters as horizyn_query_reaction: `score` is the
+        cosine similarity (0-1) in Horizyn's shared reaction/protein embedding
+        space, omitting `screening_set` searches Horizyn's current default protein
+        set, and every filter argument may be omitted.
+        """
+        body = {"sequence": sequence, "top_k": top_k, "page": page, "page_size": page_size}
+        filters = {
+            "screening_set": screening_set,
+            "ec_include": ec_include,
+            "ec_exclude": ec_exclude,
+            "tc_include": tc_include,
+            "tc_exclude": tc_exclude,
+            "cofactor_include": cofactor_include,
+            "cofactor_exclude": cofactor_exclude,
+            "min_length": min_length,
+            "max_length": max_length,
+            "cluster_by": cluster_by,
+            "cluster_size": cluster_size,
+            "fetch_metadata": fetch_metadata,
+        }
+        return await self._query("/query/protein", body, filters)
+
+    async def horizyn_protein_details(self, protein_id: str, screening_set: str) -> dict[str, Any]:
+        """Get full details for one protein from a screening set.
+
+        `protein_id` is the `id` returned by a query and `screening_set` is the
+        `screening_set` echoed in that query's response (the id is only meaningful
+        within its set). Returns name, organism, EC numbers, cofactors, expression
+        score, full sequence, length, AlphaFold id and any extra metadata.
+        """
+        return await self.request(
+            "GET", f"/protein/{_path_identifier(protein_id)}/details", params={"screening_set": screening_set}
+        )
+
+    async def horizyn_dock_submit(
+        self,
+        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,
+        screening_set: str | None = None,
+    ) -> dict[str, Any]:
+        """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.
+        """
+        if bool(protein_sequence) == bool(boltz_yaml):
+            raise ToolError("Provide exactly one of protein_sequence or boltz_yaml")
+
+        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,
+            "screening_set": screening_set,
+        }
+        return await self.request("POST", "/dock/submit", json={k: v for k, v in body.items() if v is not None})
+
+    async def horizyn_dock_waters_submit(
+        self,
+        structure_path: str,
+        selection: str | None = None,
+        label: str | None = None,
+    ) -> dict[str, Any]:
+        """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 await self.request("POST", "/dock/waters/submit", json={k: v for k, v in body.items() if v is not None})
+
+    async def horizyn_dock_status(self, job_id: str) -> dict[str, Any]:
+        """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 await self.request("GET", f"/dock/{_path_identifier(job_id)}/status")
+
+    async def horizyn_dock_result(self, job_id: str, output_dir: str | None = None) -> dict[str, Any]:
+        """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 = await self.request("GET", f"/dock/{_path_identifier(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
+
+    async def horizyn_list_docking_jobs(self, limit: int = 20) -> dict[str, Any]:
+        """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 await self.request("GET", "/me/docking-jobs", params={"limit": limit})
+
+    async def horizyn_health(self) -> dict[str, Any]:
+        """Check upstream readiness. Remote MCP still authenticates the caller."""
+        return await self.request("GET", "/healthz")
+
+    async def horizyn_detailed_health(self) -> dict[str, Any]:
+        """Inspect Horizyn GPU, registry, EFS and default-set readiness."""
+        return await self.request("GET", "/healthz/detailed")
+
+    async def horizyn_validate_reaction(self, smiles: str) -> dict[str, Any]:
+        """Validate atom/charge balance and parsing of reaction SMILES."""
+        return await self.request("POST", "/validate/reaction", json={"smiles": smiles})
+
+    async def horizyn_validate_compound(self, smiles: str) -> dict[str, Any]:
+        """Validate and canonicalize a compound SMILES."""
+        return await self.request("POST", "/validate/compound", json={"smiles": smiles})
+
+    async def horizyn_list_cofactors(self) -> dict[str, Any]:
+        """List annotated cofactors and protein counts available for filtering."""
+        return await self.request("GET", "/cofactors")
+
+    async def horizyn_get_query_history(self, limit: int = 20) -> dict[str, Any]:
+        """Return recent queries for the caller, with limit between 1 and 1000."""
+        return await self.request("GET", "/me/query-history", params={"limit": limit})
+
+    async def horizyn_get_protein_references(self, protein_id: str, screening_set: str) -> dict[str, Any]:
+        """Fetch UniProt literature and structure references for a protein hit."""
+        return await self.request(
+            "GET", f"/protein/{_path_identifier(protein_id)}/references", params={"screening_set": screening_set}
+        )
+
+    async def horizyn_submit_ordered_waters(
+        self,
+        structure_data: str,
+        structure_format: Literal[".pdb", ".cif", "pdb", "cif"],
+        selection: Literal["confidence", "density"] = "confidence",
+        label: str | None = None,
+    ) -> dict[str, Any]:
+        """Submit ordered-water placement from inline PDB/mmCIF. Costs GPU time; ask first.
+
+        Remote callers must send the structure itself, never a server filesystem path.
+        """
+        body = {
+            "structure_data": structure_data,
+            "structure_format": structure_format,
+            "selection": selection,
+            "label": label,
+        }
+        return await self.request("POST", "/dock/waters/submit", json={k: v for k, v in body.items() if v is not None})