Send defaulted union tags in request bodies

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[

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

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

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"})