Code / libs/huleedu_nlp_shared/src/huleedu_nlp_shared/feature_contracts/schema.py
libs/huleedu_nlp_shared/src/huleedu_nlp_shared/feature_contracts/schema.py
1
"""Feature contract schema helpers.
2
3
Purpose:
4
Build ordered feature names and schema snapshots from a shared registry
5
while keeping extraction schemas separate from downstream model recipes.
6
7
Relationships:
8
- Used by current and future ML/offload/ENG adapters to validate feature
9
names and ordering without importing heavy extractor runtimes.
10
- Separates feature schema hash, input contract hash, runtime fingerprint,
11
and artifact digest for cache-safe artifacts.
12
"""
13
14
from __future__ import annotations
15
16
from huleedu_nlp_shared.feature_contracts.models import (
17
FeatureFamilySpecV1,
18
FeatureLane,
19
FeatureRegistryV1,
20
FeatureSchemaSnapshotV1,
21
)
22
from huleedu_nlp_shared.feature_contracts.provenance import schema_hash_for_parts, stable_sha256
23
24
25
def generated_vector_feature_names(
26
*,
27
name_prefix: str,
28
dim: int,
29
index_width: int,
30
) -> tuple[str, ...]:
31
"""Return generated vector names such as `embedding_000`."""
32
33
if dim <= 0:
34
raise ValueError("generated vector dim must be positive")
35
if index_width <= 0:
36
raise ValueError("generated vector index_width must be positive")
37
return tuple(f"{name_prefix}{index:0{index_width}d}" for index in range(dim))
38
39
40
def feature_names_for(
41
registry: FeatureRegistryV1,
42
feature_set_id: str,
43
*,
44
embedding_dim: int,
45
) -> tuple[str, ...]:
46
"""Return ordered feature names for a feature set and embedding dimension."""
47
48
feature_set = registry.feature_set(feature_set_id)
49
names: list[str] = []
50
embedding_feature_count = 0
51
for family_id in feature_set.family_ids:
52
family = registry.family(family_id)
53
family_names = _family_feature_names(family, embedding_dim=embedding_dim)
54
if family.generated_vector is not None:
55
embedding_feature_count += len(family_names)
56
names.extend(family_names)
57
_validate_embedding_dimension(
58
embedding_dim=embedding_dim,
59
embedding_feature_count=embedding_feature_count,
60
)
61
return tuple(names)
62
63
64
def scalar_feature_names_for(
65
registry: FeatureRegistryV1,
66
feature_set_id: str,
67
) -> tuple[str, ...]:
68
"""Return scalar feature names for a registered feature set."""
69
70
feature_set = registry.feature_set(feature_set_id)
71
names: list[str] = []
72
for family_id in feature_set.family_ids:
73
family = registry.family(family_id)
74
names.extend(feature.name for feature in family.features)
75
return tuple(names)
76
77
78
def validate_scorer_feature_set(registry: FeatureRegistryV1, feature_set_id: str) -> None:
79
"""Fail when a feature set includes a non-scorer family."""
80
81
feature_set = registry.feature_set(feature_set_id)
82
if feature_set.lane != FeatureLane.SCORER_MODEL_GRADE:
83
raise ValueError(f"Feature set is not scorer model-grade: {feature_set_id}")
84
for family_id in feature_set.family_ids:
85
family = registry.family(family_id)
86
if family.lane != FeatureLane.SCORER_MODEL_GRADE:
87
raise ValueError(f"Feature family is not scorer model-grade: {family_id}")
88
89
90
def schema_hash_for_feature_set(
91
registry: FeatureRegistryV1,
92
feature_set_id: str,
93
*,
94
embedding_dim: int,
95
) -> str:
96
"""Hash a feature set for the requested generated-vector dimension."""
97
98
feature_set = registry.feature_set(feature_set_id)
99
families = tuple(registry.family(family_id) for family_id in feature_set.family_ids)
100
return schema_hash_for_parts(
101
registry_id=registry.registry_id,
102
registry_version=registry.registry_version,
103
feature_set_id=feature_set.feature_set_id,
104
feature_set_version=feature_set.feature_set_version,
105
family_parts=tuple(
106
_family_schema_hash_part(family, embedding_dim=embedding_dim) for family in families
107
),
108
ordered_feature_names=feature_names_for(
109
registry,
110
feature_set_id,
111
embedding_dim=embedding_dim,
112
),
113
)
114
115
116
def build_schema_snapshot(
117
registry: FeatureRegistryV1,
118
feature_set_id: str,
119
*,
120
embedding_dim: int,
121
runtime_fingerprint: str,
122
artifact_digest: str,
123
) -> FeatureSchemaSnapshotV1:
124
"""Build a persisted schema snapshot for a feature matrix artifact."""
125
126
return materialize_feature_set(
127
registry,
128
feature_set_id,
129
embedding_dim=embedding_dim,
130
runtime_fingerprint=runtime_fingerprint,
131
artifact_digest=artifact_digest,
132
)
133
134
135
def materialize_feature_set(
136
registry: FeatureRegistryV1,
137
feature_set_id: str,
138
*,
139
embedding_dim: int,
140
runtime_fingerprint: str,
141
artifact_digest: str,
142
) -> FeatureSchemaSnapshotV1:
143
"""Materialize one concrete feature-set schema snapshot."""
144
145
feature_set = registry.feature_set(feature_set_id)
146
ordered_feature_names = feature_names_for(registry, feature_set_id, embedding_dim=embedding_dim)
147
scalar_feature_names = scalar_feature_names_for(registry, feature_set_id)
148
embedding_feature_names = _embedding_feature_names_for(
149
registry,
150
feature_set_id,
151
embedding_dim=embedding_dim,
152
)
153
_validate_embedding_dimension(
154
embedding_dim=embedding_dim,
155
embedding_feature_count=len(embedding_feature_names),
156
)
157
schema_hash = schema_hash_for_feature_set(
158
registry,
159
feature_set_id,
160
embedding_dim=embedding_dim,
161
)
162
return FeatureSchemaSnapshotV1(
163
snapshot_id=stable_sha256(
164
{
165
"registry_id": registry.registry_id,
166
"registry_version": registry.registry_version,
167
"feature_set_id": feature_set.feature_set_id,
168
"feature_set_version": feature_set.feature_set_version,
169
"embedding_dim": embedding_dim,
170
"ordered_feature_names": list(ordered_feature_names),
171
"schema_hash": schema_hash,
172
"input_contract_hash": registry.input_contract_hash,
173
"runtime_fingerprint": runtime_fingerprint,
174
"artifact_digest": artifact_digest,
175
}
176
),
177
registry_id=registry.registry_id,
178
registry_version=registry.registry_version,
179
feature_set_id=feature_set.feature_set_id,
180
feature_set_version=feature_set.feature_set_version,
181
embedding_dim=embedding_dim,
182
ordered_feature_names=ordered_feature_names,
183
scalar_feature_names=scalar_feature_names,
184
embedding_feature_names=embedding_feature_names,
185
schema_hash=schema_hash,
186
input_contract_hash=registry.input_contract_hash,
187
runtime_fingerprint=runtime_fingerprint,
188
artifact_digest=artifact_digest,
189
)
190
191
192
def validate_schema_snapshot(
193
registry: FeatureRegistryV1, snapshot: FeatureSchemaSnapshotV1
194
) -> None:
195
"""Fail closed when a schema snapshot no longer matches the registry recipe."""
196
197
feature_set = registry.feature_set(snapshot.feature_set_id)
198
if snapshot.registry_id != registry.registry_id:
199
raise ValueError("schema snapshot registry_id does not match registry")
200
if snapshot.registry_version != registry.registry_version:
201
raise ValueError("schema snapshot registry_version does not match registry")
202
if snapshot.feature_set_version != feature_set.feature_set_version:
203
raise ValueError("schema snapshot feature_set_version does not match registry")
204
expected = materialize_feature_set(
205
registry,
206
snapshot.feature_set_id,
207
embedding_dim=snapshot.embedding_dim,
208
runtime_fingerprint=snapshot.runtime_fingerprint,
209
artifact_digest=snapshot.artifact_digest,
210
)
211
if snapshot.ordered_feature_names != expected.ordered_feature_names:
212
raise ValueError("schema snapshot ordered_feature_names do not match registry")
213
if snapshot.scalar_feature_names != expected.scalar_feature_names:
214
raise ValueError("schema snapshot scalar_feature_names do not match registry")
215
if snapshot.embedding_feature_names != expected.embedding_feature_names:
216
raise ValueError("schema snapshot embedding_feature_names do not match registry")
217
if snapshot.schema_hash != expected.schema_hash:
218
raise ValueError("schema snapshot schema_hash does not match registry")
219
if snapshot.input_contract_hash != expected.input_contract_hash:
220
raise ValueError("schema snapshot input_contract_hash does not match registry")
221
if snapshot.snapshot_id != expected.snapshot_id:
222
raise ValueError("schema snapshot snapshot_id does not match registry")
223
224
225
def _embedding_feature_names_for(
226
registry: FeatureRegistryV1,
227
feature_set_id: str,
228
*,
229
embedding_dim: int,
230
) -> tuple[str, ...]:
231
feature_set = registry.feature_set(feature_set_id)
232
names: list[str] = []
233
for family_id in feature_set.family_ids:
234
family = registry.family(family_id)
235
if family.generated_vector is not None:
236
names.extend(_family_feature_names(family, embedding_dim=embedding_dim))
237
return tuple(names)
238
239
240
def _validate_embedding_dimension(
241
*,
242
embedding_dim: int,
243
embedding_feature_count: int,
244
) -> None:
245
if embedding_dim < 0:
246
raise ValueError("embedding_dim must be non-negative")
247
if embedding_feature_count == 0:
248
if embedding_dim != 0:
249
raise ValueError("embedding_dim must be 0 for feature sets without generated vectors")
250
return
251
if embedding_feature_count != embedding_dim:
252
raise ValueError("embedding_feature_names length must match embedding_dim")
253
254
255
def _family_feature_names(
256
family: FeatureFamilySpecV1,
257
*,
258
embedding_dim: int,
259
) -> tuple[str, ...]:
260
if family.generated_vector is None:
261
return tuple(feature.name for feature in family.features)
262
return generated_vector_feature_names(
263
name_prefix=family.generated_vector.name_prefix,
264
dim=embedding_dim,
265
index_width=family.generated_vector.index_width,
266
)
267
268
269
def _family_schema_hash_part(
270
family: FeatureFamilySpecV1,
271
*,
272
embedding_dim: int,
273
) -> dict[str, object]:
274
if family.generated_vector is not None:
275
vector = family.generated_vector
276
return {
277
"family_id": family.family_id,
278
"family_version": family.family_version,
279
"value_kind": family.value_kind,
280
"meaning_version": vector.meaning_version,
281
"vector_id": vector.vector_id,
282
"name_prefix": vector.name_prefix,
283
"dim": embedding_dim,
284
"default_dim": vector.default_dim,
285
"index_width": vector.index_width,
286
"dtype": vector.dtype,
287
"model_name": vector.model_name,
288
"model_revision": vector.model_revision,
289
"tokenizer_revision": vector.tokenizer_revision,
290
"pooling": vector.pooling,
291
"max_length": vector.max_length,
292
"normalization_policy": vector.normalization_policy,
293
"truncate_policy": vector.truncate_policy,
294
"runtime_requirements": [
295
runtime.model_dump(mode="json", exclude_none=True)
296
for runtime in family.runtime_requirements
297
],
298
}
299
return {
300
"family_id": family.family_id,
301
"family_version": family.family_version,
302
"value_kind": family.value_kind,
303
"features": [
304
{
305
"name": feature.name,
306
"dtype": feature.dtype,
307
"unit": feature.unit,
308
"order": feature.order,
309
"input_role_ids": list(feature.input_role_ids),
310
"denominator_role_id": feature.denominator_role_id,
311
"meaning_version": feature.meaning_version,
312
}
313
for feature in family.features
314
],
315
"runtime_requirements": [
316
runtime.model_dump(mode="json", exclude_none=True)
317
for runtime in family.runtime_requirements
318
],
319
}