feat: (DDOS-6506) ensure all write paths apply tags across the client

[maurocolella]

to do:

• remove Datasets
• remove tag and billing from every run -- pending comments from @andrewmochalskyy

[Copilot]

Pull request overview

This PR standardizes how metadata tags are propagated through the Python client’s write paths, including entity CRUD (ligands/proteins/projects/datasets) and tool execution creation, and updates tests + the mock server to match the new shapes.

Changes:

• Add/standardize tags (jsonb object) support across entity create/update flows (ligands, proteins, projects, datasets) and ensure Ligand/Protein sync/register forward and patch tags.
• Add execution-level tagging support via client.tag plus per-call tag, and introduce client.billing_tag plus per-call billing propagation to executions.
• Update tests and mock server dataset routes to support dataset tag storage as a jsonb object (e.g. {"dataset_tags": [...]}).

Reviewed changes

Copilot reviewed 26 out of 26 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test_entities.py	Adds LV1 coverage asserting ligand tag persistence on create/update/register/sync.
tests/test_datasets.py	Updates assertions to the new dataset tags jsonb shape and tag filter behavior.
tests/test_client.py	Adds unit coverage for execution payload tag/billing propagation and override semantics.
tests/test_billing_tag_end_to_end.py	Re-enables LV2 end-to-end billing tag verification using client.tag.
tests/mock_server/routers/data_platform.py	Updates dataset admin endpoints to accept nested admin set bundles and tag filtering with jsonb tags.
src/platform/projects.py	Adds tags support on create and introduces Projects.update() supporting tags/notes.
src/platform/executions.py	Extends Executions.create() to send tag/billing from per-call args or client defaults.
src/platform/entities.py	Adds tags support to ligand/protein create and updates ligand/protein update tag type to jsonb dict.
src/platform/datasets.py	Refactors dataset create to nested admin set bundle and coerces tag lists into jsonb object shape.
src/platform/client.py	Adds billing_tag to client init, storage, and __repr__.
src/drug_discovery/tool_execution.py	Routes execution submission through shared _create_execution() helper.
src/drug_discovery/system_prep.py	Adds tag/billing parameters and applies them to execution create.
src/drug_discovery/structures/protein.py	Forwards Entity.tags on create and patches tags on sync when matching an existing row.
src/drug_discovery/structures/ligand.py	Forwards Entity.tags on create, includes tags in row serialization, and patches tags on sync for existing rows.
src/drug_discovery/structures/entity.py	Adds Entity.tags field and documents tag behavior for sync/register.
src/drug_discovery/rbfe.py	Uses _create_execution() helper for consistent tag/billing propagation.
src/drug_discovery/protonation.py	Adds tool_version attribute and supports per-call tag/billing overrides via _execution_tags().
src/drug_discovery/pocket_finder.py	Applies execution tag/billing overrides in sync paths and uses _create_execution().
src/drug_discovery/patent.py	Uses _create_execution() helper for consistent execution submission.
src/drug_discovery/molprops.py	Adds tag/billing plumbing to molprops execution creation and sets tool_version attribute.
src/drug_discovery/konnektor.py	Supports per-call tag/billing via _execution_tags() and uses _create_execution().
src/drug_discovery/execution.py	Introduces _execution_tags() context manager and _create_execution() helper to centralize tag/billing behavior.
src/drug_discovery/execution_mixins.py	Allows start() callers to pass tag/billing and applies overrides during submission.
src/drug_discovery/docking.py	Applies per-call tag/billing overrides and uses _create_execution() for async submission path.
src/drug_discovery/constrained_docking.py	Adds per-call tag/billing and applies them to execution create.
src/drug_discovery/abfe.py	Uses _create_execution() helper for consistent execution submission.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[andrewmochalskyy]

Code review

Found 3 issues (1 bug, 2 tied to tag-management consistency / SDK pass-through):

1. changelog silently dropped — Datasets.create() accepts and documents changelog, but the refactored body calls _build_dataset_admin_set(...) without it (helper has no such param), so it never reaches the API. On main it was sent as set_dict["changelog"]. Caller gets no error, silent data loss. (Same applies to summary, now only a fallback for description.)

do-dd-client/src/platform/datasets.py

Lines 210 to 233 in ec5aceb

	sample_rows: Cached first rows for preview.
	changelog: Description of what changed in this version.
	Returns:
	Dictionary containing the created dataset record.
	"""
	self._require_superuser()
	set_dict = _build_dataset_admin_set(
	name=name,
	file_path=file_path,
	dataset_key=dataset_key,
	dataset_version=dataset_version,
	summary=summary,
	description=description,
	source_url=source_url,
	source_name=source_name,
	tags=tags,
	compound_count=compound_count,
	file_size_bytes=file_size_bytes,
	dataset_schema=dataset_schema,
	sample_rows=sample_rows,
	)
	return self._c.post_json(

2. Dead tag path in molprops (tag-management duplication) — Molprops.run() wraps the body in with self._execution_tags(tag=tag, billing=billing):, but the actual create goes through run_molprops_combined() → client.executions.create(...) directly, never _create_execution() (the only reader of those overrides). The context manager is dead code here; tag/billing only work via the explicit kwargs. Drop the wrapper or route through _create_execution.

do-dd-client/src/drug_discovery/molprops.py

Lines 257 to 263 in ec5aceb

	lig.id = str(idx)
	ligand_set = LigandSet(ligands=self._ligands)
	with self._execution_tags(tag=tag, billing=billing):
	if quote:
	_unused_rows, raw = run_molprops_combined(
	ligand_set=ligand_set,

3. tags type changed list[str] → dict with no coercion (SDK pass-through risk) — update_ligand/update_protein changed the tags type. There's no beartype on these, so an existing caller (e.g. Balto) passing tags=["x"] won't error but will persist a list where the rest of the system now expects a jsonb object. datasets.create coerces list → {"dataset_tags": [...]}; entities/projects don't. Coerce consistently or keep list back-compat.

do-dd-client/src/platform/entities.py

Lines 590 to 606 in ec5aceb

	def update_ligand(
	self,
	id: str,
	*,
	smiles: str | None = None,
	project_id: str | None = None,
	name: str | None = None,
	mol_file: str | None = None,
	formal_charge: int | None = None,
	hbond_donor_count: int | None = None,
	hbond_acceptor_count: int | None = None,
	rotatable_bond_count: int | None = None,
	tpsa: float | None = None,
	molecular_weight: float | None = None,
	variant_name_tag: str | None = None,
	tags: dict[str, Any] | None = None,
	) -> dict:

Nice to have (non-blocking):

4. start() accepts tag/billing only via **kwargs (popped at the body) and documents them, but they aren't in the signature — every other tool method declares them explicitly. Make them real params for discoverability/type-checking.

do-dd-client/src/drug_discovery/execution_mixins.py

Lines 126 to 131 in ec5aceb

	)
	resolved_amount = 0 if quote else approve_amount
	call_tag = kwargs.pop("tag", None)
	call_billing = kwargs.pop("billing", None)
	with self._execution_tags(tag=call_tag, billing=call_billing):
	self._start_impl(approve_amount=resolved_amount, **kwargs)

5. billing_tag docstring says ... / DO_TAGS, but nothing in the client reads a DO_TAGS env var — drop the mention or wire it up.

do-dd-client/src/platform/client.py

Lines 424 to 428 in ec5aceb

	record: Whether to record tool execution responses for testing.
	tag: Optional billing tag applied to all tool executions (``tag`` field).
	billing_tag: Optional runtime billing tag (``billing`` field / ``DO_TAGS``).
	Distinct from :attr:`billing`, the billing API wrapper.
	_app: Internal app identifier. Part of the singleton cache key.

6. _create_execution's tool_key/tool_version/timeout params are never passed by any caller — dead surface, can be removed.

do-dd-client/src/drug_discovery/execution.py

Lines 138 to 145 in ec5aceb

	def _create_execution(
	self,
	*,
	data: dict[str, Any],
	tool_key: str | None = None,
	tool_version: str | None = None,
	timeout: float | None = None,
	) -> dict[str, Any]:

(Also: projects.update() is new and has no test.)

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[sg-s]

Addressed in a2929be: `Datasets.create(changelog=...)` now forwards the value via `datasetMeta.changelog` in the admin `set` bundle (replacing the pre-refactor flat `set.changelog` field). Mock flatten + `test_create_dataset` updated.

[Copilot]

Pull request overview

Copilot reviewed 26 out of 26 changed files in this pull request and generated 2 comments.

[sg-s]

lots of admin-only code here that i would remove IMO. also @maurocolella plz see pattern of injecting tag and billing (tag?) into every tool class -- this is not sustainable. can we only put it in the client?

[Copilot]

Pull request overview

Copilot reviewed 30 out of 30 changed files in this pull request and generated 1 comment.

[andrewmochalskyy]

Approving — all review items addressed, and merge_entity_tags now auto-stamps app/session provenance on direct entity creates (caller keys win), which is the real fix for the source-filter gap on non-tool paths. Good cleanup (Datasets removed, per-run tag/billing collapsed to client config, coercion unified, tests/test_tags.py added).

One heads-up before/at merge — breaking change: entity tags is now dict-only and raises TypeError on a list (was list[str] on main for update_ligand/update_protein). No callers pass list tags in this repo or the sibling platform repos, so it's safe here — but any out-of-workspace consumer (e.g. a Balto build) passing tags=[...] will now throw. Worth a quick ping to the Balto team so they send a dict.

do-dd-client/src/platform/tags.py

Lines 36 to 38 in b7f39df

	if tags is not None and not isinstance(tags, dict):
	raise TypeError("tags must be a dict or None")

Non-blocking: the SDK now provably writes the {app, session} shape, but the true end-to-end (UI source filter returns a directly-added ligand/protein) is still a QA confirmation, not covered by unit tests here.