From a337ae6ea0aaed57c74757b463a9fbf399cd20a0 Mon Sep 17 00:00:00 2001 From: Vincent <0426vincent@gmail.com> Date: Mon, 9 Nov 2026 10:19:56 -0800 Subject: [PATCH] Send defaulted union tags in request bodies The spec gives each tagged-union member's tag field (kind/typ) a default, so the generated models default it too. Request bodies are dumped with exclude_unset, which dropped a tag the caller left at its default, and the server rejects the body: 400 Unable to extract tag using discriminator 'kind' [type=union_tag_not_found] For example, CreateModelRequest(source=ModelArchiveSource(name=..., deployment=...)) sent no kind. Postprocessing now marks each defaulted tag as set on every instance, so it is always sent, matching the Go SDK's From* helpers. Other defaulted fields are still left to the server. --- baseten/client/managementapi/_models.py | 42 +++++++++++++++++++++++++ baseten/client/modelconfig/_models.py | 9 ++++++ scripts/apigen/postprocess.py | 40 +++++++++++++++++++++++ tests/client/test_management.py | 33 +++++++++++++++++++ 4 files changed, 124 insertions(+) diff --git a/baseten/client/managementapi/_models.py b/baseten/client/managementapi/_models.py index b5f378f..9d678f3 100644 --- a/baseten/client/managementapi/_models.py +++ b/baseten/client/managementapi/_models.py @@ -25,6 +25,9 @@ class BasetenLatestCheckpointConfig(BaseModel): "baseten_latest_checkpoint" ) + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("typ",)) + class BasetenNamedCheckpointConfig(BaseModel): checkpoint_name: Annotated[ @@ -48,6 +51,9 @@ class BasetenNamedCheckpointConfig(BaseModel): "baseten_named_checkpoint" ) + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("typ",)) + class CreateTrainingJobCacheConfig(BaseModel): enable_legacy_hf_mount: Annotated[ @@ -166,6 +172,9 @@ class LoopsCheckpointConfig(BaseModel): ] = "trainer" typ: Annotated[Literal["loops_checkpoint"], Field(title="Typ")] = "loops_checkpoint" + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("typ",)) + class TrussUserEnv(BaseModel): truss_client_version: Annotated[str | None, Field(title="Truss Client Version")] = ( @@ -4195,6 +4204,9 @@ class LibraryListingSource(BaseModel): ), ] = None + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("kind",)) + class AutoscalingSettings(BaseModel): min_replica: Annotated[ @@ -4541,6 +4553,9 @@ class DeploymentPatchOpConfig(BaseModel): str, Field(description="Config file path within the source.", title="Path") ] = "config.yaml" + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("type",)) + class DeploymentPatchOpEnvVar(BaseModel): type: Annotated[Literal["environment_variable"], Field(title="Type")] = ( @@ -4564,6 +4579,9 @@ class DeploymentPatchOpEnvVar(BaseModel): ), ] = None + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("type",)) + class DeploymentPatchOpExternalData(BaseModel): type: Annotated[Literal["external_data"], Field(title="Type")] = "external_data" @@ -4579,6 +4597,9 @@ class DeploymentPatchOpExternalData(BaseModel): Field(description="The single external data item descriptor.", title="Item"), ] + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("type",)) + class DeploymentPatchOpModelCode(BaseModel): type: Annotated[Literal["model_code"], Field(title="Type")] = "model_code" @@ -4616,6 +4637,9 @@ class DeploymentPatchOpModelCode(BaseModel): ), ] = False + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("type",)) + class DeploymentPatchOpPackage(BaseModel): type: Annotated[Literal["package"], Field(title="Type")] = "package" @@ -4647,6 +4671,9 @@ class DeploymentPatchOpPackage(BaseModel): ), ] = None + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("type",)) + class DeploymentPatchOpPythonRequirement(BaseModel): type: Annotated[Literal["python_requirement"], Field(title="Type")] = ( @@ -4667,6 +4694,9 @@ class DeploymentPatchOpPythonRequirement(BaseModel): ), ] + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("type",)) + class DeploymentPatchPoint(BaseModel): content_hashes: Annotated[ @@ -6995,6 +7025,9 @@ class LoopsCheckpointFilesResponse(BaseModel): class LoopsCheckpointS3Source(BaseModel): kind: Annotated[Literal["s3"], Field(title="Kind")] = "s3" + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("kind",)) + class LoopsCheckpointVolumeSource(BaseModel): kind: Annotated[Literal["volume"], Field(title="Kind")] = "volume" @@ -7013,6 +7046,9 @@ class LoopsCheckpointVolumeSource(BaseModel): ), ] + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("kind",)) + class LoopsCheckpointSourceResponse(BaseModel): source: Annotated[ @@ -16507,6 +16543,9 @@ class ModelArchiveSource(BaseModel): ), ] = False + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("kind",)) + class CreateModelRequest(BaseModel): source: Annotated[ @@ -16536,6 +16575,9 @@ class DeploymentArchiveSource(BaseModel): ), ] = None + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("kind",)) + class CreateModelDeploymentRequest(BaseModel): source: Annotated[ diff --git a/baseten/client/modelconfig/_models.py b/baseten/client/modelconfig/_models.py index 64ca410..fe85dc2 100644 --- a/baseten/client/modelconfig/_models.py +++ b/baseten/client/modelconfig/_models.py @@ -215,10 +215,16 @@ class FabricRequirement(BaseModel): class GRPCOptions(BaseModel): kind: Annotated[Literal["grpc"], Field(title="Kind")] = "grpc" + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("kind",)) + class HTTPOptions(BaseModel): kind: Annotated[Literal["http"], Field(title="Kind")] = "http" + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("kind",)) + class HealthChecks(BaseModel): model_config = ConfigDict( @@ -409,6 +415,9 @@ class WebsocketOptions(BaseModel): float | None, Field(title="Ping Timeout Seconds") ] = None + def model_post_init(self, context: Any, /) -> None: + self.__pydantic_fields_set__.update(("kind",)) + class NumDraftTokens(RootModel[int]): root: Annotated[int | None, Field(ge=1, title="Num Draft Tokens")] = None diff --git a/scripts/apigen/postprocess.py b/scripts/apigen/postprocess.py index 5cab701..faf0c5c 100644 --- a/scripts/apigen/postprocess.py +++ b/scripts/apigen/postprocess.py @@ -46,9 +46,49 @@ def postprocess_models(src: str) -> str: src = _allow_population_by_field_name(src) src = _validate_literal_model_defaults(src) src = _open_literals(src) + src = _send_defaulted_tags(src) return _open_tagged_unions(src) +def _send_defaulted_tags(src: str) -> str: + # A union member's tag field defaults to its only value, so callers need + # not pass it. Request bodies are dumped with exclude_unset, which would + # then drop the tag and leave the server unable to pick the member, so each + # such tag is marked set on every instance. + tree = ast.parse(src) + discriminators = _discriminator_fields(tree) + lines = src.splitlines(keepends=True) + + for cls in reversed(tree.body): + if not isinstance(cls, ast.ClassDef): + continue + tags = [ + stmt.target.id + for stmt in cls.body + if isinstance(stmt, ast.AnnAssign) + and isinstance(stmt.target, ast.Name) + and stmt.value is not None + and (cls.name, stmt.target.id) in discriminators + ] + if not tags: + continue + if any( + isinstance(stmt, ast.FunctionDef) and stmt.name == "model_post_init" + for stmt in cls.body + ): + raise ValueError(f"{cls.name} already defines model_post_init") + assert cls.end_lineno is not None + indent = " " * cls.body[0].col_offset + added = ", ".join(repr(tag) for tag in tags) + lines.insert( + cls.end_lineno, + f"\n{indent}def model_post_init(self, context: Any, /) -> None:\n" + f"{indent} self.__pydantic_fields_set__.update(({added},))\n", + ) + + return _ensure_import("".join(lines), "typing", ["Any"]) + + def _open_tagged_unions(src: str) -> str: # A tagged union rejects a tag outside its members, so a variant the server # adds later would fail validation of the whole response. Each one is kept diff --git a/tests/client/test_management.py b/tests/client/test_management.py index 3535c7e..7e3caa0 100644 --- a/tests/client/test_management.py +++ b/tests/client/test_management.py @@ -207,6 +207,39 @@ async def test_post_with_body() -> None: await client.close() +def test_defaulted_union_tag_sent() -> None: + # The server needs the tag to pick the union member, so it is sent though + # the caller left it at its default. + fake = FakeTransport(400, {"detail": "stop here"}) + client = make_sync_client(fake) + + with pytest.raises(baseten.client.managementapi.ResponseError): + client.api.post_models( + request=baseten.client.managementapi.CreateModelRequest( + source=baseten.client.managementapi.ModelArchiveSource( + name="my-model", + deployment=baseten.client.managementapi.DeploymentArchivePayload( + config={} + ), + ), + ), + ) + + body = json.loads(fake.capture.body) + assert body["source"]["kind"] == "model_archive" + # Other defaulted fields are still left to the server. + assert body["source"]["deployment"] == {"config": {}} + client.close() + + +def test_defaulted_union_tag_set_on_validated_response() -> None: + source = baseten.client.managementapi.ModelArchiveSource.model_validate( + {"name": "my-model", "deployment": {"config": {}}} + ) + + assert "kind" in source.model_fields_set + + @pytest.mark.asyncio async def test_response_error() -> None: fake = FakeTransport(500, {"detail": "boom"})