diff --git a/CHANGELOG.md b/CHANGELOG.md index 1577303..d30808a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,35 @@ # Changelog +## 1.1.0 — 2026-09-28 + +Full coverage of the current `/api/v1` (Agent platform epic, cognichem#458; covers the Python-client follow-up cognichem#269). Additive: existing 1.0 calls keep working. + +### Added + +- **Workflows** (`client.workflows`): `validate`, `estimate`, curated `templates` / `template`, catalog `node_ports`. + - `workflows.runs`: `create`, `list`, `get(include_spec=)`, `artifacts`, `resume`, `cancel`, `delete`, `wait`, `run`. `wait` stops on `completed | failed | cancelled | paused`. + - `workflows.definitions`: CRUD, JSON `download`, link `share` / `unshare`, `get_shared` / `fork_shared`, `versions` / `get_version` / `restore_version`. +- **Artifacts** (`client.artifacts`): `list` (filters, paging), `usage`, `upload` (`cognichem_zip` or `raw_file` multipart), `get` (manifest), `download` (whole zip or one `record=`), `delete`. `artifact_ref(id, port)` builds `$artifact` bindings. +- **Completion events** (`client.events`): `list` and a long-polling, de-duplicating `listen` generator. +- **Wallet** (`client.wallet.balance`), **science lookups** (`client.lookup`: PubChem, ChEMBL, Open Targets, UniProt, RCSB PDB, PubChem properties), and the **reference KG** (`client.reference`: `search`, `get`, `neighborhood`). +- **CogniChem Assistant chat** (`client.chat`, JWT only): sessions, message history, turn `estimate`, SSE `stream` / collected `send`, `stop`, thread `spend_cap` / `continue_spend_cap`, and run `proposals` (`list`, `get`, `approve`, `reject`). +- Jobs: `estimate`; `list` pagination and filters (`limit`, `offset`, `status`, `job_type`, `q`, `sort`) with full `items` rows; `status().result_artifact_id`. +- Inference / utils: `cancel`. +- API keys: `create(scopes=, expires_at=, spend_ceiling_usd=, allow_structure_search=)`, `update(name=, allow_structure_search=)`; new key metadata fields. +- `client.health()` / `client.ready()`. +- Errors: `BadRequestError` (400), `PaymentRequiredError` (402), `GoneError` (410), `PayloadTooLargeError` (413), `ServerError` (5xx); `retryability` attribute; `code` is now the problem type slug (e.g. `missing-idempotency-key`, `session-spend-cap`). +- Catalog constants: 45 submittable job types (was 23), utilities `molecule_standardize`, `druglikeness_filter`, `coolprop`, workflow run statuses. + +### Fixed + +- Mutating calls (`jobs.submit`, `submit_many`, `inference.submit*`, `utils.submit`, run / definition / upload writes) now send an auto-generated `Idempotency-Key` when none is given. The API rejects API-key mutations without one (400 `missing-idempotency-key`), so 1.0 submits failed with API keys. +- `api_keys.delete` returns `None`; it raised on the API's 204 response. + +### Changed + +- `cognichem_client.types` is now a package (same import path; every model is still importable from it). +- `scripts/sync_routes.py` checks keys **and paths** against backend `ROUTE_MAP` (`scripts/route_map.txt`). + ## 1.0.0 — 2026-07-27 Clean-break redesign of the public SDK against the current CogniChem `/api/v1` compute API. diff --git a/README.md b/README.md index 1961b20..c948e42 100644 --- a/README.md +++ b/README.md @@ -55,7 +55,9 @@ if status.status == "completed": | Mode | How | |------|-----| | **API key** (recommended for scripts) | `CogniChem(api_key=...)` → `X-Api-Key` header | -| **JWT** (API key management) | `client.auth.login_email(email, password)` then `client.api_keys.*` | +| **JWT** (API key management, Assistant chat) | `client.auth.login_email(email, password)` or `CogniChem(access_token=...)` | + +API keys carry `read` and/or `write` scopes. Submitting jobs, workflow runs, or uploads needs `write`. API-key mutations must send an `Idempotency-Key`; the client generates one per call unless you pass `idempotency_key=` (pass your own to make retries safe). Environment variables: @@ -67,28 +69,97 @@ Create and manage keys in your [CogniChem account](https://cognichem.com). Colle ## API surface -| Resource | Examples | -|----------|----------| -| `client.jobs` | `submit`, `submit_many`, `wait`, `run`, `result`, `result_many`, `info`, `list`, `cancel`, `delete` | -| `client.inference` | `submit`, `submit_batch`, `wait`, `run`, `result`, `list`, `delete` | +| Resource | Methods | +|----------|---------| +| `client.jobs` | `submit`, `submit_many`, `estimate`, `wait`, `run`, `result`, `result_many`, `info`, `status`, `list` (paginated, filters), `cancel`, `delete` | +| `client.workflows` | `validate`, `estimate`, `templates`, `template`, `node_ports` | +| `client.workflows.runs` | `create`, `list`, `get`, `artifacts`, `resume`, `cancel`, `delete`, `wait`, `run` | +| `client.workflows.definitions` | `create`, `list`, `get`, `update`, `delete`, `download`, `share`, `unshare`, `get_shared`, `fork_shared`, `versions`, `get_version`, `restore_version` | +| `client.artifacts` | `list`, `usage`, `upload`, `get`, `download` (whole zip or one `record=`), `delete` | +| `client.events` | `list`, `listen` (long-poll completion events for jobs and workflow runs) | +| `client.wallet` | `balance` | +| `client.lookup` | `pubchem`, `chembl`, `targets` (Open Targets), `uniprot`, `pdb`, `properties` | +| `client.reference` | `search`, `get`, `neighborhood` (literature / method papers) | +| `client.chat` (JWT) | `estimate`, `stream`, `send`, `stop`, `spend_cap`, `continue_spend_cap`, plus `sessions.*` and `proposals.*` | +| `client.inference` | `submit`, `submit_batch`, `wait`, `run`, `result`, `list`, `cancel`, `delete` | | `client.inference.models.mpnn` | `public()`, `user(name=None)`, `delete(name)` | -| `client.utils` | `submit`, `wait`, `run`, `result`, `list`, `delete` | +| `client.utils` | `submit`, `wait`, `run`, `result`, `list`, `cancel`, `delete` | | `client.usage` | `limits(month=None)` | -| `client.auth` / `client.api_keys` | login, check, refresh, logout; list/create/rotate/delete keys | +| `client.auth` / `client.api_keys` | login, check, refresh, logout; list/create/get/update/rename/rotate/delete keys | +| `client` | `health()`, `ready()` | + +Async twin: `AsyncCogniChem` with the same method names (`await client.jobs.run(...)`, `async for event in client.chat.stream(...)`). + +## Workflows + +```python +from cognichem_client import CogniChem, artifact_ref + +client = CogniChem.from_env() + +spec = client.workflows.template("smiles-embed-dock").spec +target = client.artifacts.upload("target.pdb", data_kind="protein_structure", data_format="pdb") +params = { + "ligands": ["CCO", "c1ccccc1O"], + "target_structure": artifact_ref(target.id, "structure"), +} + +assert client.workflows.validate(spec, params=params).valid +estimate = client.workflows.estimate(spec, params=params) # USD hold at your tier + +run = client.workflows.runs.run("dock-demo", spec, params=params, max_run_cost=estimate.total * 1.5) +# run.status is "completed" (or "paused" if it hit max_run_cost; call runs.resume) +for item in client.workflows.runs.artifacts(run.id).items: + print(item.step_key, item.port, item.artifact_id, item.record_count) +``` -Optional `idempotency_key=` on mutating POSTs (`jobs`, `inference`, `utils` submit). +`runs.wait` stops on `completed`, `failed`, `cancelled`, or `paused` (a paused run needs you to resume it). `runs.run` raises `ProcessFailedError` / `ProcessCancelledError` on failure. + +Job and run outputs are durable **artifacts**: `client.jobs.status(pid).result_artifact_id` points at one. Inspect the manifest with `client.artifacts.get(...)` and fetch single records with `download(..., record=...)` instead of whole zips. Bind artifacts into workflow params or job payloads with `artifact_ref(id, port)`. + +## Completion events + +```python +for event in client.events.listen(kind="workflow_run"): + print(event.subject_id, event.status) # re-read the run for detail +``` + +## CogniChem Assistant (JWT) + +```python +client = CogniChem(access_token=token) +session = client.chat.sessions.create("Docking plan") +reply = client.chat.send(session.id, "Dock aspirin against COX-1", reasoning_effort="low") +print(reply.content, reply.billed_usd) +for proposal in client.chat.proposals.list(session.id).items: + client.chat.proposals.approve(session.id, proposal.id, proposal.estimate_usd) +``` -Async twin: `AsyncCogniChem` with the same method names (`await client.jobs.run(...)`). +Each turn holds wallet funds at your tier's token rate and bills the metered tokens. A thread at its spend cap raises `ConflictError` (`code == "session-spend-cap"`); `continue_spend_cap` raises the cap. ## Errors -API failures raise subclasses of `CogniChemError` parsed from RFC 7807 problem details (`detail`, `title`, `status`, `type`, `request_id`, `errors` on 422). +API failures raise subclasses of `CogniChemError` parsed from RFC 7807 problem details (`detail`, `title`, `status`, `type`, `request_id`, `retryability`, `errors`). `code` is the problem type slug (for example `missing-idempotency-key`, `insufficient-scope`, `spend-ceiling-exceeded`, `insufficient-wallet`, `session-spend-cap`, `estimate-changed`), and `body` holds any extra fields. + +| Status | Exception | +|--------|-----------| +| 400 | `BadRequestError` | +| 401 | `AuthenticationError` | +| 402 | `PaymentRequiredError` | +| 403 | `ForbiddenError` | +| 404 | `NotFoundError` | +| 409 | `ConflictError` | +| 410 | `GoneError` | +| 413 | `PayloadTooLargeError` | +| 422 | `ValidationError` (WorkflowSpec diagnostics in `errors`) | +| 429 | `RateLimitError` | +| 5xx | `ServerError` | Polling helpers raise `PollTimeoutError`, `ProcessFailedError`, or `ProcessCancelledError` when appropriate. ## Examples -See [`examples/`](examples/) for scripts covering utilities, inference, and job submission. +See [`examples/`](examples/) for scripts covering utilities, inference, job submission, workflow templates, artifact records, science lookups, and the Assistant. Live OpenAPI docs: [https://api.cognichem.com/docs](https://api.cognichem.com/docs) @@ -102,6 +173,8 @@ ruff check src tests mypy pytest python scripts/sync_routes.py +# refresh the ROUTE_MAP snapshot from a monorepo checkout: +python scripts/sync_routes.py --backend-config ../cognichem/apps/back-end-api/src/core/config.py --write-snapshot ``` ## Migrating from 0.1.x diff --git a/examples/artifact_records.py b/examples/artifact_records.py new file mode 100644 index 0000000..84d5c3e --- /dev/null +++ b/examples/artifact_records.py @@ -0,0 +1,45 @@ +"""Inspect a job's result artifact and download one record instead of the zip.""" + +from __future__ import annotations + +import os +import sys + +from cognichem_client import CogniChem + + +def main(process_id: str) -> None: + client = CogniChem(api_key=os.environ["COGNICHEM_API_KEY"]) + + status = client.jobs.status(process_id) + if not status.result_artifact_id: + print("no stored artifact for", process_id, status.status) + return + + artifact = client.artifacts.get(status.result_artifact_id) + for port, spec in artifact.manifest.get("ports", {}).items(): + print(port, spec.get("kind"), spec.get("format"), spec.get("record_count")) + + usage = client.artifacts.usage() + print(f"storage {usage.used_bytes}/{usage.quota_bytes} bytes") + + ports = artifact.manifest.get("ports", {}) + first = next( + ( + (port, record["id"]) + for port, spec in ports.items() + if port != "archive" + for record in spec.get("records", []) + ), + None, + ) + if first: + port, record_id = first + result = client.artifacts.download( + artifact.id, record=record_id, port=port, save_path="." + ) + print("saved", result.filename, len(result.content), "bytes") + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/examples/assistant_chat.py b/examples/assistant_chat.py new file mode 100644 index 0000000..e3bd0a3 --- /dev/null +++ b/examples/assistant_chat.py @@ -0,0 +1,36 @@ +"""Ask the CogniChem Assistant a question and review any run proposals. + +Chat is JWT-only: set ``COGNICHEM_ACCESS_TOKEN`` (or sign in on a client with +``client.auth.login_email(email, password)``, which stores the token). +""" + +from __future__ import annotations + +import os + +from cognichem_client import CogniChem + + +def main() -> None: + client = CogniChem(access_token=os.environ["COGNICHEM_ACCESS_TOKEN"]) + + session = client.chat.sessions.create("Docking plan") + question = "How should I dock aspirin against COX-1? Estimate the cost." + quote = client.chat.estimate(session.id, question) + print(f"turn hold ${quote.hold_usd:.4f} ({quote.reasoning_effort})") + + for event in client.chat.stream(session.id, question): + if event.event == "token": + print(event.data["delta"], end="", flush=True) + elif event.event == "done": + print(f"\nbilled ${event.data['billed_usd']:.4f}") + + for proposal in client.chat.proposals.list(session.id).items: + print(proposal.id, proposal.kind, proposal.status, proposal.estimate_usd) + # To run it: client.chat.proposals.approve( + # session.id, proposal.id, proposal.estimate_usd + # ) + + +if __name__ == "__main__": + main() diff --git a/examples/science_lookup.py b/examples/science_lookup.py new file mode 100644 index 0000000..7ea2130 --- /dev/null +++ b/examples/science_lookup.py @@ -0,0 +1,30 @@ +"""Resolve a compound, look up target activities, and find method papers.""" + +from __future__ import annotations + +import os + +from cognichem_client import CogniChem + + +def main() -> None: + client = CogniChem(api_key=os.environ["COGNICHEM_API_KEY"]) + + compounds = client.lookup.pubchem(["aspirin"]) + for compound in compounds.compounds: + print(compound.name, compound.status, compound.cid, compound.smiles) + + activities = client.lookup.chembl("target_activities", ["EGFR"], limit=5) + for row in activities.results: + print(row.status, row.record_id, row.title, row.url) + print(activities.attribution) + + papers = client.reference.search("protonation state pKa", paper_set="method") + for paper in papers.papers: + print(paper.year, paper.title, paper.url) + + print("wallet", client.wallet.balance()) + + +if __name__ == "__main__": + main() diff --git a/examples/workflow_template_run.py b/examples/workflow_template_run.py new file mode 100644 index 0000000..af16f01 --- /dev/null +++ b/examples/workflow_template_run.py @@ -0,0 +1,49 @@ +"""Run a curated workflow template: upload a target, validate, estimate, run.""" + +from __future__ import annotations + +import os +import sys +import uuid + +from cognichem_client import CogniChem, artifact_ref + + +def main(pdb_path: str) -> None: + client = CogniChem(api_key=os.environ["COGNICHEM_API_KEY"]) + + for template in client.workflows.templates(): + print(f"{template.id:32} {template.complexity:9} {template.summary}") + spec = client.workflows.template("smiles-embed-dock").spec + + # Structure params take an $artifact ref, so upload the PDB first. + target = client.artifacts.upload( + pdb_path, data_kind="protein_structure", data_format="pdb" + ) + port = next(iter(target.manifest.get("ports", {})), "structure") + params = { + "ligands": ["CCO", "c1ccccc1O"], + "target_structure": artifact_ref(target.id, port), + } + + check = client.workflows.validate(spec, params=params) + if not check.valid: + for error in check.errors: + print("invalid:", error.code, error.message) + return + estimate = client.workflows.estimate(spec, params=params) + print(f"estimated hold ${estimate.total:.4f} at tier {estimate.tier}") + + run = client.workflows.runs.run( + f"example-dock-{uuid.uuid4().hex[:8]}", + spec, + params=params, + max_run_cost=estimate.total * 1.5, + ) + print("run", run.id, run.status, f"charged ${run.charged_amount:.4f}") + for item in client.workflows.runs.artifacts(run.id).items: + print(" ", item.step_key, item.port, item.artifact_id, item.record_count) + + +if __name__ == "__main__": + main(sys.argv[1]) diff --git a/pyproject.toml b/pyproject.toml index 473dc3a..c2c788b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "cognichem-client" -version = "1.0.0" +version = "1.1.0" description = "Python client for the CogniChem API" authors = [ { name = "Travis Kessler", email = "travis@cognichem.com" } diff --git a/scripts/route_map.txt b/scripts/route_map.txt new file mode 100644 index 0000000..a10a83c --- /dev/null +++ b/scripts/route_map.txt @@ -0,0 +1,69 @@ +# Backend Settings.ROUTE_MAP: key path +auth_login_email /api/v1/auth/login-email +auth_check /api/v1/auth/check +auth_refresh /api/v1/auth/refresh +auth_logout /api/v1/auth/logout +auth_api_keys /api/v1/auth/api-keys +auth_api_keys_item /api/v1/auth/api-keys/{key_id} +auth_api_keys_rotate /api/v1/auth/api-keys/{key_id}/rotate +jobs_submit /api/v1/jobs/submit +jobs_submit_multiple /api/v1/jobs/submit-multiple +jobs_estimate /api/v1/jobs/estimate +jobs_info /api/v1/jobs/info +jobs_list /api/v1/jobs/list +jobs_status /api/v1/jobs/status +jobs_result /api/v1/jobs/result +jobs_result_multiple /api/v1/jobs/result-multiple +jobs_cancel /api/v1/jobs/cancel +jobs_delete /api/v1/jobs/delete +artifacts_list /api/v1/artifacts +artifacts_upload /api/v1/artifacts +artifacts_usage /api/v1/artifacts/usage +artifacts_get /api/v1/artifacts/{artifact_id} +artifacts_download /api/v1/artifacts/{artifact_id}/download +artifacts_download_record /api/v1/artifacts/{artifact_id}/download +artifacts_delete /api/v1/artifacts/{artifact_id} +workflows_validate /api/v1/workflows/validate +workflows_estimate /api/v1/workflows/estimate +workflows_definitions_create /api/v1/workflows/definitions +workflows_definitions_list /api/v1/workflows/definitions +workflows_definitions_get /api/v1/workflows/definitions/{definition_id} +workflows_definitions_update /api/v1/workflows/definitions/{definition_id} +workflows_definitions_delete /api/v1/workflows/definitions/{definition_id} +workflows_definitions_download /api/v1/workflows/definitions/{definition_id}/download +workflows_runs_create /api/v1/workflows/runs +workflows_runs_list /api/v1/workflows/runs +workflows_runs_get /api/v1/workflows/runs/{run_id} +workflows_runs_artifacts /api/v1/workflows/runs/{run_id}/artifacts +workflows_runs_resume /api/v1/workflows/runs/{run_id}/resume +workflows_runs_cancel /api/v1/workflows/runs/{run_id}/cancel +workflows_runs_delete /api/v1/workflows/runs/{run_id} +inference_submit /api/v1/inference/submit +inference_submit_batch /api/v1/inference/submit-batch +inference_list /api/v1/inference/list +inference_status /api/v1/inference/status +inference_result /api/v1/inference/result +inference_delete /api/v1/inference/delete +inference_cancel /api/v1/inference/cancel +mpnn_user_models /api/v1/user-models/mpnn +mpnn_public_models /api/v1/public-models/mpnn +utils_submit /api/v1/utils/submit +utils_list /api/v1/utils/list +utils_status /api/v1/utils/status +utils_result /api/v1/utils/result +utils_delete /api/v1/utils/delete +utils_cancel /api/v1/utils/cancel +usage_limits /api/v1/usage-limits +events /api/v1/events +chat_sessions /api/v1/chat/sessions +chat_session /api/v1/chat/sessions/{session_id} +chat_messages /api/v1/chat/sessions/{session_id}/messages +chat_estimate /api/v1/chat/sessions/{session_id}/estimate +chat_turns /api/v1/chat/sessions/{session_id}/turns +chat_turn_stop /api/v1/chat/sessions/{session_id}/turns/{turn_id}/stop +chat_spend_cap /api/v1/chat/sessions/{session_id}/spend-cap +chat_spend_cap_continue /api/v1/chat/sessions/{session_id}/spend-cap/continue +chat_proposals /api/v1/chat/sessions/{session_id}/proposals +chat_proposal /api/v1/chat/sessions/{session_id}/proposals/{proposal_id} +chat_proposal_approve /api/v1/chat/sessions/{session_id}/proposals/{proposal_id}/approve +chat_proposal_reject /api/v1/chat/sessions/{session_id}/proposals/{proposal_id}/reject diff --git a/scripts/route_map_keys.txt b/scripts/route_map_keys.txt deleted file mode 100644 index 9db720f..0000000 --- a/scripts/route_map_keys.txt +++ /dev/null @@ -1,30 +0,0 @@ -auth_login_email -auth_check -auth_refresh -auth_logout -auth_api_keys -auth_api_keys_item -auth_api_keys_rotate -jobs_submit -jobs_submit_multiple -jobs_info -jobs_list -jobs_status -jobs_result -jobs_result_multiple -jobs_cancel -jobs_delete -inference_submit -inference_submit_batch -inference_list -inference_status -inference_result -inference_delete -mpnn_user_models -mpnn_public_models -utils_submit -utils_list -utils_status -utils_result -utils_delete -usage_limits diff --git a/scripts/sync_routes.py b/scripts/sync_routes.py index a56623b..f1247fc 100644 --- a/scripts/sync_routes.py +++ b/scripts/sync_routes.py @@ -1,13 +1,21 @@ #!/usr/bin/env python3 -"""Assert client ROUTES keys match a backend ROUTE_MAP snapshot. +"""Assert client ROUTES cover the backend ROUTE_MAP with identical paths. + +Every backend ``Settings.ROUTE_MAP`` key must exist in the client's +``ROUTES`` with the same path. The client may carry extra routes that +ROUTE_MAP does not list (health, templates, lookup, …); those are reported +but do not fail the check. Usage: - # Against a checked-in snapshot (default): + # Against the checked-in snapshot (default): python scripts/sync_routes.py - # Against a live monorepo path: + # Against a live monorepo checkout: python scripts/sync_routes.py --backend-config ../cognichem/apps/back-end-api/src/core/config.py -""" + + # Refresh the snapshot from a monorepo checkout: + python scripts/sync_routes.py --backend-config --write-snapshot +""" # noqa: E501 from __future__ import annotations @@ -19,104 +27,130 @@ ROOT = Path(__file__).resolve().parents[1] CLIENT_CONSTANTS = ROOT / "src" / "cognichem_client" / "constants.py" -SNAPSHOT = ROOT / "scripts" / "route_map_keys.txt" - -EXPECTED_KEYS = [ - "auth_login_email", - "auth_check", - "auth_refresh", - "auth_logout", - "auth_api_keys", - "auth_api_keys_item", - "auth_api_keys_rotate", - "jobs_submit", - "jobs_submit_multiple", - "jobs_info", - "jobs_list", - "jobs_status", - "jobs_result", - "jobs_result_multiple", - "jobs_cancel", - "jobs_delete", - "inference_submit", - "inference_submit_batch", - "inference_list", - "inference_status", - "inference_result", - "inference_delete", - "mpnn_user_models", - "mpnn_public_models", - "utils_submit", - "utils_list", - "utils_status", - "utils_result", - "utils_delete", - "usage_limits", -] - - -def client_route_keys() -> set[str]: +SNAPSHOT = ROOT / "scripts" / "route_map.txt" +API_V1_STR = "/api/v1" + +# ``"jobs_submit": f"{b}/jobs/submit"`` or the parenthesized multi-line form. +_BACKEND_ENTRY = re.compile(r'"([a-z0-9_]+)":\s*\(?\s*f"\{b\}([^"]*)"') + + +def client_routes() -> dict[str, str]: + """Return the client's ROUTES as ``{key: path}``.""" tree = ast.parse(CLIENT_CONSTANTS.read_text()) for node in tree.body: - if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name): - if node.target.id == "ROUTES" and isinstance(node.value, ast.Dict): - keys: set[str] = set() - for key in node.value.keys: - if isinstance(key, ast.Constant) and isinstance(key.value, str): - keys.add(key.value) - return keys + if ( + isinstance(node, ast.AnnAssign) + and isinstance(node.target, ast.Name) + and node.target.id == "ROUTES" + and isinstance(node.value, ast.Dict) + ): + routes: dict[str, str] = {} + for key, value in zip(node.value.keys, node.value.values, strict=True): + assert isinstance(key, ast.Constant) and isinstance(key.value, str) + routes[key.value] = _render(value) + return routes raise RuntimeError("ROUTES dict not found in constants.py") -def keys_from_backend_config(path: Path) -> set[str]: +def _render(value: ast.expr) -> str: + """Render a ROUTES value (literal or ``f"{API_V1_STR}…"``) to a path.""" + if isinstance(value, ast.Constant) and isinstance(value.value, str): + return value.value + if isinstance(value, ast.JoinedStr): + parts: list[str] = [] + for part in value.values: + if isinstance(part, ast.Constant): + parts.append(str(part.value)) + elif ( + isinstance(part, ast.FormattedValue) + and isinstance(part.value, ast.Name) + and part.value.id == "API_V1_STR" + ): + parts.append(API_V1_STR) + else: + raise RuntimeError(f"Unsupported ROUTES value: {ast.dump(value)}") + return "".join(parts) + raise RuntimeError(f"Unsupported ROUTES value: {ast.dump(value)}") + + +def backend_routes(path: Path) -> dict[str, str]: + """Parse ``Settings.ROUTE_MAP`` from the backend config source.""" text = path.read_text() - # Match string keys in ROUTE_MAP return dict literals: "jobs_submit": - return set(re.findall(r'"([a-z0-9_]+)":\s*f?"?\{?b\}?', text)) & set(EXPECTED_KEYS) + start = text.find("def ROUTE_MAP") + if start == -1: + return {} + body = text[start:] + end = body.find("\n def ", 1) + body = body if end == -1 else body[:end] + return { + key: API_V1_STR + suffix.replace("{{", "{").replace("}}", "}") + for key, suffix in _BACKEND_ENTRY.findall(body) + } + + +def read_snapshot() -> dict[str, str]: + """Read ``route_map.txt`` (``key path`` per line).""" + routes: dict[str, str] = {} + for line in SNAPSHOT.read_text().splitlines(): + line = line.strip() + if line and not line.startswith("#"): + key, path = line.split(maxsplit=1) + routes[key] = path + return routes def main() -> int: + """Run the drift check and return a process exit code.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--backend-config", type=Path, - help="Optional path to cognichem apps/back-end-api/src/core/config.py", + help="Path to cognichem apps/back-end-api/src/core/config.py", ) parser.add_argument( "--write-snapshot", action="store_true", - help="Write scripts/route_map_keys.txt from EXPECTED_KEYS", + help="Write scripts/route_map.txt from --backend-config", ) args = parser.parse_args() - if args.write_snapshot: - SNAPSHOT.write_text("\n".join(EXPECTED_KEYS) + "\n") - print(f"Wrote {SNAPSHOT}") - return 0 - - client_keys = client_route_keys() - expected = set(EXPECTED_KEYS) if args.backend_config: - expected = keys_from_backend_config(args.backend_config) + expected = backend_routes(args.backend_config) if not expected: - print("Failed to parse backend ROUTE_MAP keys", file=sys.stderr) + print("Failed to parse backend ROUTE_MAP", file=sys.stderr) return 1 - elif SNAPSHOT.exists(): - expected = { - line.strip() - for line in SNAPSHOT.read_text().splitlines() - if line.strip() and not line.startswith("#") - } - - missing = expected - client_keys - extra = client_keys - expected - if missing or extra: - if missing: - print(f"Missing in client ROUTES: {sorted(missing)}", file=sys.stderr) - if extra: - print(f"Extra in client ROUTES: {sorted(extra)}", file=sys.stderr) + elif args.write_snapshot: + print("--write-snapshot needs --backend-config", file=sys.stderr) return 1 + else: + expected = read_snapshot() + + if args.write_snapshot: + lines = ["# Backend Settings.ROUTE_MAP: key path"] + lines += [f"{key} {path}" for key, path in expected.items()] + SNAPSHOT.write_text("\n".join(lines) + "\n") + print(f"Wrote {SNAPSHOT} ({len(expected)} routes)") + return 0 - print(f"OK: {len(client_keys)} route keys match contract") + client = client_routes() + missing = sorted(set(expected) - set(client)) + changed = sorted( + key for key in set(expected) & set(client) if expected[key] != client[key] + ) + extra = sorted(set(client) - set(expected)) + for key in missing: + print(f"Missing in client ROUTES: {key} {expected[key]}", file=sys.stderr) + for key in changed: + print( + f"Path mismatch for {key}: backend {expected[key]} != client {client[key]}", + file=sys.stderr, + ) + if missing or changed: + return 1 + print( + f"OK: {len(expected)} backend routes match; " + f"{len(extra)} client-only routes ({', '.join(extra)})" + ) return 0 diff --git a/src/cognichem_client/__init__.py b/src/cognichem_client/__init__.py index c9570a5..c38fbdd 100644 --- a/src/cognichem_client/__init__.py +++ b/src/cognichem_client/__init__.py @@ -1,7 +1,8 @@ """Public package entry for the CogniChem Python client. Re-exports the sync and async clients, common response types, catalog -constants, and exception classes used by application code. +constants, and exception classes used by application code. Every response +model is also importable from :mod:`cognichem_client.types`. """ from cognichem_client.async_client import AsyncCogniChem @@ -12,48 +13,68 @@ MODEL_TYPES, RESOURCES, UTILITY_TYPES, + WORKFLOW_RUN_STATUSES, + WORKFLOW_RUN_STOP_STATUSES, ) from cognichem_client.errors import ( AuthenticationError, + BadRequestError, CogniChemError, ConflictError, ForbiddenError, + GoneError, NotFoundError, + PayloadTooLargeError, + PaymentRequiredError, PollTimeoutError, ProcessCancelledError, ProcessFailedError, RateLimitError, + ServerError, ValidationError, ) +from cognichem_client.resources.artifacts import artifact_ref from cognichem_client.types import ( + Artifact, BinaryResult, DataResult, JobSubmitRequest, ProcessStatus, + WorkflowRun, ) __all__ = [ + "Artifact", "AsyncCogniChem", "AuthenticationError", + "BadRequestError", "BinaryResult", "CogniChem", "CogniChemError", "ConflictError", "DataResult", "ForbiddenError", + "GoneError", "JOB_TERMINAL_STATUSES", "JOB_TYPES", "JobSubmitRequest", "MODEL_TYPES", "NotFoundError", + "PayloadTooLargeError", + "PaymentRequiredError", "PollTimeoutError", "ProcessCancelledError", "ProcessFailedError", "ProcessStatus", "RESOURCES", "RateLimitError", + "ServerError", "UTILITY_TYPES", "ValidationError", + "WORKFLOW_RUN_STATUSES", + "WORKFLOW_RUN_STOP_STATUSES", + "WorkflowRun", + "artifact_ref", ] -__version__ = "1.0.0" +__version__ = "1.1.0" diff --git a/src/cognichem_client/_http.py b/src/cognichem_client/_http.py index 9874154..cba94a6 100644 --- a/src/cognichem_client/_http.py +++ b/src/cognichem_client/_http.py @@ -2,7 +2,10 @@ from __future__ import annotations +import json as _json import re +import uuid +from collections.abc import AsyncIterator, Iterable, Iterator from typing import Any from urllib.parse import urljoin @@ -10,7 +13,7 @@ from cognichem_client.constants import DEFAULT_BASE_URL, DEFAULT_HTTP_TIMEOUT from cognichem_client.errors import raise_for_problem -from cognichem_client.types import BinaryResult +from cognichem_client.types import BinaryResult, ServerSentEvent _FILENAME_RE = re.compile(r'filename\*?=(?:UTF-8\'\')?"?([^";]+)"?', re.IGNORECASE) @@ -36,6 +39,92 @@ def _parse_filename(content_disposition: str | None) -> str | None: return match.group(1).strip() +def ensure_idempotency_key(key: str | None = None) -> str: + """Return ``key``, or a fresh random ``Idempotency-Key`` when omitted. + + API-key callers must send ``Idempotency-Key`` on every mutating POST + (400 ``missing-idempotency-key`` otherwise), so resources default to a + new UUID4 per call. Pass your own key to make retries safe. + + Parameters + ---------- + key : str or None, optional + Caller-supplied key. + + Returns + ------- + str + ``key`` when given, otherwise a new UUID4 string. + """ + return key or str(uuid.uuid4()) + + +class _SseParser: + """Incremental ``text/event-stream`` parser (``event:`` / ``data:`` frames).""" + + def __init__(self) -> None: + """Start with an empty frame.""" + self._event = "message" + self._data: list[str] = [] + + def feed(self, line: str) -> ServerSentEvent | None: + """Consume one line and return an event when a frame completes. + + Parameters + ---------- + line : str + One line of the stream without its line terminator. + + Returns + ------- + ServerSentEvent or None + The completed event on a blank line, otherwise ``None``. + """ + if not line: + if not self._data: + self._event = "message" + return None + raw = "\n".join(self._data) + try: + data: Any = _json.loads(raw) + except ValueError: + data = raw + event = ServerSentEvent(event=self._event, data=data) + self._event = "message" + self._data = [] + return event + if line.startswith(":"): + return None + field, _, value = line.partition(":") + value = value[1:] if value.startswith(" ") else value + if field == "event": + self._event = value + elif field == "data": + self._data.append(value) + return None + + def iter_events(self, lines: Iterable[str]) -> Iterator[ServerSentEvent]: + """Yield events parsed from ``lines`` (flushes a trailing frame). + + Parameters + ---------- + lines : iterable of str + Stream lines without terminators. + + Yields + ------ + ServerSentEvent + Parsed events in stream order. + """ + for line in lines: + event = self.feed(line) + if event is not None: + yield event + event = self.feed("") + if event is not None: + yield event + + class _BaseHttp: """Shared configuration and header helpers for HTTP clients. @@ -300,6 +389,7 @@ def request( params: dict[str, Any] | None = None, json: Any = None, data: Any = None, + files: Any = None, idempotency_key: str | None = None, content_type: str | None = "application/json", prefer_bearer: bool = False, @@ -319,6 +409,9 @@ def request( JSON body (encoded by httpx). data : any, optional Form or raw body (e.g. OAuth2 password form). + files : any, optional + Multipart files (httpx ``files=``). Pass ``content_type=None`` so + httpx sets the multipart boundary. idempotency_key : str or None, optional Optional ``Idempotency-Key`` header. content_type : str or None, optional @@ -345,6 +438,7 @@ def request( params=params, json=json, data=data, + files=files, headers=self._headers( idempotency_key=idempotency_key, content_type=content_type, @@ -358,6 +452,59 @@ def request( return response.json() return self._binary_result(response) + def stream_events( + self, + method: str, + path: str, + *, + json: Any = None, + idempotency_key: str | None = None, + prefer_bearer: bool = False, + timeout: float | None = None, + ) -> Iterator[ServerSentEvent]: + """Send a request and yield its ``text/event-stream`` events. + + Parameters + ---------- + method : str + HTTP method. + path : str + API path under the configured base URL. + json : any, optional + JSON body. + idempotency_key : str or None, optional + Optional ``Idempotency-Key`` header. + prefer_bearer : bool, optional + Prefer Bearer auth over API key when both are available. + timeout : float or None, optional + Read timeout for the stream; defaults to the client timeout. + + Yields + ------ + ServerSentEvent + Parsed events in stream order. + + Raises + ------ + CogniChemError + When the response status is not a success (before any event). + """ + headers = self._headers( + idempotency_key=idempotency_key, prefer_bearer=prefer_bearer + ) + headers["Accept"] = "text/event-stream" + with self._client.stream( + method, + self.url(path), + json=json, + headers=headers, + timeout=timeout if timeout is not None else self.timeout, + ) as response: + if not response.is_success: + response.read() + self._raise_if_error(response) + yield from _SseParser().iter_events(response.iter_lines()) + class AsyncHttpClient(_BaseHttp): """Asynchronous httpx transport used by async resource classes. @@ -447,6 +594,7 @@ async def request( params: dict[str, Any] | None = None, json: Any = None, data: Any = None, + files: Any = None, idempotency_key: str | None = None, content_type: str | None = "application/json", prefer_bearer: bool = False, @@ -466,6 +614,9 @@ async def request( JSON body (encoded by httpx). data : any, optional Form or raw body. + files : any, optional + Multipart files (httpx ``files=``). Pass ``content_type=None`` so + httpx sets the multipart boundary. idempotency_key : str or None, optional Optional ``Idempotency-Key`` header. content_type : str or None, optional @@ -492,6 +643,7 @@ async def request( params=params, json=json, data=data, + files=files, headers=self._headers( idempotency_key=idempotency_key, content_type=content_type, @@ -504,3 +656,63 @@ async def request( if expect_json: return response.json() return self._binary_result(response) + + async def stream_events( + self, + method: str, + path: str, + *, + json: Any = None, + idempotency_key: str | None = None, + prefer_bearer: bool = False, + timeout: float | None = None, + ) -> AsyncIterator[ServerSentEvent]: + """Send a request and yield its ``text/event-stream`` events. + + Parameters + ---------- + method : str + HTTP method. + path : str + API path under the configured base URL. + json : any, optional + JSON body. + idempotency_key : str or None, optional + Optional ``Idempotency-Key`` header. + prefer_bearer : bool, optional + Prefer Bearer auth over API key when both are available. + timeout : float or None, optional + Read timeout for the stream; defaults to the client timeout. + + Yields + ------ + ServerSentEvent + Parsed events in stream order. + + Raises + ------ + CogniChemError + When the response status is not a success (before any event). + """ + headers = self._headers( + idempotency_key=idempotency_key, prefer_bearer=prefer_bearer + ) + headers["Accept"] = "text/event-stream" + async with self._client.stream( + method, + self.url(path), + json=json, + headers=headers, + timeout=timeout if timeout is not None else self.timeout, + ) as response: + if not response.is_success: + await response.aread() + self._raise_if_error(response) + parser = _SseParser() + async for line in response.aiter_lines(): + event = parser.feed(line) + if event is not None: + yield event + event = parser.feed("") + if event is not None: + yield event diff --git a/src/cognichem_client/async_client.py b/src/cognichem_client/async_client.py index 547b3ab..8da9b67 100644 --- a/src/cognichem_client/async_client.py +++ b/src/cognichem_client/async_client.py @@ -5,13 +5,21 @@ import os from cognichem_client._http import AsyncHttpClient -from cognichem_client.constants import DEFAULT_BASE_URL, DEFAULT_HTTP_TIMEOUT +from cognichem_client.constants import DEFAULT_BASE_URL, DEFAULT_HTTP_TIMEOUT, ROUTES from cognichem_client.resources.api_keys import AsyncApiKeysResource +from cognichem_client.resources.artifacts import AsyncArtifactsResource from cognichem_client.resources.auth import AsyncAuthResource +from cognichem_client.resources.chat import AsyncChatResource +from cognichem_client.resources.events import AsyncEventsResource from cognichem_client.resources.inference import AsyncInferenceResource from cognichem_client.resources.jobs import AsyncJobsResource +from cognichem_client.resources.lookup import AsyncLookupResource +from cognichem_client.resources.reference import AsyncReferenceResource from cognichem_client.resources.usage import AsyncUsageResource from cognichem_client.resources.utils import AsyncUtilsResource +from cognichem_client.resources.wallet import AsyncWalletResource +from cognichem_client.resources.workflows import AsyncWorkflowsResource +from cognichem_client.types import HealthResponse class AsyncCogniChem: @@ -43,6 +51,22 @@ class AsyncCogniChem: Short-running utility endpoints. usage : AsyncUsageResource Usage limit queries. + workflows : AsyncWorkflowsResource + Workflow validate / estimate / templates, plus ``runs`` and + ``definitions``. + artifacts : AsyncArtifactsResource + Durable job / workflow artifacts (list, inspect, upload, download). + events : AsyncEventsResource + Job and workflow run completion events feed. + wallet : AsyncWalletResource + Wallet balance. + lookup : AsyncLookupResource + Science database lookups (PubChem, ChEMBL, Open Targets, UniProt, + RCSB PDB). + reference : AsyncReferenceResource + Reference-KG literature search. + chat : AsyncChatResource + CogniChem Assistant chat (JWT only). """ def __init__( @@ -78,6 +102,13 @@ def __init__( self.inference = AsyncInferenceResource(self._http) self.utils = AsyncUtilsResource(self._http) self.usage = AsyncUsageResource(self._http) + self.workflows = AsyncWorkflowsResource(self._http) + self.artifacts = AsyncArtifactsResource(self._http) + self.events = AsyncEventsResource(self._http) + self.wallet = AsyncWalletResource(self._http) + self.lookup = AsyncLookupResource(self._http) + self.reference = AsyncReferenceResource(self._http) + self.chat = AsyncChatResource(self._http) @classmethod def from_env( @@ -118,6 +149,33 @@ def from_env( timeout=timeout, ) + async def health(self) -> HealthResponse: + """Check API liveness (``GET /health``, no auth needed). + + Returns + ------- + HealthResponse + ``status`` is ``ok`` when the API is up. + """ + data = await self._http.request("GET", ROUTES["health"]) + return HealthResponse.model_validate(data) + + async def ready(self) -> HealthResponse: + """Check API readiness (``GET /ready``, no auth needed). + + Returns + ------- + HealthResponse + ``status`` is ``ready`` when dependencies are reachable. + + Raises + ------ + ServerError + When the API is not ready (503). + """ + data = await self._http.request("GET", ROUTES["ready"]) + return HealthResponse.model_validate(data) + async def aclose(self) -> None: """Close the underlying async HTTP client and free connections. diff --git a/src/cognichem_client/client.py b/src/cognichem_client/client.py index 7713b26..311e699 100644 --- a/src/cognichem_client/client.py +++ b/src/cognichem_client/client.py @@ -5,13 +5,21 @@ import os from cognichem_client._http import HttpClient -from cognichem_client.constants import DEFAULT_BASE_URL, DEFAULT_HTTP_TIMEOUT +from cognichem_client.constants import DEFAULT_BASE_URL, DEFAULT_HTTP_TIMEOUT, ROUTES from cognichem_client.resources.api_keys import ApiKeysResource +from cognichem_client.resources.artifacts import ArtifactsResource from cognichem_client.resources.auth import AuthResource +from cognichem_client.resources.chat import ChatResource +from cognichem_client.resources.events import EventsResource from cognichem_client.resources.inference import InferenceResource from cognichem_client.resources.jobs import JobsResource +from cognichem_client.resources.lookup import LookupResource +from cognichem_client.resources.reference import ReferenceResource from cognichem_client.resources.usage import UsageResource from cognichem_client.resources.utils import UtilsResource +from cognichem_client.resources.wallet import WalletResource +from cognichem_client.resources.workflows import WorkflowsResource +from cognichem_client.types import HealthResponse class CogniChem: @@ -44,6 +52,22 @@ class CogniChem: Short-running utility endpoints. usage : UsageResource Usage limit queries. + workflows : WorkflowsResource + Workflow validate / estimate / templates, plus ``runs`` and + ``definitions``. + artifacts : ArtifactsResource + Durable job / workflow artifacts (list, inspect, upload, download). + events : EventsResource + Job and workflow run completion events feed. + wallet : WalletResource + Wallet balance. + lookup : LookupResource + Science database lookups (PubChem, ChEMBL, Open Targets, UniProt, + RCSB PDB). + reference : ReferenceResource + Reference-KG literature search. + chat : ChatResource + CogniChem Assistant chat (JWT only). """ def __init__( @@ -79,6 +103,13 @@ def __init__( self.inference = InferenceResource(self._http) self.utils = UtilsResource(self._http) self.usage = UsageResource(self._http) + self.workflows = WorkflowsResource(self._http) + self.artifacts = ArtifactsResource(self._http) + self.events = EventsResource(self._http) + self.wallet = WalletResource(self._http) + self.lookup = LookupResource(self._http) + self.reference = ReferenceResource(self._http) + self.chat = ChatResource(self._http) @classmethod def from_env( @@ -119,6 +150,33 @@ def from_env( timeout=timeout, ) + def health(self) -> HealthResponse: + """Check API liveness (``GET /health``, no auth needed). + + Returns + ------- + HealthResponse + ``status`` is ``ok`` when the API is up. + """ + data = self._http.request("GET", ROUTES["health"]) + return HealthResponse.model_validate(data) + + def ready(self) -> HealthResponse: + """Check API readiness (``GET /ready``, no auth needed). + + Returns + ------- + HealthResponse + ``status`` is ``ready`` when dependencies are reachable. + + Raises + ------ + ServerError + When the API is not ready (503). + """ + data = self._http.request("GET", ROUTES["ready"]) + return HealthResponse.model_validate(data) + def close(self) -> None: """Close the underlying HTTP client and free connections. diff --git a/src/cognichem_client/constants.py b/src/cognichem_client/constants.py index 1bca0e9..4e47cc2 100644 --- a/src/cognichem_client/constants.py +++ b/src/cognichem_client/constants.py @@ -7,9 +7,11 @@ API_V1_STR : str Versioned API path prefix (``/api/v1``). ROUTES : dict of str to str - Named endpoint path map mirroring back-end ``Settings.ROUTE_MAP``. + Named endpoint path map mirroring back-end ``Settings.ROUTE_MAP``, plus + routes that map does not list. JobType : typing.Literal - Allowed long-running job type string literals. + Allowed long-running job type string literals (catalog job types that can + be submitted on their own). JOB_TYPES : frozenset of str Runtime set of supported job type strings. UtilityType : typing.Literal @@ -24,6 +26,23 @@ Allowed compute resource tier string literals. RESOURCES : frozenset of str Runtime set of supported resource tier strings. +WorkflowRunStatus : typing.Literal + Workflow run status string literals. +WORKFLOW_RUN_STATUSES : frozenset of str + Runtime set of workflow run statuses. +WORKFLOW_RUN_STOP_STATUSES : frozenset of str + Statuses where :meth:`WorkflowRunsResource.wait` stops polling + (``completed``, ``failed``, ``cancelled``, ``paused``). +ReferencePaperSet : typing.Literal + Reference-KG paper sets (``method``, ``used_cognichem``, ``general``). +ReferenceEdgeKind : typing.Literal + Reference-KG edge kinds. +EventKind : typing.Literal + Completion event subject kinds (``job``, ``workflow_run``). +ReasoningEffort : typing.Literal + CogniChem Assistant reasoning levels. +DEFAULT_ASSISTANT_MODEL_ID : str + Catalog model id of the CogniChem Assistant. JOB_TERMINAL_STATUSES : frozenset of str Terminal statuses for long-running jobs (``completed``, ``error``, ``cancelled``). @@ -37,6 +56,12 @@ Default seconds between inference/utility status polls. DEFAULT_FAST_TIMEOUT : float Default maximum seconds to wait for inference/utility completion. +DEFAULT_WORKFLOW_POLL_INTERVAL : float + Default seconds between workflow run polls. +DEFAULT_WORKFLOW_TIMEOUT : float + Default maximum seconds to wait for a workflow run to stop. +MAX_EVENTS_WAIT : float + Longest ``wait`` (seconds) ``GET /events`` accepts. DEFAULT_HTTP_TIMEOUT : float Default per-request HTTP timeout in seconds. """ @@ -48,8 +73,12 @@ DEFAULT_BASE_URL: Final = "https://api.cognichem.com" API_V1_STR: Final = "/api/v1" -# Mirrors apps/back-end-api Settings.ROUTE_MAP +# Mirrors apps/back-end-api Settings.ROUTE_MAP (plus routes it does not list: +# health/ready, templates, node-ports, definition share/versions, lookup, +# reference, and wallet). ROUTES: Final[dict[str, str]] = { + "health": "/health", + "ready": "/ready", "auth_login_email": f"{API_V1_STR}/auth/login-email", "auth_check": f"{API_V1_STR}/auth/check", "auth_refresh": f"{API_V1_STR}/auth/refresh", @@ -59,6 +88,7 @@ "auth_api_keys_rotate": f"{API_V1_STR}/auth/api-keys/{{key_id}}/rotate", "jobs_submit": f"{API_V1_STR}/jobs/submit", "jobs_submit_multiple": f"{API_V1_STR}/jobs/submit-multiple", + "jobs_estimate": f"{API_V1_STR}/jobs/estimate", "jobs_info": f"{API_V1_STR}/jobs/info", "jobs_list": f"{API_V1_STR}/jobs/list", "jobs_status": f"{API_V1_STR}/jobs/status", @@ -66,12 +96,64 @@ "jobs_result_multiple": f"{API_V1_STR}/jobs/result-multiple", "jobs_cancel": f"{API_V1_STR}/jobs/cancel", "jobs_delete": f"{API_V1_STR}/jobs/delete", + "artifacts_list": f"{API_V1_STR}/artifacts", + "artifacts_upload": f"{API_V1_STR}/artifacts", + "artifacts_usage": f"{API_V1_STR}/artifacts/usage", + "artifacts_get": f"{API_V1_STR}/artifacts/{{artifact_id}}", + "artifacts_download": f"{API_V1_STR}/artifacts/{{artifact_id}}/download", + "artifacts_download_record": f"{API_V1_STR}/artifacts/{{artifact_id}}/download", + "artifacts_delete": f"{API_V1_STR}/artifacts/{{artifact_id}}", + "workflows_validate": f"{API_V1_STR}/workflows/validate", + "workflows_estimate": f"{API_V1_STR}/workflows/estimate", + "workflows_templates": f"{API_V1_STR}/workflows/templates", + "workflows_template": f"{API_V1_STR}/workflows/templates/{{template_id}}", + "workflows_node_ports": f"{API_V1_STR}/workflows/node-ports", + "workflows_definitions_create": f"{API_V1_STR}/workflows/definitions", + "workflows_definitions_list": f"{API_V1_STR}/workflows/definitions", + "workflows_definitions_get": ( + f"{API_V1_STR}/workflows/definitions/{{definition_id}}" + ), + "workflows_definitions_update": ( + f"{API_V1_STR}/workflows/definitions/{{definition_id}}" + ), + "workflows_definitions_delete": ( + f"{API_V1_STR}/workflows/definitions/{{definition_id}}" + ), + "workflows_definitions_download": ( + f"{API_V1_STR}/workflows/definitions/{{definition_id}}/download" + ), + "workflows_definitions_share": ( + f"{API_V1_STR}/workflows/definitions/{{definition_id}}/share" + ), + "workflows_definitions_versions": ( + f"{API_V1_STR}/workflows/definitions/{{definition_id}}/versions" + ), + "workflows_definitions_version": ( + f"{API_V1_STR}/workflows/definitions/{{definition_id}}/versions/{{version}}" + ), + "workflows_definitions_version_restore": ( + f"{API_V1_STR}/workflows/definitions/{{definition_id}}/versions/{{version}}/restore" + ), + "workflows_definitions_shared": ( + f"{API_V1_STR}/workflows/definitions/shared/{{token}}" + ), + "workflows_definitions_shared_fork": ( + f"{API_V1_STR}/workflows/definitions/shared/{{token}}/fork" + ), + "workflows_runs_create": f"{API_V1_STR}/workflows/runs", + "workflows_runs_list": f"{API_V1_STR}/workflows/runs", + "workflows_runs_get": f"{API_V1_STR}/workflows/runs/{{run_id}}", + "workflows_runs_artifacts": f"{API_V1_STR}/workflows/runs/{{run_id}}/artifacts", + "workflows_runs_resume": f"{API_V1_STR}/workflows/runs/{{run_id}}/resume", + "workflows_runs_cancel": f"{API_V1_STR}/workflows/runs/{{run_id}}/cancel", + "workflows_runs_delete": f"{API_V1_STR}/workflows/runs/{{run_id}}", "inference_submit": f"{API_V1_STR}/inference/submit", "inference_submit_batch": f"{API_V1_STR}/inference/submit-batch", "inference_list": f"{API_V1_STR}/inference/list", "inference_status": f"{API_V1_STR}/inference/status", "inference_result": f"{API_V1_STR}/inference/result", "inference_delete": f"{API_V1_STR}/inference/delete", + "inference_cancel": f"{API_V1_STR}/inference/cancel", "mpnn_user_models": f"{API_V1_STR}/user-models/mpnn", "mpnn_public_models": f"{API_V1_STR}/public-models/mpnn", "utils_submit": f"{API_V1_STR}/utils/submit", @@ -79,65 +161,160 @@ "utils_status": f"{API_V1_STR}/utils/status", "utils_result": f"{API_V1_STR}/utils/result", "utils_delete": f"{API_V1_STR}/utils/delete", + "utils_cancel": f"{API_V1_STR}/utils/cancel", "usage_limits": f"{API_V1_STR}/usage-limits", + "wallet": f"{API_V1_STR}/wallet", + "events": f"{API_V1_STR}/events", + "lookup_pubchem": f"{API_V1_STR}/lookup/pubchem", + "lookup_chembl": f"{API_V1_STR}/lookup/chembl", + "lookup_targets": f"{API_V1_STR}/lookup/targets", + "lookup_uniprot": f"{API_V1_STR}/lookup/uniprot", + "lookup_pdb": f"{API_V1_STR}/lookup/pdb", + "lookup_properties": f"{API_V1_STR}/lookup/properties", + "reference_papers": f"{API_V1_STR}/reference/papers", + "reference_paper": f"{API_V1_STR}/reference/papers/{{paper_id}}", + "reference_neighborhood": f"{API_V1_STR}/reference/neighborhood", + "chat_sessions": f"{API_V1_STR}/chat/sessions", + "chat_session": f"{API_V1_STR}/chat/sessions/{{session_id}}", + "chat_messages": f"{API_V1_STR}/chat/sessions/{{session_id}}/messages", + "chat_estimate": f"{API_V1_STR}/chat/sessions/{{session_id}}/estimate", + "chat_turns": f"{API_V1_STR}/chat/sessions/{{session_id}}/turns", + "chat_turn_stop": ( + f"{API_V1_STR}/chat/sessions/{{session_id}}/turns/{{turn_id}}/stop" + ), + "chat_spend_cap": f"{API_V1_STR}/chat/sessions/{{session_id}}/spend-cap", + "chat_spend_cap_continue": ( + f"{API_V1_STR}/chat/sessions/{{session_id}}/spend-cap/continue" + ), + "chat_proposals": f"{API_V1_STR}/chat/sessions/{{session_id}}/proposals", + "chat_proposal": ( + f"{API_V1_STR}/chat/sessions/{{session_id}}/proposals/{{proposal_id}}" + ), + "chat_proposal_approve": ( + f"{API_V1_STR}/chat/sessions/{{session_id}}/proposals/{{proposal_id}}/approve" + ), + "chat_proposal_reject": ( + f"{API_V1_STR}/chat/sessions/{{session_id}}/proposals/{{proposal_id}}/reject" + ), } +# Mirrors packages/job-catalog JOB_TYPES minus workflow-only helpers +# (``workflow-adapter`` / ``workflow-filter``). JobType = Literal[ + "admet-predict", "autodockvina", "boltz2", "boltzgen", "cantera", + "conformer-ensemble", "convert-batch", + "coolprop", "diffdock", "drugflow", + "druglikeness-filter", "esmfold2", "evo2", + "fingerprint-cluster", + "freebindcraft", "freewilson", "gnina", + "group-contribution", + "md-analyze", + "mm-gbsa", + "mmp-analysis", "molecule-rmsd-matrix", + "molecule-standardize", + "nasa-thermo-fit", "opendde", "openmm-md", "padel-descriptor", + "phreeqc", + "plip-profile", + "pocket-detect", + "protein-prepare", + "reaction-enumerate", + "retrosynthesis", "rfantibody-finetune", "rfantibody-pipeline", "rfantibody-proteinmpnn", "rfantibody-rfdiffusion", "rfantibody-tcr-predict", + "rmg-mechanism", + "scaffold-analyze", + "surfdock", "surfmap", "thompsonsampling", "train-mpnn", + "vle-flash", ] JOB_TYPES: Final[frozenset[str]] = frozenset( { + "admet-predict", "autodockvina", "boltz2", "boltzgen", "cantera", + "conformer-ensemble", "convert-batch", + "coolprop", "diffdock", "drugflow", + "druglikeness-filter", "esmfold2", "evo2", + "fingerprint-cluster", + "freebindcraft", "freewilson", "gnina", + "group-contribution", + "md-analyze", + "mm-gbsa", + "mmp-analysis", "molecule-rmsd-matrix", + "molecule-standardize", + "nasa-thermo-fit", "opendde", "openmm-md", "padel-descriptor", + "phreeqc", + "plip-profile", + "pocket-detect", + "protein-prepare", + "reaction-enumerate", + "retrosynthesis", "rfantibody-finetune", "rfantibody-pipeline", "rfantibody-proteinmpnn", "rfantibody-rfdiffusion", "rfantibody-tcr-predict", + "rmg-mechanism", + "scaffold-analyze", + "surfdock", "surfmap", "thompsonsampling", "train-mpnn", + "vle-flash", } ) -UtilityType = Literal["convert", "molecule_rmsd"] -UTILITY_TYPES: Final[frozenset[str]] = frozenset({"convert", "molecule_rmsd"}) +UtilityType = Literal[ + "convert", + "molecule_rmsd", + "molecule_standardize", + "druglikeness_filter", + "coolprop", +] + +UTILITY_TYPES: Final[frozenset[str]] = frozenset( + { + "convert", + "molecule_rmsd", + "molecule_standardize", + "druglikeness_filter", + "coolprop", + } +) ModelType = Literal["mpnn"] MODEL_TYPES: Final[frozenset[str]] = frozenset({"mpnn"}) @@ -181,8 +358,30 @@ {"completed", "error", "cancelled"} ) +WorkflowRunStatus = Literal[ + "pending", "running", "paused", "completed", "failed", "cancelled" +] +WORKFLOW_RUN_STATUSES: Final[frozenset[str]] = frozenset( + {"pending", "running", "paused", "completed", "failed", "cancelled"} +) +# ``paused`` needs a user action (resume or a higher ``max_run_cost``), so +# waiting stops there too (mirrors cognichem-mcp WORKFLOW_RUN_STOP_STATUSES). +WORKFLOW_RUN_STOP_STATUSES: Final[frozenset[str]] = frozenset( + {"completed", "failed", "cancelled", "paused"} +) + +ReferencePaperSet = Literal["method", "used_cognichem", "general"] +ReferenceEdgeKind = Literal["cites", "method_of", "uses_tool", "compares", "describes"] +EventKind = Literal["job", "workflow_run"] +ReasoningEffort = Literal["low", "medium", "high"] + +DEFAULT_ASSISTANT_MODEL_ID: Final = "cognichem-assistant" + DEFAULT_JOB_POLL_INTERVAL: Final = 5.0 DEFAULT_JOB_TIMEOUT: Final = 3600.0 DEFAULT_FAST_POLL_INTERVAL: Final = 0.3 DEFAULT_FAST_TIMEOUT: Final = 300.0 +DEFAULT_WORKFLOW_POLL_INTERVAL: Final = 10.0 +DEFAULT_WORKFLOW_TIMEOUT: Final = 6 * 3600.0 +MAX_EVENTS_WAIT: Final = 25.0 DEFAULT_HTTP_TIMEOUT: Final = 60.0 diff --git a/src/cognichem_client/errors.py b/src/cognichem_client/errors.py index 6403e06..86daac2 100644 --- a/src/cognichem_client/errors.py +++ b/src/cognichem_client/errors.py @@ -28,6 +28,9 @@ class CogniChemError(Exception): Stable client-side error code (e.g. ``api-key-limit-exceeded``). body : any, optional Raw parsed response body for debugging. + retryability : str or None, optional + Server retry class: ``validation``, ``auth``, ``quota``, ``conflict``, + ``upstream``, ``timeout``, or ``terminal``. Attributes ---------- @@ -48,7 +51,11 @@ class CogniChemError(Exception): code : str Stable client-side error code. body : any - Raw response body, if any. + Raw response body, if any. Problem extensions (for example + ``cap_usd`` on ``session-spend-cap``, or ``proposal`` on + ``estimate-changed``) live here. + retryability : str or None + Server retry class, if any. """ def __init__( @@ -63,6 +70,7 @@ def __init__( errors: list[Any] | None = None, code: str = "unknown", body: Any = None, + retryability: str | None = None, ) -> None: """Store problem-detail fields on the exception instance. @@ -86,6 +94,8 @@ def __init__( Stable client-side error code. body : any, optional Raw parsed response body. + retryability : str or None, optional + Server retry class. """ super().__init__(message) self.message = message @@ -97,6 +107,18 @@ def __init__( self.errors = errors self.code = code self.body = body + self.retryability = retryability + + +class BadRequestError(CogniChemError): + """Raised for HTTP 400 Bad Request responses. + + Notes + ----- + Inherits attributes from :class:`CogniChemError`. Operational failures + such as ``missing-idempotency-key`` or ``insufficient_wallet`` on + workflow run create arrive as 400. + """ class AuthenticationError(CogniChemError): @@ -108,6 +130,16 @@ class AuthenticationError(CogniChemError): """ +class PaymentRequiredError(CogniChemError): + """Raised for HTTP 402 Payment Required responses. + + Notes + ----- + The wallet cannot cover the hold (``insufficient-wallet``). Inherits + attributes from :class:`CogniChemError`. + """ + + class ForbiddenError(CogniChemError): """Raised for HTTP 403 Forbidden responses. @@ -136,6 +168,28 @@ class ConflictError(CogniChemError): """ +class GoneError(CogniChemError): + """Raised for HTTP 410 Gone responses. + + Notes + ----- + Examples: an expired Assistant proposal (``proposal-expired``) or a + ``?record=`` download whose parent artifact expired. Inherits attributes + from :class:`CogniChemError`. + """ + + +class PayloadTooLargeError(CogniChemError): + """Raised for HTTP 413 Content Too Large responses. + + Notes + ----- + Examples: an upload over the size cap or storage quota, or a WorkflowSpec + over the byte cap (Diagnostic list on ``errors``). Inherits attributes + from :class:`CogniChemError`. + """ + + class ValidationError(CogniChemError): """Raised for HTTP 422 Unprocessable Entity responses. @@ -155,6 +209,16 @@ class RateLimitError(CogniChemError): """ +class ServerError(CogniChemError): + """Raised for HTTP 5xx responses. + + Notes + ----- + Usually safe to retry (``retryability`` is ``upstream`` or ``timeout``). + Inherits attributes from :class:`CogniChemError`. + """ + + class PollTimeoutError(CogniChemError): """Raised when a polling wait exceeds the configured timeout. @@ -167,6 +231,8 @@ class PollTimeoutError(CogniChemError): class ProcessFailedError(CogniChemError): """Raised when a job, inference, or utility ends in ``error`` status. + Also raised when a workflow run ends in ``failed`` status. + Notes ----- Inherits attributes from :class:`CogniChemError`. @@ -174,7 +240,7 @@ class ProcessFailedError(CogniChemError): class ProcessCancelledError(CogniChemError): - """Raised when a job ends in ``cancelled`` status. + """Raised when a job or workflow run ends in ``cancelled`` status. Notes ----- @@ -182,7 +248,7 @@ class ProcessCancelledError(CogniChemError): """ -_API_KEY_LIMIT_SUFFIX = "/problems/api-key-limit-exceeded" +_PROBLEMS_SEGMENT = "/problems/" _DUPLICATE_NAME_DETAIL = "An API key with this name already exists" @@ -250,15 +316,19 @@ def _error_code(status: int, body: dict[str, Any], message: str) -> str: Returns ------- str - Client error code such as ``api-key-limit-exceeded``, - ``duplicate-name``, or ``unknown``. + Client error code: ``duplicate-name``, the problem ``type`` slug + (e.g. ``api-key-limit-exceeded``, ``missing-idempotency-key``, + ``session-spend-cap``), or ``unknown`` for generic ``http-`` + problems. """ - raw_type = body.get("type") - problem_type = raw_type if isinstance(raw_type, str) else "" - if status == 403 and problem_type.endswith(_API_KEY_LIMIT_SUFFIX): - return "api-key-limit-exceeded" if status == 409 and message == _DUPLICATE_NAME_DETAIL: return "duplicate-name" + raw_type = body.get("type") + problem_type = raw_type if isinstance(raw_type, str) else "" + if _PROBLEMS_SEGMENT in problem_type: + slug = problem_type.rsplit(_PROBLEMS_SEGMENT, 1)[1].strip("/") + if slug and not slug.startswith("http-"): + return slug return "unknown" @@ -275,18 +345,28 @@ def _exception_class(status: int) -> type[CogniChemError]: type of CogniChemError Exception class to raise. """ + if status == 400: + return BadRequestError if status == 401: return AuthenticationError + if status == 402: + return PaymentRequiredError if status == 403: return ForbiddenError if status == 404: return NotFoundError if status == 409: return ConflictError + if status == 410: + return GoneError + if status == 413: + return PayloadTooLargeError if status == 422: return ValidationError if status == 429: return RateLimitError + if status >= 500: + return ServerError return CogniChemError @@ -324,6 +404,9 @@ def raise_for_problem( request_id = ( body.get("request_id") if isinstance(body.get("request_id"), str) else None ) + retryability = ( + body.get("retryability") if isinstance(body.get("retryability"), str) else None + ) errors = body.get("errors") if isinstance(body.get("errors"), list) else None if status == 422 and errors is None and isinstance(body.get("detail"), list): errors = body["detail"] @@ -339,4 +422,5 @@ def raise_for_problem( errors=errors, code=code, body=body, + retryability=retryability, ) diff --git a/src/cognichem_client/polling.py b/src/cognichem_client/polling.py index b7ff3cf..f4aa7db 100644 --- a/src/cognichem_client/polling.py +++ b/src/cognichem_client/polling.py @@ -1,10 +1,11 @@ -"""Shared polling helpers for long-running processes.""" +"""Shared polling helpers for long-running processes and workflow runs.""" from __future__ import annotations import asyncio import time from collections.abc import Awaitable, Callable +from typing import Protocol, TypeVar from cognichem_client.constants import POLLABLE_TERMINAL_STATUSES from cognichem_client.errors import ( @@ -15,6 +16,21 @@ from cognichem_client.types import ProcessStatus +class _HasStatus(Protocol): + """Anything with a string ``status`` (process status, workflow run).""" + + @property + def status(self) -> str: + """Current status string.""" + ... + + +S = TypeVar("S", bound=_HasStatus) + +_PROCESS_FAILED = frozenset({"error"}) +_RUN_FAILED = frozenset({"failed"}) + + def wait_for_terminal( status_fn: Callable[[], ProcessStatus], *, @@ -53,20 +69,14 @@ def wait_for_terminal( ProcessCancelledError If ``raise_on_failure`` is true and status is ``cancelled``. """ - deadline = time.monotonic() + timeout - while True: - status = status_fn() - if status.status in terminal_statuses: - if raise_on_failure: - _raise_if_failed(status) - return status - if time.monotonic() >= deadline: - raise PollTimeoutError( - f"Timed out after {timeout}s waiting for process " - f"{status.process_id} (last status={status.status!r})", - status=None, - ) - time.sleep(poll_interval) + return poll_until( + status_fn, + poll_interval=poll_interval, + timeout=timeout, + stop_statuses=terminal_statuses, + subject=lambda s: f"process {s.process_id}", + failed_statuses=_PROCESS_FAILED if raise_on_failure else None, + ) async def await_for_terminal( @@ -107,44 +117,193 @@ async def await_for_terminal( ProcessCancelledError If ``raise_on_failure`` is true and status is ``cancelled``. """ + return await apoll_until( + status_fn, + poll_interval=poll_interval, + timeout=timeout, + stop_statuses=terminal_statuses, + subject=lambda s: f"process {s.process_id}", + failed_statuses=_PROCESS_FAILED if raise_on_failure else None, + ) + + +def poll_until( + fetch: Callable[[], S], + *, + poll_interval: float, + timeout: float, + stop_statuses: frozenset[str], + subject: Callable[[S], str], + failed_statuses: frozenset[str] | None = None, +) -> S: + """Poll ``fetch`` until its ``status`` is in ``stop_statuses``. + + Parameters + ---------- + fetch : callable + Zero-argument callable returning an object with a ``status``. + poll_interval : float + Seconds to sleep between polls. + timeout : float + Maximum seconds to wait before raising :class:`PollTimeoutError`. + stop_statuses : frozenset of str + Statuses that end the wait. + subject : callable + Describes the polled object for error messages. + failed_statuses : frozenset of str or None, optional + When set, raise :class:`ProcessFailedError` for these statuses and + :class:`ProcessCancelledError` for ``cancelled``. + + Returns + ------- + object + The last fetched object (its status is in ``stop_statuses``). + + Raises + ------ + PollTimeoutError + If ``timeout`` elapses first. + ProcessFailedError + If ``failed_statuses`` is set and the final status is in it. + ProcessCancelledError + If ``failed_statuses`` is set and the final status is ``cancelled``. + """ + deadline = time.monotonic() + timeout + while True: + current = fetch() + if current.status in stop_statuses: + if failed_statuses is not None: + _raise_if_failed(current, subject(current), failed_statuses) + return current + if time.monotonic() >= deadline: + raise _timeout(timeout, subject(current), current.status) + time.sleep(poll_interval) + + +async def apoll_until( + fetch: Callable[[], Awaitable[S]], + *, + poll_interval: float, + timeout: float, + stop_statuses: frozenset[str], + subject: Callable[[S], str], + failed_statuses: frozenset[str] | None = None, +) -> S: + """Asynchronously poll ``fetch`` until its ``status`` is in ``stop_statuses``. + + Parameters + ---------- + fetch : callable + Zero-argument async callable returning an object with a ``status``. + poll_interval : float + Seconds to sleep between polls. + timeout : float + Maximum seconds to wait before raising :class:`PollTimeoutError`. + stop_statuses : frozenset of str + Statuses that end the wait. + subject : callable + Describes the polled object for error messages. + failed_statuses : frozenset of str or None, optional + When set, raise :class:`ProcessFailedError` for these statuses and + :class:`ProcessCancelledError` for ``cancelled``. + + Returns + ------- + object + The last fetched object (its status is in ``stop_statuses``). + + Raises + ------ + PollTimeoutError + If ``timeout`` elapses first. + ProcessFailedError + If ``failed_statuses`` is set and the final status is in it. + ProcessCancelledError + If ``failed_statuses`` is set and the final status is ``cancelled``. + """ deadline = time.monotonic() + timeout while True: - status = await status_fn() - if status.status in terminal_statuses: - if raise_on_failure: - _raise_if_failed(status) - return status + current = await fetch() + if current.status in stop_statuses: + if failed_statuses is not None: + _raise_if_failed(current, subject(current), failed_statuses) + return current if time.monotonic() >= deadline: - raise PollTimeoutError( - f"Timed out after {timeout}s waiting for process " - f"{status.process_id} (last status={status.status!r})", - status=None, - ) + raise _timeout(timeout, subject(current), current.status) await asyncio.sleep(poll_interval) -def _raise_if_failed(status: ProcessStatus) -> None: - """Raise for failed or cancelled terminal statuses. +def run_failure_statuses(raise_on_failure: bool) -> frozenset[str] | None: + """Return the workflow-run failure set when ``raise_on_failure`` is set. + + Parameters + ---------- + raise_on_failure : bool + Whether the caller wants failures raised. + + Returns + ------- + frozenset of str or None + ``{"failed"}`` or ``None``. + """ + return _RUN_FAILED if raise_on_failure else None + + +def _timeout(timeout: float, subject: str, status: str) -> PollTimeoutError: + """Build the timeout error raised by the poll loops. + + Parameters + ---------- + timeout : float + Configured timeout in seconds. + subject : str + Description of the polled object. + status : str + Last observed status. + + Returns + ------- + PollTimeoutError + Error to raise. + """ + return PollTimeoutError( + f"Timed out after {timeout}s waiting for {subject} (last status={status!r})", + status=None, + ) + + +def _raise_if_failed( + current: _HasStatus, subject: str, failed_statuses: frozenset[str] +) -> None: + """Raise for failed or cancelled final statuses. Parameters ---------- - status : ProcessStatus - Terminal status payload. + current : object + Final polled object. + subject : str + Description of the polled object. + failed_statuses : frozenset of str + Statuses that count as failure. Raises ------ ProcessFailedError - When ``status.status`` is ``error``. + When the status is in ``failed_statuses``. ProcessCancelledError - When ``status.status`` is ``cancelled``. + When the status is ``cancelled``. """ - if status.status == "error": + message = getattr(current, "message", None) or getattr( + current, "status_message", None + ) + label = subject[:1].upper() + subject[1:] + if current.status in failed_statuses: raise ProcessFailedError( - status.message or f"Process {status.process_id} failed", - detail=status.message, + message or f"{label} failed", + detail=message, ) - if status.status == "cancelled": + if current.status == "cancelled": raise ProcessCancelledError( - status.message or f"Process {status.process_id} was cancelled", - detail=status.message, + message or f"{label} was cancelled", + detail=message, ) diff --git a/src/cognichem_client/resources/__init__.py b/src/cognichem_client/resources/__init__.py index 675027b..de8223a 100644 --- a/src/cognichem_client/resources/__init__.py +++ b/src/cognichem_client/resources/__init__.py @@ -1,23 +1,38 @@ """API resource modules for sync CogniChem clients. Exposes resource classes used by :class:`~cognichem_client.client.CogniChem` -(auth, API keys, jobs, inference, utilities, and usage limits). Async +(auth, API keys, jobs, inference, utilities, usage limits, workflows, +artifacts, events, wallet, lookups, reference KG, and Assistant chat). Async counterparts live alongside these modules and are wired by :class:`~cognichem_client.async_client.AsyncCogniChem`. """ from cognichem_client.resources.api_keys import ApiKeysResource +from cognichem_client.resources.artifacts import ArtifactsResource from cognichem_client.resources.auth import AuthResource +from cognichem_client.resources.chat import ChatResource +from cognichem_client.resources.events import EventsResource from cognichem_client.resources.inference import InferenceResource from cognichem_client.resources.jobs import JobsResource +from cognichem_client.resources.lookup import LookupResource +from cognichem_client.resources.reference import ReferenceResource from cognichem_client.resources.usage import UsageResource from cognichem_client.resources.utils import UtilsResource +from cognichem_client.resources.wallet import WalletResource +from cognichem_client.resources.workflows import WorkflowsResource __all__ = [ "ApiKeysResource", + "ArtifactsResource", "AuthResource", + "ChatResource", + "EventsResource", "InferenceResource", "JobsResource", + "LookupResource", + "ReferenceResource", "UsageResource", "UtilsResource", + "WalletResource", + "WorkflowsResource", ] diff --git a/src/cognichem_client/resources/api_keys.py b/src/cognichem_client/resources/api_keys.py index f2dec6e..ff36b95 100644 --- a/src/cognichem_client/resources/api_keys.py +++ b/src/cognichem_client/resources/api_keys.py @@ -2,16 +2,93 @@ from __future__ import annotations +from collections.abc import Sequence +from datetime import datetime +from typing import Any + from cognichem_client._http import AsyncHttpClient, HttpClient from cognichem_client.constants import ROUTES from cognichem_client.types import ( ApiKeyCreatedResponse, ApiKeyListItem, ApiKeyListResponse, - MessageResponse, ) +def _create_body( + name: str, + *, + scopes: Sequence[str] | None, + expires_at: datetime | str | None, + spend_ceiling_usd: float | None, + allow_structure_search: bool, +) -> dict[str, Any]: + """Build the ``POST /auth/api-keys`` body, omitting unset options. + + Parameters + ---------- + name : str + Display name. + scopes : sequence of str or None + Scopes (subset of ``read``, ``write``). + expires_at : datetime or str or None + Expiry instant. + spend_ceiling_usd : float or None + Per-key USD cap. + allow_structure_search : bool + ChEMBL structure-search opt-in. + + Returns + ------- + dict + JSON body. + """ + body: dict[str, Any] = {"name": name} + if scopes is not None: + body["scopes"] = list(scopes) + if expires_at is not None: + body["expires_at"] = ( + expires_at.isoformat() if isinstance(expires_at, datetime) else expires_at + ) + if spend_ceiling_usd is not None: + body["spend_ceiling_usd"] = spend_ceiling_usd + if allow_structure_search: + body["allow_structure_search"] = True + return body + + +def _update_body( + name: str | None, allow_structure_search: bool | None +) -> dict[str, Any]: + """Build the ``PATCH /auth/api-keys/{key_id}`` body. + + Parameters + ---------- + name : str or None + New display name. + allow_structure_search : bool or None + New structure-search opt-in. + + Returns + ------- + dict + JSON body with only the fields to change. + + Raises + ------ + ValueError + If neither field is given. + """ + body: dict[str, Any] = {} + if name is not None: + body["name"] = name + if allow_structure_search is not None: + body["allow_structure_search"] = allow_structure_search + if not body: + raise ValueError("Provide name and/or allow_structure_search") + return body + + class ApiKeysResource: """Synchronous API key management under ``/auth/api-keys``. @@ -41,18 +118,34 @@ def list(self) -> ApiKeyListResponse: ApiKeyListResponse Keys, limit, and count. """ - data = self._http.request( - "GET", ROUTES["auth_api_keys"], prefer_bearer=True - ) + data = self._http.request("GET", ROUTES["auth_api_keys"], prefer_bearer=True) return ApiKeyListResponse.model_validate(data) - def create(self, name: str) -> ApiKeyCreatedResponse: + def create( + self, + name: str, + *, + scopes: Sequence[str] | None = None, + expires_at: datetime | str | None = None, + spend_ceiling_usd: float | None = None, + allow_structure_search: bool = False, + ) -> ApiKeyCreatedResponse: """Create a named API key. Parameters ---------- name : str - Display name (unique per user). + Display name (unique per user, 1-64 characters). + scopes : sequence of str or None, optional + Subset of ``read`` and ``write`` (default both). Submitting jobs + or workflow runs needs ``write``. + expires_at : datetime or str or None, optional + Instant after which the key is rejected. + spend_ceiling_usd : float or None, optional + Per-key USD spend cap (``None`` means wallet-only). + allow_structure_search : bool, optional + Let this key send SMILES to ChEMBL similarity / substructure + search via ``/lookup/chembl`` (off by default). Returns ------- @@ -62,7 +155,13 @@ def create(self, name: str) -> ApiKeyCreatedResponse: data = self._http.request( "POST", ROUTES["auth_api_keys"], - json={"name": name}, + json=_create_body( + name, + scopes=scopes, + expires_at=expires_at, + spend_ceiling_usd=spend_ceiling_usd, + allow_structure_search=allow_structure_search, + ), prefer_bearer=True, ) return ApiKeyCreatedResponse.model_validate(data) @@ -87,48 +186,78 @@ def get(self, key_id: str) -> ApiKeyListItem: ) return ApiKeyListItem.model_validate(data) - def rename(self, key_id: str, name: str) -> ApiKeyListItem: - """Rename an API key. + def update( + self, + key_id: str, + *, + name: str | None = None, + allow_structure_search: bool | None = None, + ) -> ApiKeyListItem: + """Rename an API key and/or change its structure-search opt-in. Parameters ---------- key_id : str Credential UUID. - name : str + name : str or None, optional New display name. + allow_structure_search : bool or None, optional + Let the key send SMILES to ChEMBL similarity / substructure + search. Returns ------- ApiKeyListItem Updated key metadata. + + Raises + ------ + ValueError + If neither ``name`` nor ``allow_structure_search`` is given. """ data = self._http.request( "PATCH", ROUTES["auth_api_keys_item"].format(key_id=key_id), - json={"name": name}, + json=_update_body(name, allow_structure_search), prefer_bearer=True, ) return ApiKeyListItem.model_validate(data) - def delete(self, key_id: str) -> MessageResponse: - """Revoke (hard-delete) an API key. + def rename(self, key_id: str, name: str) -> ApiKeyListItem: + """Rename an API key (shorthand for :meth:`update`). Parameters ---------- key_id : str Credential UUID. + name : str + New display name. Returns ------- - MessageResponse - Server confirmation message. + ApiKeyListItem + Updated key metadata. """ - data = self._http.request( + return self.update(key_id, name=name) + + def delete(self, key_id: str) -> None: + """Revoke an API key. + + Parameters + ---------- + key_id : str + Credential UUID. + + Returns + ------- + None + The API answers 204 No Content. + """ + self._http.request( "DELETE", ROUTES["auth_api_keys_item"].format(key_id=key_id), prefer_bearer=True, ) - return MessageResponse.model_validate(data) def rotate(self, key_id: str) -> ApiKeyCreatedResponse: """Rotate an API key and return the new secret once. @@ -185,13 +314,31 @@ async def list(self) -> ApiKeyListResponse: ) return ApiKeyListResponse.model_validate(data) - async def create(self, name: str) -> ApiKeyCreatedResponse: + async def create( + self, + name: str, + *, + scopes: Sequence[str] | None = None, + expires_at: datetime | str | None = None, + spend_ceiling_usd: float | None = None, + allow_structure_search: bool = False, + ) -> ApiKeyCreatedResponse: """Create a named API key. Parameters ---------- name : str - Display name (unique per user). + Display name (unique per user, 1-64 characters). + scopes : sequence of str or None, optional + Subset of ``read`` and ``write`` (default both). Submitting jobs + or workflow runs needs ``write``. + expires_at : datetime or str or None, optional + Instant after which the key is rejected. + spend_ceiling_usd : float or None, optional + Per-key USD spend cap (``None`` means wallet-only). + allow_structure_search : bool, optional + Let this key send SMILES to ChEMBL similarity / substructure + search via ``/lookup/chembl`` (off by default). Returns ------- @@ -201,7 +348,13 @@ async def create(self, name: str) -> ApiKeyCreatedResponse: data = await self._http.request( "POST", ROUTES["auth_api_keys"], - json={"name": name}, + json=_create_body( + name, + scopes=scopes, + expires_at=expires_at, + spend_ceiling_usd=spend_ceiling_usd, + allow_structure_search=allow_structure_search, + ), prefer_bearer=True, ) return ApiKeyCreatedResponse.model_validate(data) @@ -226,48 +379,78 @@ async def get(self, key_id: str) -> ApiKeyListItem: ) return ApiKeyListItem.model_validate(data) - async def rename(self, key_id: str, name: str) -> ApiKeyListItem: - """Rename an API key. + async def update( + self, + key_id: str, + *, + name: str | None = None, + allow_structure_search: bool | None = None, + ) -> ApiKeyListItem: + """Rename an API key and/or change its structure-search opt-in. Parameters ---------- key_id : str Credential UUID. - name : str + name : str or None, optional New display name. + allow_structure_search : bool or None, optional + Let the key send SMILES to ChEMBL similarity / substructure + search. Returns ------- ApiKeyListItem Updated key metadata. + + Raises + ------ + ValueError + If neither ``name`` nor ``allow_structure_search`` is given. """ data = await self._http.request( "PATCH", ROUTES["auth_api_keys_item"].format(key_id=key_id), - json={"name": name}, + json=_update_body(name, allow_structure_search), prefer_bearer=True, ) return ApiKeyListItem.model_validate(data) - async def delete(self, key_id: str) -> MessageResponse: - """Revoke (hard-delete) an API key. + async def rename(self, key_id: str, name: str) -> ApiKeyListItem: + """Rename an API key (shorthand for :meth:`update`). Parameters ---------- key_id : str Credential UUID. + name : str + New display name. Returns ------- - MessageResponse - Server confirmation message. + ApiKeyListItem + Updated key metadata. """ - data = await self._http.request( + return await self.update(key_id, name=name) + + async def delete(self, key_id: str) -> None: + """Revoke an API key. + + Parameters + ---------- + key_id : str + Credential UUID. + + Returns + ------- + None + The API answers 204 No Content. + """ + await self._http.request( "DELETE", ROUTES["auth_api_keys_item"].format(key_id=key_id), prefer_bearer=True, ) - return MessageResponse.model_validate(data) async def rotate(self, key_id: str) -> ApiKeyCreatedResponse: """Rotate an API key and return the new secret once. diff --git a/src/cognichem_client/resources/artifacts.py b/src/cognichem_client/resources/artifacts.py new file mode 100644 index 0000000..2b31c69 --- /dev/null +++ b/src/cognichem_client/resources/artifacts.py @@ -0,0 +1,554 @@ +"""Durable artifact resource clients under ``/artifacts``.""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Any, Literal + +from cognichem_client._http import AsyncHttpClient, HttpClient, ensure_idempotency_key +from cognichem_client.constants import ROUTES +from cognichem_client.resources.jobs import _maybe_save +from cognichem_client.types import ( + Artifact, + ArtifactListResponse, + ArtifactUsage, + BinaryResult, + MessageResponse, +) + +UploadMode = Literal["cognichem_zip", "raw_file"] + + +def artifact_ref( + artifact_id: str, port: str, *, records: Sequence[str] | None = None +) -> dict[str, Any]: + """Build an ``$artifact`` reference for workflow params or job payloads. + + Parameters + ---------- + artifact_id : str + Artifact identifier. + port : str + Manifest port to bind (e.g. ``structure``, ``poses``). + records : sequence of str or None, optional + Keep only these manifest record ids. + + Returns + ------- + dict + ``{"$artifact": {"id": ..., "port": ..., "records"?: [...]}}``. + + Examples + -------- + >>> artifact_ref("art-1", "structure") + {'$artifact': {'id': 'art-1', 'port': 'structure'}} + """ + ref: dict[str, Any] = {"id": artifact_id, "port": port} + if records is not None: + ref["records"] = list(records) + return {"$artifact": ref} + + +def _list_params( + *, + limit: int | None, + offset: int | None, + data_kind: str | None, + data_format: str | None, + run_id: str | None, + include_archive: bool, +) -> dict[str, Any] | None: + """Build ``GET /artifacts`` query parameters, dropping unset values. + + Parameters + ---------- + limit, offset : int or None + Pagination. + data_kind, data_format, run_id : str or None + Filters. + include_archive : bool + Whether to include ``archive`` ports. + + Returns + ------- + dict or None + Query parameters, or ``None`` when nothing is set. + """ + params: dict[str, Any] = { + "limit": limit, + "offset": offset, + "data_kind": data_kind, + "data_format": data_format, + "run_id": run_id, + } + if include_archive: + params["include_archive"] = "true" + return {k: v for k, v in params.items() if v is not None} or None + + +def _upload_parts( + file: str | Path | bytes, + *, + filename: str | None, + mode: UploadMode | None, + data_kind: str | None, + data_format: str | None, +) -> tuple[dict[str, Any], dict[str, str]]: + """Build multipart ``files`` and form ``data`` for ``POST /artifacts``. + + Parameters + ---------- + file : str or Path or bytes + File path or raw bytes. + filename : str or None + Upload filename (defaults to the path name, or ``upload`` / ``.zip``). + mode : {"cognichem_zip", "raw_file"} or None + Upload mode; inferred when ``None`` (``raw_file`` when ``data_kind`` + is set, otherwise ``cognichem_zip``). + data_kind, data_format : str or None + Required for ``raw_file`` (e.g. ``protein_structure`` / ``pdb``). + + Returns + ------- + tuple of (dict, dict) + httpx ``files`` and ``data`` arguments. + + Raises + ------ + ValueError + If ``raw_file`` is missing ``data_kind`` or ``data_format``. + """ + if isinstance(file, bytes): + content = file + name = filename + else: + path = Path(file) + content = path.read_bytes() + name = filename or path.name + resolved_mode = mode or ("raw_file" if data_kind else "cognichem_zip") + if resolved_mode == "raw_file" and not (data_kind and data_format): + raise ValueError("raw_file uploads need data_kind and data_format") + name = name or ("result.zip" if resolved_mode == "cognichem_zip" else "upload") + data: dict[str, str] = {"mode": resolved_mode} + if data_kind: + data["data_kind"] = data_kind + if data_format: + data["data_format"] = data_format + return {"file": (name, content)}, data + + +class ArtifactsResource: + """Synchronous artifact endpoints under ``/artifacts``. + + Artifacts are durable job / workflow outputs described by a manifest of + ports and records. Bind them into workflow params or job payloads with + :func:`artifact_ref`, and fetch single records with :meth:`download` + ``record=`` instead of whole zips. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def list( + self, + *, + limit: int | None = None, + offset: int | None = None, + data_kind: str | None = None, + data_format: str | None = None, + run_id: str | None = None, + include_archive: bool = False, + ) -> ArtifactListResponse: + """List owned artifacts, one item per chainable manifest port. + + Parameters + ---------- + limit : int or None, optional + Artifacts per page (1-100). + offset : int or None, optional + Page offset (use ``next_offset`` from the previous page). + data_kind : str or None, optional + Keep only ports of this kind (e.g. ``molecule_set``). + data_format : str or None, optional + Keep only ports of this format (e.g. ``sdf``). + run_id : str or None, optional + Keep only artifacts from this workflow run. + include_archive : bool, optional + Also list whole-zip ``archive`` ports. + + Returns + ------- + ArtifactListResponse + Page of artifact ports. + """ + data = self._http.request( + "GET", + ROUTES["artifacts_list"], + params=_list_params( + limit=limit, + offset=offset, + data_kind=data_kind, + data_format=data_format, + run_id=run_id, + include_archive=include_archive, + ), + ) + return ArtifactListResponse.model_validate(data) + + def usage(self) -> ArtifactUsage: + """Return artifact storage usage against the tier quota. + + Returns + ------- + ArtifactUsage + Used bytes, quota, artifact count, and tier. + """ + data = self._http.request("GET", ROUTES["artifacts_usage"]) + return ArtifactUsage.model_validate(data) + + def upload( + self, + file: str | Path | bytes, + *, + mode: UploadMode | None = None, + data_kind: str | None = None, + data_format: str | None = None, + filename: str | None = None, + idempotency_key: str | None = None, + ) -> Artifact: + """Upload a library artifact. + + Parameters + ---------- + file : str or Path or bytes + File path or raw bytes. + mode : {"cognichem_zip", "raw_file"} or None, optional + ``cognichem_zip`` for a CogniChem ``result.zip``; ``raw_file`` for + an allowlisted PDB / SDF / SMILES file. Inferred when omitted + (``raw_file`` if ``data_kind`` is given). + data_kind : str or None, optional + Required for ``raw_file`` (e.g. ``protein_structure``). + data_format : str or None, optional + Required for ``raw_file`` (e.g. ``pdb``). + filename : str or None, optional + Upload filename (defaults to the path name). + idempotency_key : str or None, optional + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). + + Returns + ------- + Artifact + The stored artifact and its manifest. + + Raises + ------ + ValueError + If ``raw_file`` is missing ``data_kind`` or ``data_format``. + """ + files, form = _upload_parts( + file, + filename=filename, + mode=mode, + data_kind=data_kind, + data_format=data_format, + ) + data = self._http.request( + "POST", + ROUTES["artifacts_upload"], + data=form, + files=files, + content_type=None, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return Artifact.model_validate(data) + + def get(self, artifact_id: str) -> Artifact: + """Fetch metadata and the parsed manifest of an owned artifact. + + Parameters + ---------- + artifact_id : str + Artifact identifier. + + Returns + ------- + Artifact + Metadata plus manifest (ports and records). + """ + data = self._http.request( + "GET", ROUTES["artifacts_get"].format(artifact_id=artifact_id) + ) + return Artifact.model_validate(data) + + def download( + self, + artifact_id: str, + *, + record: str | None = None, + port: str | None = None, + save_path: str | Path | None = None, + ) -> BinaryResult: + """Download a whole artifact zip, or one manifest record. + + Parameters + ---------- + artifact_id : str + Artifact identifier. + record : str or None, optional + Manifest record id to fetch alone. Required for virtual artifacts. + port : str or None, optional + Port that scopes ``record`` when the id is ambiguous. + save_path : str or Path or None, optional + Optional file or directory path to write the bytes. + + Returns + ------- + BinaryResult + Content bytes plus filename and media type. + """ + params = {k: v for k, v in {"record": record, "port": port}.items() if v} + result = self._http.request( + "GET", + ROUTES["artifacts_download"].format(artifact_id=artifact_id), + params=params or None, + expect_json=False, + ) + assert isinstance(result, BinaryResult) + return _maybe_save(result, save_path) + + def delete(self, artifact_id: str) -> MessageResponse: + """Delete an owned artifact and its stored files. + + Parameters + ---------- + artifact_id : str + Artifact identifier. + + Returns + ------- + MessageResponse + Server confirmation message. + """ + data = self._http.request( + "DELETE", ROUTES["artifacts_delete"].format(artifact_id=artifact_id) + ) + return MessageResponse.model_validate(data) + + +class AsyncArtifactsResource: + """Asynchronous artifact endpoints under ``/artifacts``. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def list( + self, + *, + limit: int | None = None, + offset: int | None = None, + data_kind: str | None = None, + data_format: str | None = None, + run_id: str | None = None, + include_archive: bool = False, + ) -> ArtifactListResponse: + """List owned artifacts, one item per chainable manifest port. + + Parameters + ---------- + limit : int or None, optional + Artifacts per page (1-100). + offset : int or None, optional + Page offset (use ``next_offset`` from the previous page). + data_kind : str or None, optional + Keep only ports of this kind (e.g. ``molecule_set``). + data_format : str or None, optional + Keep only ports of this format (e.g. ``sdf``). + run_id : str or None, optional + Keep only artifacts from this workflow run. + include_archive : bool, optional + Also list whole-zip ``archive`` ports. + + Returns + ------- + ArtifactListResponse + Page of artifact ports. + """ + data = await self._http.request( + "GET", + ROUTES["artifacts_list"], + params=_list_params( + limit=limit, + offset=offset, + data_kind=data_kind, + data_format=data_format, + run_id=run_id, + include_archive=include_archive, + ), + ) + return ArtifactListResponse.model_validate(data) + + async def usage(self) -> ArtifactUsage: + """Return artifact storage usage against the tier quota. + + Returns + ------- + ArtifactUsage + Used bytes, quota, artifact count, and tier. + """ + data = await self._http.request("GET", ROUTES["artifacts_usage"]) + return ArtifactUsage.model_validate(data) + + async def upload( + self, + file: str | Path | bytes, + *, + mode: UploadMode | None = None, + data_kind: str | None = None, + data_format: str | None = None, + filename: str | None = None, + idempotency_key: str | None = None, + ) -> Artifact: + """Upload a library artifact. + + Parameters + ---------- + file : str or Path or bytes + File path or raw bytes. + mode : {"cognichem_zip", "raw_file"} or None, optional + Upload mode; inferred when omitted (``raw_file`` if ``data_kind`` + is given). + data_kind : str or None, optional + Required for ``raw_file`` (e.g. ``protein_structure``). + data_format : str or None, optional + Required for ``raw_file`` (e.g. ``pdb``). + filename : str or None, optional + Upload filename (defaults to the path name). + idempotency_key : str or None, optional + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). + + Returns + ------- + Artifact + The stored artifact and its manifest. + + Raises + ------ + ValueError + If ``raw_file`` is missing ``data_kind`` or ``data_format``. + """ + files, form = _upload_parts( + file, + filename=filename, + mode=mode, + data_kind=data_kind, + data_format=data_format, + ) + data = await self._http.request( + "POST", + ROUTES["artifacts_upload"], + data=form, + files=files, + content_type=None, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return Artifact.model_validate(data) + + async def get(self, artifact_id: str) -> Artifact: + """Fetch metadata and the parsed manifest of an owned artifact. + + Parameters + ---------- + artifact_id : str + Artifact identifier. + + Returns + ------- + Artifact + Metadata plus manifest (ports and records). + """ + data = await self._http.request( + "GET", ROUTES["artifacts_get"].format(artifact_id=artifact_id) + ) + return Artifact.model_validate(data) + + async def download( + self, + artifact_id: str, + *, + record: str | None = None, + port: str | None = None, + save_path: str | Path | None = None, + ) -> BinaryResult: + """Download a whole artifact zip, or one manifest record. + + Parameters + ---------- + artifact_id : str + Artifact identifier. + record : str or None, optional + Manifest record id to fetch alone. Required for virtual artifacts. + port : str or None, optional + Port that scopes ``record`` when the id is ambiguous. + save_path : str or Path or None, optional + Optional file or directory path to write the bytes. + + Returns + ------- + BinaryResult + Content bytes plus filename and media type. + """ + params = {k: v for k, v in {"record": record, "port": port}.items() if v} + result = await self._http.request( + "GET", + ROUTES["artifacts_download"].format(artifact_id=artifact_id), + params=params or None, + expect_json=False, + ) + assert isinstance(result, BinaryResult) + return _maybe_save(result, save_path) + + async def delete(self, artifact_id: str) -> MessageResponse: + """Delete an owned artifact and its stored files. + + Parameters + ---------- + artifact_id : str + Artifact identifier. + + Returns + ------- + MessageResponse + Server confirmation message. + """ + data = await self._http.request( + "DELETE", ROUTES["artifacts_delete"].format(artifact_id=artifact_id) + ) + return MessageResponse.model_validate(data) diff --git a/src/cognichem_client/resources/chat.py b/src/cognichem_client/resources/chat.py new file mode 100644 index 0000000..1045040 --- /dev/null +++ b/src/cognichem_client/resources/chat.py @@ -0,0 +1,1213 @@ +"""CogniChem Assistant chat resource clients under ``/chat`` (JWT only).""" + +from __future__ import annotations + +from collections.abc import AsyncIterator, Iterator, Sequence +from typing import Any + +from cognichem_client._http import AsyncHttpClient, HttpClient, ensure_idempotency_key +from cognichem_client.constants import ( + DEFAULT_ASSISTANT_MODEL_ID, + ROUTES, + ReasoningEffort, +) +from cognichem_client.errors import CogniChemError +from cognichem_client.types import ( + ChatEstimate, + ChatMessageList, + ChatRunProposal, + ChatRunProposalList, + ChatSession, + ChatSessionList, + ChatSpendCap, + ChatTurnResult, + ServerSentEvent, +) + +# Turns end within the server's 180 s deadline, but tool rounds can leave +# the stream quiet for a while; allow longer than the default read timeout. +_TURN_READ_TIMEOUT = 240.0 + + +def _turn_body( + message: str, + reasoning_effort: ReasoningEffort | None, + artifact_ids: Sequence[str], + model_id: str, +) -> dict[str, Any]: + """Build a turn / estimate body. + + Parameters + ---------- + message : str + User message. + reasoning_effort : str or None + Reasoning level. + artifact_ids : sequence of str + Attached artifacts. + model_id : str + Catalog model id. + + Returns + ------- + dict + JSON body. + """ + body: dict[str, Any] = { + "message": message, + "model_id": model_id, + "artifact_ids": list(artifact_ids), + } + if reasoning_effort is not None: + body["reasoning_effort"] = reasoning_effort + return body + + +class _TurnCollector: + """Fold turn events into a :class:`ChatTurnResult`.""" + + def __init__(self) -> None: + """Start an empty result.""" + self._result = ChatTurnResult() + self._parts: list[str] = [] + + def add(self, event: ServerSentEvent) -> None: + """Apply one event. + + Parameters + ---------- + event : ServerSentEvent + Turn event. + """ + data = event.data if isinstance(event.data, dict) else {} + result = self._result + if event.event == "hold": + result.turn_id = data.get("turn_id") + result.hold_usd = data.get("hold_usd") + result.reasoning_effort = data.get("reasoning_effort") + elif event.event == "token": + self._parts.append(str(data.get("delta") or "")) + elif event.event == "tool_call": + result.tool_calls.append(data) + elif event.event == "citation": + result.citations.append(data) + elif event.event == "proposal": + result.proposals.append(data) + elif event.event == "done": + result.done = data + result.turn_id = data.get("turn_id") or result.turn_id + result.billed_usd = data.get("billed_usd") + result.status = data.get("status") + if isinstance(data.get("citations"), list): + result.citations = list(data["citations"]) + elif event.event == "error": + result.error = str(data.get("detail") or event.data) + + def finish(self) -> ChatTurnResult: + """Return the collected result. + + Returns + ------- + ChatTurnResult + Collected turn. + + Raises + ------ + CogniChemError + If the stream carried an ``error`` event. + """ + self._result.content = "".join(self._parts) + if self._result.error is not None: + raise CogniChemError( + self._result.error, + detail=self._result.error, + code="assistant-turn-error", + body=self._result.model_dump(), + ) + return self._result + + +class ChatSessionsResource: + """Synchronous Assistant thread endpoints under ``/chat/sessions``. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def list(self) -> ChatSessionList: + """List the caller's Assistant threads. + + Returns + ------- + ChatSessionList + Threads. + """ + data = self._http.request("GET", ROUTES["chat_sessions"], prefer_bearer=True) + return ChatSessionList.model_validate(data) + + def create(self, title: str | None = None) -> ChatSession: + """Create an Assistant thread. + + Parameters + ---------- + title : str or None, optional + Thread title (up to 200 characters). + + Returns + ------- + ChatSession + The new thread. + """ + data = self._http.request( + "POST", + ROUTES["chat_sessions"], + json={"title": title}, + prefer_bearer=True, + ) + return ChatSession.model_validate(data) + + def get(self, session_id: str) -> ChatSession: + """Fetch one Assistant thread. + + Parameters + ---------- + session_id : str + Session identifier. + + Returns + ------- + ChatSession + The thread. + """ + data = self._http.request( + "GET", + ROUTES["chat_session"].format(session_id=session_id), + prefer_bearer=True, + ) + return ChatSession.model_validate(data) + + def update( + self, + session_id: str, + *, + title: str | None = None, + allow_structure_search: bool | None = None, + ) -> ChatSession: + """Rename a thread and/or change its structure-search opt-in. + + Parameters + ---------- + session_id : str + Session identifier. + title : str or None, optional + New title. + allow_structure_search : bool or None, optional + Let the Assistant send SMILES to external databases for + similarity / substructure search in this thread. + + Returns + ------- + ChatSession + The updated thread. + + Raises + ------ + ValueError + If neither field is given. + """ + body: dict[str, Any] = {} + if title is not None: + body["title"] = title + if allow_structure_search is not None: + body["allow_structure_search"] = allow_structure_search + if not body: + raise ValueError("Provide title and/or allow_structure_search") + data = self._http.request( + "PATCH", + ROUTES["chat_session"].format(session_id=session_id), + json=body, + prefer_bearer=True, + ) + return ChatSession.model_validate(data) + + def delete(self, session_id: str) -> None: + """Delete a thread with its messages. + + Parameters + ---------- + session_id : str + Session identifier. + + Returns + ------- + None + The API answers 204 No Content. + """ + self._http.request( + "DELETE", + ROUTES["chat_session"].format(session_id=session_id), + prefer_bearer=True, + ) + + def messages(self, session_id: str) -> ChatMessageList: + """Return a thread's message history. + + Parameters + ---------- + session_id : str + Session identifier. + + Returns + ------- + ChatMessageList + Messages, oldest first. + """ + data = self._http.request( + "GET", + ROUTES["chat_messages"].format(session_id=session_id), + prefer_bearer=True, + ) + return ChatMessageList.model_validate(data) + + +class ChatProposalsResource: + """Synchronous Assistant run proposal (plan card) endpoints. + + The Assistant only proposes jobs and workflow runs; nothing runs until + you approve the proposal at its server estimate. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def list(self, session_id: str) -> ChatRunProposalList: + """List the proposals in a thread. + + Parameters + ---------- + session_id : str + Session identifier. + + Returns + ------- + ChatRunProposalList + Plan cards. + """ + data = self._http.request( + "GET", + ROUTES["chat_proposals"].format(session_id=session_id), + prefer_bearer=True, + ) + return ChatRunProposalList.model_validate(data) + + def get( + self, session_id: str, proposal_id: str, *, include_spec: bool = False + ) -> ChatRunProposal: + """Fetch one proposal. + + Parameters + ---------- + session_id : str + Session identifier. + proposal_id : str + Proposal identifier. + include_spec : bool, optional + Include the frozen WorkflowSpec (``workflow_spec``). + + Returns + ------- + ChatRunProposal + The plan card. + """ + data = self._http.request( + "GET", + ROUTES["chat_proposal"].format( + session_id=session_id, proposal_id=proposal_id + ), + params={"include_spec": "true"} if include_spec else None, + prefer_bearer=True, + ) + return ChatRunProposal.model_validate(data) + + def approve( + self, + session_id: str, + proposal_id: str, + accepted_estimate_usd: float, + *, + max_run_cost_usd: float | None = None, + idempotency_key: str | None = None, + ) -> ChatRunProposal: + """Approve a proposal: re-validate, re-estimate, then enqueue the run. + + Parameters + ---------- + session_id : str + Session identifier. + proposal_id : str + Proposal identifier. + accepted_estimate_usd : float + The ``estimate_usd`` you reviewed. If the server estimate has + risen, nothing runs and a + :class:`~cognichem_client.errors.ConflictError` with code + ``estimate-changed`` carries the new ``proposal`` in ``body``. + max_run_cost_usd : float or None, optional + Workflow spend cap (defaults to the estimate total). + idempotency_key : str or None, optional + Idempotency key (required by the API; generated when omitted). + + Returns + ------- + ChatRunProposal + The proposal with ``job_id`` or ``workflow_run_id`` once submitted. + """ + data = self._http.request( + "POST", + ROUTES["chat_proposal_approve"].format( + session_id=session_id, proposal_id=proposal_id + ), + json={ + "accepted_estimate_usd": accepted_estimate_usd, + "max_run_cost_usd": max_run_cost_usd, + }, + idempotency_key=ensure_idempotency_key(idempotency_key), + prefer_bearer=True, + ) + return ChatRunProposal.model_validate(data) + + def reject(self, session_id: str, proposal_id: str) -> ChatRunProposal: + """Dismiss a pending proposal (no spend). + + Parameters + ---------- + session_id : str + Session identifier. + proposal_id : str + Proposal identifier. + + Returns + ------- + ChatRunProposal + The rejected proposal. + """ + data = self._http.request( + "POST", + ROUTES["chat_proposal_reject"].format( + session_id=session_id, proposal_id=proposal_id + ), + prefer_bearer=True, + ) + return ChatRunProposal.model_validate(data) + + +class ChatResource: + """Synchronous CogniChem Assistant chat endpoints under ``/chat``. + + Chat is JWT-only: sign in with ``client.auth.login_email`` (or pass + ``access_token``) first. Each turn holds wallet funds at your tier's + token rate, debits the metered tokens, and releases the rest. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + + Attributes + ---------- + sessions : ChatSessionsResource + Thread endpoints. + proposals : ChatProposalsResource + Run proposal endpoints. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport and nested resources. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + self.sessions = ChatSessionsResource(http) + self.proposals = ChatProposalsResource(http) + + def estimate( + self, + session_id: str, + message: str, + *, + reasoning_effort: ReasoningEffort | None = None, + artifact_ids: Sequence[str] = (), + model_id: str = DEFAULT_ASSISTANT_MODEL_ID, + ) -> ChatEstimate: + """Quote the wallet hold for a turn without starting it. + + Parameters + ---------- + session_id : str + Session identifier. + message : str + User message (up to 16,000 characters). + reasoning_effort : {"low", "medium", "high"} or None, optional + Reasoning level (catalog default ``low``). + artifact_ids : sequence of str, optional + Artifacts to attach. + model_id : str, optional + Catalog model id. + + Returns + ------- + ChatEstimate + Hold in USD and the token assumptions. + """ + data = self._http.request( + "POST", + ROUTES["chat_estimate"].format(session_id=session_id), + json=_turn_body(message, reasoning_effort, artifact_ids, model_id), + prefer_bearer=True, + ) + return ChatEstimate.model_validate(data) + + def stream( + self, + session_id: str, + message: str, + *, + reasoning_effort: ReasoningEffort | None = None, + artifact_ids: Sequence[str] = (), + model_id: str = DEFAULT_ASSISTANT_MODEL_ID, + ) -> Iterator[ServerSentEvent]: + """Start a turn and yield its server-sent events as they arrive. + + Events: ``hold`` (``turn_id``, ``hold_usd``), ``token`` (``delta``), + ``tool_call``, ``citation``, ``proposal``, ``done`` (``billed_usd``, + spend fields), and ``error`` (``detail``). + + Parameters + ---------- + session_id : str + Session identifier. + message : str + User message (up to 16,000 characters). + reasoning_effort : {"low", "medium", "high"} or None, optional + Reasoning level (catalog default ``low``). + artifact_ids : sequence of str, optional + Artifacts to attach (bound as ``$artifact`` in proposals). + model_id : str, optional + Catalog model id. + + Yields + ------ + ServerSentEvent + Turn events in order. + + Raises + ------ + PaymentRequiredError + If the wallet cannot cover the hold. + ConflictError + If the thread is at its spend cap (code ``session-spend-cap``; + call :meth:`continue_spend_cap` to raise it). + """ + return self._http.stream_events( + "POST", + ROUTES["chat_turns"].format(session_id=session_id), + json=_turn_body(message, reasoning_effort, artifact_ids, model_id), + prefer_bearer=True, + timeout=_TURN_READ_TIMEOUT, + ) + + def send( + self, + session_id: str, + message: str, + *, + reasoning_effort: ReasoningEffort | None = None, + artifact_ids: Sequence[str] = (), + model_id: str = DEFAULT_ASSISTANT_MODEL_ID, + ) -> ChatTurnResult: + """Run one turn to completion and return the collected reply. + + Parameters + ---------- + session_id : str + Session identifier. + message : str + User message (up to 16,000 characters). + reasoning_effort : {"low", "medium", "high"} or None, optional + Reasoning level (catalog default ``low``). + artifact_ids : sequence of str, optional + Artifacts to attach. + model_id : str, optional + Catalog model id. + + Returns + ------- + ChatTurnResult + Reply text, spend, tool calls, citations, and proposals. + + Raises + ------ + CogniChemError + If the stream ends with an ``error`` event (code + ``assistant-turn-error``), or on HTTP errors as in :meth:`stream`. + """ + collector = _TurnCollector() + for event in self.stream( + session_id, + message, + reasoning_effort=reasoning_effort, + artifact_ids=artifact_ids, + model_id=model_id, + ): + collector.add(event) + return collector.finish() + + def stop(self, session_id: str, turn_id: str) -> dict[str, str]: + """Interrupt an in-flight turn; metered tokens bill, the rest releases. + + Parameters + ---------- + session_id : str + Session identifier. + turn_id : str + Turn identifier (from the ``hold`` event). + + Returns + ------- + dict + ``{"status": "stopping", "turn_id": ...}``. + """ + data = self._http.request( + "POST", + ROUTES["chat_turn_stop"].format(session_id=session_id, turn_id=turn_id), + prefer_bearer=True, + ) + assert isinstance(data, dict) + return data + + def spend_cap( + self, session_id: str, *, model_id: str = DEFAULT_ASSISTANT_MODEL_ID + ) -> ChatSpendCap: + """Return a thread's Assistant spend and spend cap. + + Parameters + ---------- + session_id : str + Session identifier. + model_id : str, optional + Catalog model id. + + Returns + ------- + ChatSpendCap + Spent, held, cap, and step (USD). + """ + data = self._http.request( + "GET", + ROUTES["chat_spend_cap"].format(session_id=session_id), + params={"model_id": model_id}, + prefer_bearer=True, + ) + return ChatSpendCap.model_validate(data) + + def continue_spend_cap( + self, + session_id: str, + current_cap_usd: float, + *, + model_id: str = DEFAULT_ASSISTANT_MODEL_ID, + ) -> ChatSpendCap: + """Raise a thread's spend cap by one step (the user's Continue). + + Parameters + ---------- + session_id : str + Session identifier. + current_cap_usd : float + The cap you saw. If it has already moved, nothing changes and + ``raised`` is false. + model_id : str, optional + Catalog model id. + + Returns + ------- + ChatSpendCap + State after the request. + """ + data = self._http.request( + "POST", + ROUTES["chat_spend_cap_continue"].format(session_id=session_id), + json={"current_cap_usd": current_cap_usd, "model_id": model_id}, + prefer_bearer=True, + ) + return ChatSpendCap.model_validate(data) + + +class AsyncChatSessionsResource: + """Asynchronous Assistant thread endpoints under ``/chat/sessions``. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def list(self) -> ChatSessionList: + """List the caller's Assistant threads. + + Returns + ------- + ChatSessionList + Threads. + """ + data = await self._http.request( + "GET", ROUTES["chat_sessions"], prefer_bearer=True + ) + return ChatSessionList.model_validate(data) + + async def create(self, title: str | None = None) -> ChatSession: + """Create an Assistant thread. + + Parameters + ---------- + title : str or None, optional + Thread title (up to 200 characters). + + Returns + ------- + ChatSession + The new thread. + """ + data = await self._http.request( + "POST", + ROUTES["chat_sessions"], + json={"title": title}, + prefer_bearer=True, + ) + return ChatSession.model_validate(data) + + async def get(self, session_id: str) -> ChatSession: + """Fetch one Assistant thread. + + Parameters + ---------- + session_id : str + Session identifier. + + Returns + ------- + ChatSession + The thread. + """ + data = await self._http.request( + "GET", + ROUTES["chat_session"].format(session_id=session_id), + prefer_bearer=True, + ) + return ChatSession.model_validate(data) + + async def update( + self, + session_id: str, + *, + title: str | None = None, + allow_structure_search: bool | None = None, + ) -> ChatSession: + """Rename a thread and/or change its structure-search opt-in. + + Parameters + ---------- + session_id : str + Session identifier. + title : str or None, optional + New title. + allow_structure_search : bool or None, optional + Let the Assistant send SMILES to external databases for + similarity / substructure search in this thread. + + Returns + ------- + ChatSession + The updated thread. + + Raises + ------ + ValueError + If neither field is given. + """ + body: dict[str, Any] = {} + if title is not None: + body["title"] = title + if allow_structure_search is not None: + body["allow_structure_search"] = allow_structure_search + if not body: + raise ValueError("Provide title and/or allow_structure_search") + data = await self._http.request( + "PATCH", + ROUTES["chat_session"].format(session_id=session_id), + json=body, + prefer_bearer=True, + ) + return ChatSession.model_validate(data) + + async def delete(self, session_id: str) -> None: + """Delete a thread with its messages. + + Parameters + ---------- + session_id : str + Session identifier. + + Returns + ------- + None + The API answers 204 No Content. + """ + await self._http.request( + "DELETE", + ROUTES["chat_session"].format(session_id=session_id), + prefer_bearer=True, + ) + + async def messages(self, session_id: str) -> ChatMessageList: + """Return a thread's message history. + + Parameters + ---------- + session_id : str + Session identifier. + + Returns + ------- + ChatMessageList + Messages, oldest first. + """ + data = await self._http.request( + "GET", + ROUTES["chat_messages"].format(session_id=session_id), + prefer_bearer=True, + ) + return ChatMessageList.model_validate(data) + + +class AsyncChatProposalsResource: + """Asynchronous Assistant run proposal (plan card) endpoints. + + The Assistant only proposes jobs and workflow runs; nothing runs until + you approve the proposal at its server estimate. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def list(self, session_id: str) -> ChatRunProposalList: + """List the proposals in a thread. + + Parameters + ---------- + session_id : str + Session identifier. + + Returns + ------- + ChatRunProposalList + Plan cards. + """ + data = await self._http.request( + "GET", + ROUTES["chat_proposals"].format(session_id=session_id), + prefer_bearer=True, + ) + return ChatRunProposalList.model_validate(data) + + async def get( + self, session_id: str, proposal_id: str, *, include_spec: bool = False + ) -> ChatRunProposal: + """Fetch one proposal. + + Parameters + ---------- + session_id : str + Session identifier. + proposal_id : str + Proposal identifier. + include_spec : bool, optional + Include the frozen WorkflowSpec (``workflow_spec``). + + Returns + ------- + ChatRunProposal + The plan card. + """ + data = await self._http.request( + "GET", + ROUTES["chat_proposal"].format( + session_id=session_id, proposal_id=proposal_id + ), + params={"include_spec": "true"} if include_spec else None, + prefer_bearer=True, + ) + return ChatRunProposal.model_validate(data) + + async def approve( + self, + session_id: str, + proposal_id: str, + accepted_estimate_usd: float, + *, + max_run_cost_usd: float | None = None, + idempotency_key: str | None = None, + ) -> ChatRunProposal: + """Approve a proposal: re-validate, re-estimate, then enqueue the run. + + Parameters + ---------- + session_id : str + Session identifier. + proposal_id : str + Proposal identifier. + accepted_estimate_usd : float + The ``estimate_usd`` you reviewed. If the server estimate has + risen, nothing runs and a + :class:`~cognichem_client.errors.ConflictError` with code + ``estimate-changed`` carries the new ``proposal`` in ``body``. + max_run_cost_usd : float or None, optional + Workflow spend cap (defaults to the estimate total). + idempotency_key : str or None, optional + Idempotency key (required by the API; generated when omitted). + + Returns + ------- + ChatRunProposal + The proposal with ``job_id`` or ``workflow_run_id`` once submitted. + """ + data = await self._http.request( + "POST", + ROUTES["chat_proposal_approve"].format( + session_id=session_id, proposal_id=proposal_id + ), + json={ + "accepted_estimate_usd": accepted_estimate_usd, + "max_run_cost_usd": max_run_cost_usd, + }, + idempotency_key=ensure_idempotency_key(idempotency_key), + prefer_bearer=True, + ) + return ChatRunProposal.model_validate(data) + + async def reject(self, session_id: str, proposal_id: str) -> ChatRunProposal: + """Dismiss a pending proposal (no spend). + + Parameters + ---------- + session_id : str + Session identifier. + proposal_id : str + Proposal identifier. + + Returns + ------- + ChatRunProposal + The rejected proposal. + """ + data = await self._http.request( + "POST", + ROUTES["chat_proposal_reject"].format( + session_id=session_id, proposal_id=proposal_id + ), + prefer_bearer=True, + ) + return ChatRunProposal.model_validate(data) + + +class AsyncChatResource: + """Asynchronous CogniChem Assistant chat endpoints under ``/chat``. + + Chat is JWT-only: sign in with ``client.auth.login_email`` (or pass + ``access_token``) first. Each turn holds wallet funds at your tier's + token rate, debits the metered tokens, and releases the rest. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + + Attributes + ---------- + sessions : AsyncChatSessionsResource + Thread endpoints. + proposals : AsyncChatProposalsResource + Run proposal endpoints. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport and nested resources. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + self.sessions = AsyncChatSessionsResource(http) + self.proposals = AsyncChatProposalsResource(http) + + async def estimate( + self, + session_id: str, + message: str, + *, + reasoning_effort: ReasoningEffort | None = None, + artifact_ids: Sequence[str] = (), + model_id: str = DEFAULT_ASSISTANT_MODEL_ID, + ) -> ChatEstimate: + """Quote the wallet hold for a turn without starting it. + + Parameters + ---------- + session_id : str + Session identifier. + message : str + User message (up to 16,000 characters). + reasoning_effort : {"low", "medium", "high"} or None, optional + Reasoning level (catalog default ``low``). + artifact_ids : sequence of str, optional + Artifacts to attach. + model_id : str, optional + Catalog model id. + + Returns + ------- + ChatEstimate + Hold in USD and the token assumptions. + """ + data = await self._http.request( + "POST", + ROUTES["chat_estimate"].format(session_id=session_id), + json=_turn_body(message, reasoning_effort, artifact_ids, model_id), + prefer_bearer=True, + ) + return ChatEstimate.model_validate(data) + + def stream( + self, + session_id: str, + message: str, + *, + reasoning_effort: ReasoningEffort | None = None, + artifact_ids: Sequence[str] = (), + model_id: str = DEFAULT_ASSISTANT_MODEL_ID, + ) -> AsyncIterator[ServerSentEvent]: + """Start a turn and yield its server-sent events as they arrive. + + Events: ``hold`` (``turn_id``, ``hold_usd``), ``token`` (``delta``), + ``tool_call``, ``citation``, ``proposal``, ``done`` (``billed_usd``, + spend fields), and ``error`` (``detail``). + + Parameters + ---------- + session_id : str + Session identifier. + message : str + User message (up to 16,000 characters). + reasoning_effort : {"low", "medium", "high"} or None, optional + Reasoning level (catalog default ``low``). + artifact_ids : sequence of str, optional + Artifacts to attach (bound as ``$artifact`` in proposals). + model_id : str, optional + Catalog model id. + + Yields + ------ + ServerSentEvent + Turn events in order. + + Raises + ------ + PaymentRequiredError + If the wallet cannot cover the hold. + ConflictError + If the thread is at its spend cap (code ``session-spend-cap``; + call :meth:`continue_spend_cap` to raise it). + """ + return self._http.stream_events( + "POST", + ROUTES["chat_turns"].format(session_id=session_id), + json=_turn_body(message, reasoning_effort, artifact_ids, model_id), + prefer_bearer=True, + timeout=_TURN_READ_TIMEOUT, + ) + + async def send( + self, + session_id: str, + message: str, + *, + reasoning_effort: ReasoningEffort | None = None, + artifact_ids: Sequence[str] = (), + model_id: str = DEFAULT_ASSISTANT_MODEL_ID, + ) -> ChatTurnResult: + """Run one turn to completion and return the collected reply. + + Parameters + ---------- + session_id : str + Session identifier. + message : str + User message (up to 16,000 characters). + reasoning_effort : {"low", "medium", "high"} or None, optional + Reasoning level (catalog default ``low``). + artifact_ids : sequence of str, optional + Artifacts to attach. + model_id : str, optional + Catalog model id. + + Returns + ------- + ChatTurnResult + Reply text, spend, tool calls, citations, and proposals. + + Raises + ------ + CogniChemError + If the stream ends with an ``error`` event (code + ``assistant-turn-error``), or on HTTP errors as in :meth:`stream`. + """ + collector = _TurnCollector() + async for event in self.stream( + session_id, + message, + reasoning_effort=reasoning_effort, + artifact_ids=artifact_ids, + model_id=model_id, + ): + collector.add(event) + return collector.finish() + + async def stop(self, session_id: str, turn_id: str) -> dict[str, str]: + """Interrupt an in-flight turn; metered tokens bill, the rest releases. + + Parameters + ---------- + session_id : str + Session identifier. + turn_id : str + Turn identifier (from the ``hold`` event). + + Returns + ------- + dict + ``{"status": "stopping", "turn_id": ...}``. + """ + data = await self._http.request( + "POST", + ROUTES["chat_turn_stop"].format(session_id=session_id, turn_id=turn_id), + prefer_bearer=True, + ) + assert isinstance(data, dict) + return data + + async def spend_cap( + self, session_id: str, *, model_id: str = DEFAULT_ASSISTANT_MODEL_ID + ) -> ChatSpendCap: + """Return a thread's Assistant spend and spend cap. + + Parameters + ---------- + session_id : str + Session identifier. + model_id : str, optional + Catalog model id. + + Returns + ------- + ChatSpendCap + Spent, held, cap, and step (USD). + """ + data = await self._http.request( + "GET", + ROUTES["chat_spend_cap"].format(session_id=session_id), + params={"model_id": model_id}, + prefer_bearer=True, + ) + return ChatSpendCap.model_validate(data) + + async def continue_spend_cap( + self, + session_id: str, + current_cap_usd: float, + *, + model_id: str = DEFAULT_ASSISTANT_MODEL_ID, + ) -> ChatSpendCap: + """Raise a thread's spend cap by one step (the user's Continue). + + Parameters + ---------- + session_id : str + Session identifier. + current_cap_usd : float + The cap you saw. If it has already moved, nothing changes and + ``raised`` is false. + model_id : str, optional + Catalog model id. + + Returns + ------- + ChatSpendCap + State after the request. + """ + data = await self._http.request( + "POST", + ROUTES["chat_spend_cap_continue"].format(session_id=session_id), + json={"current_cap_usd": current_cap_usd, "model_id": model_id}, + prefer_bearer=True, + ) + return ChatSpendCap.model_validate(data) diff --git a/src/cognichem_client/resources/events.py b/src/cognichem_client/resources/events.py new file mode 100644 index 0000000..20e0fee --- /dev/null +++ b/src/cognichem_client/resources/events.py @@ -0,0 +1,280 @@ +"""Run completion events resource clients (``GET /events``).""" + +from __future__ import annotations + +from collections import deque +from collections.abc import AsyncIterator, Iterator +from typing import Any + +from cognichem_client._http import AsyncHttpClient, HttpClient +from cognichem_client.constants import MAX_EVENTS_WAIT, ROUTES, EventKind +from cognichem_client.types import RunEvent, RunEventList + +# Recent event ids remembered by ``listen`` to drop at-least-once repeats. +_DEDUPE_WINDOW = 1000 + + +def _params( + *, + after: str | None, + kind: EventKind | None, + subject_id: str | None, + limit: int, + wait: float, +) -> dict[str, Any]: + """Build ``GET /events`` query parameters, dropping unset values. + + Parameters + ---------- + after : str or None + Cursor. + kind : str or None + Subject kind filter. + subject_id : str or None + Subject filter. + limit : int + Page size. + wait : float + Long-poll seconds. + + Returns + ------- + dict + Query parameters. + """ + params: dict[str, Any] = {"limit": limit} + if wait: + params["wait"] = wait + for key, value in (("after", after), ("kind", kind), ("subject_id", subject_id)): + if value is not None: + params[key] = value + return params + + +class _Deduper: + """Bounded memory of seen ``event_id`` values.""" + + def __init__(self) -> None: + """Start with no seen ids.""" + self._order: deque[int] = deque() + self._seen: set[int] = set() + + def first_time(self, event_id: int) -> bool: + """Record ``event_id`` and report whether it is new. + + Parameters + ---------- + event_id : int + Event dedupe key. + + Returns + ------- + bool + ``True`` the first time an id is seen. + """ + if event_id in self._seen: + return False + self._seen.add(event_id) + self._order.append(event_id) + if len(self._order) > _DEDUPE_WINDOW: + self._seen.discard(self._order.popleft()) + return True + + +class EventsResource: + """Synchronous run completion events feed. + + Each event is one stop transition of a job or workflow run, retained for + 7 days. Delivery is at-least-once and gap-free; dedupe on ``event_id``. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def list( + self, + *, + after: str | None = None, + kind: EventKind | None = None, + subject_id: str | None = None, + limit: int = 50, + wait: float = 0, + ) -> RunEventList: + """Fetch one page of events after a cursor, oldest first. + + Parameters + ---------- + after : str or None, optional + ``next_cursor`` from a previous page; omit for the oldest + retained event. + kind : {"job", "workflow_run"} or None, optional + Keep only one subject kind. + subject_id : str or None, optional + Keep only events for one job or run. + limit : int, optional + Page size (1-100). + wait : float, optional + Seconds (up to 25) to hold the request open when there are no + events yet. + + Returns + ------- + RunEventList + Events and ``next_cursor``. + """ + data = self._http.request( + "GET", + ROUTES["events"], + params=_params( + after=after, kind=kind, subject_id=subject_id, limit=limit, wait=wait + ), + ) + return RunEventList.model_validate(data) + + def listen( + self, + *, + after: str | None = None, + kind: EventKind | None = None, + subject_id: str | None = None, + wait: float = MAX_EVENTS_WAIT, + ) -> Iterator[RunEvent]: + """Yield events forever, long-polling and following the cursor. + + Repeats from at-least-once delivery are dropped. Stop by breaking out + of the loop. + + Parameters + ---------- + after : str or None, optional + Cursor to start after; omit for the oldest retained event. + kind : {"job", "workflow_run"} or None, optional + Keep only one subject kind. + subject_id : str or None, optional + Keep only events for one job or run. + wait : float, optional + Long-poll seconds per request (up to 25). + + Yields + ------ + RunEvent + Events in order. + """ + seen = _Deduper() + cursor = after + while True: + page = self.list(after=cursor, kind=kind, subject_id=subject_id, wait=wait) + for event in page.events: + if seen.first_time(event.event_id): + yield event + if page.next_cursor: + cursor = page.next_cursor + + +class AsyncEventsResource: + """Asynchronous run completion events feed. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def list( + self, + *, + after: str | None = None, + kind: EventKind | None = None, + subject_id: str | None = None, + limit: int = 50, + wait: float = 0, + ) -> RunEventList: + """Fetch one page of events after a cursor, oldest first. + + Parameters + ---------- + after : str or None, optional + ``next_cursor`` from a previous page. + kind : {"job", "workflow_run"} or None, optional + Keep only one subject kind. + subject_id : str or None, optional + Keep only events for one job or run. + limit : int, optional + Page size (1-100). + wait : float, optional + Seconds (up to 25) to hold the request open when there are no + events yet. + + Returns + ------- + RunEventList + Events and ``next_cursor``. + """ + data = await self._http.request( + "GET", + ROUTES["events"], + params=_params( + after=after, kind=kind, subject_id=subject_id, limit=limit, wait=wait + ), + ) + return RunEventList.model_validate(data) + + async def listen( + self, + *, + after: str | None = None, + kind: EventKind | None = None, + subject_id: str | None = None, + wait: float = MAX_EVENTS_WAIT, + ) -> AsyncIterator[RunEvent]: + """Yield events forever, long-polling and following the cursor. + + Parameters + ---------- + after : str or None, optional + Cursor to start after. + kind : {"job", "workflow_run"} or None, optional + Keep only one subject kind. + subject_id : str or None, optional + Keep only events for one job or run. + wait : float, optional + Long-poll seconds per request (up to 25). + + Yields + ------ + RunEvent + Events in order, without repeats. + """ + seen = _Deduper() + cursor = after + while True: + page = await self.list( + after=cursor, kind=kind, subject_id=subject_id, wait=wait + ) + for event in page.events: + if seen.first_time(event.event_id): + yield event + if page.next_cursor: + cursor = page.next_cursor diff --git a/src/cognichem_client/resources/inference/__init__.py b/src/cognichem_client/resources/inference/__init__.py index 4ca9282..f138852 100644 --- a/src/cognichem_client/resources/inference/__init__.py +++ b/src/cognichem_client/resources/inference/__init__.py @@ -4,7 +4,11 @@ from typing import Any -from cognichem_client._http import AsyncHttpClient, HttpClient +from cognichem_client._http import ( + AsyncHttpClient, + HttpClient, + ensure_idempotency_key, +) from cognichem_client.constants import ( DEFAULT_FAST_POLL_INTERVAL, DEFAULT_FAST_TIMEOUT, @@ -70,7 +74,8 @@ def submit( payload : dict Model-specific request payload. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -85,7 +90,7 @@ def submit( "model_name": model_name, "payload": payload, }, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitResponse.model_validate(data) @@ -102,7 +107,8 @@ def submit_batch( jobs : list of InferenceSubmitRequest or dict Inference specifications to submit. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -119,7 +125,7 @@ def submit_batch( "POST", ROUTES["inference_submit_batch"], json=body, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitMultipleResponse.model_validate(data) @@ -190,6 +196,26 @@ def delete(self, process_id: str) -> MessageResponse: ) return MessageResponse.model_validate(data) + def cancel(self, process_id: str) -> MessageResponse: + """Request cancellation of a queued or running inference job. + + Parameters + ---------- + process_id : str + Process identifier. + + Returns + ------- + MessageResponse + Server confirmation message. + """ + data = self._http.request( + "POST", + ROUTES["inference_cancel"], + params={"process_id": process_id}, + ) + return MessageResponse.model_validate(data) + def wait( self, process_id: str, @@ -255,7 +281,8 @@ def run( payload : dict Model-specific request payload. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). poll_interval : float, optional Seconds between status polls. timeout : float, optional @@ -335,7 +362,8 @@ async def submit( payload : dict Model-specific request payload. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -350,7 +378,7 @@ async def submit( "model_name": model_name, "payload": payload, }, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitResponse.model_validate(data) @@ -367,7 +395,8 @@ async def submit_batch( jobs : list of InferenceSubmitRequest or dict Inference specifications to submit. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -384,7 +413,7 @@ async def submit_batch( "POST", ROUTES["inference_submit_batch"], json=body, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitMultipleResponse.model_validate(data) @@ -455,6 +484,26 @@ async def delete(self, process_id: str) -> MessageResponse: ) return MessageResponse.model_validate(data) + async def cancel(self, process_id: str) -> MessageResponse: + """Request cancellation of a queued or running inference job. + + Parameters + ---------- + process_id : str + Process identifier. + + Returns + ------- + MessageResponse + Server confirmation message. + """ + data = await self._http.request( + "POST", + ROUTES["inference_cancel"], + params={"process_id": process_id}, + ) + return MessageResponse.model_validate(data) + async def wait( self, process_id: str, @@ -520,7 +569,8 @@ async def run( payload : dict Model-specific request payload. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). poll_interval : float, optional Seconds between status polls. timeout : float, optional diff --git a/src/cognichem_client/resources/jobs.py b/src/cognichem_client/resources/jobs.py index 5110082..fb9553d 100644 --- a/src/cognichem_client/resources/jobs.py +++ b/src/cognichem_client/resources/jobs.py @@ -6,7 +6,7 @@ from pathlib import Path from typing import Any -from cognichem_client._http import AsyncHttpClient, HttpClient +from cognichem_client._http import AsyncHttpClient, HttpClient, ensure_idempotency_key from cognichem_client.constants import ( DEFAULT_JOB_POLL_INTERVAL, DEFAULT_JOB_TIMEOUT, @@ -16,6 +16,7 @@ from cognichem_client.polling import await_for_terminal, wait_for_terminal from cognichem_client.types import ( BinaryResult, + JobEstimateResponse, JobInfoResponse, JobSubmitMultipleResponse, JobSubmitRequest, @@ -52,6 +53,44 @@ def _maybe_save(result: BinaryResult, save_path: str | Path | None) -> BinaryRes return result +def _list_params( + *, + limit: int | None, + offset: int | None, + status: str | Sequence[str] | None, + job_type: str | None, + q: str | None, + sort: str | None, +) -> dict[str, Any] | None: + """Build ``GET /jobs/list`` query parameters, dropping unset values. + + Parameters + ---------- + limit, offset : int or None + Pagination. + status : str or sequence of str or None + Status filter (joined with commas). + job_type, q, sort : str or None + Other filters. + + Returns + ------- + dict or None + Query parameters, or ``None`` when nothing is set. + """ + if status is not None and not isinstance(status, str): + status = ",".join(status) + params = { + "limit": limit, + "offset": offset, + "status": status, + "job_type": job_type, + "q": q, + "sort": sort, + } + return {k: v for k, v in params.items() if v is not None} or None + + class JobsResource: """Synchronous long-running job endpoints under ``/jobs``. @@ -93,7 +132,8 @@ def submit( resource : str, optional Compute resource tier (default ``default``). idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -109,7 +149,7 @@ def submit( "payload": payload, "resource": resource, }, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitResponse.model_validate(data) @@ -126,7 +166,8 @@ def submit_many( jobs : sequence of JobSubmitRequest or dict Job specifications to submit. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -142,19 +183,92 @@ def submit_many( "POST", ROUTES["jobs_submit_multiple"], json=body, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitMultipleResponse.model_validate(data) - def list(self) -> ListJobsResponse: - """List jobs belonging to the authenticated caller. + def estimate( + self, + job_type: str, + payload: dict[str, Any], + *, + resource: str = "default", + ) -> JobEstimateResponse: + """Estimate the wallet hold for one job at the caller's tier. + + Uses the same catalog reservation math as submit. Incomplete payloads + may still return a scaled estimate (``assumptions["validated"]`` is + false). + + Parameters + ---------- + job_type : str + Catalog job type. + payload : dict + Type-specific request payload. + resource : str, optional + Compute resource tier (default ``default``). + + Returns + ------- + JobEstimateResponse + Estimated cost, tier, rate, and assumptions. + """ + data = self._http.request( + "POST", + ROUTES["jobs_estimate"], + json={"job_type": job_type, "payload": payload, "resource": resource}, + ) + return JobEstimateResponse.model_validate(data) + + def list( + self, + *, + limit: int | None = None, + offset: int | None = None, + status: str | Sequence[str] | None = None, + job_type: str | None = None, + q: str | None = None, + sort: str | None = None, + ) -> ListJobsResponse: + """List standalone jobs belonging to the authenticated caller. + + Workflow step jobs are excluded. Results are paginated. + + Parameters + ---------- + limit : int or None, optional + Page size (1-100). + offset : int or None, optional + Page offset. + status : str or sequence of str or None, optional + Keep only these statuses (e.g. ``["queued", "running"]``). + job_type : str or None, optional + Keep only this catalog job type. + q : str or None, optional + Case-insensitive job name search. + sort : str or None, optional + ``created_at``, ``finished_at``, ``runtime_seconds``, or + ``job_name``; prefix ``-`` for descending (default + ``-created_at``). Returns ------- ListJobsResponse - Job summaries for the current account. + Page of jobs (``items`` carries full rows). """ - data = self._http.request("GET", ROUTES["jobs_list"]) + data = self._http.request( + "GET", + ROUTES["jobs_list"], + params=_list_params( + limit=limit, + offset=offset, + status=status, + job_type=job_type, + q=q, + sort=sort, + ), + ) return ListJobsResponse.model_validate(data) def info(self, process_id: str) -> JobInfoResponse: @@ -357,7 +471,8 @@ def run( resource : str, optional Compute resource tier (default ``default``). idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). poll_interval : float, optional Seconds between status polls. timeout : float, optional @@ -437,7 +552,8 @@ async def submit( resource : str, optional Compute resource tier (default ``default``). idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -453,7 +569,7 @@ async def submit( "payload": payload, "resource": resource, }, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitResponse.model_validate(data) @@ -470,7 +586,8 @@ async def submit_many( jobs : sequence of JobSubmitRequest or dict Job specifications to submit. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -486,19 +603,92 @@ async def submit_many( "POST", ROUTES["jobs_submit_multiple"], json=body, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitMultipleResponse.model_validate(data) - async def list(self) -> ListJobsResponse: - """List jobs belonging to the authenticated caller. + async def estimate( + self, + job_type: str, + payload: dict[str, Any], + *, + resource: str = "default", + ) -> JobEstimateResponse: + """Estimate the wallet hold for one job at the caller's tier. + + Uses the same catalog reservation math as submit. Incomplete payloads + may still return a scaled estimate (``assumptions["validated"]`` is + false). + + Parameters + ---------- + job_type : str + Catalog job type. + payload : dict + Type-specific request payload. + resource : str, optional + Compute resource tier (default ``default``). + + Returns + ------- + JobEstimateResponse + Estimated cost, tier, rate, and assumptions. + """ + data = await self._http.request( + "POST", + ROUTES["jobs_estimate"], + json={"job_type": job_type, "payload": payload, "resource": resource}, + ) + return JobEstimateResponse.model_validate(data) + + async def list( + self, + *, + limit: int | None = None, + offset: int | None = None, + status: str | Sequence[str] | None = None, + job_type: str | None = None, + q: str | None = None, + sort: str | None = None, + ) -> ListJobsResponse: + """List standalone jobs belonging to the authenticated caller. + + Workflow step jobs are excluded. Results are paginated. + + Parameters + ---------- + limit : int or None, optional + Page size (1-100). + offset : int or None, optional + Page offset. + status : str or sequence of str or None, optional + Keep only these statuses (e.g. ``["queued", "running"]``). + job_type : str or None, optional + Keep only this catalog job type. + q : str or None, optional + Case-insensitive job name search. + sort : str or None, optional + ``created_at``, ``finished_at``, ``runtime_seconds``, or + ``job_name``; prefix ``-`` for descending (default + ``-created_at``). Returns ------- ListJobsResponse - Job summaries for the current account. + Page of jobs (``items`` carries full rows). """ - data = await self._http.request("GET", ROUTES["jobs_list"]) + data = await self._http.request( + "GET", + ROUTES["jobs_list"], + params=_list_params( + limit=limit, + offset=offset, + status=status, + job_type=job_type, + q=q, + sort=sort, + ), + ) return ListJobsResponse.model_validate(data) async def info(self, process_id: str) -> JobInfoResponse: @@ -701,7 +891,8 @@ async def run( resource : str, optional Compute resource tier (default ``default``). idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). poll_interval : float, optional Seconds between status polls. timeout : float, optional diff --git a/src/cognichem_client/resources/lookup.py b/src/cognichem_client/resources/lookup.py new file mode 100644 index 0000000..40ad0a5 --- /dev/null +++ b/src/cognichem_client/resources/lookup.py @@ -0,0 +1,481 @@ +"""Science database lookup resource clients under ``/lookup``.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any, Literal + +from cognichem_client._http import AsyncHttpClient, HttpClient +from cognichem_client.constants import ROUTES +from cognichem_client.types import LookupResponse, PubchemLookupResponse + +ChemblOperation = Literal[ + "molecule", + "target", + "molecule_activities", + "target_activities", + "similarity", + "substructure", +] +TargetsOperation = Literal[ + "target_diseases", "disease_targets", "target_drugs", "disease_drugs" +] +EntryOrSearch = Literal["entry", "search"] +PropertyKey = Literal[ + "boiling_point", + "melting_point", + "density", + "vapor_pressure", + "vapor_density", + "heat_of_vaporization", + "heat_of_combustion", + "viscosity", + "surface_tension", + "solubility", + "logp", + "henrys_law_constant", + "flash_point", + "autoignition_temperature", + "refractive_index", + "dissociation_constants", + "other", +] + + +def _body(operation: str, queries: Sequence[str], **options: Any) -> dict[str, Any]: + """Build a lookup request body, omitting unset options. + + Parameters + ---------- + operation : str + Lookup operation. + queries : sequence of str + Queries (a bare string is treated as one query). + **options : any + Optional fields (``limit``, ``threshold``, ``organism``, …). + + Returns + ------- + dict + JSON body. + """ + items = [queries] if isinstance(queries, str) else list(queries) + body: dict[str, Any] = {"operation": operation, "queries": items} + body.update({k: v for k, v in options.items() if v is not None}) + return body + + +class LookupResource: + """Synchronous science database lookups under ``/lookup``. + + Results carry source links and licenses; show ``attribution`` with any + data you display. Names and ids are sent upstream; SMILES are sent only + by ChEMBL ``similarity`` / ``substructure`` and only for API keys with + ``allow_structure_search`` (otherwise rows come back with status + ``structure_egress_not_allowed``). + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def pubchem(self, names: Sequence[str]) -> PubchemLookupResponse: + """Resolve compound names to PubChem CIDs and SMILES. + + Parameters + ---------- + names : sequence of str + Up to 20 compound names (names only; structures are refused). + + Returns + ------- + PubchemLookupResponse + One row per name, plus attribution. + """ + data = self._http.request( + "POST", + ROUTES["lookup_pubchem"], + json={"names": [names] if isinstance(names, str) else list(names)}, + ) + return PubchemLookupResponse.model_validate(data) + + def chembl( + self, + operation: ChemblOperation, + queries: Sequence[str], + *, + limit: int | None = None, + threshold: int | None = None, + ) -> LookupResponse: + """Query ChEMBL molecules, targets, activities, or structure search. + + Parameters + ---------- + operation : str + ``molecule``, ``target``, ``molecule_activities``, + ``target_activities``, ``similarity``, or ``substructure``. + queries : sequence of str + Up to 20 ChEMBL ids or names; SMILES for ``similarity`` / + ``substructure``. + limit : int or None, optional + Rows per query (1-25, default 10). + threshold : int or None, optional + ``similarity`` only: percent similarity (40-100, default 70). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = self._http.request( + "POST", + ROUTES["lookup_chembl"], + json=_body(operation, queries, limit=limit, threshold=threshold), + ) + return LookupResponse.model_validate(data) + + def targets( + self, + operation: TargetsOperation, + queries: Sequence[str], + *, + limit: int | None = None, + ) -> LookupResponse: + """Query Open Targets target–disease–drug associations. + + Parameters + ---------- + operation : str + ``target_diseases``, ``disease_targets``, ``target_drugs``, or + ``disease_drugs``. + queries : sequence of str + Up to 20 target symbols, disease names, or ids. + limit : int or None, optional + Rows per query (1-25, default 10). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = self._http.request( + "POST", + ROUTES["lookup_targets"], + json=_body(operation, queries, limit=limit), + ) + return LookupResponse.model_validate(data) + + def uniprot( + self, + operation: EntryOrSearch, + queries: Sequence[str], + *, + organism: str | None = None, + ) -> LookupResponse: + """Fetch UniProt entries or search proteins. + + Parameters + ---------- + operation : {"entry", "search"} + ``entry`` for accessions, ``search`` for names / genes. + queries : sequence of str + Up to 20 accessions or search terms. + organism : str or None, optional + ``search`` only: organism name or NCBI taxon id (``9606`` = human). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = self._http.request( + "POST", + ROUTES["lookup_uniprot"], + json=_body(operation, queries, organism=organism), + ) + return LookupResponse.model_validate(data) + + def pdb( + self, + operation: EntryOrSearch, + queries: Sequence[str], + *, + limit: int | None = None, + ) -> LookupResponse: + """Fetch RCSB PDB entries or search structures. + + Parameters + ---------- + operation : {"entry", "search"} + ``entry`` for PDB ids, ``search`` for text queries. + queries : sequence of str + Up to 20 PDB ids or search terms. + limit : int or None, optional + Rows per query (1-25, default 10). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = self._http.request( + "POST", + ROUTES["lookup_pdb"], + json=_body(operation, queries, limit=limit), + ) + return LookupResponse.model_validate(data) + + def properties( + self, + queries: Sequence[str], + *, + properties: Sequence[PropertyKey] | None = None, + limit: int | None = None, + ) -> LookupResponse: + """Fetch PubChem experimental properties (boiling point, density, …). + + Parameters + ---------- + queries : sequence of str + Up to 20 PubChem CIDs or compound names. + properties : sequence of str or None, optional + Property keys (e.g. ``boiling_point``, ``vapor_pressure``, + ``logp``); defaults to a common thermophysical set. + limit : int or None, optional + Values per property (1-10, default 3). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = self._http.request( + "POST", + ROUTES["lookup_properties"], + json=_body( + "experimental", + queries, + properties=list(properties) if properties is not None else None, + limit=limit, + ), + ) + return LookupResponse.model_validate(data) + + +class AsyncLookupResource: + """Asynchronous science database lookups under ``/lookup``. + + Results carry source links and licenses; show ``attribution`` with any + data you display. Names and ids are sent upstream; SMILES are sent only + by ChEMBL ``similarity`` / ``substructure`` and only for API keys with + ``allow_structure_search`` (otherwise rows come back with status + ``structure_egress_not_allowed``). + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def pubchem(self, names: Sequence[str]) -> PubchemLookupResponse: + """Resolve compound names to PubChem CIDs and SMILES. + + Parameters + ---------- + names : sequence of str + Up to 20 compound names (names only; structures are refused). + + Returns + ------- + PubchemLookupResponse + One row per name, plus attribution. + """ + data = await self._http.request( + "POST", + ROUTES["lookup_pubchem"], + json={"names": [names] if isinstance(names, str) else list(names)}, + ) + return PubchemLookupResponse.model_validate(data) + + async def chembl( + self, + operation: ChemblOperation, + queries: Sequence[str], + *, + limit: int | None = None, + threshold: int | None = None, + ) -> LookupResponse: + """Query ChEMBL molecules, targets, activities, or structure search. + + Parameters + ---------- + operation : str + ``molecule``, ``target``, ``molecule_activities``, + ``target_activities``, ``similarity``, or ``substructure``. + queries : sequence of str + Up to 20 ChEMBL ids or names; SMILES for ``similarity`` / + ``substructure``. + limit : int or None, optional + Rows per query (1-25, default 10). + threshold : int or None, optional + ``similarity`` only: percent similarity (40-100, default 70). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = await self._http.request( + "POST", + ROUTES["lookup_chembl"], + json=_body(operation, queries, limit=limit, threshold=threshold), + ) + return LookupResponse.model_validate(data) + + async def targets( + self, + operation: TargetsOperation, + queries: Sequence[str], + *, + limit: int | None = None, + ) -> LookupResponse: + """Query Open Targets target–disease–drug associations. + + Parameters + ---------- + operation : str + ``target_diseases``, ``disease_targets``, ``target_drugs``, or + ``disease_drugs``. + queries : sequence of str + Up to 20 target symbols, disease names, or ids. + limit : int or None, optional + Rows per query (1-25, default 10). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = await self._http.request( + "POST", + ROUTES["lookup_targets"], + json=_body(operation, queries, limit=limit), + ) + return LookupResponse.model_validate(data) + + async def uniprot( + self, + operation: EntryOrSearch, + queries: Sequence[str], + *, + organism: str | None = None, + ) -> LookupResponse: + """Fetch UniProt entries or search proteins. + + Parameters + ---------- + operation : {"entry", "search"} + ``entry`` for accessions, ``search`` for names / genes. + queries : sequence of str + Up to 20 accessions or search terms. + organism : str or None, optional + ``search`` only: organism name or NCBI taxon id (``9606`` = human). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = await self._http.request( + "POST", + ROUTES["lookup_uniprot"], + json=_body(operation, queries, organism=organism), + ) + return LookupResponse.model_validate(data) + + async def pdb( + self, + operation: EntryOrSearch, + queries: Sequence[str], + *, + limit: int | None = None, + ) -> LookupResponse: + """Fetch RCSB PDB entries or search structures. + + Parameters + ---------- + operation : {"entry", "search"} + ``entry`` for PDB ids, ``search`` for text queries. + queries : sequence of str + Up to 20 PDB ids or search terms. + limit : int or None, optional + Rows per query (1-25, default 10). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = await self._http.request( + "POST", + ROUTES["lookup_pdb"], + json=_body(operation, queries, limit=limit), + ) + return LookupResponse.model_validate(data) + + async def properties( + self, + queries: Sequence[str], + *, + properties: Sequence[PropertyKey] | None = None, + limit: int | None = None, + ) -> LookupResponse: + """Fetch PubChem experimental properties (boiling point, density, …). + + Parameters + ---------- + queries : sequence of str + Up to 20 PubChem CIDs or compound names. + properties : sequence of str or None, optional + Property keys (e.g. ``boiling_point``, ``vapor_pressure``, + ``logp``); defaults to a common thermophysical set. + limit : int or None, optional + Values per property (1-10, default 3). + + Returns + ------- + LookupResponse + Results in input order, plus attribution. + """ + data = await self._http.request( + "POST", + ROUTES["lookup_properties"], + json=_body( + "experimental", + queries, + properties=list(properties) if properties is not None else None, + limit=limit, + ), + ) + return LookupResponse.model_validate(data) diff --git a/src/cognichem_client/resources/reference.py b/src/cognichem_client/resources/reference.py new file mode 100644 index 0000000..9eccdad --- /dev/null +++ b/src/cognichem_client/resources/reference.py @@ -0,0 +1,235 @@ +"""Reference-KG (literature) resource clients under ``/reference``.""" + +from __future__ import annotations + +from typing import Any + +from cognichem_client._http import AsyncHttpClient, HttpClient +from cognichem_client.constants import ROUTES, ReferenceEdgeKind, ReferencePaperSet +from cognichem_client.types import ( + ReferenceNeighborhood, + ReferencePaperDetail, + ReferenceSearchResponse, +) + + +class ReferenceResource: + """Synchronous reference-KG (literature) endpoints under ``/reference``. + + The reference KG holds method papers for every catalog job type, papers + that use CogniChem tools, and curated general papers. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def search( + self, + q: str, + *, + paper_set: ReferencePaperSet | None = None, + limit: int = 10, + ) -> ReferenceSearchResponse: + """Search reference papers (hybrid semantic + lexical when available). + + Parameters + ---------- + q : str + Search text (1-200 characters). + paper_set : {"method", "used_cognichem", "general"} or None, optional + Keep only one paper set. + limit : int, optional + Hits to return (1-50). + + Returns + ------- + ReferenceSearchResponse + Ranked papers and provenance. + """ + params: dict[str, Any] = {"q": q, "limit": limit} + if paper_set is not None: + params["paper_set"] = paper_set + data = self._http.request("GET", ROUTES["reference_papers"], params=params) + return ReferenceSearchResponse.model_validate(data) + + def get(self, paper_id: str) -> ReferencePaperDetail: + """Fetch one paper with its links (and abstract, when licensed). + + Parameters + ---------- + paper_id : str + Paper id (``paper:…``). + + Returns + ------- + ReferencePaperDetail + Paper, edges, citing count, and provenance. + """ + data = self._http.request( + "GET", ROUTES["reference_paper"].format(paper_id=paper_id) + ) + return ReferencePaperDetail.model_validate(data) + + def neighborhood( + self, + *, + node_id: str | None = None, + paper_id: str | None = None, + kind: ReferenceEdgeKind | None = None, + limit: int = 25, + ) -> ReferenceNeighborhood: + """Return papers linked to a Product KG node or to a paper. + + Parameters + ---------- + node_id : str or None, optional + Product KG node (e.g. ``job_type:protein-prepare``). + paper_id : str or None, optional + Paper id. Give ``node_id`` or ``paper_id``. + kind : str or None, optional + Keep only one edge kind (``cites``, ``method_of``, ``uses_tool``, + ``compares``, ``describes``). + limit : int, optional + Edges to return (1-50). + + Returns + ------- + ReferenceNeighborhood + Center, edges, papers, and provenance. + """ + params: dict[str, Any] = {"limit": limit} + for key, value in ( + ("node_id", node_id), + ("paper_id", paper_id), + ("kind", kind), + ): + if value is not None: + params[key] = value + data = self._http.request( + "GET", ROUTES["reference_neighborhood"], params=params + ) + return ReferenceNeighborhood.model_validate(data) + + +class AsyncReferenceResource: + """Asynchronous reference-KG (literature) endpoints under ``/reference``. + + The reference KG holds method papers for every catalog job type, papers + that use CogniChem tools, and curated general papers. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def search( + self, + q: str, + *, + paper_set: ReferencePaperSet | None = None, + limit: int = 10, + ) -> ReferenceSearchResponse: + """Search reference papers (hybrid semantic + lexical when available). + + Parameters + ---------- + q : str + Search text (1-200 characters). + paper_set : {"method", "used_cognichem", "general"} or None, optional + Keep only one paper set. + limit : int, optional + Hits to return (1-50). + + Returns + ------- + ReferenceSearchResponse + Ranked papers and provenance. + """ + params: dict[str, Any] = {"q": q, "limit": limit} + if paper_set is not None: + params["paper_set"] = paper_set + data = await self._http.request( + "GET", ROUTES["reference_papers"], params=params + ) + return ReferenceSearchResponse.model_validate(data) + + async def get(self, paper_id: str) -> ReferencePaperDetail: + """Fetch one paper with its links (and abstract, when licensed). + + Parameters + ---------- + paper_id : str + Paper id (``paper:…``). + + Returns + ------- + ReferencePaperDetail + Paper, edges, citing count, and provenance. + """ + data = await self._http.request( + "GET", ROUTES["reference_paper"].format(paper_id=paper_id) + ) + return ReferencePaperDetail.model_validate(data) + + async def neighborhood( + self, + *, + node_id: str | None = None, + paper_id: str | None = None, + kind: ReferenceEdgeKind | None = None, + limit: int = 25, + ) -> ReferenceNeighborhood: + """Return papers linked to a Product KG node or to a paper. + + Parameters + ---------- + node_id : str or None, optional + Product KG node (e.g. ``job_type:protein-prepare``). + paper_id : str or None, optional + Paper id. Give ``node_id`` or ``paper_id``. + kind : str or None, optional + Keep only one edge kind (``cites``, ``method_of``, ``uses_tool``, + ``compares``, ``describes``). + limit : int, optional + Edges to return (1-50). + + Returns + ------- + ReferenceNeighborhood + Center, edges, papers, and provenance. + """ + params: dict[str, Any] = {"limit": limit} + for key, value in ( + ("node_id", node_id), + ("paper_id", paper_id), + ("kind", kind), + ): + if value is not None: + params[key] = value + data = await self._http.request( + "GET", ROUTES["reference_neighborhood"], params=params + ) + return ReferenceNeighborhood.model_validate(data) diff --git a/src/cognichem_client/resources/utils.py b/src/cognichem_client/resources/utils.py index c6d5cc1..747dca2 100644 --- a/src/cognichem_client/resources/utils.py +++ b/src/cognichem_client/resources/utils.py @@ -4,7 +4,11 @@ from typing import Any -from cognichem_client._http import AsyncHttpClient, HttpClient +from cognichem_client._http import ( + AsyncHttpClient, + HttpClient, + ensure_idempotency_key, +) from cognichem_client.constants import ( DEFAULT_FAST_POLL_INTERVAL, DEFAULT_FAST_TIMEOUT, @@ -55,7 +59,8 @@ def submit( payload : dict Type-specific request payload. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -66,7 +71,7 @@ def submit( "POST", ROUTES["utils_submit"], json={"utility_type": utility_type, "payload": payload}, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitResponse.model_validate(data) @@ -135,6 +140,26 @@ def delete(self, process_id: str) -> MessageResponse: ) return MessageResponse.model_validate(data) + def cancel(self, process_id: str) -> MessageResponse: + """Request cancellation of a queued or running utility process. + + Parameters + ---------- + process_id : str + Process identifier. + + Returns + ------- + MessageResponse + Server confirmation message. + """ + data = self._http.request( + "POST", + ROUTES["utils_cancel"], + params={"process_id": process_id}, + ) + return MessageResponse.model_validate(data) + def wait( self, process_id: str, @@ -197,7 +222,8 @@ def run( payload : dict Type-specific request payload. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). poll_interval : float, optional Seconds between status polls. timeout : float, optional @@ -265,7 +291,8 @@ async def submit( payload : dict Type-specific request payload. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). Returns ------- @@ -276,7 +303,7 @@ async def submit( "POST", ROUTES["utils_submit"], json={"utility_type": utility_type, "payload": payload}, - idempotency_key=idempotency_key, + idempotency_key=ensure_idempotency_key(idempotency_key), ) return JobSubmitResponse.model_validate(data) @@ -345,6 +372,26 @@ async def delete(self, process_id: str) -> MessageResponse: ) return MessageResponse.model_validate(data) + async def cancel(self, process_id: str) -> MessageResponse: + """Request cancellation of a queued or running utility process. + + Parameters + ---------- + process_id : str + Process identifier. + + Returns + ------- + MessageResponse + Server confirmation message. + """ + data = await self._http.request( + "POST", + ROUTES["utils_cancel"], + params={"process_id": process_id}, + ) + return MessageResponse.model_validate(data) + async def wait( self, process_id: str, @@ -407,7 +454,8 @@ async def run( payload : dict Type-specific request payload. idempotency_key : str or None, optional - Optional idempotency key for safe retries. + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). poll_interval : float, optional Seconds between status polls. timeout : float, optional diff --git a/src/cognichem_client/resources/wallet.py b/src/cognichem_client/resources/wallet.py new file mode 100644 index 0000000..7450f2c --- /dev/null +++ b/src/cognichem_client/resources/wallet.py @@ -0,0 +1,69 @@ +"""Wallet resource clients (``GET /wallet``).""" + +from __future__ import annotations + +from cognichem_client._http import AsyncHttpClient, HttpClient +from cognichem_client.constants import ROUTES +from cognichem_client.types import WalletBalance + + +class WalletResource: + """Synchronous wallet endpoint (``GET /wallet``). + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def balance(self) -> WalletBalance: + """Return the wallet balance and USD held by active reservations. + + Returns + ------- + WalletBalance + ``balance`` and ``active_reserved`` in USD. + """ + data = self._http.request("GET", ROUTES["wallet"]) + return WalletBalance.model_validate(data) + + +class AsyncWalletResource: + """Asynchronous wallet endpoint (``GET /wallet``). + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def balance(self) -> WalletBalance: + """Return the wallet balance and USD held by active reservations. + + Returns + ------- + WalletBalance + ``balance`` and ``active_reserved`` in USD. + """ + data = await self._http.request("GET", ROUTES["wallet"]) + return WalletBalance.model_validate(data) diff --git a/src/cognichem_client/resources/workflows.py b/src/cognichem_client/resources/workflows.py new file mode 100644 index 0000000..295595c --- /dev/null +++ b/src/cognichem_client/resources/workflows.py @@ -0,0 +1,1775 @@ +"""Workflow resource clients under ``/workflows`` (validate, estimate, catalog, +runs, and saved definitions).""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from cognichem_client._http import AsyncHttpClient, HttpClient, ensure_idempotency_key +from cognichem_client.constants import ( + DEFAULT_WORKFLOW_POLL_INTERVAL, + DEFAULT_WORKFLOW_TIMEOUT, + ROUTES, + WORKFLOW_RUN_STOP_STATUSES, + WorkflowRunStatus, +) +from cognichem_client.polling import apoll_until, poll_until, run_failure_statuses +from cognichem_client.resources.jobs import _maybe_save +from cognichem_client.types import ( + BinaryResult, + WorkflowDefinition, + WorkflowDefinitionList, + WorkflowDefinitionVersion, + WorkflowDefinitionVersionList, + WorkflowEstimate, + WorkflowNodePorts, + WorkflowRun, + WorkflowRunArtifacts, + WorkflowRunCancelResponse, + WorkflowRunList, + WorkflowTemplate, + WorkflowTemplateSummary, + WorkflowValidation, +) + + +def _drop_none(**values: Any) -> dict[str, Any]: + """Return values without None entries. + + Parameters + ---------- + **values : any + Candidate query or body fields. + + Returns + ------- + dict + Fields that are set. + """ + return {k: v for k, v in values.items() if v is not None} + + +def _spec_body(spec: dict[str, Any], params: dict[str, Any] | None) -> dict[str, Any]: + """Build a validate / estimate body. + + Parameters + ---------- + spec : dict + WorkflowSpec. + params : dict or None + Runtime param values. + + Returns + ------- + dict + JSON body. + """ + body: dict[str, Any] = {"spec": spec} + if params is not None: + body["params"] = params + return body + + +class WorkflowRunsResource: + """Synchronous workflow run endpoints under ``/workflows/runs``. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def create( + self, + run_name: str, + spec: dict[str, Any], + *, + params: dict[str, Any] | None = None, + max_run_cost: float | None = None, + definition_id: str | None = None, + idempotency_key: str | None = None, + ) -> WorkflowRun: + """Create (enqueue) a workflow run. + + The server validates the spec, estimates it at the caller's tier, and + opens a wallet hold before any step starts. + + Parameters + ---------- + run_name : str + Unique-per-user run name. + spec : dict + WorkflowSpec (frozen on the run). + params : dict or None, optional + Runtime param values the spec declares. + max_run_cost : float or None, optional + Spend cap; the run pauses instead of exceeding it. + definition_id : str or None, optional + Saved definition (``wfd-…``) the spec came from. + idempotency_key : str or None, optional + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). + + Returns + ------- + WorkflowRun + The new run with its step rows and ``estimated_cost``. + """ + body: dict[str, Any] = { + "run_name": run_name, + "spec": spec, + "params": params or {}, + } + if max_run_cost is not None: + body["max_run_cost"] = max_run_cost + if definition_id is not None: + body["definition_id"] = definition_id + data = self._http.request( + "POST", + ROUTES["workflows_runs_create"], + json=body, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowRun.model_validate(data) + + def list( + self, + *, + limit: int | None = None, + offset: int | None = None, + status: WorkflowRunStatus | None = None, + ) -> WorkflowRunList: + """List the caller's workflow runs. + + Parameters + ---------- + limit : int or None, optional + Page size (1-100). + offset : int or None, optional + Page offset. + status : str or None, optional + Keep only runs with this status. + + Returns + ------- + WorkflowRunList + Page of runs. + """ + data = self._http.request( + "GET", + ROUTES["workflows_runs_list"], + params=_drop_none(limit=limit, offset=offset, status=status), + ) + return WorkflowRunList.model_validate(data) + + def get(self, run_id: str, *, include_spec: bool = False) -> WorkflowRun: + """Fetch a workflow run and its steps. + + Parameters + ---------- + run_id : str + Run identifier. + include_spec : bool, optional + Also return the frozen ``spec`` and ``definition_id``. + + Returns + ------- + WorkflowRun + Run status, costs, and step rows. + """ + data = self._http.request( + "GET", + ROUTES["workflows_runs_get"].format(run_id=run_id), + params={"include_spec": 1} if include_spec else None, + ) + return WorkflowRun.model_validate(data) + + def artifacts(self, run_id: str) -> WorkflowRunArtifacts: + """List artifact ports produced by a workflow run. + + Parameters + ---------- + run_id : str + Run identifier. + + Returns + ------- + WorkflowRunArtifacts + Artifact ports by step. + """ + data = self._http.request( + "GET", ROUTES["workflows_runs_artifacts"].format(run_id=run_id) + ) + return WorkflowRunArtifacts.model_validate(data) + + def resume(self, run_id: str) -> WorkflowRun: + """Resume a paused workflow run. + + Parameters + ---------- + run_id : str + Run identifier. + + Returns + ------- + WorkflowRun + The resumed run. + """ + data = self._http.request( + "POST", ROUTES["workflows_runs_resume"].format(run_id=run_id) + ) + return WorkflowRun.model_validate(data) + + def cancel(self, run_id: str) -> WorkflowRunCancelResponse: + """Cancel a workflow run and its active step jobs. + + Parameters + ---------- + run_id : str + Run identifier. + + Returns + ------- + WorkflowRunCancelResponse + Final status and cancelled job ids. + """ + data = self._http.request( + "POST", ROUTES["workflows_runs_cancel"].format(run_id=run_id) + ) + return WorkflowRunCancelResponse.model_validate(data) + + def delete(self, run_id: str) -> None: + """Delete a workflow run record. + + Artifacts stay in storage (usable as ``$artifact`` inputs until they + expire) and linked jobs are kept. + + Parameters + ---------- + run_id : str + Run identifier. + + Returns + ------- + None + The API answers 204 No Content. + """ + self._http.request( + "DELETE", ROUTES["workflows_runs_delete"].format(run_id=run_id) + ) + + def wait( + self, + run_id: str, + *, + poll_interval: float = DEFAULT_WORKFLOW_POLL_INTERVAL, + timeout: float = DEFAULT_WORKFLOW_TIMEOUT, + raise_on_failure: bool = False, + ) -> WorkflowRun: + """Poll until the run stops: completed, failed, cancelled, or paused. + + A ``paused`` run needs your action (:meth:`resume`, usually after + raising ``max_run_cost``), so waiting stops there too. + + Parameters + ---------- + run_id : str + Run identifier. + poll_interval : float, optional + Seconds between polls. + timeout : float, optional + Maximum seconds to wait before raising + :class:`~cognichem_client.errors.PollTimeoutError`. + raise_on_failure : bool, optional + When ``True``, raise if the run ends ``failed`` or ``cancelled``. + + Returns + ------- + WorkflowRun + The run at its stop status. + + Raises + ------ + PollTimeoutError + If ``timeout`` elapses first. + ProcessFailedError + If ``raise_on_failure`` is true and the run failed. + ProcessCancelledError + If ``raise_on_failure`` is true and the run was cancelled. + """ + return poll_until( + lambda: self.get(run_id), + poll_interval=poll_interval, + timeout=timeout, + stop_statuses=WORKFLOW_RUN_STOP_STATUSES, + subject=lambda run: f"workflow run {run.id}", + failed_statuses=run_failure_statuses(raise_on_failure), + ) + + def run( + self, + run_name: str, + spec: dict[str, Any], + *, + params: dict[str, Any] | None = None, + max_run_cost: float | None = None, + definition_id: str | None = None, + idempotency_key: str | None = None, + poll_interval: float = DEFAULT_WORKFLOW_POLL_INTERVAL, + timeout: float = DEFAULT_WORKFLOW_TIMEOUT, + ) -> WorkflowRun: + """Create a workflow run and wait until it stops. + + Parameters + ---------- + run_name : str + Unique-per-user run name. + spec : dict + WorkflowSpec. + params : dict or None, optional + Runtime param values. + max_run_cost : float or None, optional + Spend cap. + definition_id : str or None, optional + Saved definition the spec came from. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + poll_interval : float, optional + Seconds between polls. + timeout : float, optional + Maximum seconds to wait. + + Returns + ------- + WorkflowRun + The run at ``completed`` or ``paused``. List its outputs with + :meth:`artifacts`. + + Raises + ------ + PollTimeoutError + If ``timeout`` elapses first. + ProcessFailedError + If the run fails. + ProcessCancelledError + If the run is cancelled. + """ + created = self.create( + run_name, + spec, + params=params, + max_run_cost=max_run_cost, + definition_id=definition_id, + idempotency_key=idempotency_key, + ) + return self.wait( + created.id, + poll_interval=poll_interval, + timeout=timeout, + raise_on_failure=True, + ) + + +class WorkflowDefinitionsResource: + """Synchronous saved-definition endpoints under ``/workflows/definitions``. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + + def create( + self, + name: str, + spec: dict[str, Any], + *, + description: str = "", + graph_layout: dict[str, Any] | None = None, + spec_version: int = 1, + idempotency_key: str | None = None, + ) -> WorkflowDefinition: + """Save a workflow definition. + + Parameters + ---------- + name : str + Unique-per-user name. + spec : dict + WorkflowSpec (a top-level ``graph_layout`` moves to its own column). + description : str, optional + Description. + graph_layout : dict or None, optional + Builder layout (wins over one inside ``spec``). + spec_version : int, optional + Spec schema version. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + The saved definition. + """ + body: dict[str, Any] = { + "name": name, + "description": description, + "spec": spec, + "spec_version": spec_version, + } + if graph_layout is not None: + body["graph_layout"] = graph_layout + data = self._http.request( + "POST", + ROUTES["workflows_definitions_create"], + json=body, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + def list( + self, *, limit: int | None = None, offset: int | None = None + ) -> WorkflowDefinitionList: + """List saved definitions (without specs). + + Parameters + ---------- + limit : int or None, optional + Page size. + offset : int or None, optional + Page offset. + + Returns + ------- + WorkflowDefinitionList + Page of definitions. + """ + data = self._http.request( + "GET", + ROUTES["workflows_definitions_list"], + params=_drop_none(limit=limit, offset=offset), + ) + return WorkflowDefinitionList.model_validate(data) + + def get(self, definition_id: str) -> WorkflowDefinition: + """Fetch a saved definition with its spec. + + Parameters + ---------- + definition_id : str + Definition identifier (``wfd-…``). + + Returns + ------- + WorkflowDefinition + The definition. + """ + data = self._http.request( + "GET", + ROUTES["workflows_definitions_get"].format(definition_id=definition_id), + ) + return WorkflowDefinition.model_validate(data) + + def update( + self, + definition_id: str, + *, + name: str | None = None, + description: str | None = None, + spec: dict[str, Any] | None = None, + graph_layout: dict[str, Any] | None = None, + spec_version: int | None = None, + idempotency_key: str | None = None, + ) -> WorkflowDefinition: + """Update a saved definition (only the given fields change). + + Parameters + ---------- + definition_id : str + Definition identifier. + name, description : str or None, optional + New name or description. + spec, graph_layout : dict or None, optional + New spec or builder layout (a spec change adds a version). + spec_version : int or None, optional + New spec schema version. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + The updated definition. + + Raises + ------ + ValueError + If no field is given. + """ + body = _drop_none( + name=name, + description=description, + spec=spec, + graph_layout=graph_layout, + spec_version=spec_version, + ) + if not body: + raise ValueError("Provide at least one field to update") + data = self._http.request( + "PATCH", + ROUTES["workflows_definitions_update"].format(definition_id=definition_id), + json=body, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + def delete(self, definition_id: str) -> None: + """Delete a saved definition. + + Parameters + ---------- + definition_id : str + Definition identifier. + + Returns + ------- + None + The API answers 204 No Content. + """ + self._http.request( + "DELETE", + ROUTES["workflows_definitions_delete"].format(definition_id=definition_id), + ) + + def download( + self, + definition_id: str, + *, + include_layout: bool = False, + save_path: str | Path | None = None, + ) -> BinaryResult: + """Download a definition as WorkflowSpec JSON. + + Parameters + ---------- + definition_id : str + Definition identifier. + include_layout : bool, optional + Merge the builder ``graph_layout`` into the export. + save_path : str or Path or None, optional + Optional file or directory path to write the JSON. + + Returns + ------- + BinaryResult + JSON bytes plus filename. + """ + result = self._http.request( + "GET", + ROUTES["workflows_definitions_download"].format( + definition_id=definition_id + ), + params={"include_layout": 1} if include_layout else None, + expect_json=False, + ) + assert isinstance(result, BinaryResult) + return _maybe_save(result, save_path) + + def share( + self, definition_id: str, *, idempotency_key: str | None = None + ) -> WorkflowDefinition: + """Turn on link sharing and return the ``share_token``. + + Parameters + ---------- + definition_id : str + Definition identifier. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + The definition with ``visibility="link"`` and ``share_token``. + """ + data = self._http.request( + "POST", + ROUTES["workflows_definitions_share"].format(definition_id=definition_id), + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + def unshare(self, definition_id: str) -> WorkflowDefinition: + """Revoke link sharing. + + Parameters + ---------- + definition_id : str + Definition identifier. + + Returns + ------- + WorkflowDefinition + The definition with ``visibility="private"``. + """ + data = self._http.request( + "DELETE", + ROUTES["workflows_definitions_share"].format(definition_id=definition_id), + ) + return WorkflowDefinition.model_validate(data) + + def get_shared(self, token: str) -> dict[str, Any]: + """Fetch someone's shared definition export by share token. + + Parameters + ---------- + token : str + Share token. + + Returns + ------- + dict + Export (spec, plus ``graph_layout`` when present). + """ + data = self._http.request( + "GET", ROUTES["workflows_definitions_shared"].format(token=token) + ) + assert isinstance(data, dict) + return data + + def fork_shared( + self, + token: str, + name: str, + *, + description: str = "", + idempotency_key: str | None = None, + ) -> WorkflowDefinition: + """Copy a shared definition into your library. + + Parameters + ---------- + token : str + Share token. + name : str + Name for your copy. + description : str, optional + Description for your copy. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + Your new definition. + """ + data = self._http.request( + "POST", + ROUTES["workflows_definitions_shared_fork"].format(token=token), + json={"name": name, "description": description}, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + def versions( + self, + definition_id: str, + *, + limit: int | None = None, + offset: int | None = None, + ) -> WorkflowDefinitionVersionList: + """List saved versions of a definition. + + Parameters + ---------- + definition_id : str + Definition identifier. + limit : int or None, optional + Page size. + offset : int or None, optional + Page offset. + + Returns + ------- + WorkflowDefinitionVersionList + Page of version metadata. + """ + data = self._http.request( + "GET", + ROUTES["workflows_definitions_versions"].format( + definition_id=definition_id + ), + params=_drop_none(limit=limit, offset=offset), + ) + return WorkflowDefinitionVersionList.model_validate(data) + + def get_version( + self, definition_id: str, version: int + ) -> WorkflowDefinitionVersion: + """Fetch one version snapshot of a definition. + + Parameters + ---------- + definition_id : str + Definition identifier. + version : int + Version number. + + Returns + ------- + WorkflowDefinitionVersion + Spec and layout at that version. + """ + data = self._http.request( + "GET", + ROUTES["workflows_definitions_version"].format( + definition_id=definition_id, version=version + ), + ) + return WorkflowDefinitionVersion.model_validate(data) + + def restore_version( + self, + definition_id: str, + version: int, + *, + idempotency_key: str | None = None, + ) -> WorkflowDefinition: + """Restore a definition to an earlier version. + + Parameters + ---------- + definition_id : str + Definition identifier. + version : int + Version number to restore. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + The definition after the restore. + """ + data = self._http.request( + "POST", + ROUTES["workflows_definitions_version_restore"].format( + definition_id=definition_id, version=version + ), + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + +class WorkflowsResource: + """Synchronous workflow endpoints under ``/workflows``. + + Validate and estimate WorkflowSpecs, browse curated templates and catalog + ports, and manage runs (:attr:`runs`) and saved definitions + (:attr:`definitions`). + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + + Attributes + ---------- + runs : WorkflowRunsResource + Workflow run endpoints. + definitions : WorkflowDefinitionsResource + Saved definition endpoints. + """ + + def __init__(self, http: HttpClient) -> None: + """Attach the shared HTTP transport and nested resources. + + Parameters + ---------- + http : HttpClient + Shared HTTP transport. + """ + self._http = http + self.runs = WorkflowRunsResource(http) + self.definitions = WorkflowDefinitionsResource(http) + + def validate( + self, spec: dict[str, Any], *, params: dict[str, Any] | None = None + ) -> WorkflowValidation: + """Validate a WorkflowSpec without running it. + + An invalid spec is not an error: check ``valid`` and ``errors``. + + Parameters + ---------- + spec : dict + WorkflowSpec. + params : dict or None, optional + Runtime param values; supply them to check for missing or unknown + params. + + Returns + ------- + WorkflowValidation + Validity, diagnostics, and step order. + """ + data = self._http.request( + "POST", + ROUTES["workflows_validate"], + json=_spec_body(spec, params), + ) + return WorkflowValidation.model_validate(data) + + def estimate( + self, spec: dict[str, Any], *, params: dict[str, Any] | None = None + ) -> WorkflowEstimate: + """Estimate the wallet hold for a WorkflowSpec at the caller's tier. + + Same formula as run create; nothing is charged. + + Parameters + ---------- + spec : dict + WorkflowSpec. + params : dict or None, optional + Runtime param values (scale per-step reservations). + + Returns + ------- + WorkflowEstimate + Total, per-step costs, and assumptions. + + Raises + ------ + ValidationError + If the spec is invalid (Diagnostic list on ``errors``). + """ + data = self._http.request( + "POST", + ROUTES["workflows_estimate"], + json=_spec_body(spec, params), + ) + return WorkflowEstimate.model_validate(data) + + def templates(self) -> list[WorkflowTemplateSummary]: + """List curated workflow templates in gallery order. + + Returns + ------- + list of WorkflowTemplateSummary + Template outlines (fetch a full spec with :meth:`template`). + """ + data = self._http.request("GET", ROUTES["workflows_templates"]) + return [ + WorkflowTemplateSummary.model_validate(item) + for item in data.get("templates", []) + ] + + def template(self, template_id: str) -> WorkflowTemplate: + """Fetch one curated template with its full WorkflowSpec. + + Parameters + ---------- + template_id : str + Template identifier. + + Returns + ------- + WorkflowTemplate + Template metadata and ``spec``. + """ + data = self._http.request( + "GET", ROUTES["workflows_template"].format(template_id=template_id) + ) + return WorkflowTemplate.model_validate(data) + + def node_ports(self, job_type: str) -> WorkflowNodePorts: + """Return the catalog workflow input / output ports of a job type. + + Parameters + ---------- + job_type : str + Catalog job type. + + Returns + ------- + WorkflowNodePorts + Input and output port definitions. + """ + data = self._http.request( + "GET", ROUTES["workflows_node_ports"], params={"job_type": job_type} + ) + return WorkflowNodePorts.model_validate(data) + + +class AsyncWorkflowRunsResource: + """Asynchronous workflow run endpoints under ``/workflows/runs``. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def create( + self, + run_name: str, + spec: dict[str, Any], + *, + params: dict[str, Any] | None = None, + max_run_cost: float | None = None, + definition_id: str | None = None, + idempotency_key: str | None = None, + ) -> WorkflowRun: + """Create (enqueue) a workflow run. + + The server validates the spec, estimates it at the caller's tier, and + opens a wallet hold before any step starts. + + Parameters + ---------- + run_name : str + Unique-per-user run name. + spec : dict + WorkflowSpec (frozen on the run). + params : dict or None, optional + Runtime param values the spec declares. + max_run_cost : float or None, optional + Spend cap; the run pauses instead of exceeding it. + definition_id : str or None, optional + Saved definition (``wfd-…``) the spec came from. + idempotency_key : str or None, optional + Idempotency key for safe retries. A random key is generated when + omitted (API-key callers must send one). + + Returns + ------- + WorkflowRun + The new run with its step rows and ``estimated_cost``. + """ + body: dict[str, Any] = { + "run_name": run_name, + "spec": spec, + "params": params or {}, + } + if max_run_cost is not None: + body["max_run_cost"] = max_run_cost + if definition_id is not None: + body["definition_id"] = definition_id + data = await self._http.request( + "POST", + ROUTES["workflows_runs_create"], + json=body, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowRun.model_validate(data) + + async def list( + self, + *, + limit: int | None = None, + offset: int | None = None, + status: WorkflowRunStatus | None = None, + ) -> WorkflowRunList: + """List the caller's workflow runs. + + Parameters + ---------- + limit : int or None, optional + Page size (1-100). + offset : int or None, optional + Page offset. + status : str or None, optional + Keep only runs with this status. + + Returns + ------- + WorkflowRunList + Page of runs. + """ + data = await self._http.request( + "GET", + ROUTES["workflows_runs_list"], + params=_drop_none(limit=limit, offset=offset, status=status), + ) + return WorkflowRunList.model_validate(data) + + async def get(self, run_id: str, *, include_spec: bool = False) -> WorkflowRun: + """Fetch a workflow run and its steps. + + Parameters + ---------- + run_id : str + Run identifier. + include_spec : bool, optional + Also return the frozen ``spec`` and ``definition_id``. + + Returns + ------- + WorkflowRun + Run status, costs, and step rows. + """ + data = await self._http.request( + "GET", + ROUTES["workflows_runs_get"].format(run_id=run_id), + params={"include_spec": 1} if include_spec else None, + ) + return WorkflowRun.model_validate(data) + + async def artifacts(self, run_id: str) -> WorkflowRunArtifacts: + """List artifact ports produced by a workflow run. + + Parameters + ---------- + run_id : str + Run identifier. + + Returns + ------- + WorkflowRunArtifacts + Artifact ports by step. + """ + data = await self._http.request( + "GET", ROUTES["workflows_runs_artifacts"].format(run_id=run_id) + ) + return WorkflowRunArtifacts.model_validate(data) + + async def resume(self, run_id: str) -> WorkflowRun: + """Resume a paused workflow run. + + Parameters + ---------- + run_id : str + Run identifier. + + Returns + ------- + WorkflowRun + The resumed run. + """ + data = await self._http.request( + "POST", ROUTES["workflows_runs_resume"].format(run_id=run_id) + ) + return WorkflowRun.model_validate(data) + + async def cancel(self, run_id: str) -> WorkflowRunCancelResponse: + """Cancel a workflow run and its active step jobs. + + Parameters + ---------- + run_id : str + Run identifier. + + Returns + ------- + WorkflowRunCancelResponse + Final status and cancelled job ids. + """ + data = await self._http.request( + "POST", ROUTES["workflows_runs_cancel"].format(run_id=run_id) + ) + return WorkflowRunCancelResponse.model_validate(data) + + async def delete(self, run_id: str) -> None: + """Delete a workflow run record. + + Artifacts stay in storage (usable as ``$artifact`` inputs until they + expire) and linked jobs are kept. + + Parameters + ---------- + run_id : str + Run identifier. + + Returns + ------- + None + The API answers 204 No Content. + """ + await self._http.request( + "DELETE", ROUTES["workflows_runs_delete"].format(run_id=run_id) + ) + + async def wait( + self, + run_id: str, + *, + poll_interval: float = DEFAULT_WORKFLOW_POLL_INTERVAL, + timeout: float = DEFAULT_WORKFLOW_TIMEOUT, + raise_on_failure: bool = False, + ) -> WorkflowRun: + """Poll until the run stops: completed, failed, cancelled, or paused. + + A ``paused`` run needs your action (:meth:`resume`, usually after + raising ``max_run_cost``), so waiting stops there too. + + Parameters + ---------- + run_id : str + Run identifier. + poll_interval : float, optional + Seconds between polls. + timeout : float, optional + Maximum seconds to wait before raising + :class:`~cognichem_client.errors.PollTimeoutError`. + raise_on_failure : bool, optional + When ``True``, raise if the run ends ``failed`` or ``cancelled``. + + Returns + ------- + WorkflowRun + The run at its stop status. + + Raises + ------ + PollTimeoutError + If ``timeout`` elapses first. + ProcessFailedError + If ``raise_on_failure`` is true and the run failed. + ProcessCancelledError + If ``raise_on_failure`` is true and the run was cancelled. + """ + return await apoll_until( + lambda: self.get(run_id), + poll_interval=poll_interval, + timeout=timeout, + stop_statuses=WORKFLOW_RUN_STOP_STATUSES, + subject=lambda run: f"workflow run {run.id}", + failed_statuses=run_failure_statuses(raise_on_failure), + ) + + async def run( + self, + run_name: str, + spec: dict[str, Any], + *, + params: dict[str, Any] | None = None, + max_run_cost: float | None = None, + definition_id: str | None = None, + idempotency_key: str | None = None, + poll_interval: float = DEFAULT_WORKFLOW_POLL_INTERVAL, + timeout: float = DEFAULT_WORKFLOW_TIMEOUT, + ) -> WorkflowRun: + """Create a workflow run and wait until it stops. + + Parameters + ---------- + run_name : str + Unique-per-user run name. + spec : dict + WorkflowSpec. + params : dict or None, optional + Runtime param values. + max_run_cost : float or None, optional + Spend cap. + definition_id : str or None, optional + Saved definition the spec came from. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + poll_interval : float, optional + Seconds between polls. + timeout : float, optional + Maximum seconds to wait. + + Returns + ------- + WorkflowRun + The run at ``completed`` or ``paused``. List its outputs with + :meth:`artifacts`. + + Raises + ------ + PollTimeoutError + If ``timeout`` elapses first. + ProcessFailedError + If the run fails. + ProcessCancelledError + If the run is cancelled. + """ + created = await self.create( + run_name, + spec, + params=params, + max_run_cost=max_run_cost, + definition_id=definition_id, + idempotency_key=idempotency_key, + ) + return await self.wait( + created.id, + poll_interval=poll_interval, + timeout=timeout, + raise_on_failure=True, + ) + + +class AsyncWorkflowDefinitionsResource: + """Asynchronous saved-definition endpoints under ``/workflows/definitions``. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + + async def create( + self, + name: str, + spec: dict[str, Any], + *, + description: str = "", + graph_layout: dict[str, Any] | None = None, + spec_version: int = 1, + idempotency_key: str | None = None, + ) -> WorkflowDefinition: + """Save a workflow definition. + + Parameters + ---------- + name : str + Unique-per-user name. + spec : dict + WorkflowSpec (a top-level ``graph_layout`` moves to its own column). + description : str, optional + Description. + graph_layout : dict or None, optional + Builder layout (wins over one inside ``spec``). + spec_version : int, optional + Spec schema version. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + The saved definition. + """ + body: dict[str, Any] = { + "name": name, + "description": description, + "spec": spec, + "spec_version": spec_version, + } + if graph_layout is not None: + body["graph_layout"] = graph_layout + data = await self._http.request( + "POST", + ROUTES["workflows_definitions_create"], + json=body, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + async def list( + self, *, limit: int | None = None, offset: int | None = None + ) -> WorkflowDefinitionList: + """List saved definitions (without specs). + + Parameters + ---------- + limit : int or None, optional + Page size. + offset : int or None, optional + Page offset. + + Returns + ------- + WorkflowDefinitionList + Page of definitions. + """ + data = await self._http.request( + "GET", + ROUTES["workflows_definitions_list"], + params=_drop_none(limit=limit, offset=offset), + ) + return WorkflowDefinitionList.model_validate(data) + + async def get(self, definition_id: str) -> WorkflowDefinition: + """Fetch a saved definition with its spec. + + Parameters + ---------- + definition_id : str + Definition identifier (``wfd-…``). + + Returns + ------- + WorkflowDefinition + The definition. + """ + data = await self._http.request( + "GET", + ROUTES["workflows_definitions_get"].format(definition_id=definition_id), + ) + return WorkflowDefinition.model_validate(data) + + async def update( + self, + definition_id: str, + *, + name: str | None = None, + description: str | None = None, + spec: dict[str, Any] | None = None, + graph_layout: dict[str, Any] | None = None, + spec_version: int | None = None, + idempotency_key: str | None = None, + ) -> WorkflowDefinition: + """Update a saved definition (only the given fields change). + + Parameters + ---------- + definition_id : str + Definition identifier. + name, description : str or None, optional + New name or description. + spec, graph_layout : dict or None, optional + New spec or builder layout (a spec change adds a version). + spec_version : int or None, optional + New spec schema version. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + The updated definition. + + Raises + ------ + ValueError + If no field is given. + """ + body = _drop_none( + name=name, + description=description, + spec=spec, + graph_layout=graph_layout, + spec_version=spec_version, + ) + if not body: + raise ValueError("Provide at least one field to update") + data = await self._http.request( + "PATCH", + ROUTES["workflows_definitions_update"].format(definition_id=definition_id), + json=body, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + async def delete(self, definition_id: str) -> None: + """Delete a saved definition. + + Parameters + ---------- + definition_id : str + Definition identifier. + + Returns + ------- + None + The API answers 204 No Content. + """ + await self._http.request( + "DELETE", + ROUTES["workflows_definitions_delete"].format(definition_id=definition_id), + ) + + async def download( + self, + definition_id: str, + *, + include_layout: bool = False, + save_path: str | Path | None = None, + ) -> BinaryResult: + """Download a definition as WorkflowSpec JSON. + + Parameters + ---------- + definition_id : str + Definition identifier. + include_layout : bool, optional + Merge the builder ``graph_layout`` into the export. + save_path : str or Path or None, optional + Optional file or directory path to write the JSON. + + Returns + ------- + BinaryResult + JSON bytes plus filename. + """ + result = await self._http.request( + "GET", + ROUTES["workflows_definitions_download"].format( + definition_id=definition_id + ), + params={"include_layout": 1} if include_layout else None, + expect_json=False, + ) + assert isinstance(result, BinaryResult) + return _maybe_save(result, save_path) + + async def share( + self, definition_id: str, *, idempotency_key: str | None = None + ) -> WorkflowDefinition: + """Turn on link sharing and return the ``share_token``. + + Parameters + ---------- + definition_id : str + Definition identifier. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + The definition with ``visibility="link"`` and ``share_token``. + """ + data = await self._http.request( + "POST", + ROUTES["workflows_definitions_share"].format(definition_id=definition_id), + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + async def unshare(self, definition_id: str) -> WorkflowDefinition: + """Revoke link sharing. + + Parameters + ---------- + definition_id : str + Definition identifier. + + Returns + ------- + WorkflowDefinition + The definition with ``visibility="private"``. + """ + data = await self._http.request( + "DELETE", + ROUTES["workflows_definitions_share"].format(definition_id=definition_id), + ) + return WorkflowDefinition.model_validate(data) + + async def get_shared(self, token: str) -> dict[str, Any]: + """Fetch someone's shared definition export by share token. + + Parameters + ---------- + token : str + Share token. + + Returns + ------- + dict + Export (spec, plus ``graph_layout`` when present). + """ + data = await self._http.request( + "GET", ROUTES["workflows_definitions_shared"].format(token=token) + ) + assert isinstance(data, dict) + return data + + async def fork_shared( + self, + token: str, + name: str, + *, + description: str = "", + idempotency_key: str | None = None, + ) -> WorkflowDefinition: + """Copy a shared definition into your library. + + Parameters + ---------- + token : str + Share token. + name : str + Name for your copy. + description : str, optional + Description for your copy. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + Your new definition. + """ + data = await self._http.request( + "POST", + ROUTES["workflows_definitions_shared_fork"].format(token=token), + json={"name": name, "description": description}, + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + async def versions( + self, + definition_id: str, + *, + limit: int | None = None, + offset: int | None = None, + ) -> WorkflowDefinitionVersionList: + """List saved versions of a definition. + + Parameters + ---------- + definition_id : str + Definition identifier. + limit : int or None, optional + Page size. + offset : int or None, optional + Page offset. + + Returns + ------- + WorkflowDefinitionVersionList + Page of version metadata. + """ + data = await self._http.request( + "GET", + ROUTES["workflows_definitions_versions"].format( + definition_id=definition_id + ), + params=_drop_none(limit=limit, offset=offset), + ) + return WorkflowDefinitionVersionList.model_validate(data) + + async def get_version( + self, definition_id: str, version: int + ) -> WorkflowDefinitionVersion: + """Fetch one version snapshot of a definition. + + Parameters + ---------- + definition_id : str + Definition identifier. + version : int + Version number. + + Returns + ------- + WorkflowDefinitionVersion + Spec and layout at that version. + """ + data = await self._http.request( + "GET", + ROUTES["workflows_definitions_version"].format( + definition_id=definition_id, version=version + ), + ) + return WorkflowDefinitionVersion.model_validate(data) + + async def restore_version( + self, + definition_id: str, + version: int, + *, + idempotency_key: str | None = None, + ) -> WorkflowDefinition: + """Restore a definition to an earlier version. + + Parameters + ---------- + definition_id : str + Definition identifier. + version : int + Version number to restore. + idempotency_key : str or None, optional + Idempotency key for safe retries (generated when omitted). + + Returns + ------- + WorkflowDefinition + The definition after the restore. + """ + data = await self._http.request( + "POST", + ROUTES["workflows_definitions_version_restore"].format( + definition_id=definition_id, version=version + ), + idempotency_key=ensure_idempotency_key(idempotency_key), + ) + return WorkflowDefinition.model_validate(data) + + +class AsyncWorkflowsResource: + """Asynchronous workflow endpoints under ``/workflows``. + + Validate and estimate WorkflowSpecs, browse curated templates and catalog + ports, and manage runs (:attr:`runs`) and saved definitions + (:attr:`definitions`). + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + + Attributes + ---------- + runs : AsyncWorkflowRunsResource + Workflow run endpoints. + definitions : AsyncWorkflowDefinitionsResource + Saved definition endpoints. + """ + + def __init__(self, http: AsyncHttpClient) -> None: + """Attach the shared async HTTP transport and nested resources. + + Parameters + ---------- + http : AsyncHttpClient + Shared async HTTP transport. + """ + self._http = http + self.runs = AsyncWorkflowRunsResource(http) + self.definitions = AsyncWorkflowDefinitionsResource(http) + + async def validate( + self, spec: dict[str, Any], *, params: dict[str, Any] | None = None + ) -> WorkflowValidation: + """Validate a WorkflowSpec without running it. + + An invalid spec is not an error: check ``valid`` and ``errors``. + + Parameters + ---------- + spec : dict + WorkflowSpec. + params : dict or None, optional + Runtime param values; supply them to check for missing or unknown + params. + + Returns + ------- + WorkflowValidation + Validity, diagnostics, and step order. + """ + data = await self._http.request( + "POST", + ROUTES["workflows_validate"], + json=_spec_body(spec, params), + ) + return WorkflowValidation.model_validate(data) + + async def estimate( + self, spec: dict[str, Any], *, params: dict[str, Any] | None = None + ) -> WorkflowEstimate: + """Estimate the wallet hold for a WorkflowSpec at the caller's tier. + + Same formula as run create; nothing is charged. + + Parameters + ---------- + spec : dict + WorkflowSpec. + params : dict or None, optional + Runtime param values (scale per-step reservations). + + Returns + ------- + WorkflowEstimate + Total, per-step costs, and assumptions. + + Raises + ------ + ValidationError + If the spec is invalid (Diagnostic list on ``errors``). + """ + data = await self._http.request( + "POST", + ROUTES["workflows_estimate"], + json=_spec_body(spec, params), + ) + return WorkflowEstimate.model_validate(data) + + async def templates(self) -> list[WorkflowTemplateSummary]: + """List curated workflow templates in gallery order. + + Returns + ------- + list of WorkflowTemplateSummary + Template outlines (fetch a full spec with :meth:`template`). + """ + data = await self._http.request("GET", ROUTES["workflows_templates"]) + return [ + WorkflowTemplateSummary.model_validate(item) + for item in data.get("templates", []) + ] + + async def template(self, template_id: str) -> WorkflowTemplate: + """Fetch one curated template with its full WorkflowSpec. + + Parameters + ---------- + template_id : str + Template identifier. + + Returns + ------- + WorkflowTemplate + Template metadata and ``spec``. + """ + data = await self._http.request( + "GET", ROUTES["workflows_template"].format(template_id=template_id) + ) + return WorkflowTemplate.model_validate(data) + + async def node_ports(self, job_type: str) -> WorkflowNodePorts: + """Return the catalog workflow input / output ports of a job type. + + Parameters + ---------- + job_type : str + Catalog job type. + + Returns + ------- + WorkflowNodePorts + Input and output port definitions. + """ + data = await self._http.request( + "GET", ROUTES["workflows_node_ports"], params={"job_type": job_type} + ) + return WorkflowNodePorts.model_validate(data) diff --git a/src/cognichem_client/types/__init__.py b/src/cognichem_client/types/__init__.py new file mode 100644 index 0000000..bdcf1c3 --- /dev/null +++ b/src/cognichem_client/types/__init__.py @@ -0,0 +1,158 @@ +"""Pydantic models for CogniChem API request/response envelopes. + +Every model is importable from ``cognichem_client.types``. +""" + +from cognichem_client.types.artifacts import ( + Artifact, + ArtifactListItem, + ArtifactListResponse, + ArtifactUsage, +) +from cognichem_client.types.chat import ( + ChatCitation, + ChatEstimate, + ChatMessage, + ChatMessageList, + ChatRunProposal, + ChatRunProposalList, + ChatSession, + ChatSessionList, + ChatSpendCap, + ChatTurnResult, +) +from cognichem_client.types.common import ( + ApiKeyCreatedResponse, + ApiKeyListItem, + ApiKeyListResponse, + AuthCheckResponse, + BinaryResult, + DataResult, + HealthResponse, + InferenceSubmitRequest, + JobEstimateResponse, + JobInfoResponse, + JobListItem, + JobSubmitMultipleResponse, + JobSubmitRequest, + JobSubmitResponse, + ListJobsResponse, + ListProcessIdsResponse, + MessageResponse, + ProcessStatus, + ServerSentEvent, + TokenResponse, + UsageLimitsResponse, + WalletBalance, +) +from cognichem_client.types.events import ( + RunEvent, + RunEventList, +) +from cognichem_client.types.science import ( + LookupItem, + LookupResponse, + PubchemCompound, + PubchemLookupResponse, + ReferenceEdge, + ReferenceNeighborhood, + ReferencePaper, + ReferencePaperDetail, + ReferenceProvenance, + ReferenceSearchResponse, +) +from cognichem_client.types.workflows import ( + Diagnostic, + WorkflowDefinition, + WorkflowDefinitionList, + WorkflowDefinitionListItem, + WorkflowDefinitionVersion, + WorkflowDefinitionVersionList, + WorkflowDefinitionVersionListItem, + WorkflowEstimate, + WorkflowNodePorts, + WorkflowRun, + WorkflowRunArtifact, + WorkflowRunArtifacts, + WorkflowRunCancelResponse, + WorkflowRunList, + WorkflowRunListItem, + WorkflowRunStep, + WorkflowStepCost, + WorkflowTemplate, + WorkflowTemplateStep, + WorkflowTemplateSummary, + WorkflowValidation, +) + +__all__ = [ + "ApiKeyCreatedResponse", + "ApiKeyListItem", + "ApiKeyListResponse", + "Artifact", + "ArtifactListItem", + "ArtifactListResponse", + "ArtifactUsage", + "AuthCheckResponse", + "BinaryResult", + "ChatCitation", + "ChatEstimate", + "ChatMessage", + "ChatMessageList", + "ChatRunProposal", + "ChatRunProposalList", + "ChatSession", + "ChatSessionList", + "ChatSpendCap", + "ChatTurnResult", + "DataResult", + "Diagnostic", + "HealthResponse", + "InferenceSubmitRequest", + "JobEstimateResponse", + "JobInfoResponse", + "JobListItem", + "JobSubmitMultipleResponse", + "JobSubmitRequest", + "JobSubmitResponse", + "ListJobsResponse", + "ListProcessIdsResponse", + "LookupItem", + "LookupResponse", + "MessageResponse", + "ProcessStatus", + "PubchemCompound", + "PubchemLookupResponse", + "ReferenceEdge", + "ReferenceNeighborhood", + "ReferencePaper", + "ReferencePaperDetail", + "ReferenceProvenance", + "ReferenceSearchResponse", + "RunEvent", + "RunEventList", + "ServerSentEvent", + "TokenResponse", + "UsageLimitsResponse", + "WalletBalance", + "WorkflowDefinition", + "WorkflowDefinitionList", + "WorkflowDefinitionListItem", + "WorkflowDefinitionVersion", + "WorkflowDefinitionVersionList", + "WorkflowDefinitionVersionListItem", + "WorkflowEstimate", + "WorkflowNodePorts", + "WorkflowRun", + "WorkflowRunArtifact", + "WorkflowRunArtifacts", + "WorkflowRunCancelResponse", + "WorkflowRunList", + "WorkflowRunListItem", + "WorkflowRunStep", + "WorkflowStepCost", + "WorkflowTemplate", + "WorkflowTemplateStep", + "WorkflowTemplateSummary", + "WorkflowValidation", +] diff --git a/src/cognichem_client/types/artifacts.py b/src/cognichem_client/types/artifacts.py new file mode 100644 index 0000000..a017465 --- /dev/null +++ b/src/cognichem_client/types/artifacts.py @@ -0,0 +1,123 @@ +"""Models for durable artifacts under ``/artifacts``.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + + +class ArtifactListItem(BaseModel): + """One chainable manifest port of an owned artifact. + + Attributes + ---------- + artifact_id : str + Artifact identifier (``art-…``); bind it as ``$artifact`` in a + WorkflowSpec. + port : str + Manifest port name (``archive`` only with ``include_archive=True``). + data_kind : str + Port data kind (e.g. ``molecule_set``, ``protein_structure``). + data_format : str + Port data format (e.g. ``smiles``, ``pdb``). + record_count : int + Records on this port. + size_bytes : int + Stored size in bytes. + is_virtual : bool + Whether the artifact is a view over a parent artifact. + created_at, expires_at : str or None + ISO-8601 timestamps. + job_id, run_id, step_id : str or None + Producing job, workflow run, and step, when known. + job_name, job_type, run_name : str or None + Display labels for the producer. + """ + + artifact_id: str + port: str + data_kind: str + data_format: str + record_count: int = 0 + size_bytes: int = 0 + is_virtual: bool = False + created_at: str | None = None + expires_at: str | None = None + job_id: str | None = None + run_id: str | None = None + step_id: str | None = None + job_name: str | None = None + job_type: str | None = None + run_name: str | None = None + + +class ArtifactListResponse(BaseModel): + """Page of artifact ports. + + Paging is over artifacts, not items: every port of one artifact lands on + the same page. + + Attributes + ---------- + items : list of ArtifactListItem + Port rows. + has_more : bool + Whether another page may exist. + next_offset : int or None + ``offset`` for the next page, when ``has_more``. + """ + + items: list[ArtifactListItem] = Field(default_factory=list) + has_more: bool = False + next_offset: int | None = None + + +class ArtifactUsage(BaseModel): + """Artifact storage usage against the tier quota. + + Attributes + ---------- + used_bytes : int + Bytes stored. + quota_bytes : int + Tier quota in bytes. + artifact_count : int + Number of stored artifacts. + tier : int + Caller subscription tier. + """ + + used_bytes: int + quota_bytes: int + artifact_count: int + tier: int + + +class Artifact(BaseModel): + """Owner-visible artifact metadata plus its parsed manifest. + + Attributes + ---------- + id : str + Artifact identifier. + job_id : str or None + Producing job, when known. + size_bytes : int + Stored size in bytes. + content_type : str + Media type of the stored object. + expires_at, created_at : str or None + ISO-8601 timestamps. + manifest : dict + Manifest with ``ports`` → ``records`` (ids, files, metrics). Pass a + record ``id`` to :meth:`ArtifactsResource.download` as ``record``. + """ + + id: str + job_id: str | None = None + size_bytes: int = 0 + content_type: str = "application/octet-stream" + expires_at: str | None = None + created_at: str | None = None + manifest: dict[str, Any] = Field(default_factory=dict) diff --git a/src/cognichem_client/types/chat.py b/src/cognichem_client/types/chat.py new file mode 100644 index 0000000..28deabc --- /dev/null +++ b/src/cognichem_client/types/chat.py @@ -0,0 +1,295 @@ +"""Models for CogniChem Assistant chat (``/chat``; JWT only).""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + + +class ChatSession(BaseModel): + """One Assistant thread. + + Attributes + ---------- + id : str + Session identifier. + title : str or None + Thread title. + created_at, updated_at : str + ISO-8601 timestamps. + allow_structure_search : bool + Whether the Assistant may send SMILES to external databases for + similarity / substructure search in this thread. + """ + + id: str + title: str | None = None + created_at: str + updated_at: str + allow_structure_search: bool = False + + +class ChatSessionList(BaseModel): + """The caller's Assistant threads. + + Attributes + ---------- + items : list of ChatSession + Threads. + """ + + items: list[ChatSession] = Field(default_factory=list) + + +class ChatMessage(BaseModel): + """One persisted chat message. + + Attributes + ---------- + id : str + Message identifier. + session_id : str + Owning thread. + turn_id : str or None + Turn (wallet hold) the message belongs to. + role : str + ``user``, ``assistant``, ``tool``, or ``system``. + content : str or None + Message text. + model_id : str or None + Catalog model id. + prompt_tokens, completion_tokens, cached_tokens, embedding_tokens : int or None + Metered token counts. + billed_usd : float or None + USD debited for the turn. + reasoning_effort : str or None + Level the assistant turn was billed at. + citations, tool_calls : list of dict or None + Citation chips and tool call records. + artifact_ids, job_ids, workflow_run_ids : list of str or None + Linked resources. + created_at : str + ISO-8601 timestamp. + """ + + id: str + session_id: str + turn_id: str | None = None + role: str + content: str | None = None + model_id: str | None = None + prompt_tokens: int | None = None + completion_tokens: int | None = None + cached_tokens: int | None = None + embedding_tokens: int | None = None + billed_usd: float | None = None + reasoning_effort: str | None = None + citations: list[dict[str, Any]] | None = None + tool_calls: list[dict[str, Any]] | None = None + artifact_ids: list[str] | None = None + job_ids: list[str] | None = None + workflow_run_ids: list[str] | None = None + created_at: str + + +class ChatMessageList(BaseModel): + """History for one thread. + + Attributes + ---------- + items : list of ChatMessage + Messages, oldest first. + """ + + items: list[ChatMessage] = Field(default_factory=list) + + +class ChatEstimate(BaseModel): + """Server quote for one Assistant turn's wallet hold. + + Attributes + ---------- + model_id : str + Catalog model id. + reasoning_effort : str + Reasoning level priced. + tier : int + Caller subscription tier. + prompt_tokens : int + Estimated prompt tokens (message plus history window). + max_completion_tokens : int + Completion token ceiling. + hold_usd : float + USD the turn would hold. + """ + + model_id: str + reasoning_effort: str + tier: int + prompt_tokens: int + max_completion_tokens: int + hold_usd: float + + +class ChatSpendCap(BaseModel): + """Assistant spend and spend cap of one thread (USD). + + Attributes + ---------- + spent_usd : float + Debited turns in this thread. + held_usd : float + Holds of turns still in flight. + cap_usd : float + Current thread spend cap. + step_usd : float + How much one Continue raises the cap. + raised : bool + Continue only: whether this request raised the cap. + """ + + spent_usd: float + held_usd: float + cap_usd: float + step_usd: float + raised: bool = False + + +class ChatCitation(BaseModel): + """Citation chip on a proposal. + + Attributes + ---------- + title : str or None + Link title. + url : str + Public URL. + """ + + title: str | None = None + url: str + + +class ChatRunProposal(BaseModel): + """Assistant run proposal (plan card) awaiting the user's decision. + + Attributes + ---------- + id : str + Proposal identifier. + session_id : str + Owning thread. + turn_id, message_id : str or None + Turn and message that proposed it. + kind : str + ``job`` or ``workflow``. + status : str + ``pending``, ``approving``, ``submitted``, ``rejected``, ``expired``, + or ``failed``. + summary, display_name : str or None + Plan text. + job_type, resource : str or None + Catalog job type and resource for job proposals. + estimate_usd : float + Server estimate; pass it back as ``accepted_estimate_usd`` to approve. + estimate_breakdown : dict or None + Per-step estimate for workflows. + tier : int or None + Tier the estimate used. + accepted_estimate_usd, max_run_cost_usd : float or None + Values recorded at approval. + citations : list of ChatCitation + Supporting links. + failure_code, last_failure_code : str or None + Failure details. + job_id, workflow_run_id : str or None + Enqueued run once submitted. + expires_at, created_at : str + ISO-8601 timestamps. + payload_preview, params_preview, spec_outline : dict or None + Previews of what would run. + workflow_spec : dict or None + Frozen WorkflowSpec (``include_spec=True`` only). + """ + + id: str + session_id: str + turn_id: str | None = None + message_id: str | None = None + kind: str + status: str + summary: str | None = None + display_name: str | None = None + job_type: str | None = None + resource: str | None = None + estimate_usd: float + estimate_breakdown: dict[str, Any] | None = None + tier: int | None = None + accepted_estimate_usd: float | None = None + max_run_cost_usd: float | None = None + citations: list[ChatCitation] = Field(default_factory=list) + failure_code: str | None = None + last_failure_code: str | None = None + job_id: str | None = None + workflow_run_id: str | None = None + expires_at: str + created_at: str + payload_preview: dict[str, Any] | None = None + params_preview: dict[str, Any] | None = None + spec_outline: dict[str, Any] | None = None + workflow_spec: dict[str, Any] | None = None + + +class ChatRunProposalList(BaseModel): + """Proposals for one thread. + + Attributes + ---------- + items : list of ChatRunProposal + Plan cards. + """ + + items: list[ChatRunProposal] = Field(default_factory=list) + + +class ChatTurnResult(BaseModel): + """Collected outcome of one streamed Assistant turn. + + Attributes + ---------- + turn_id : str or None + Turn identifier (from the ``hold`` event). + content : str + Concatenated ``token`` deltas. + hold_usd : float or None + Wallet hold opened for the turn. + billed_usd : float or None + USD debited (from the ``done`` event). + reasoning_effort : str or None + Level the turn ran at. + status : str or None + ``done`` status when the turn was stopped or settled early. + tool_calls : list of dict + ``tool_call`` events (``name``, ``id``). + citations : list of dict + Citation events and ``done`` citations. + proposals : list of dict + ``proposal`` events (plan cards to approve or reject). + done : dict or None + Raw ``done`` payload (spend fields included). + error : str or None + Detail from an ``error`` event, if any. + """ + + turn_id: str | None = None + content: str = "" + hold_usd: float | None = None + billed_usd: float | None = None + reasoning_effort: str | None = None + status: str | None = None + tool_calls: list[dict[str, Any]] = Field(default_factory=list) + citations: list[dict[str, Any]] = Field(default_factory=list) + proposals: list[dict[str, Any]] = Field(default_factory=list) + done: dict[str, Any] | None = None + error: str | None = None diff --git a/src/cognichem_client/types.py b/src/cognichem_client/types/common.py similarity index 57% rename from src/cognichem_client/types.py rename to src/cognichem_client/types/common.py index e5e4832..0568a63 100644 --- a/src/cognichem_client/types.py +++ b/src/cognichem_client/types/common.py @@ -1,4 +1,4 @@ -"""Pydantic models for CogniChem API request/response envelopes.""" +"""Core request/response models: auth, API keys, jobs, inference, utilities.""" from __future__ import annotations @@ -44,13 +44,24 @@ class ApiKeyListItem(BaseModel): key_prefix : str or None Non-secret prefix for display. scopes : list of str - Reserved scopes (not enforced yet). + Enforced scopes (``read`` and/or ``write``). created_at : datetime Creation timestamp. last_used_at : datetime or None Last successful authentication time, if any. rotated_at : datetime or None Last rotation time, if any. + expires_at : datetime or None + When the key stops authenticating, if set. + revoked_at : datetime or None + Soft-revoke timestamp, if any. + spend_ceiling_usd : float or None + Per-key USD spend cap (``None`` means wallet-only). + spend_accrued_usd : float or None + USD accrued in the current spend window. + allow_structure_search : bool + Whether the key may send SMILES to ChEMBL similarity / substructure + search via ``/lookup/chembl``. """ id: str @@ -60,6 +71,11 @@ class ApiKeyListItem(BaseModel): created_at: datetime last_used_at: datetime | None = None rotated_at: datetime | None = None + expires_at: datetime | None = None + revoked_at: datetime | None = None + spend_ceiling_usd: float | None = None + spend_accrued_usd: float | None = None + allow_structure_search: bool = False class ApiKeyListResponse(BaseModel): @@ -137,19 +153,106 @@ class JobSubmitMultipleResponse(BaseModel): process_ids: list[str] = Field(default_factory=list) +class JobEstimateResponse(BaseModel): + """Wallet-reservation estimate for one job at the caller's tier. + + Attributes + ---------- + cost : float + Soft-hold estimate in USD. + tier : int + Caller subscription tier used for the rate. + resource : str + Resolved compute resource. + job_type : str + Catalog job type. + expected_runtime_sec : float + Catalog runtime assumption in seconds. + rate_per_sec : float + USD per second at ``tier`` and ``resource``. + assumptions : dict + Estimate basis (``validated`` is false for incomplete payloads). + """ + + cost: float + tier: int + resource: str + job_type: str + expected_runtime_sec: float + rate_per_sec: float + assumptions: dict[str, Any] = Field(default_factory=dict) + + +class JobListItem(BaseModel): + """One Job Queue row. + + Attributes + ---------- + process_id : str + Job process identifier. + job_name : str + User-chosen job name. + job_type : str + Catalog job type. + resource : str + Compute resource. + status : str + ``queued``, ``dispatched``, ``running``, ``completed``, ``error``, or + ``cancelled``. + message : str + Latest status message. + created_at, started_at, finished_at : str or None + ISO-8601 lifecycle timestamps. + runtime_seconds : float + Billed runtime so far. + expected_billing_cost : float or None + Reservation estimate (USD). + charged_amount : float or None + Final charge (USD) once billed. + result_artifact_id : str or None + Durable result artifact id, when stored. + """ + + process_id: str + job_name: str + job_type: str + resource: str + status: str + message: str = "" + created_at: str | None = None + started_at: str | None = None + finished_at: str | None = None + runtime_seconds: float = 0.0 + expected_billing_cost: float | None = None + charged_amount: float | None = None + result_artifact_id: str | None = None + + class ListJobsResponse(BaseModel): - """Response from listing user jobs. + """Paginated standalone jobs (workflow step jobs are excluded). Attributes ---------- job_names : list of str - Job names. + Job names (parallel to ``job_pids``). job_pids : list of str Corresponding process IDs. + items : list of JobListItem + Full rows for the same page. + total : int + Total matching jobs. + limit : int + Page size. + offset : int + Page offset. """ job_names: list[str] = Field(default_factory=list) job_pids: list[str] = Field(default_factory=list) + items: list[JobListItem] = Field(default_factory=list) + total: int = 0 + limit: int = 100 + offset: int = 0 class ProcessStatus(BaseModel): @@ -163,11 +266,15 @@ class ProcessStatus(BaseModel): Current status string (e.g. ``queued``, ``running``, ``completed``). message : str or None Optional status message. + result_artifact_id : str or None + Durable result artifact id for jobs, when stored. Inspect it with + ``client.artifacts`` instead of downloading the whole zip. """ process_id: str status: str message: str | None = None + result_artifact_id: str | None = None class JobInfoResponse(BaseModel): @@ -283,3 +390,47 @@ class AuthCheckResponse(BaseModel): model_config = ConfigDict(extra="allow") message: str | None = None + + +class HealthResponse(BaseModel): + """Response from ``GET /health`` or ``GET /ready``. + + Attributes + ---------- + status : str + ``ok`` (health) or ``ready`` (readiness). + """ + + model_config = ConfigDict(extra="allow") + + status: str + + +class WalletBalance(BaseModel): + """Wallet balance and USD reserved by in-flight jobs, runs and turns. + + Attributes + ---------- + balance : float + Wallet balance in USD. + active_reserved : float + USD currently held by active reservations. + """ + + balance: float + active_reserved: float + + +class ServerSentEvent(BaseModel): + """One ``text/event-stream`` frame. + + Attributes + ---------- + event : str + Event name (``message`` when the frame had no ``event:`` line). + data : any + Parsed JSON ``data``, or the raw string when it is not JSON. + """ + + event: str + data: Any = None diff --git a/src/cognichem_client/types/events.py b/src/cognichem_client/types/events.py new file mode 100644 index 0000000..b76e435 --- /dev/null +++ b/src/cognichem_client/types/events.py @@ -0,0 +1,56 @@ +"""Models for the run completion events feed (``GET /events``).""" + +from __future__ import annotations + +from datetime import datetime + +from pydantic import BaseModel, Field + + +class RunEvent(BaseModel): + """One stop transition of a job or workflow run. + + Events are thin: re-read ``client.jobs.status`` or + ``client.workflows.runs.get`` for detail. + + Attributes + ---------- + event_id : int + Dedupe key (delivery is at-least-once). + cursor : str + Opaque cursor positioned at this event. + kind : str + ``job`` or ``workflow_run``. + subject_id : str + Job process id or workflow run id. + status : str + Jobs: ``completed``, ``error``, ``cancelled``. Workflow runs: + ``completed``, ``failed``, ``cancelled``, ``paused``. + workflow_run_id : str or None + Owning run when the job is a workflow step. + occurred_at : datetime + When the transition happened. + """ + + event_id: int + cursor: str + kind: str + subject_id: str + status: str + workflow_run_id: str | None = None + occurred_at: datetime + + +class RunEventList(BaseModel): + """A page of events after the ``after`` cursor, oldest first. + + Attributes + ---------- + events : list of RunEvent + Events in order. + next_cursor : str or None + Pass as ``after`` on the next call (unchanged when no events). + """ + + events: list[RunEvent] = Field(default_factory=list) + next_cursor: str | None = None diff --git a/src/cognichem_client/types/science.py b/src/cognichem_client/types/science.py new file mode 100644 index 0000000..a5b5ebd --- /dev/null +++ b/src/cognichem_client/types/science.py @@ -0,0 +1,277 @@ +"""Models for science lookups (``/lookup/*``) and the reference KG (``/reference``).""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + + +class LookupItem(BaseModel): + """One lookup result. + + Attributes + ---------- + source : str + Upstream source (``ChEMBL``, ``Open Targets``, ``UniProt``, ``RCSB PDB``, + ``PubChem``). + operation : str + Operation that produced the row. + query : str or None + Input query the row answers. + status : str + ``found``, ``ambiguous``, ``not_found``, ``error``, or + ``structure_egress_not_allowed``. + record_id : str or None + Source record id. + title : str or None + Record title. + url : str or None + Public source link (cite it). + release : str or None + Source release, when known. + license : str + Data license. + data : any + Source-specific payload. + """ + + source: str + operation: str + query: str | None = None + status: str + record_id: str | None = None + title: str | None = None + url: str | None = None + release: str | None = None + license: str = "" + data: Any = None + + +class LookupResponse(BaseModel): + """Lookup results in input order plus source attribution. + + Attributes + ---------- + results : list of LookupItem + One or more rows per query. + attribution : str + Attribution text to show with the data. + """ + + results: list[LookupItem] = Field(default_factory=list) + attribution: str = "" + + +class PubchemCompound(BaseModel): + """One PubChem name resolution. + + Attributes + ---------- + name : str + Queried compound name. + status : str + ``found``, ``ambiguous``, ``not_found``, or ``error``. + cid : int or None + PubChem CID. + title : str or None + PubChem title. + smiles : str or None + Resolved SMILES. + url : str or None + PubChem compound page. + other_cids : list of int or None + Other candidate CIDs when ambiguous. + source : str + Always ``PubChem``. + license : str + Data license. + """ + + name: str + status: str + cid: int | None = None + title: str | None = None + smiles: str | None = None + url: str | None = None + other_cids: list[int] | None = None + source: str = "PubChem" + license: str = "" + + +class PubchemLookupResponse(BaseModel): + """Resolved compounds in input order plus source attribution. + + Attributes + ---------- + compounds : list of PubchemCompound + One row per name. + attribution : str + Attribution text to show with the data. + """ + + compounds: list[PubchemCompound] = Field(default_factory=list) + attribution: str = "" + + +class ReferencePaper(BaseModel): + """One paper in the reference KG (public metadata and links). + + Attributes + ---------- + id : str + Paper id (``paper:…``). + title : str + Title. + year : int or None + Publication year. + venue : str or None + Journal or venue. + authors : list of str + First authors (capped). + authors_total : int + Total author count. + doi : str or None + DOI. + url : str or None + Public landing link. + oa_url : str or None + Open-access link, when known. + paper_set : str + ``method``, ``used_cognichem``, or ``general``. + abstract : str or None + Abstract (only on get, and only when licensed CC0 / CC BY). + abstract_license : str or None + ``CC0`` or ``CC BY`` when an abstract is stored. + abstract_license_url : str or None + Creative Commons page for ``abstract_license``. + abstract_source : str or None + ``Europe PMC`` when an abstract is stored. + abstract_notice : str or None + How the stored abstract differs from the original. + source : str + Ingest source. + retrieved_at : str or None + ISO-8601 retrieval time. + """ + + id: str + title: str + year: int | None = None + venue: str | None = None + authors: list[str] = Field(default_factory=list) + authors_total: int = 0 + doi: str | None = None + url: str | None = None + oa_url: str | None = None + paper_set: str + abstract: str | None = None + abstract_license: str | None = None + abstract_license_url: str | None = None + abstract_source: str | None = None + abstract_notice: str | None = None + source: str = "" + retrieved_at: str | None = None + + +class ReferenceEdge(BaseModel): + """A paper → paper or paper → Product KG node link. + + Attributes + ---------- + src_paper_id : str + Source paper id. + dst_kind : str + ``paper`` or ``product_node``. + dst_id : str + Destination id (e.g. ``job_type:protein-prepare``). + kind : str + ``cites``, ``method_of``, ``uses_tool``, ``compares``, or + ``describes``. + """ + + src_paper_id: str + dst_kind: str + dst_id: str + kind: str + + +class ReferenceProvenance(BaseModel): + """Where a reference answer came from. + + Attributes + ---------- + source : str + Always ``cognichem_reference_kg``. + entity_ids : list of str + Entity ids used in the answer. + outbound : list of dict + Public links to cite. + """ + + source: str = "cognichem_reference_kg" + entity_ids: list[str] = Field(default_factory=list) + outbound: list[dict[str, str]] = Field(default_factory=list) + + +class ReferenceSearchResponse(BaseModel): + """Ranked paper hits from ``GET /reference/papers``. + + Attributes + ---------- + query : str + Search text. + mode : str + ``hybrid`` when the query embedding was used, else ``lexical``. + papers : list of ReferencePaper + Ranked hits. + provenance : ReferenceProvenance + Answer provenance. + """ + + query: str + mode: str = "lexical" + papers: list[ReferencePaper] = Field(default_factory=list) + provenance: ReferenceProvenance = Field(default_factory=ReferenceProvenance) + + +class ReferencePaperDetail(BaseModel): + """One paper plus its links, from ``GET /reference/papers/{paper_id}``. + + Attributes + ---------- + paper : ReferencePaper + The paper. + edges : list of ReferenceEdge + Outgoing links. + cited_by_count : int + Reference-KG papers that cite this one. + provenance : ReferenceProvenance + Answer provenance. + """ + + paper: ReferencePaper + edges: list[ReferenceEdge] = Field(default_factory=list) + cited_by_count: int = 0 + provenance: ReferenceProvenance = Field(default_factory=ReferenceProvenance) + + +class ReferenceNeighborhood(BaseModel): + """Papers linked to a Product KG node or a paper. + + Attributes + ---------- + center : str + Node or paper id the neighborhood is centered on. + edges : list of ReferenceEdge + Links around the center. + papers : list of ReferencePaper + Linked papers. + provenance : ReferenceProvenance + Answer provenance. + """ + + center: str + edges: list[ReferenceEdge] = Field(default_factory=list) + papers: list[ReferencePaper] = Field(default_factory=list) + provenance: ReferenceProvenance = Field(default_factory=ReferenceProvenance) diff --git a/src/cognichem_client/types/workflows.py b/src/cognichem_client/types/workflows.py new file mode 100644 index 0000000..08cf868 --- /dev/null +++ b/src/cognichem_client/types/workflows.py @@ -0,0 +1,557 @@ +"""Models for workflow validate / estimate / catalog, runs, and definitions.""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + + +class Diagnostic(BaseModel): + """WorkflowSpec diagnostic (validation error or warning). + + Attributes + ---------- + code : str + Stable diagnostic code (e.g. ``cycle_detected``, ``missing_param``). + message : str + Human-readable explanation. + step_id, port, path : str or None + Where the problem is, when known. + """ + + code: str + message: str + step_id: str | None = None + port: str | None = None + path: str | None = None + + +class WorkflowValidation(BaseModel): + """Result of ``POST /workflows/validate``. + + Attributes + ---------- + valid : bool + Whether the spec is valid. + errors : list of Diagnostic + Blocking problems (empty when ``valid``). + warnings : list of Diagnostic + Non-blocking notes. + order : list of str + Topological step order (empty when invalid). + step_count : int + Number of steps in the spec. + """ + + valid: bool + errors: list[Diagnostic] = Field(default_factory=list) + warnings: list[Diagnostic] = Field(default_factory=list) + order: list[str] = Field(default_factory=list) + step_count: int = 0 + + +class WorkflowStepCost(BaseModel): + """Per-step cost line of a workflow estimate. + + Attributes + ---------- + step_id : str + Step identifier. + cost : float + Soft-hold USD for the step. + """ + + step_id: str + cost: float + + +class WorkflowEstimate(BaseModel): + """Upper-bound soft-hold estimate for a WorkflowSpec. + + Attributes + ---------- + total : float + Total USD hold. + per_step : list of WorkflowStepCost + Cost per step. + tier : int + Caller subscription tier. + assumptions : list of dict + Per-step estimate basis (``reason``, ``multiplier``, ``unit_cost``, + fan-out caps, unresolved binds, …). + timing_assumptions : list of dict + Coarse concurrency advisories (not dollars, not an ETA). + warnings : list of Diagnostic + Estimate warnings (e.g. ``storage_quota_soft``). + """ + + total: float + per_step: list[WorkflowStepCost] = Field(default_factory=list) + tier: int + assumptions: list[dict[str, Any]] = Field(default_factory=list) + timing_assumptions: list[dict[str, Any]] = Field(default_factory=list) + warnings: list[Diagnostic] = Field(default_factory=list) + + +class WorkflowTemplateStep(BaseModel): + """One step outline of a curated template. + + Attributes + ---------- + id : str + Step identifier. + type : str + Step type (``job``, ``filter``, …). + job_type : str or None + Catalog job type for job steps. + op : str or None + Operation for non-job steps. + """ + + id: str + type: str + job_type: str | None = None + op: str | None = None + + +class WorkflowTemplateSummary(BaseModel): + """Curated template outline (no spec body). + + Attributes + ---------- + id : str + Template identifier. + name : str + Display name. + summary : str + One-line description. + complexity : str + ``simple``, ``moderate``, or ``advanced``. + steps : list of WorkflowTemplateStep + Step outline. + params : list of str + Runtime params the template declares. + """ + + id: str + name: str + summary: str + complexity: str + steps: list[WorkflowTemplateStep] = Field(default_factory=list) + params: list[str] = Field(default_factory=list) + + +class WorkflowTemplate(BaseModel): + """Curated template with its full WorkflowSpec. + + Attributes + ---------- + id : str + Template identifier. + name : str + Display name. + summary : str + One-line description. + complexity : str + Template complexity. + spec : dict + WorkflowSpec to validate, estimate, or run. + """ + + id: str + name: str + summary: str + complexity: str + spec: dict[str, Any] + + +class WorkflowNodePorts(BaseModel): + """Catalog workflow input / output ports for one job type. + + Attributes + ---------- + job_type : str + Catalog job type. + inputs : dict + Input ports (kind, formats, payload field, cardinality, …). + outputs : dict + Output ports (kind, format, metrics, resolver). + """ + + job_type: str + inputs: dict[str, Any] = Field(default_factory=dict) + outputs: dict[str, Any] = Field(default_factory=dict) + + +class WorkflowRunStep(BaseModel): + """One step row of a workflow run. + + Attributes + ---------- + id : str + Step row identifier. + step_key : str + Spec step id. + step_index : int + Fan-out index. + status : str + Step status. + job_type : str or None + Catalog job type. + job_id : str or None + Linked job process id. + status_message : str + Latest status message. + """ + + id: str + step_key: str + step_index: int = 0 + status: str + job_type: str | None = None + job_id: str | None = None + status_message: str = "" + + +class WorkflowRun(BaseModel): + """Workflow run with step summary (create / get / resume). + + Attributes + ---------- + id : str + Run identifier. + run_name : str + User-chosen run name. + status : str + ``pending``, ``running``, ``paused``, ``completed``, ``failed``, or + ``cancelled``. + status_message : str + Latest status message. + step_counts : dict + Step counts by status. + remaining_hold, estimated_cost, charged_amount : float + Wallet amounts in USD. + max_run_cost : float or None + Spend cap; the run pauses when it would be exceeded. + steps : list of WorkflowRunStep + Step rows. + created_at, updated_at, finished_at : str or None + ISO-8601 timestamps. + spec : dict or None + Frozen spec (only with ``include_spec=True``). + definition_id : str or None + Saved definition, when the run came from one (``include_spec=True``). + warnings : list of dict + Create-time warnings (e.g. ``storage_quota_soft``). + """ + + id: str + run_name: str + status: str + status_message: str = "" + step_counts: dict[str, Any] = Field(default_factory=dict) + remaining_hold: float = 0.0 + estimated_cost: float = 0.0 + charged_amount: float = 0.0 + max_run_cost: float | None = None + steps: list[WorkflowRunStep] = Field(default_factory=list) + created_at: str | None = None + updated_at: str | None = None + finished_at: str | None = None + spec: dict[str, Any] | None = None + definition_id: str | None = None + warnings: list[dict[str, Any]] = Field(default_factory=list) + + +class WorkflowRunListItem(BaseModel): + """Workflow run list row (no spec, params, or steps). + + Attributes + ---------- + id : str + Run identifier. + run_name : str + Run name. + status : str + Run status. + status_message : str + Latest status message. + step_counts : dict + Step counts by status. + estimated_cost, charged_amount, remaining_hold : float + Wallet amounts in USD. + max_run_cost : float or None + Spend cap. + created_at, updated_at, finished_at : str or None + ISO-8601 timestamps. + definition_id : str or None + Saved definition, if any. + """ + + id: str + run_name: str + status: str + status_message: str = "" + step_counts: dict[str, Any] = Field(default_factory=dict) + estimated_cost: float = 0.0 + charged_amount: float = 0.0 + remaining_hold: float = 0.0 + max_run_cost: float | None = None + created_at: str | None = None + updated_at: str | None = None + finished_at: str | None = None + definition_id: str | None = None + + +class WorkflowRunList(BaseModel): + """Paginated workflow runs. + + Attributes + ---------- + items : list of WorkflowRunListItem + Page rows. + total : int + Total matching runs. + limit : int + Page size. + offset : int + Page offset. + """ + + items: list[WorkflowRunListItem] = Field(default_factory=list) + total: int = 0 + limit: int = 0 + offset: int = 0 + + +class WorkflowRunArtifact(BaseModel): + """One artifact port produced by a workflow run. + + Attributes + ---------- + artifact_id : str + Artifact identifier. + step_key : str or None + Producing spec step id. + step_index : int or None + Fan-out index. + port : str + Manifest port. + data_kind, data_format : str + Port kind and format. + record_count, size_bytes : int + Records and stored size. + is_virtual : bool + Whether the artifact is a view over a parent. + expires_at : str or None + ISO-8601 expiry. + job_id : str or None + Producing job. + """ + + artifact_id: str + step_key: str | None = None + step_index: int | None = None + port: str + data_kind: str + data_format: str + record_count: int = 0 + size_bytes: int = 0 + is_virtual: bool = False + expires_at: str | None = None + job_id: str | None = None + + +class WorkflowRunArtifacts(BaseModel): + """Run-scoped artifact list. + + Attributes + ---------- + items : list of WorkflowRunArtifact + Artifact ports. + """ + + items: list[WorkflowRunArtifact] = Field(default_factory=list) + + +class WorkflowRunCancelResponse(BaseModel): + """Response from ``POST /workflows/runs/{run_id}/cancel``. + + Attributes + ---------- + run_id : str + Run identifier. + status : str + Run status after the request. + cancelled_job_ids : list of str + Step jobs that were cancelled. + message : str + Server message. + """ + + run_id: str + status: str + cancelled_job_ids: list[str] = Field(default_factory=list) + message: str = "" + + +class WorkflowDefinition(BaseModel): + """Saved workflow definition (``wfd-…``). + + Attributes + ---------- + id : str + Definition identifier. + user_id : str + Owner. + name : str + Unique per-user name. + description : str + Description. + spec : dict + WorkflowSpec. + spec_version : int + Spec schema version. + graph_layout : dict or None + Builder layout. + visibility : str + ``private`` or ``link``. + share_token : str or None + Share token (owner responses when ``visibility`` is ``link``). + created_at, updated_at : str or None + ISO-8601 timestamps. + """ + + id: str + user_id: str + name: str + description: str = "" + spec: dict[str, Any] = Field(default_factory=dict) + spec_version: int = 1 + graph_layout: dict[str, Any] | None = None + visibility: str = "private" + share_token: str | None = None + created_at: str | None = None + updated_at: str | None = None + + +class WorkflowDefinitionListItem(BaseModel): + """Definition list row (no spec or share token). + + Attributes + ---------- + id : str + Definition identifier. + user_id : str + Owner. + name : str + Name. + description : str + Description. + spec_version : int + Spec schema version. + graph_layout : dict or None + Builder layout. + visibility : str + ``private`` or ``link``. + created_at, updated_at : str or None + ISO-8601 timestamps. + """ + + id: str + user_id: str + name: str + description: str = "" + spec_version: int = 1 + graph_layout: dict[str, Any] | None = None + visibility: str = "private" + created_at: str | None = None + updated_at: str | None = None + + +class WorkflowDefinitionList(BaseModel): + """Paginated workflow definitions. + + Attributes + ---------- + items : list of WorkflowDefinitionListItem + Page rows. + total : int + Total definitions. + limit : int + Page size. + offset : int + Page offset. + """ + + items: list[WorkflowDefinitionListItem] = Field(default_factory=list) + total: int = 0 + limit: int = 0 + offset: int = 0 + + +class WorkflowDefinitionVersionListItem(BaseModel): + """Definition version metadata. + + Attributes + ---------- + version : int + Version number. + created_at : str or None + ISO-8601 timestamp. + created_by_user_id : str or None + Author. + """ + + version: int + created_at: str | None = None + created_by_user_id: str | None = None + + +class WorkflowDefinitionVersionList(BaseModel): + """Paginated definition versions. + + Attributes + ---------- + items : list of WorkflowDefinitionVersionListItem + Page rows. + total : int + Total versions. + limit : int + Page size. + offset : int + Page offset. + """ + + items: list[WorkflowDefinitionVersionListItem] = Field(default_factory=list) + total: int = 0 + limit: int = 0 + offset: int = 0 + + +class WorkflowDefinitionVersion(BaseModel): + """Full definition version snapshot. + + Attributes + ---------- + definition_id : str + Definition identifier. + version : int + Version number. + spec : dict + WorkflowSpec at this version. + graph_layout : dict or None + Builder layout at this version. + created_at : str or None + ISO-8601 timestamp. + created_by_user_id : str or None + Author. + """ + + definition_id: str + version: int + spec: dict[str, Any] = Field(default_factory=dict) + graph_layout: dict[str, Any] | None = None + created_at: str | None = None + created_by_user_id: str | None = None diff --git a/tests/test_artifacts_jobs.py b/tests/test_artifacts_jobs.py new file mode 100644 index 0000000..6ba9d1e --- /dev/null +++ b/tests/test_artifacts_jobs.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import json + +import httpx +import pytest +import respx + +from cognichem_client import AsyncCogniChem, CogniChem + +BASE = "https://api.test.cognichem.com" +API = f"{BASE}/api/v1" + +ARTIFACT = { + "id": "art-1", + "job_id": "job-1", + "size_bytes": 10, + "content_type": "application/zip", + "manifest": {"ports": {"poses": {"records": [{"id": "lig-1"}]}}}, +} + + +@respx.mock +def test_artifacts_list_usage_get_delete() -> None: + listing = respx.get(f"{API}/artifacts").mock( + return_value=httpx.Response( + 200, + json={ + "items": [ + { + "artifact_id": "art-1", + "port": "poses", + "data_kind": "docking_poses", + "data_format": "pdbqt", + "record_count": 1, + } + ], + "has_more": True, + "next_offset": 25, + }, + ) + ) + respx.get(f"{API}/artifacts/usage").mock( + return_value=httpx.Response( + 200, + json={ + "used_bytes": 1, + "quota_bytes": 10, + "artifact_count": 1, + "tier": 2, + }, + ) + ) + respx.get(f"{API}/artifacts/art-1").mock( + return_value=httpx.Response(200, json=ARTIFACT) + ) + respx.delete(f"{API}/artifacts/art-1").mock( + return_value=httpx.Response(200, json={"message": "Artifact art-1 deleted"}) + ) + client = CogniChem(api_key="k", base_url=BASE) + + page = client.artifacts.list(data_kind="docking_poses", include_archive=True) + assert page.items[0].artifact_id == "art-1" + assert page.next_offset == 25 + params = listing.calls.last.request.url.params + assert params["data_kind"] == "docking_poses" + assert params["include_archive"] == "true" + assert "limit" not in params + + assert client.artifacts.usage().tier == 2 + assert client.artifacts.get("art-1").manifest["ports"]["poses"] + assert "deleted" in client.artifacts.delete("art-1").message + + +@respx.mock +def test_artifact_record_download(tmp_path) -> None: # type: ignore[no-untyped-def] + route = respx.get(f"{API}/artifacts/art-1/download").mock( + return_value=httpx.Response( + 200, + content=b"ATOM", + headers={"Content-Disposition": "attachment; filename=lig-1.pdbqt"}, + ) + ) + client = CogniChem(api_key="k", base_url=BASE) + result = client.artifacts.download( + "art-1", record="lig-1", port="poses", save_path=tmp_path + ) + assert result.content == b"ATOM" + assert (tmp_path / "lig-1.pdbqt").read_bytes() == b"ATOM" + assert route.calls.last.request.url.params["record"] == "lig-1" + assert route.calls.last.request.url.params["port"] == "poses" + + +@respx.mock +def test_artifact_upload_raw_file(tmp_path) -> None: # type: ignore[no-untyped-def] + route = respx.post(f"{API}/artifacts").mock( + return_value=httpx.Response(201, json=ARTIFACT) + ) + pdb = tmp_path / "target.pdb" + pdb.write_text("ATOM 1 N\n") + client = CogniChem(api_key="k", base_url=BASE) + artifact = client.artifacts.upload( + pdb, data_kind="protein_structure", data_format="pdb" + ) + assert artifact.id == "art-1" + request = route.calls.last.request + assert request.headers["Content-Type"].startswith("multipart/form-data") + assert request.headers["Idempotency-Key"] + body = request.content + assert b'name="mode"' in body and b"raw_file" in body + assert b'filename="target.pdb"' in body + assert b"protein_structure" in body + + with pytest.raises(ValueError): + client.artifacts.upload(b"x", mode="raw_file") + + +@respx.mock +def test_job_estimate_list_filters_and_status_artifact() -> None: + respx.post(f"{API}/jobs/estimate").mock( + return_value=httpx.Response( + 200, + json={ + "cost": 0.12, + "tier": 1, + "resource": "cpu", + "job_type": "autodockvina", + "expected_runtime_sec": 600.0, + "rate_per_sec": 0.0002, + "assumptions": {"validated": True}, + }, + ) + ) + listing = respx.get(f"{API}/jobs/list").mock( + return_value=httpx.Response( + 200, + json={ + "job_names": ["a"], + "job_pids": ["job-a"], + "items": [ + { + "process_id": "job-a", + "job_name": "a", + "job_type": "autodockvina", + "resource": "cpu", + "status": "running", + } + ], + "total": 1, + "limit": 20, + "offset": 0, + }, + ) + ) + respx.get(f"{API}/jobs/status").mock( + return_value=httpx.Response( + 200, + json={ + "process_id": "job-a", + "status": "completed", + "result_artifact_id": "art-1", + }, + ) + ) + client = CogniChem(api_key="k", base_url=BASE) + estimate = client.jobs.estimate("autodockvina", {"x": 1}, resource="cpu") + assert estimate.cost == 0.12 + + page = client.jobs.list(limit=20, status=["queued", "running"], q="dock") + assert page.items[0].status == "running" + params = listing.calls.last.request.url.params + assert params["status"] == "queued,running" + assert params["q"] == "dock" + assert "sort" not in params + + assert client.jobs.status("job-a").result_artifact_id == "art-1" + + +@respx.mock +def test_submit_generates_idempotency_key_per_call() -> None: + route = respx.post(f"{API}/jobs/submit").mock( + return_value=httpx.Response(200, json={"process_id": "job-1"}) + ) + client = CogniChem(api_key="k", base_url=BASE) + client.jobs.submit("a", "padel-descriptor", {}) + client.jobs.submit("b", "padel-descriptor", {}) + keys = [call.request.headers["Idempotency-Key"] for call in route.calls] + assert all(keys) and keys[0] != keys[1] + + +@respx.mock +def test_inference_and_utils_cancel() -> None: + inference = respx.post(f"{API}/inference/cancel").mock( + return_value=httpx.Response(200, json={"message": "cancelled"}) + ) + utils = respx.post(f"{API}/utils/cancel").mock( + return_value=httpx.Response(200, json={"message": "cancelled"}) + ) + client = CogniChem(api_key="k", base_url=BASE) + assert client.inference.cancel("inf-1").message == "cancelled" + assert client.utils.cancel("util-1").message == "cancelled" + assert inference.calls.last.request.url.params["process_id"] == "inf-1" + assert utils.calls.last.request.url.params["process_id"] == "util-1" + + +@respx.mock +def test_api_keys_create_update_delete() -> None: + key = { + "id": "key-1", + "name": "ci", + "scopes": ["read"], + "created_at": "2026-09-01T00:00:00Z", + "allow_structure_search": True, + "spend_ceiling_usd": 5.0, + } + create = respx.post(f"{API}/auth/api-keys").mock( + return_value=httpx.Response(201, json={**key, "api_key": "secret"}) + ) + update = respx.patch(f"{API}/auth/api-keys/key-1").mock( + return_value=httpx.Response(200, json=key) + ) + respx.delete(f"{API}/auth/api-keys/key-1").mock(return_value=httpx.Response(204)) + client = CogniChem(access_token="jwt", base_url=BASE) + created = client.api_keys.create( + "ci", + scopes=["read"], + spend_ceiling_usd=5.0, + allow_structure_search=True, + expires_at="2027-01-01T00:00:00Z", + ) + assert created.api_key == "secret" + assert created.allow_structure_search is True + assert json.loads(create.calls.last.request.content) == { + "name": "ci", + "scopes": ["read"], + "expires_at": "2027-01-01T00:00:00Z", + "spend_ceiling_usd": 5.0, + "allow_structure_search": True, + } + assert create.calls.last.request.headers["Authorization"] == "Bearer jwt" + + client.api_keys.update("key-1", allow_structure_search=False) + assert json.loads(update.calls.last.request.content) == { + "allow_structure_search": False + } + client.api_keys.rename("key-1", "prod") + assert json.loads(update.calls.last.request.content) == {"name": "prod"} + assert client.api_keys.delete("key-1") is None + + +@respx.mock +def test_health_and_ready() -> None: + respx.get(f"{BASE}/health").mock( + return_value=httpx.Response(200, json={"status": "ok"}) + ) + respx.get(f"{BASE}/ready").mock( + return_value=httpx.Response(200, json={"status": "ready"}) + ) + client = CogniChem(api_key="k", base_url=BASE) + assert client.health().status == "ok" + assert client.ready().status == "ready" + + +@respx.mock +async def test_async_artifact_upload_bytes() -> None: + route = respx.post(f"{API}/artifacts").mock( + return_value=httpx.Response(201, json=ARTIFACT) + ) + async with AsyncCogniChem(api_key="k", base_url=BASE) as client: + artifact = await client.artifacts.upload(b"PK\x03\x04") + assert artifact.id == "art-1" + body = route.calls.last.request.content + assert b"cognichem_zip" in body and b'filename="result.zip"' in body diff --git a/tests/test_errors.py b/tests/test_errors.py index 1f9a048..1f5e4ce 100644 --- a/tests/test_errors.py +++ b/tests/test_errors.py @@ -4,9 +4,15 @@ from cognichem_client.errors import ( AuthenticationError, + BadRequestError, + CogniChemError, ConflictError, ForbiddenError, + GoneError, + PayloadTooLargeError, + PaymentRequiredError, RateLimitError, + ServerError, ValidationError, raise_for_problem, ) @@ -67,3 +73,40 @@ def test_validation_errors_list() -> None: def test_rate_limit() -> None: with pytest.raises(RateLimitError): raise_for_problem(429, {"detail": "Slow down"}) + + +@pytest.mark.parametrize( + ("status", "cls"), + [ + (400, BadRequestError), + (402, PaymentRequiredError), + (410, GoneError), + (413, PayloadTooLargeError), + (503, ServerError), + ], +) +def test_status_classes(status: int, cls: type[CogniChemError]) -> None: + with pytest.raises(cls): + raise_for_problem(status, {"detail": "nope"}) + + +def test_problem_type_slug_code_and_retryability() -> None: + with pytest.raises(BadRequestError) as exc: + raise_for_problem( + 400, + { + "type": "/api/v1/problems/missing-idempotency-key", + "title": "Missing Idempotency-Key", + "status": 400, + "detail": "Idempotency-Key is required for API-key mutations", + "retryability": "validation", + }, + ) + assert exc.value.code == "missing-idempotency-key" + assert exc.value.retryability == "validation" + + +def test_generic_http_problem_code_is_unknown() -> None: + with pytest.raises(CogniChemError) as exc: + raise_for_problem(500, {"type": "/api/v1/problems/http-500", "detail": "x"}) + assert exc.value.code == "unknown" diff --git a/tests/test_science_events_chat.py b/tests/test_science_events_chat.py new file mode 100644 index 0000000..0c54261 --- /dev/null +++ b/tests/test_science_events_chat.py @@ -0,0 +1,415 @@ +from __future__ import annotations + +import json + +import httpx +import pytest +import respx + +from cognichem_client import AsyncCogniChem, CogniChem +from cognichem_client.errors import ( + CogniChemError, + ConflictError, + PaymentRequiredError, +) + +BASE = "https://api.test.cognichem.com" +API = f"{BASE}/api/v1" + + +def _lookup(source: str) -> dict[str, object]: + return { + "results": [ + { + "source": source, + "operation": "entry", + "query": "q", + "status": "found", + "license": "CC BY 4.0", + "data": {"x": 1}, + } + ], + "attribution": f"Data from {source}", + } + + +@respx.mock +def test_lookups_send_expected_bodies() -> None: + pubchem = respx.post(f"{API}/lookup/pubchem").mock( + return_value=httpx.Response( + 200, + json={ + "compounds": [ + { + "name": "aspirin", + "status": "found", + "cid": 2244, + "smiles": "CC(=O)OC1=CC=CC=C1C(=O)O", + "license": "Public domain", + } + ], + "attribution": "Data from PubChem", + }, + ) + ) + chembl = respx.post(f"{API}/lookup/chembl").mock( + return_value=httpx.Response(200, json=_lookup("ChEMBL")) + ) + targets = respx.post(f"{API}/lookup/targets").mock( + return_value=httpx.Response(200, json=_lookup("Open Targets")) + ) + uniprot = respx.post(f"{API}/lookup/uniprot").mock( + return_value=httpx.Response(200, json=_lookup("UniProt")) + ) + pdb = respx.post(f"{API}/lookup/pdb").mock( + return_value=httpx.Response(200, json=_lookup("RCSB PDB")) + ) + props = respx.post(f"{API}/lookup/properties").mock( + return_value=httpx.Response(200, json=_lookup("PubChem")) + ) + client = CogniChem(api_key="k", base_url=BASE) + + assert client.lookup.pubchem("aspirin").compounds[0].cid == 2244 + assert json.loads(pubchem.calls.last.request.content) == {"names": ["aspirin"]} + + result = client.lookup.chembl("similarity", ["CCO"], limit=5, threshold=80) + assert result.results[0].source == "ChEMBL" + assert json.loads(chembl.calls.last.request.content) == { + "operation": "similarity", + "queries": ["CCO"], + "limit": 5, + "threshold": 80, + } + + client.lookup.targets("target_diseases", "EGFR") + assert json.loads(targets.calls.last.request.content) == { + "operation": "target_diseases", + "queries": ["EGFR"], + } + client.lookup.uniprot("search", ["EGFR"], organism="9606") + assert json.loads(uniprot.calls.last.request.content)["organism"] == "9606" + client.lookup.pdb("entry", ["1M17"]) + assert json.loads(pdb.calls.last.request.content)["queries"] == ["1M17"] + client.lookup.properties(["methanol"], properties=["boiling_point"], limit=2) + assert json.loads(props.calls.last.request.content) == { + "operation": "experimental", + "queries": ["methanol"], + "properties": ["boiling_point"], + "limit": 2, + } + + +@respx.mock +def test_wallet_and_reference() -> None: + respx.get(f"{API}/wallet").mock( + return_value=httpx.Response(200, json={"balance": 12.5, "active_reserved": 1}) + ) + paper = {"id": "paper:propka", "title": "PROPKA3", "paper_set": "method"} + search = respx.get(f"{API}/reference/papers").mock( + return_value=httpx.Response( + 200, json={"query": "pKa", "mode": "hybrid", "papers": [paper]} + ) + ) + respx.get(f"{API}/reference/papers/paper:propka").mock( + return_value=httpx.Response( + 200, + json={ + "paper": paper, + "edges": [ + { + "src_paper_id": "paper:propka", + "dst_kind": "product_node", + "dst_id": "job_type:protein-prepare", + "kind": "method_of", + } + ], + "cited_by_count": 3, + "provenance": {"entity_ids": ["paper:propka"]}, + }, + ) + ) + hood = respx.get(f"{API}/reference/neighborhood").mock( + return_value=httpx.Response( + 200, + json={"center": "job_type:protein-prepare", "papers": [paper]}, + ) + ) + client = CogniChem(api_key="k", base_url=BASE) + assert client.wallet.balance().balance == 12.5 + + hits = client.reference.search("pKa", paper_set="method", limit=5) + assert hits.mode == "hybrid" and hits.papers[0].id == "paper:propka" + assert search.calls.last.request.url.params["paper_set"] == "method" + detail = client.reference.get("paper:propka") + assert detail.edges[0].kind == "method_of" and detail.cited_by_count == 3 + client.reference.neighborhood(node_id="job_type:protein-prepare") + params = hood.calls.last.request.url.params + assert params["node_id"] == "job_type:protein-prepare" + assert "paper_id" not in params + + +def _event(event_id: int) -> dict[str, object]: + return { + "event_id": event_id, + "cursor": f"c{event_id}", + "kind": "job", + "subject_id": f"job-{event_id}", + "status": "completed", + "occurred_at": "2026-09-27T12:00:00Z", + } + + +@respx.mock +def test_events_list_and_listen_dedupes() -> None: + pages = iter( + [ + {"events": [_event(1), _event(2)], "next_cursor": "c2"}, + {"events": [_event(2), _event(3)], "next_cursor": "c3"}, + ] + ) + route = respx.get(f"{API}/events").mock( + side_effect=lambda request: httpx.Response(200, json=next(pages)) + ) + client = CogniChem(api_key="k", base_url=BASE) + received = [] + for event in client.events.listen(kind="job", wait=1): + received.append(event.event_id) + if len(received) == 3: + break + assert received == [1, 2, 3] + first, second = (call.request.url.params for call in route.calls) + assert "after" not in first and first["kind"] == "job" and first["wait"] == "1" + assert second["after"] == "c2" + + +def _sse(*frames: tuple[str, dict[str, object]]) -> bytes: + return "".join( + f"event: {name}\ndata: {json.dumps(data)}\n\n" for name, data in frames + ).encode() + + +TURN = _sse( + ("hold", {"turn_id": "t-1", "hold_usd": 0.05, "reasoning_effort": "low"}), + ("token", {"delta": "Use "}), + ("tool_call", {"name": "kg_product_how_to_run", "id": "call-1"}), + ("token", {"delta": "autodockvina."}), + ("proposal", {"id": "prop-1", "estimate_usd": 0.12}), + ( + "done", + { + "turn_id": "t-1", + "billed_usd": 0.01, + "held_usd": 0.05, + "citations": [{"url": "https://app.cognichem.com/tool/autodockvina"}], + }, + ), +) + + +@respx.mock +def test_chat_send_collects_stream_with_bearer() -> None: + route = respx.post(f"{API}/chat/sessions/s-1/turns").mock( + return_value=httpx.Response( + 200, content=TURN, headers={"Content-Type": "text/event-stream"} + ) + ) + client = CogniChem(api_key="k", access_token="jwt", base_url=BASE) + result = client.chat.send("s-1", "How do I dock?", reasoning_effort="medium") + assert result.turn_id == "t-1" + assert result.content == "Use autodockvina." + assert result.billed_usd == 0.01 + assert result.tool_calls[0]["name"] == "kg_product_how_to_run" + assert result.proposals[0]["id"] == "prop-1" + assert result.citations[0]["url"].endswith("autodockvina") + request = route.calls.last.request + assert request.headers["Authorization"] == "Bearer jwt" + assert "X-Api-Key" not in request.headers + assert json.loads(request.content) == { + "message": "How do I dock?", + "model_id": "cognichem-assistant", + "artifact_ids": [], + "reasoning_effort": "medium", + } + + +@respx.mock +def test_chat_stream_errors() -> None: + respx.post(f"{API}/chat/sessions/s-1/turns").mock( + return_value=httpx.Response( + 402, + json={ + "type": "/api/v1/problems/insufficient-wallet", + "status": 402, + "detail": "Top up", + }, + ) + ) + respx.post(f"{API}/chat/sessions/s-2/turns").mock( + return_value=httpx.Response( + 409, + json={ + "type": "/api/v1/problems/session-spend-cap", + "status": 409, + "detail": "Cap reached", + "cap_usd": 1.0, + "step_usd": 1.0, + }, + ) + ) + respx.post(f"{API}/chat/sessions/s-3/turns").mock( + return_value=httpx.Response( + 200, content=_sse(("error", {"detail": "Assistant turn failed"})) + ) + ) + client = CogniChem(access_token="jwt", base_url=BASE) + with pytest.raises(PaymentRequiredError) as paid: + client.chat.send("s-1", "hi") + assert paid.value.code == "insufficient-wallet" + with pytest.raises(ConflictError) as cap: + list(client.chat.stream("s-2", "hi")) + assert cap.value.code == "session-spend-cap" + assert cap.value.body["cap_usd"] == 1.0 + with pytest.raises(CogniChemError) as failed: + client.chat.send("s-3", "hi") + assert failed.value.code == "assistant-turn-error" + + +@respx.mock +def test_chat_sessions_spend_cap_and_proposals() -> None: + session = { + "id": "s-1", + "title": "Docking", + "created_at": "2026-09-27T00:00:00Z", + "updated_at": "2026-09-27T00:00:00Z", + } + proposal = { + "id": "p-1", + "session_id": "s-1", + "kind": "job", + "status": "pending", + "estimate_usd": 0.12, + "expires_at": "2026-09-28T00:00:00Z", + "created_at": "2026-09-27T00:00:00Z", + } + respx.post(f"{API}/chat/sessions").mock( + return_value=httpx.Response(201, json=session) + ) + respx.get(f"{API}/chat/sessions").mock( + return_value=httpx.Response(200, json={"items": [session]}) + ) + patch = respx.patch(f"{API}/chat/sessions/s-1").mock( + return_value=httpx.Response( + 200, json={**session, "allow_structure_search": True} + ) + ) + respx.delete(f"{API}/chat/sessions/s-1").mock(return_value=httpx.Response(204)) + respx.get(f"{API}/chat/sessions/s-1/messages").mock( + return_value=httpx.Response( + 200, + json={ + "items": [ + { + "id": "m-1", + "session_id": "s-1", + "role": "user", + "content": "hi", + "created_at": "2026-09-27T00:00:00Z", + } + ] + }, + ) + ) + respx.post(f"{API}/chat/sessions/s-1/estimate").mock( + return_value=httpx.Response( + 200, + json={ + "model_id": "cognichem-assistant", + "reasoning_effort": "low", + "tier": 1, + "prompt_tokens": 100, + "max_completion_tokens": 2000, + "hold_usd": 0.03, + }, + ) + ) + respx.get(f"{API}/chat/sessions/s-1/spend-cap").mock( + return_value=httpx.Response( + 200, + json={"spent_usd": 1, "held_usd": 0, "cap_usd": 1, "step_usd": 1}, + ) + ) + cont = respx.post(f"{API}/chat/sessions/s-1/spend-cap/continue").mock( + return_value=httpx.Response( + 200, + json={ + "spent_usd": 1, + "held_usd": 0, + "cap_usd": 2, + "step_usd": 1, + "raised": True, + }, + ) + ) + respx.post(f"{API}/chat/sessions/s-1/turns/t-1/stop").mock( + return_value=httpx.Response(200, json={"status": "stopping", "turn_id": "t-1"}) + ) + respx.get(f"{API}/chat/sessions/s-1/proposals").mock( + return_value=httpx.Response(200, json={"items": [proposal]}) + ) + get_prop = respx.get(f"{API}/chat/sessions/s-1/proposals/p-1").mock( + return_value=httpx.Response(200, json=proposal) + ) + approve = respx.post(f"{API}/chat/sessions/s-1/proposals/p-1/approve").mock( + return_value=httpx.Response( + 200, json={**proposal, "status": "submitted", "job_id": "job-9"} + ) + ) + respx.post(f"{API}/chat/sessions/s-1/proposals/p-1/reject").mock( + return_value=httpx.Response(200, json={**proposal, "status": "rejected"}) + ) + client = CogniChem(access_token="jwt", base_url=BASE) + chat = client.chat + + assert chat.sessions.create("Docking").id == "s-1" + assert chat.sessions.list().items[0].title == "Docking" + assert chat.sessions.update("s-1", allow_structure_search=True) + assert json.loads(patch.calls.last.request.content) == { + "allow_structure_search": True + } + assert chat.sessions.messages("s-1").items[0].role == "user" + assert chat.sessions.delete("s-1") is None + assert chat.estimate("s-1", "hi").hold_usd == 0.03 + assert chat.spend_cap("s-1").cap_usd == 1 + assert chat.continue_spend_cap("s-1", 1.0).raised is True + assert json.loads(cont.calls.last.request.content)["current_cap_usd"] == 1.0 + assert chat.stop("s-1", "t-1")["status"] == "stopping" + assert chat.proposals.list("s-1").items[0].id == "p-1" + chat.proposals.get("s-1", "p-1", include_spec=True) + assert get_prop.calls.last.request.url.params["include_spec"] == "true" + approved = chat.proposals.approve("s-1", "p-1", 0.12) + assert approved.job_id == "job-9" + assert approve.calls.last.request.headers["Idempotency-Key"] + assert json.loads(approve.calls.last.request.content) == { + "accepted_estimate_usd": 0.12, + "max_run_cost_usd": None, + } + assert chat.proposals.reject("s-1", "p-1").status == "rejected" + + +@respx.mock +async def test_async_chat_stream_and_events() -> None: + respx.post(f"{API}/chat/sessions/s-1/turns").mock( + return_value=httpx.Response(200, content=TURN) + ) + respx.get(f"{API}/events").mock( + return_value=httpx.Response( + 200, json={"events": [_event(7)], "next_cursor": "c7"} + ) + ) + async with AsyncCogniChem(access_token="jwt", base_url=BASE) as client: + names = [event.event async for event in client.chat.stream("s-1", "hi")] + result = await client.chat.send("s-1", "hi") + page = await client.events.list(after="c6") + assert names == ["hold", "token", "tool_call", "token", "proposal", "done"] + assert result.content == "Use autodockvina." + assert page.events[0].event_id == 7 diff --git a/tests/test_workflows.py b/tests/test_workflows.py new file mode 100644 index 0000000..9fe8f84 --- /dev/null +++ b/tests/test_workflows.py @@ -0,0 +1,370 @@ +from __future__ import annotations + +import json + +import httpx +import pytest +import respx + +from cognichem_client import AsyncCogniChem, CogniChem +from cognichem_client.errors import ProcessFailedError, ValidationError + +BASE = "https://api.test.cognichem.com" +API = f"{BASE}/api/v1" + +SPEC = { + "spec_version": 1, + "name": "vina-boltz-linear", + "params": {}, + "steps": [ + {"id": "dock", "type": "job", "job_type": "autodockvina", "inputs": {}}, + ], +} + + +def _run(status: str, **extra: object) -> dict[str, object]: + return { + "id": "run-1", + "run_name": "demo", + "status": status, + "steps": [ + {"id": "s1", "step_key": "dock", "status": status, "job_id": "job-1"} + ], + **extra, + } + + +@respx.mock +def test_validate_and_estimate() -> None: + validate = respx.post(f"{API}/workflows/validate").mock( + return_value=httpx.Response( + 200, + json={ + "valid": False, + "errors": [ + {"code": "cycle_detected", "message": "cycle", "step_id": "dock"} + ], + "warnings": [], + "order": [], + "step_count": 1, + }, + ) + ) + respx.post(f"{API}/workflows/estimate").mock( + return_value=httpx.Response( + 200, + json={ + "total": 1.25, + "per_step": [{"step_id": "dock", "cost": 1.25}], + "tier": 1, + "assumptions": [{"step_id": "dock", "reason": "static"}], + }, + ) + ) + client = CogniChem(api_key="k", base_url=BASE) + + result = client.workflows.validate(SPEC, params={"ligands": "CCO"}) + assert result.valid is False + assert result.errors[0].code == "cycle_detected" + body = json.loads(validate.calls.last.request.content) + assert body == {"spec": SPEC, "params": {"ligands": "CCO"}} + + estimate = client.workflows.estimate(SPEC) + assert estimate.total == 1.25 + assert estimate.per_step[0].step_id == "dock" + + +@respx.mock +def test_estimate_diagnostics_raise_validation_error() -> None: + respx.post(f"{API}/workflows/estimate").mock( + return_value=httpx.Response( + 422, + json={ + "type": "/api/v1/problems/http-422", + "title": "Unprocessable Content", + "status": 422, + "detail": "1 validation error", + "errors": [{"code": "estimate_failed", "message": "nope"}], + "retryability": "validation", + }, + ) + ) + client = CogniChem(api_key="k", base_url=BASE) + with pytest.raises(ValidationError) as exc: + client.workflows.estimate(SPEC) + assert exc.value.errors == [{"code": "estimate_failed", "message": "nope"}] + assert exc.value.retryability == "validation" + + +@respx.mock +def test_templates_and_node_ports() -> None: + respx.get(f"{API}/workflows/templates").mock( + return_value=httpx.Response( + 200, + json={ + "templates": [ + { + "id": "smiles-embed-dock", + "name": "SMILES to dock", + "summary": "Embed then dock", + "complexity": "simple", + "steps": [{"id": "dock", "type": "job"}], + "params": ["ligands"], + } + ] + }, + ) + ) + respx.get(f"{API}/workflows/templates/smiles-embed-dock").mock( + return_value=httpx.Response( + 200, + json={ + "id": "smiles-embed-dock", + "name": "SMILES to dock", + "summary": "Embed then dock", + "complexity": "simple", + "spec": SPEC, + }, + ) + ) + ports = respx.get(f"{API}/workflows/node-ports").mock( + return_value=httpx.Response( + 200, + json={"job_type": "padel-descriptor", "inputs": {"m": {}}, "outputs": {}}, + ) + ) + client = CogniChem(api_key="k", base_url=BASE) + templates = client.workflows.templates() + assert [t.id for t in templates] == ["smiles-embed-dock"] + assert client.workflows.template("smiles-embed-dock").spec == SPEC + assert client.workflows.node_ports("padel-descriptor").inputs == {"m": {}} + assert ports.calls.last.request.url.params["job_type"] == "padel-descriptor" + + +@respx.mock +def test_run_create_sends_generated_idempotency_key() -> None: + route = respx.post(f"{API}/workflows/runs").mock( + return_value=httpx.Response(200, json=_run("pending", estimated_cost=2.5)) + ) + client = CogniChem(api_key="k", base_url=BASE) + run = client.workflows.runs.create( + "demo", SPEC, params={"ligands": "CCO"}, max_run_cost=5.0 + ) + assert run.estimated_cost == 2.5 + request = route.calls.last.request + assert request.headers["Idempotency-Key"] + assert json.loads(request.content) == { + "run_name": "demo", + "spec": SPEC, + "params": {"ligands": "CCO"}, + "max_run_cost": 5.0, + } + + client.workflows.runs.create("demo", SPEC, idempotency_key="mine") + assert route.calls.last.request.headers["Idempotency-Key"] == "mine" + + +@respx.mock +def test_run_list_get_artifacts_resume_cancel_delete() -> None: + listing = respx.get(f"{API}/workflows/runs").mock( + return_value=httpx.Response( + 200, + json={ + "items": [{"id": "run-1", "run_name": "demo", "status": "running"}], + "total": 1, + "limit": 10, + "offset": 0, + }, + ) + ) + get = respx.get(f"{API}/workflows/runs/run-1").mock( + return_value=httpx.Response(200, json=_run("running", spec=SPEC)) + ) + respx.get(f"{API}/workflows/runs/run-1/artifacts").mock( + return_value=httpx.Response( + 200, + json={ + "items": [ + { + "artifact_id": "art-1", + "step_key": "dock", + "port": "poses", + "data_kind": "docking_poses", + "data_format": "pdbqt", + } + ] + }, + ) + ) + respx.post(f"{API}/workflows/runs/run-1/resume").mock( + return_value=httpx.Response(200, json=_run("running")) + ) + respx.post(f"{API}/workflows/runs/run-1/cancel").mock( + return_value=httpx.Response( + 200, + json={ + "run_id": "run-1", + "status": "cancelled", + "cancelled_job_ids": ["job-1"], + "message": "cancelled", + }, + ) + ) + respx.delete(f"{API}/workflows/runs/run-1").mock(return_value=httpx.Response(204)) + client = CogniChem(api_key="k", base_url=BASE) + + page = client.workflows.runs.list(limit=10, status="running") + assert page.items[0].id == "run-1" + assert listing.calls.last.request.url.params["status"] == "running" + assert "offset" not in listing.calls.last.request.url.params + + run = client.workflows.runs.get("run-1", include_spec=True) + assert run.spec == SPEC + assert run.steps[0].step_key == "dock" + assert get.calls.last.request.url.params["include_spec"] == "1" + + assert client.workflows.runs.artifacts("run-1").items[0].artifact_id == "art-1" + assert client.workflows.runs.resume("run-1").status == "running" + assert client.workflows.runs.cancel("run-1").cancelled_job_ids == ["job-1"] + assert client.workflows.runs.delete("run-1") is None + + +@respx.mock +def test_run_waits_and_stops_on_paused() -> None: + respx.post(f"{API}/workflows/runs").mock( + return_value=httpx.Response(200, json=_run("pending")) + ) + statuses = iter([_run("running"), _run("paused", status_message="cap")]) + respx.get(f"{API}/workflows/runs/run-1").mock( + side_effect=lambda request: httpx.Response(200, json=next(statuses)) + ) + client = CogniChem(api_key="k", base_url=BASE) + run = client.workflows.runs.run("demo", SPEC, poll_interval=0.01, timeout=2.0) + assert run.status == "paused" + + +@respx.mock +def test_run_raises_when_failed() -> None: + respx.get(f"{API}/workflows/runs/run-1").mock( + return_value=httpx.Response( + 200, json=_run("failed", status_message="step dock failed") + ) + ) + client = CogniChem(api_key="k", base_url=BASE) + with pytest.raises(ProcessFailedError, match="step dock failed"): + client.workflows.runs.wait( + "run-1", poll_interval=0.01, timeout=1.0, raise_on_failure=True + ) + + +@respx.mock +def test_definitions_crud_share_versions() -> None: + definition = { + "id": "wfd-1", + "user_id": "u-1", + "name": "triage", + "description": "", + "spec": SPEC, + "spec_version": 1, + } + create = respx.post(f"{API}/workflows/definitions").mock( + return_value=httpx.Response(201, json=definition) + ) + respx.get(f"{API}/workflows/definitions").mock( + return_value=httpx.Response( + 200, + json={ + "items": [{k: v for k, v in definition.items() if k != "spec"}], + "total": 1, + "limit": 50, + "offset": 0, + }, + ) + ) + update = respx.patch(f"{API}/workflows/definitions/wfd-1").mock( + return_value=httpx.Response(200, json={**definition, "name": "v2"}) + ) + respx.delete(f"{API}/workflows/definitions/wfd-1").mock( + return_value=httpx.Response(204) + ) + respx.get(f"{API}/workflows/definitions/wfd-1/download").mock( + return_value=httpx.Response( + 200, + content=json.dumps(SPEC).encode(), + headers={"Content-Disposition": "attachment; filename=triage.json"}, + ) + ) + respx.post(f"{API}/workflows/definitions/wfd-1/share").mock( + return_value=httpx.Response( + 200, json={**definition, "visibility": "link", "share_token": "tok"} + ) + ) + respx.delete(f"{API}/workflows/definitions/wfd-1/share").mock( + return_value=httpx.Response(200, json=definition) + ) + respx.get(f"{API}/workflows/definitions/shared/tok").mock( + return_value=httpx.Response(200, json={"spec": SPEC}) + ) + fork = respx.post(f"{API}/workflows/definitions/shared/tok/fork").mock( + return_value=httpx.Response(201, json={**definition, "id": "wfd-2"}) + ) + respx.get(f"{API}/workflows/definitions/wfd-1/versions").mock( + return_value=httpx.Response( + 200, + json={"items": [{"version": 2}], "total": 1, "limit": 50, "offset": 0}, + ) + ) + respx.get(f"{API}/workflows/definitions/wfd-1/versions/1").mock( + return_value=httpx.Response( + 200, json={"definition_id": "wfd-1", "version": 1, "spec": SPEC} + ) + ) + respx.post(f"{API}/workflows/definitions/wfd-1/versions/1/restore").mock( + return_value=httpx.Response(200, json=definition) + ) + client = CogniChem(api_key="k", base_url=BASE) + defs = client.workflows.definitions + + assert defs.create("triage", SPEC, graph_layout={"nodes": {}}).id == "wfd-1" + assert json.loads(create.calls.last.request.content)["graph_layout"] == { + "nodes": {} + } + assert create.calls.last.request.headers["Idempotency-Key"] + assert defs.list().items[0].name == "triage" + assert defs.update("wfd-1", name="v2").name == "v2" + assert json.loads(update.calls.last.request.content) == {"name": "v2"} + with pytest.raises(ValueError): + defs.update("wfd-1") + assert defs.delete("wfd-1") is None + exported = defs.download("wfd-1") + assert json.loads(exported.content) == SPEC + assert exported.filename == "triage.json" + assert defs.share("wfd-1").share_token == "tok" + assert defs.unshare("wfd-1").visibility == "private" + assert defs.get_shared("tok") == {"spec": SPEC} + assert defs.fork_shared("tok", "mine").id == "wfd-2" + assert json.loads(fork.calls.last.request.content)["name"] == "mine" + assert defs.versions("wfd-1").items[0].version == 2 + assert defs.get_version("wfd-1", 1).spec == SPEC + assert defs.restore_version("wfd-1", 1).id == "wfd-1" + + +@respx.mock +async def test_async_workflow_run() -> None: + respx.post(f"{API}/workflows/runs").mock( + return_value=httpx.Response(200, json=_run("pending")) + ) + respx.get(f"{API}/workflows/runs/run-1").mock( + return_value=httpx.Response(200, json=_run("completed")) + ) + respx.post(f"{API}/workflows/validate").mock( + return_value=httpx.Response( + 200, json={"valid": True, "order": ["dock"], "step_count": 1} + ) + ) + async with AsyncCogniChem(api_key="k", base_url=BASE) as client: + assert (await client.workflows.validate(SPEC)).valid + run = await client.workflows.runs.run( + "demo", SPEC, poll_interval=0.01, timeout=1.0 + ) + assert run.status == "completed"