Skip to content

Commit d019ecf

Browse files
committed
feat(audio): agent-ready foundation — AgentReady contract mixin + residency resolver
Shared base imported by the audio stage modules to make them agent-ready: - _agent_ready.py: AgentReady mixin, StageContract/IOSpec/Gates/Role - _residency.py: input residency resolver (file/waveform/auto) - common.py: agent-ready helpers (ensure_mono/ensure_waveform_2d) Excludes the agent orchestration layer (registry/catalog/conformance/planning/agent).
1 parent 67977c8 commit d019ecf

3 files changed

Lines changed: 524 additions & 5 deletions

File tree

Lines changed: 318 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,318 @@
1+
# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from __future__ import annotations
16+
17+
from dataclasses import asdict, dataclass, field, is_dataclass
18+
from enum import Enum
19+
from typing import TYPE_CHECKING, Any, ClassVar, Literal
20+
21+
if TYPE_CHECKING:
22+
from collections.abc import Mapping
23+
24+
AudioForm = Literal["file", "waveform"]
25+
ProducedForm = Literal["tensor", "disk"]
26+
Cardinality = Literal["1:1", "1:1 nested-list", "1:N fan-out", "N:1", "filter"]
27+
Dispatch = Literal["process", "process_batch", "auto"]
28+
# How a stage handles per-item failures at runtime. "unknown" means the stage
29+
# has not declared a uniform policy (the default for most stages today).
30+
ErrorPolicy = Literal["skip", "fail", "annotate", "unknown"]
31+
32+
# Semantic role vocabulary. Because key *names* are agent-configurable
33+
# (config-knobs-only standardization: every read/written key is a ``*_key``
34+
# constructor field that an agent may rename), two stages can only be chained
35+
# reliably by matching the *role* a key plays, not its literal string. Roles are
36+
# resolved from the invariant ``*_key`` field name (see ``_roles.py``), so they
37+
# survive an agent renaming a key's value. ``"unknown"`` is the safe default and
38+
# never blocks composition.
39+
Role = Literal[
40+
"audio_filepath",
41+
"waveform",
42+
"sample_rate",
43+
"duration",
44+
"num_samples",
45+
"segments",
46+
"diar_segments",
47+
"vad_segments",
48+
"overlap_segments",
49+
"timestamps",
50+
"start",
51+
"end",
52+
"start_ms",
53+
"end_ms",
54+
"text",
55+
"pred_text",
56+
"reference_text",
57+
"words",
58+
"alignment",
59+
"speaker_id",
60+
"num_speakers",
61+
"score",
62+
"metrics",
63+
"prediction",
64+
"segment_num",
65+
"original_file",
66+
"item_id",
67+
"windows",
68+
"manifest_path",
69+
"output_path",
70+
"unknown",
71+
]
72+
73+
74+
@dataclass(frozen=True)
75+
class ParamSpec:
76+
"""A single constructor parameter an agent can set on a stage.
77+
78+
Usually derived automatically from the stage's dataclass fields (or
79+
``__init__`` signature) via
80+
:func:`nemo_curator.stages.audio._agent_registry.stage_params`, but a stage
81+
may also override/augment entries in ``StageContract.params``.
82+
"""
83+
84+
name: str
85+
type: str = "Any"
86+
default: Any = None
87+
required: bool = False
88+
choices: list[Any] | None = None # populated for Literal[...] params
89+
description: str | None = None
90+
role: Role | None = None # semantic role of the key this param configures
91+
92+
93+
@dataclass(frozen=True)
94+
class IOSpec:
95+
"""Task data keys and audio forms read or written by a stage."""
96+
97+
data_keys: list[str] = field(default_factory=list)
98+
segment_data_keys: list[str] = field(default_factory=list)
99+
accepts: list[AudioForm] = field(default_factory=list)
100+
produces: list[ProducedForm] = field(default_factory=list)
101+
102+
103+
@dataclass(frozen=True)
104+
class Gates:
105+
"""Execution gates or side effects an agent should know before wrapping a stage."""
106+
107+
writes_to_disk: bool = False
108+
requires_gpu: bool = False
109+
requires_internet_first_run: bool = False
110+
requires_ffmpeg: bool = False
111+
lifecycle_side_effects: bool = False
112+
runtime_secrets: list[str] = field(default_factory=list)
113+
# Serializability contract for sinks (distinguishes the two JSON sinks):
114+
# requires_serializable_input — this stage serializes task.data as-is (e.g.
115+
# raw json.dumps) and will fail on a resident tensor/audio blob.
116+
# sanitizes_output — this stage strips tensors/audio blobs, so anything
117+
# downstream of it is serialization-safe.
118+
requires_serializable_input: bool = False
119+
sanitizes_output: bool = False
120+
121+
122+
@dataclass(frozen=True)
123+
class SizeEnvelope:
124+
"""Coarse size and memory hints for agent planning."""
125+
126+
max_input_sec: float | None = None
127+
allowed_sample_rates: list[int] | None = None
128+
channels: Literal["mono", "stereo", "any"] = "any"
129+
memory_hint: str | None = None
130+
131+
132+
@dataclass(frozen=True)
133+
class StaticHints:
134+
"""Instance-independent hints a stage may declare for discovery/planning.
135+
136+
Lets a stage expose ``cardinality_options``, ``gates``, ``dispatch``,
137+
``error_policy``, ``description`` and ``stage_id`` *without* being
138+
instantiated (see :meth:`AgentReady.describe_static`). Declared on a class
139+
via the ``AGENT_STATIC`` ClassVar; every field defaults so declaring it is
140+
fully optional and additive.
141+
"""
142+
143+
cardinality_options: list[str] = field(default_factory=list)
144+
gates: Gates = field(default_factory=Gates)
145+
dispatch: Dispatch = "auto"
146+
error_policy: ErrorPolicy = "unknown"
147+
description: str | None = None
148+
stage_id: str | None = None
149+
150+
151+
@dataclass(frozen=True)
152+
class StageContract:
153+
"""Read-only discovery contract for an agent-ready processing stage."""
154+
155+
reads: IOSpec = field(default_factory=IOSpec)
156+
writes: IOSpec = field(default_factory=IOSpec)
157+
reads_one_of: list[IOSpec] = field(default_factory=list)
158+
metadata_reads: list[str] = field(default_factory=list)
159+
metadata_writes: list[str] = field(default_factory=list)
160+
cardinality: Cardinality = "1:1"
161+
cardinality_options: list[str] = field(default_factory=list)
162+
iteration_key: str | None = None
163+
preserves_upstream_keys: bool = True
164+
wrappable: bool = True
165+
size_envelope: SizeEnvelope = field(default_factory=SizeEnvelope)
166+
gates: Gates = field(default_factory=Gates)
167+
# Agent-facing metadata (advisory; defaults keep older describe() calls valid).
168+
stage_id: str | None = None # stable semantic id; defaults to the class name when None
169+
description: str | None = None # one-line human summary for planners/UIs
170+
params: list[ParamSpec] = field(default_factory=list) # usually auto-derived at discovery time
171+
dispatch: Dispatch = "auto" # "auto" => infer from the stage at runtime
172+
error_policy: ErrorPolicy = "unknown"
173+
# Resolved-key-value -> semantic role. Populated at discovery time by
174+
# ``_agent_registry.build_contract``; empty when a contract is built by hand.
175+
key_roles: dict[str, Role] = field(default_factory=dict)
176+
# True when the stage only implements ``process_batch`` (``process`` raises).
177+
# Auto-derived at discovery time; agents must not call ``process`` on these.
178+
batch_only: bool = False
179+
180+
def to_dict(self) -> dict[str, Any]:
181+
"""Return a JSON-safe dict of this contract (``json.dumps`` never raises)."""
182+
return {
183+
"reads": asdict(self.reads),
184+
"writes": asdict(self.writes),
185+
"reads_one_of": [asdict(spec) for spec in self.reads_one_of],
186+
"metadata_reads": list(self.metadata_reads),
187+
"metadata_writes": list(self.metadata_writes),
188+
"cardinality": self.cardinality,
189+
"cardinality_options": list(self.cardinality_options),
190+
"iteration_key": self.iteration_key,
191+
"preserves_upstream_keys": self.preserves_upstream_keys,
192+
"wrappable": self.wrappable,
193+
"size_envelope": asdict(self.size_envelope),
194+
"gates": asdict(self.gates),
195+
"stage_id": self.stage_id,
196+
"description": self.description,
197+
"params": [_paramspec_to_dict(p) for p in self.params],
198+
"dispatch": self.dispatch,
199+
"error_policy": self.error_policy,
200+
"key_roles": dict(self.key_roles),
201+
"batch_only": self.batch_only,
202+
}
203+
204+
205+
def _jsonable_default(value: Any) -> Any: # noqa: ANN401, PLR0911 (complexity accepted: one early return per JSON type family)
206+
"""Coerce an arbitrary param default to a JSON-serializable value.
207+
208+
Non-serializable objects (tensors, models, callables) become a short
209+
``"<non-serializable: Type>"`` sentinel rather than raising.
210+
"""
211+
if value is None or isinstance(value, (str, int, float, bool)):
212+
return value
213+
if isinstance(value, (list, tuple, set)):
214+
return [_jsonable_default(v) for v in value]
215+
if isinstance(value, dict):
216+
return {str(k): _jsonable_default(v) for k, v in value.items()}
217+
if isinstance(value, Enum):
218+
return value.value
219+
if is_dataclass(value) and not isinstance(value, type):
220+
try:
221+
return {k: _jsonable_default(v) for k, v in asdict(value).items()}
222+
except Exception: # noqa: BLE001
223+
return f"<non-serializable: {type(value).__name__}>"
224+
return f"<non-serializable: {type(value).__name__}>"
225+
226+
227+
def _paramspec_to_dict(p: ParamSpec) -> dict[str, Any]:
228+
return {
229+
"name": p.name,
230+
"type": p.type,
231+
"default": _jsonable_default(p.default),
232+
"required": p.required,
233+
"choices": None if p.choices is None else [_jsonable_default(c) for c in p.choices],
234+
"description": p.description,
235+
"role": p.role,
236+
}
237+
238+
239+
def _json_type(type_str: str | None) -> str | None:
240+
"""Map a rendered ParamSpec.type string to a JSON-Schema type (or None to omit)."""
241+
if not type_str:
242+
return None
243+
t = type_str.replace(" ", "").replace("|None", "")
244+
if t.startswith("Optional["):
245+
t = t[len("Optional[") : -1] if t.endswith("]") else t
246+
scalar = {"str": "string", "int": "integer", "float": "number", "bool": "boolean"}
247+
if t in scalar:
248+
return scalar[t]
249+
lowered = t.lower()
250+
if lowered.startswith(("list", "sequence", "tuple")):
251+
return "array"
252+
if lowered.startswith(("dict", "mapping")):
253+
return "object"
254+
return None
255+
256+
257+
def to_json_schema(params: list[ParamSpec]) -> dict[str, Any]:
258+
"""Build a JSON-Schema ``object`` describing a stage's configurable params.
259+
260+
Suitable as the argument schema for an agent's stage-configuration form.
261+
"""
262+
properties: dict[str, Any] = {}
263+
required: list[str] = []
264+
for p in params:
265+
schema: dict[str, Any] = {}
266+
json_type = _json_type(p.type)
267+
if json_type is not None:
268+
schema["type"] = json_type
269+
if p.choices is not None:
270+
schema["enum"] = [_jsonable_default(c) for c in p.choices]
271+
if p.default is not None:
272+
schema["default"] = _jsonable_default(p.default)
273+
if p.description:
274+
schema["description"] = p.description
275+
if p.role:
276+
schema["x-role"] = p.role
277+
properties[p.name] = schema
278+
if p.required:
279+
required.append(p.name)
280+
out: dict[str, Any] = {"type": "object", "properties": properties}
281+
if required:
282+
out["required"] = required
283+
return out
284+
285+
286+
class AgentReady:
287+
"""Mixin for stages that expose a read-only agent discovery contract."""
288+
289+
# Opt-in, instance-independent discovery hints. Annotated as ClassVar so
290+
# dataclass stages do NOT treat these as fields. All optional/additive.
291+
AGENT_STATIC: ClassVar[StaticHints | None] = None
292+
# Set True on stages whose ``process`` raises (only ``process_batch`` works).
293+
BATCH_ONLY: ClassVar[bool] = False
294+
# Rare per-field role overrides keyed by ``*_key`` field name. Consulted
295+
# before the shared ``_roles.KEY_ROLES`` table.
296+
KEY_ROLE_OVERRIDES: ClassVar[Mapping[str, Role]] = {}
297+
298+
def describe(self) -> StageContract:
299+
raise NotImplementedError
300+
301+
@classmethod
302+
def describe_static(cls) -> StageContract:
303+
"""Instance-free contract for discovery/planning.
304+
305+
Uses class defaults + ``AGENT_STATIC`` and never runs ``__init__`` side
306+
effects, so it is safe for stages with required constructor args. Prefer
307+
the instance-level :meth:`describe` (or
308+
``_agent_registry.build_contract``) when resolved key *values* are
309+
needed.
310+
"""
311+
from nemo_curator.stages.audio._agent_registry import static_contract
312+
313+
return static_contract(cls)
314+
315+
316+
# (resolve_contract was removed: dead code whose signature promised class
317+
# acceptance the body rejected. Instances: use stage.describe() / build_contract;
318+
# classes: use describe_static() / static_contract.)

0 commit comments

Comments
 (0)