From 3a7fae8d882825624bb8186e083e5faaa40aec52 Mon Sep 17 00:00:00 2001 From: tjkessler Date: Mon, 28 Sep 2026 21:30:18 -0400 Subject: [PATCH] Cover the full /api/v1 surface in cognichem-client 1.1.0. Adds workflows (validate, estimate, templates, node ports, runs, definitions), artifacts (list, usage, upload, manifest, record download), completion events, wallet, science lookups, the reference KG, and CogniChem Assistant chat, with sync and async twins. Covers the Python client follow-up (cognichem#269) from the agent platform epic (cognichem#458). Fixes API-key submits, which the API rejects without an Idempotency-Key: mutating calls now generate one per call. Fixes api_keys.delete raising on the 204 response. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 30 + README.md | 95 +- examples/artifact_records.py | 45 + examples/assistant_chat.py | 36 + examples/science_lookup.py | 30 + examples/workflow_template_run.py | 49 + pyproject.toml | 2 +- scripts/route_map.txt | 69 + scripts/route_map_keys.txt | 30 - scripts/sync_routes.py | 188 +- src/cognichem_client/__init__.py | 25 +- src/cognichem_client/_http.py | 214 +- src/cognichem_client/async_client.py | 60 +- src/cognichem_client/client.py | 60 +- src/cognichem_client/constants.py | 209 +- src/cognichem_client/errors.py | 102 +- src/cognichem_client/polling.py | 233 ++- src/cognichem_client/resources/__init__.py | 17 +- src/cognichem_client/resources/api_keys.py | 243 ++- src/cognichem_client/resources/artifacts.py | 554 +++++ src/cognichem_client/resources/chat.py | 1213 +++++++++++ src/cognichem_client/resources/events.py | 280 +++ .../resources/inference/__init__.py | 72 +- src/cognichem_client/resources/jobs.py | 229 ++- src/cognichem_client/resources/lookup.py | 481 +++++ src/cognichem_client/resources/reference.py | 235 +++ src/cognichem_client/resources/utils.py | 62 +- src/cognichem_client/resources/wallet.py | 69 + src/cognichem_client/resources/workflows.py | 1775 +++++++++++++++++ src/cognichem_client/types/__init__.py | 158 ++ src/cognichem_client/types/artifacts.py | 123 ++ src/cognichem_client/types/chat.py | 295 +++ .../{types.py => types/common.py} | 159 +- src/cognichem_client/types/events.py | 56 + src/cognichem_client/types/science.py | 277 +++ src/cognichem_client/types/workflows.py | 557 ++++++ tests/test_artifacts_jobs.py | 273 +++ tests/test_errors.py | 43 + tests/test_science_events_chat.py | 415 ++++ tests/test_workflows.py | 370 ++++ 40 files changed, 9186 insertions(+), 247 deletions(-) create mode 100644 examples/artifact_records.py create mode 100644 examples/assistant_chat.py create mode 100644 examples/science_lookup.py create mode 100644 examples/workflow_template_run.py create mode 100644 scripts/route_map.txt delete mode 100644 scripts/route_map_keys.txt create mode 100644 src/cognichem_client/resources/artifacts.py create mode 100644 src/cognichem_client/resources/chat.py create mode 100644 src/cognichem_client/resources/events.py create mode 100644 src/cognichem_client/resources/lookup.py create mode 100644 src/cognichem_client/resources/reference.py create mode 100644 src/cognichem_client/resources/wallet.py create mode 100644 src/cognichem_client/resources/workflows.py create mode 100644 src/cognichem_client/types/__init__.py create mode 100644 src/cognichem_client/types/artifacts.py create mode 100644 src/cognichem_client/types/chat.py rename src/cognichem_client/{types.py => types/common.py} (57%) create mode 100644 src/cognichem_client/types/events.py create mode 100644 src/cognichem_client/types/science.py create mode 100644 src/cognichem_client/types/workflows.py create mode 100644 tests/test_artifacts_jobs.py create mode 100644 tests/test_science_events_chat.py create mode 100644 tests/test_workflows.py 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"