Refresh API specs, add query parameter support, bump to 0.10.0

baseten/client/modelconfig/__init__.py

@@ -21,6 +21,7 @@
     DockerAuthSettings,
     DockerAuthType,
     DockerServer,
+    EgressRestrictions,
     ExternalData,
     ExternalDataItem,
     GRPCOptions,
@@ -84,6 +85,7 @@
     "DockerAuthSettings",
     "DockerAuthType",
     "DockerServer",
+    "EgressRestrictions",
     "ExternalData",
     "ExternalDataItem",
     "GRPCOptions",

baseten/client/modelconfig/_models.py

@@ -102,6 +102,28 @@ class DockerServer(BaseModel):
     ] = None
 
 
+class EgressRestrictions(BaseModel):
+    model_config = ConfigDict(
+        extra="allow",
+    )
+    ip_allow_list: Annotated[
+        list[str] | None,
+        Field(
+            description="Allowed outbound IPv4 addresses or CIDR ranges. Use null or [] alongside an equally restrictive fqdn_allow_list to block all egress.",
+            examples=[["1.1.1.1", "8.8.8.8/32"]],
+            title="Ip Allow List",
+        ),
+    ] = None
+    fqdn_allow_list: Annotated[
+        list[str] | None,
+        Field(
+            description="Allowed outbound fully-qualified domain names. Supports wildcards: '*' may appear anywhere in a label.",
+            examples=[["*.baseten.co", "huggingface.co"]],
+            title="Fqdn Allow List",
+        ),
+    ] = None
+
+
 class ExternalDataItem(BaseModel):
     model_config = ConfigDict(
         extra="allow",
@@ -690,6 +712,12 @@ class Runtime(BaseModel):
     ] = None
     health_checks: HealthChecks | None = None
     remote_ssh: RemoteSSH | None = None
+    egress_restrictions: Annotated[
+        EgressRestrictions | None,
+        Field(
+            description="Egress network restrictions for the model version. When unset, all egress is allowed (default)."
+        ),
+    ] = None
     truss_server_version_override: Annotated[
         str | None,
         Field(
@@ -752,7 +780,7 @@ class WeightsSource(BaseModel):
     source: Annotated[
         str,
         Field(
-            description="URI with scheme prefix. Use hf://, s3://, gs://, azure://, r2://, or https://. For HuggingFace, use @revision suffix (e.g., hf://owner/repo@main).",
+            description="URI with scheme prefix. Use hf://, s3://, gs://, azure://, r2://, cw://, or https://. For HuggingFace, use @revision suffix (e.g., hf://owner/repo@main).",
             min_length=1,
             title="Source",
         ),

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "baseten"
-version = "0.9.0"
+version = "0.10.0"
 description = "Baseten Python SDK"
 readme = "README.md"
 license = "MIT"

scripts/apigen/clientgen.py

@@ -26,12 +26,19 @@ class _Operation:
     success_code: int
     error_codes: dict[int, str] | None
     summary: str
+    query_ref: str
+    query_required: bool
 
 
-def _extract_operations(spec: dict) -> list[_Operation]:
-    paths = spec.get("paths", {})
+def resolve_method_names(spec: dict) -> dict[tuple[str, str], str]:
+    """Map each (path, http_method) to its resolved client method name.
 
-    # Collect raw operation data with short names (all params stripped).
+    Names are derived from the method and path, using a trailing path
+    parameter only where needed to disambiguate collisions. Shared with
+    preprocessing so injected query-parameter schemas can be named to
+    match their operation's method.
+    """
+    paths = spec.get("paths", {})
     raw: list[tuple[str, str, dict]] = []
     short_names: dict[str, int] = {}
     for path, path_item in paths.items():
@@ -44,8 +51,7 @@ def _extract_operations(spec: dict) -> list[_Operation]:
             )
             short_names[name] = short_names.get(name, 0) + 1
 
-    # Build operations, using trailing param only where needed to disambiguate.
-    ops: list[_Operation] = []
+    result: dict[tuple[str, str], str] = {}
     for path, http_method, op_data in raw:
         short = _derive_method_name(
             http_method, path, op_data, keep_trailing_param=False
@@ -56,20 +62,52 @@ def _extract_operations(spec: dict) -> list[_Operation]:
             )
         else:
             name = short
-        ops.append(
-            _Operation(
-                name=name,
-                http_method=http_method.upper(),
-                path=path,
-                path_params=_PATH_PARAM_RE.findall(path),
-                has_body="requestBody" in op_data,
-                req_body_ref=_body_schema_ref(spec, op_data),
-                resp_ref=_response_schema_ref(spec, op_data),
-                success_code=_extract_success_code(op_data, http_method, path),
-                error_codes=_error_code_map(spec, op_data),
-                summary=op_data.get("summary", ""),
+        result[(path, http_method)] = name
+    return result
+
+
+def query_request_model_name(method_name: str) -> str:
+    """Model name for an operation's injected query-parameter schema."""
+    return _snake_to_pascal(method_name) + "Request"
+
+
+def _snake_to_pascal(s: str) -> str:
+    return "".join(part.capitalize() for part in s.split("_"))
+
+
+def _extract_operations(spec: dict) -> list[_Operation]:
+    paths = spec.get("paths", {})
+    names = resolve_method_names(spec)
+
+    ops: list[_Operation] = []
+    for path, path_item in paths.items():
+        for http_method, op_data in path_item.items():
+            if http_method == "parameters" or not isinstance(op_data, dict):
+                continue
+            name = names[(path, http_method)]
+            query_params = [
+                p
+                for p in op_data.get("parameters", [])
+                if isinstance(p, dict) and p.get("in") == "query"
+            ]
+            query_ref = query_request_model_name(name) if query_params else ""
+            query_required = any(p.get("required") for p in query_params)
+            ops.append(
+                _Operation(
+                    name=name,
+                    http_method=http_method.upper(),
+                    path=path,
+                    path_params=_PATH_PARAM_RE.findall(path),
+                    has_body="requestBody" in op_data,
+                    req_body_ref=_body_schema_ref(spec, op_data),
+                    resp_ref=_response_schema_ref(spec, op_data),
+                    success_code=_extract_success_code(op_data, http_method, path),
+                    error_codes=_error_code_map(spec, op_data),
+                    summary=op_data.get("summary", ""),
+                    query_ref=query_ref,
+                    query_required=query_required,
+                )
             )
-        )
     ops.sort(key=lambda o: o.name)
     return ops
 
@@ -189,6 +227,8 @@ def _render_client(ops: list[_Operation]) -> str:
     for op in ops:
         if op.req_body_ref:
             model_imports.add(op.req_body_ref)
+        if op.query_ref:
+            model_imports.add(op.query_ref)
         if op.resp_ref:
             model_imports.add(op.resp_ref)
         for ref in (op.error_codes or {}).values():
@@ -255,6 +295,7 @@ class _ApiRequest:
     path_fmt: str
     path_args: list[str]
     body: Any
+    query: Any
     success_code: int
     error_codes: dict[int, str] | None
 """
@@ -325,10 +366,23 @@ def __init__(self, http_client: {http_cls}) -> None:
         json_body = None
         if request.body is not None:
             if isinstance(request.body, BaseModel):
-                json_body = request.body.model_dump(mode="json")
+                # Only fields the caller set are sent, so unset fields fall
+                # back to the server default rather than being reset here.
+                # An explicit None is kept, since null can mean "clear".
+                json_body = request.body.model_dump(mode="json", exclude_unset=True)
             else:
                 json_body = request.body
-        response = {aw}self._http_client.request(request.method, path, json=json_body)
+        params = None
+        if request.query is not None:
+            if isinstance(request.query, BaseModel):
+                # As above, plus dropping None: a null query parameter is
+                # meaningless and would otherwise serialize as an empty string.
+                params = request.query.model_dump(
+                    mode="json", exclude_unset=True, exclude_none=True
+                )
+            else:
+                params = request.query
+        response = {aw}self._http_client.request(request.method, path, json=json_body, params=params)
         if response.status_code != request.success_code:
 {error_dispatch}\
             raise ResponseError(status_code=response.status_code, body=response.text)
@@ -358,18 +412,36 @@ def _render_method(op: _Operation, *, is_async: bool) -> str:
     adef = "async def" if is_async else "def"
     aw = "await " if is_async else ""
 
+    # An operation carries at most one input model: a request body (any
+    # method other than GET) or query parameters (GET only). Both surface
+    # as a single keyword-only `request` argument. Query requests with no
+    # required fields are optional so callers can omit them entirely; body
+    # requests are always required so an empty body still sends `{}`.
+    if op.query_ref:
+        input_type = op.query_ref
+        input_required = op.query_required
+    elif op.has_body:
+        input_type = op.req_body_ref if op.req_body_ref else "Any"
+        input_required = True
+    else:
+        input_type = ""
+        input_required = False
+
     params = ["self"]
-    if op.path_params or op.has_body:
+    if op.path_params or input_type:
         params.append("*")
     for p in op.path_params:
         params.append(f"{p}: str")
-    if op.has_body:
-        body_type = op.req_body_ref if op.req_body_ref else "Any"
-        params.append(f"body: {body_type}")
+    if input_type:
+        if input_required:
+            params.append(f"request: {input_type}")
+        else:
+            params.append(f"request: {input_type} | None = None")
 
     ret = f" -> {op.resp_ref}" if op.resp_ref else " -> None"
     path_args = f"[{', '.join(op.path_params)}]" if op.path_params else "[]"
-    body_arg = "body" if op.has_body else "None"
+    body_arg = "request" if (input_type and not op.query_ref) else "None"
+    query_arg = "request" if op.query_ref else "None"
 
     if op.error_codes:
         codes = sorted(op.error_codes.items())
@@ -383,6 +455,7 @@ def _render_method(op: _Operation, *, is_async: bool) -> str:
         f"path_fmt={_path_fmt(op.path)!r}, "
         f"path_args={path_args}, "
         f"body={body_arg}, "
+        f"query={query_arg}, "
         f"success_code={op.success_code}, "
         f"error_codes={error_expr})"
     )

scripts/apigen/postprocess.py

@@ -18,8 +18,8 @@
 # https://github.com/koxudaxi/datamodel-code-generator/issues/2027 (closed
 # but the issue persists for this shape in 0.55.0).
 _ROOT_MODEL_BLOCK = re.compile(
-    r"(class \w+\(RootModel\[(?P<wrapped>[^\]]+)\]\):\n"
-    r"    root: )(?P<rest>.+?)(?P<eq> = None\n)",
+    r"(class \w+\(RootModel\[(?P<wrapped>[^\n]+?)\]\):\n"
+    r"    root: )(?P<rest>(?:(?!\nclass ).)+?)(?P<eq> = None\n)",
     re.DOTALL,
 )
 

scripts/apigen/preprocess.py

@@ -1,8 +1,11 @@
 """Preprocesses OpenAPI specs for code generation."""
 
+import copy
 import json
 import re
 
+from scripts.apigen.clientgen import query_request_model_name, resolve_method_names
+
 
 def preprocess_truss_config_schema(data: bytes) -> bytes:
     doc = json.loads(data)
@@ -56,6 +59,13 @@ def preprocess_spec(data: bytes) -> bytes:
     # wrapper classes.
     _hoist_component_schemas(doc)
 
+    # Synthesize a request schema per GET operation's query parameters so
+    # datamodel-code-generator emits a typed model for them (it only
+    # generates query-parameter models under the paths scope, which drags
+    # in unwanted per-operation wrappers). Injected before the V1 rename
+    # below so their $refs to enums are rewritten with everything else.
+    _inject_query_request_schemas(doc)
+
     # datamodel-code-generator generates empty BaseModel classes for
     # schemas that are bare type: object with no properties (e.g.
     # PredictInput). Adding additionalProperties makes it correctly
@@ -98,6 +108,53 @@ def _hoist_component_schemas(doc: dict) -> None:
             content["schema"] = {"$ref": f"#/components/schemas/{name}"}
 
 
+def _inject_query_request_schemas(doc: dict) -> None:
+    # Build an object schema whose properties are the operation's query
+    # parameters, named to match its client method (e.g. get_users ->
+    # GetUsersRequest). Each parameter's own schema (enum $refs, arrays,
+    # nullable wrappers, constraints) is reused verbatim so the third
+    # party types every field. GET carries only query params and every
+    # other method only a body, so this name never collides with a body.
+    schemas = doc.setdefault("components", {}).setdefault("schemas", {})
+    method_names = resolve_method_names(doc)
+
+    for path, path_item in doc.get("paths", {}).items():
+        for http_method, op in path_item.items():
+            if http_method == "parameters" or not isinstance(op, dict):
+                continue
+            query_params = [
+                p
+                for p in op.get("parameters", [])
+                if isinstance(p, dict) and p.get("in") == "query"
+            ]
+            if not query_params:
+                continue
+            if "requestBody" in op:
+                raise ValueError(
+                    f"{http_method.upper()} {path} has both a request body and "
+                    "query parameters; the client generator assumes GET carries "
+                    "only query parameters and other methods only a body"
+                )
+            name = query_request_model_name(method_names[(path, http_method)])
+            if name in schemas:
+                raise ValueError(
+                    f"injected query schema {name} collides with an existing schema"
+                )
+            properties: dict = {}
+            required: list[str] = []
+            for p in query_params:
+                schema = copy.deepcopy(p.get("schema", {}))
+                if "description" not in schema and p.get("description"):
+                    schema["description"] = p["description"]
+                properties[p["name"]] = schema
+                if p.get("required"):
+                    required.append(p["name"])
+            obj: dict = {"type": "object", "title": name, "properties": properties}
+            if required:
+                obj["required"] = required
+            schemas[name] = obj
+
+
 def _fix_bare_object_schemas(doc: dict) -> None:
     schemas = doc.get("components", {}).get("schemas", {})
     for schema in schemas.values():