Repository: bentoml/BentoML
Pull request #5217 (2025-02-17 09:37:35)

feat: support root input spec using positonal-only argument
Signed-off-by: Frost Ming <me@frostming.com>

## What does this PR address?

<!--
Thanks for sending a pull request!

Congrats for making it this far! Here's a 🍱 for you. There are still a few steps ahead.

Please make sure to read the contribution guidelines, then fill out the blanks below before requesting a code review.

Name your Pull Request with one of the following prefixes, e.g. "feat: add support for PyTorch", to indicate the type of changes proposed. This is based on the [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary).
  - feat: (new feature for the user, not a new feature for build script)
  - fix: (bug fix for the user, not a fix to a build script)
  - docs: (changes to the documentation)
  - style: (formatting, missing semicolons, etc; no production code change)
  - refactor: (refactoring production code, eg. renaming a variable)
  - perf: (code changes that improve performance)
  - test: (adding missing tests, refactoring tests; no production code change)
  - chore: (updating grunt tasks etc; no production code change)
  - build: (changes that affect the build system or external dependencies)
  - ci: (changes to configuration files and scripts)
  - revert: (reverts a previous commit)

Describe your changes in detail. Attach screenshots here if appropriate.

Once you're done with this, someone from BentoML team or community member will help review your PR (see "Who can help review?" section for potential reviewers.). If no one has reviewed your PR after a week have passed, don't hesitate to post a new comment and ping @-the same person. Notifications sometimes get lost 🥲.
-->

<!-- Remove if not applicable -->

Fixes #(issue)

## Before submitting:

<!--- Go over all the following points, and put an `x` in all the boxes that apply. -->
<!--- If you're unsure about any of these, don't hesitate to ask. We're here to help! -->
<!--- If you plan to update documentation or tests in follow-up, please note -->

- [ ] Does the Pull Request follow [Conventional Commits specification](https://www.conventionalcommits.org/en/v1.0.0/#summary) naming? Here are [GitHub's
      guide](https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/proposing-changes-to-your-work-with-pull-requests/creating-a-pull-request) on how to create a pull request.
- [ ] Does the code follow BentoML's code style, `pre-commit run -a` script has passed ([instructions](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#style-check-auto-formatting-type-checking))?
- [ ] Did you read through [contribution guidelines](https://github.com/bentoml/BentoML/blob/main/CONTRIBUTING.md#ways-to-contribute) and follow [development guidelines](https://github.com/bentoml/BentoML/blob/main/DEVELOPMENT.md#start-developing)?
- [ ] Did your changes require updates to the documentation? Have you updated
      those accordingly? Here are [documentation guidelines](https://github.com/bentoml/BentoML/tree/main/docs) and [tips on writting docs](https://github.com/bentoml/BentoML/tree/main/docs#writing-documentation).
- [ ] Did you write tests to cover your changes?



Patch:
diff --git a/src/_bentoml_impl/client/http.py b/src/_bentoml_impl/client/http.py
index bff116f9..f69b6194 100644
--- a/src/_bentoml_impl/client/http.py
+++ b/src/_bentoml_impl/client/http.py
@@ -22,6 +22,7 @@ import httpx
 from _bentoml_sdk import IODescriptor
 from _bentoml_sdk.typing_utils import is_image_type
 from bentoml import __version__
+from bentoml._internal.utils.uri import is_http_url
 from bentoml._internal.utils.uri import uri_to_path
 from bentoml.exceptions import BentoMLException
 from bentoml.exceptions import NotFound
@@ -51,10 +52,6 @@ logger = logging.getLogger("bentoml.io")
 MAX_RETRIES = 3
 
 
-def is_http_url(url: str) -> bool:
-    return urlparse(url).scheme in {"http", "https"}
-
-
 def to_async_iterable(iterable: t.Iterable[A]) -> t.AsyncIterable[A]:
     async def _gen() -> t.AsyncIterator[A]:
         for item in iterable:
@@ -222,12 +219,23 @@ class HTTPClient(AbstractClient, t.Generic[C]):
     ) -> httpx.Request:
         from opentelemetry import propagate
 
+        from _bentoml_sdk.io_models import IORootModel
+
         headers = httpx.Headers({"Content-Type": self.media_type, **headers})
         propagate.inject(headers)
         if endpoint.input_spec is not None:
             model = endpoint.input_spec.from_inputs(*args, **kwargs)
-            if model.multipart_fields and self.media_type == "application/json":
+            if (
+                not isinstance(model, IORootModel)
+                and model.multipart_fields
+                and self.media_type == "application/json"
+            ):
                 return self._build_multipart(endpoint, model, headers)
+            elif isinstance(rendered := model.model_dump(), (str, bytes)):
+                headers.update({"content-type": model.mime_type()})
+                return self.client.build_request(
+                    "POST", endpoint.route, content=rendered, headers=headers
+                )
             else:
                 payload = self.serde.serialize_model(model)
                 headers.update(payload.headers)
@@ -239,6 +247,47 @@ class HTTPClient(AbstractClient, t.Generic[C]):
                     if self.client_cls is httpx.AsyncClient
                     else payload.data,
                 )
+        assert self.media_type == "application/json", (
+            "Non-JSON request is not supported"
+        )
+        if endpoint.input.get("root_input", False):
+            if len(args) > 1 or kwargs:
+                raise TypeError("Expected one positional argument for root input")
+            if not args:
+                return self.client.build_request(
+                    "POST", endpoint.route, headers=headers
+                )
+            value = args[0]
+            passthrough = False
+            content = None
+            if "properties" in endpoint.input:
+                kwargs = value
+                args = ()
+                passthrough = True
+            elif endpoint.input.get("type") == "file":
+                file = self._get_file(value)
+                if isinstance(file, str):
+                    content = file
+                else:
+                    file_io, content_type = file[1:]
+                    content = iter(lambda: file_io.read(4096), b"")
+                    if content_type:
+                        headers.update({"content-type": content_type})
+            elif isinstance(value, (str, bytes)):
+                content = value.encode("utf-8") if isinstance(value, str) else value
+                headers.update({"content-type": "text/plain"})
+            else:
+                payload = self.serde.serialize(value, endpoint.input)
+                headers.update(payload.headers)
+                content = (
+                    to_async_iterable(payload.data)
+                    if self.client_cls is httpx.AsyncClient
+                    else payload.data
+                )
+            if not passthrough:
+                return self.client.build_request(
+                    "POST", endpoint.route, content=content, headers=headers
+                )
 
         for name, value in zip(endpoint.input["properties"], args):
             if name in kwargs:
@@ -263,7 +312,7 @@ class HTTPClient(AbstractClient, t.Generic[C]):
             and schema["items"].get("type") == "file"
             for schema in endpoint.input["properties"].values()
         )
-        if has_file and self.media_type == "application/json":
+        if has_file:
             return self._build_multipart(endpoint, kwargs, headers)
         payload = self.serde.serialize(kwargs, endpoint.input)
         headers.update(payload.headers)
@@ -292,6 +341,30 @@ class HTTPClient(AbstractClient, t.Generic[C]):
                     pass
         raise ServiceUnavailable(f"Server is not ready after {timeout} seconds")
 
+    def _get_file(self, value: t.Any) -> str | tuple[str, t.IO[bytes], str | None]:
+        if isinstance(value, str) and not is_http_url(value):
+            value = pathlib.Path(value)
+        if is_image_type(type(value)):
+            fp = getattr(value, "_fp", value.fp)
+            fname = getattr(fp, "name", None)
+            fmt = value.format.lower()
+            return (
+                pathlib.Path(fname).name if fname else f"upload-image.{fmt}",
+                fp,
+                f"image/{fmt}",
+            )
+        elif isinstance(value, pathlib.PurePath):
+            file = open(value, "rb")
+            self._opened_files.append(file)
+            return (value.name, file, mimetypes.guess_type(value)[0])
+        elif isinstance(value, str):
+            return value
+        else:
+            assert isinstance(value, t.BinaryIO)
+            filename = pathlib.Path(getattr(value, "name", "upload-file")).name
+            content_type = mimetypes.guess_type(filename)[0]
+            return (filename, value, content_type)
+
     def _build_multipart(
         self,
         endpoint: ClientEndpoint,
@@ -322,35 +395,11 @@ class HTTPClient(AbstractClient, t.Generic[C]):
                 value = [value]
 
             for v in value:
-                if isinstance(v, str) and not is_http_url(v):
-                    v = pathlib.Path(v)
-                if is_image_type(type(v)):
-                    fp = getattr(v, "_fp", v.fp)
-                    fname = getattr(fp, "name", None)
-                    fmt = v.format.lower()
-                    files.append(
-                        (
-                            name,
-                            (
-                                pathlib.Path(fname).name
-                                if fname
-                                else f"upload-image.{fmt}",
-                                fp,
-                                f"image/{fmt}",
-                            ),
-                        )
-                    )
-                elif isinstance(v, pathlib.PurePath):
-                    file = open(v, "rb")
-                    files.append((name, (v.name, file, mimetypes.guess_type(v)[0])))
-                    self._opened_files.append(file)
-                elif isinstance(v, str):
-                    data.setdefault(name, []).append(v)
+                file = self._get_file(v)
+                if isinstance(file, str):
+                    data[name] = file
                 else:
-                    assert isinstance(v, t.BinaryIO)
-                    filename = pathlib.Path(getattr(v, "name", "upload-file")).name
-                    content_type = mimetypes.guess_type(filename)[0]
-                    files.append((name, (filename, v, content_type)))
+                    files.append((name, file))
         headers.pop("content-type", None)
         return self.client.build_request(
             "POST", endpoint.route, data=data, files=files, headers=headers
diff --git a/src/_bentoml_impl/serde.py b/src/_bentoml_impl/serde.py
index 039878d7..b206e431 100644
--- a/src/_bentoml_impl/serde.py
+++ b/src/_bentoml_impl/serde.py
@@ -20,6 +20,7 @@ from _bentoml_sdk.typing_utils import is_list_type
 from _bentoml_sdk.typing_utils import is_union_type
 from _bentoml_sdk.validators import DataframeSchema
 from _bentoml_sdk.validators import TensorSchema
+from bentoml._internal.utils.uri import is_http_url
 
 if t.TYPE_CHECKING:
     from starlette.requests import Request
@@ -155,6 +156,20 @@ class JSONSerde(GenericSerde, Serde):
             )
         )
 
+    async def parse_request(self, request: Request, cls: type[T]) -> T:
+        import httpx
+
+        from _bentoml_sdk.io_models import IORootModel
+
+        body = await request.body()
+        if issubclass(cls, IORootModel) and cls.multipart_fields:
+            if is_http_url(url := body.decode("utf-8", "ignore")):
+                async with httpx.AsyncClient() as client:
+                    logger.debug("Request with URL, downloading file from %s", url)
+                    resp = await client.get(url)
+                    body = await resp.aread()
+        return self.deserialize_model(Payload((body,), metadata=request.headers), cls)
+
     def deserialize_model(self, payload: Payload, cls: type[T]) -> T:
         return cls.model_validate_json(b"".join(payload.data) or b"{}")
 
@@ -168,7 +183,8 @@ class JSONSerde(GenericSerde, Serde):
 class MultipartSerde(JSONSerde):
     media_type = "multipart/form-data"
 
-    async def ensure_file(self, obj: str | UploadFile) -> UploadFile:
+    @staticmethod
+    async def ensure_file(obj: str | UploadFile) -> UploadFile:
         import httpx
 
         if isinstance(obj, UploadFile):
diff --git a/src/_bentoml_impl/server/app.py b/src/_bentoml_impl/server/app.py
index 1aca430d..c7bdb859 100644
--- a/src/_bentoml_impl/server/app.py
+++ b/src/_bentoml_impl/server/app.py
@@ -638,6 +638,7 @@ class ServiceAppFactory(BaseAppFactory):
     async def api_endpoint(self, name: str, request: Request) -> Response:
         from _bentoml_sdk.io_models import ARGS
         from _bentoml_sdk.io_models import KWARGS
+        from _bentoml_sdk.io_models import IORootModel
         from bentoml._internal.utils import get_original_func
         from bentoml._internal.utils.http import set_cookies
 
@@ -649,14 +650,21 @@ class ServiceAppFactory(BaseAppFactory):
         method = self.service.apis[name]
         func = getattr(self._service_instance, name).local
         ctx = self.service.context
-        serde = ALL_SERDE[media_type]()
+        serde = ALL_SERDE.get(media_type, ALL_SERDE["application/json"])()
         input_data = await method.input_spec.from_http_request(request, serde)
         input_args: tuple[t.Any, ...] = ()
-        input_params = {k: getattr(input_data, k) for k in input_data.model_fields}
+        input_params: dict[str, t.Any] = {}
         if method.ctx_param is not None:
             input_params[method.ctx_param] = ctx
+        if getattr(method.input_spec, "__root_input__", False):
+            if isinstance(input_data, IORootModel):
+                input_args = t.cast(t.Tuple[t.Any], (input_data.root,))
+            else:
+                input_args = (input_data,)
+        else:
+            input_params = {k: getattr(input_data, k) for k in input_data.model_fields}
         if ARGS in input_params:
-            input_args = tuple(input_params.pop(ARGS))
+            input_args = (*input_args, input_params.pop(ARGS))
         if KWARGS in input_params:
             input_params.update(input_params.pop(KWARGS))
 
diff --git a/src/_bentoml_sdk/images.py b/src/_bentoml_sdk/images.py
index fdea64f5..6e70d8f0 100644
--- a/src/_bentoml_sdk/images.py
+++ b/src/_bentoml_sdk/images.py
@@ -22,6 +22,11 @@ from bentoml._internal.utils.pkg import get_local_bentoml_dependency
 from bentoml.exceptions import BentoMLConfigException
 from bentoml.exceptions import BentoMLException
 
+if sys.version_info >= (3, 11):
+    import tomllib
+else:
+    import tomli as tomllib
+
 logger = logging.getLogger("bentoml.build")
 
 DEFAULT_PYTHON_VERSION = f"{sys.version_info.major}.{sys.version_info.minor}"
@@ -53,6 +58,21 @@ class Image:
         self._after_pip_install = True
         return self
 
+    def pyproject_toml(self, file_path: str) -> t.Self:
+        """Add a pyproject.toml to the image. Supports chaining call.
+
+        Example:
+
+        .. code-block:: python
+
+            image = Image("debian:latest").pyproject_toml("pyproject.toml")
+        """
+        with Path(file_path).open("rb") as f:
+            pyproject_toml = tomllib.load(f)
+        dependencies = pyproject_toml.get("project", {}).get("dependencies", {})
+        self.python_packages(*dependencies)
+        return self
+
     def python_packages(self, *packages: str) -> t.Self:
         """Add python dependencies to the image. Supports chaining call.
 
diff --git a/src/_bentoml_sdk/io_models.py b/src/_bentoml_sdk/io_models.py
index 958ff309..ba10279f 100644
--- a/src/_bentoml_sdk/io_models.py
+++ b/src/_bentoml_sdk/io_models.py
@@ -93,6 +93,13 @@ class IOMixin:
     multipart_fields: ClassVar[t.List[str]]
     media_type: ClassVar[t.Optional[str]] = None
 
+    @classmethod
+    def model_json_schema(cls, *args: t.Any, **kwargs: t.Any) -> dict[str, t.Any]:
+        schema = super().model_json_schema(*args, **kwargs)
+        if getattr(cls, "__root_input__", False):
+            schema["root_input"] = True
+        return schema
+
     @classmethod
     def openapi_components(cls, name: str) -> dict[str, Schema]:
         from .service.openapi import REF_TEMPLATE
@@ -174,6 +181,16 @@ class IOMixin:
     @classmethod
     def from_inputs(cls, *args: t.Any, **kwargs: t.Any) -> IODescriptor:
         assert issubclass(cls, IODescriptor)
+        if getattr(cls, "__root_input__", False):
+            if len(args) > 1:
+                raise TypeError("Expected exactly 1 argument")
+            if len(args) < 1:
+                return cls()
+            arg = args[0]
+            if issubclass(cls, IORootModel) or not isinstance(arg, dict):
+                return cls(arg)
+            else:
+                return cls(**arg)
         model_fields = list(cls.model_fields)
         for i, arg in enumerate(args):
             if i < len(model_fields) and model_fields[i] == ARGS:
@@ -336,7 +353,8 @@ class IODescriptor(IOMixin, BaseModel):
         parameter_tuples = iter(signature.parameters.items())
         if skip_self:
             next(parameter_tuples)
-
+        positional_only_param: t.Any = None
+        positional_only_default: t.Any = signature.empty
         for name, param in parameter_tuples:
             if name in skip_names:
                 # Reserved name for context object passed in
@@ -346,6 +364,16 @@ class IODescriptor(IOMixin, BaseModel):
                 annotation = t.Any
             else:
                 annotation = eval_type_lenient(annotation, global_ns, None)
+            if param.kind == param.POSITIONAL_ONLY:
+                if positional_only_param is None:
+                    positional_only_param = annotation
+                    positional_only_default = param.default
+                    continue
+
+            if positional_only_param is not None:
+                raise TypeError(
+                    "When positional-only argument is used, no other parameters can be specified"
+                )
             if param.kind == param.VAR_KEYWORD:
                 name = KWARGS
                 annotation = t.Dict[str, t.Any]
@@ -358,6 +386,23 @@ class IODescriptor(IOMixin, BaseModel):
             fields[name] = (annotation, default)
 
         try:
+            if positional_only_param is not None:
+                content_type = None
+                if is_annotated(positional_only_param):
+                    content_type = next(
+                        (
+                            a
+                            for a in get_args(positional_only_param)
+                            if isinstance(a, ContentType)
+                        ),
+                        None,
+                    )
+                typ_ = ensure_io_descriptor(
+                    positional_only_param, positional_only_default
+                )
+                typ_.__root_input__ = True
+                typ_.media_type = content_type.content_type if content_type else None
+                return typ_
             return t.cast(
                 t.Type[IODescriptor],
                 create_model(
@@ -429,7 +474,11 @@ class IORootModel(IODescriptor, t.Generic[RootModelRootType]):
     def model_validate_json(
         cls, json_data: str | bytes | bytearray, **kwargs: t.Any
     ) -> t.Self:
-        return cls.model_validate(json.loads(json_data), **kwargs)
+        if getattr(cls, "__root_input__", False):
+            parsed = json_data
+        else:
+            parsed = json.loads(json_data)
+        return cls.model_validate(parsed, **kwargs)
 
     @classmethod
     def model_json_schema(cls, *args: t.Any, **kwargs: t.Any) -> dict[str, t.Any]:
@@ -438,10 +487,15 @@ class IORootModel(IODescriptor, t.Generic[RootModelRootType]):
             typ_ = t.Annotated[(field.annotation, *field.metadata)]
         else:
             typ_ = field.annotation
-        return TypeAdapter(typ_).json_schema(*args, **kwargs)
+        schema = TypeAdapter(typ_).json_schema(*args, **kwargs)
+        if getattr(cls, "__root_input__", False):
+            schema["root_input"] = True
+        return schema
 
 
-def ensure_io_descriptor(typ_: type) -> type[IODescriptor]:
+def ensure_io_descriptor(
+    typ_: type, root_default: t.Any = inspect.Signature.empty
+) -> type[IODescriptor]:
     from pydantic._internal._utils import lenient_issubclass
 
     type_name = getattr(typ_, "__name__", "")
@@ -453,4 +507,13 @@ def ensure_io_descriptor(typ_: type) -> type[IODescriptor]:
                 create_model(f"{type_name}IODescriptor", __base__=(IOMixin, typ_)),
             )
         return typ_
-    return IORootModel[typ_]
+
+    if root_default is inspect.Signature.empty:
+        extras = {}
+    else:
+        extras = {"root": (typ_, root_default)}
+
+    return t.cast(
+        t.Type[IODescriptor],
+        create_model(f"{type_name}IODescriptor", __base__=IORootModel[typ_], **extras),
+    )
diff --git a/src/bentoml/_internal/service/openapi/specification.py b/src/bentoml/_internal/service/openapi/specification.py
index 0604c91a..388e76e7 100644
--- a/src/bentoml/_internal/service/openapi/specification.py
+++ b/src/bentoml/_internal/service/openapi/specification.py
@@ -118,6 +118,7 @@ class Schema:
     externalDocs: t.Optional[ExternalDocumentation] = None
     example: t.Optional[t.Any] = None
     deprecated: t.Optional[bool] = None
+    root_input: bool = False
     # not yet supported: xml
 
 
diff --git a/src/bentoml/_internal/utils/uri.py b/src/bentoml/_internal/utils/uri.py
index ec9032db..7190805c 100644
--- a/src/bentoml/_internal/utils/uri.py
+++ b/src/bentoml/_internal/utils/uri.py
@@ -46,3 +46,7 @@ def uri_to_path(uri: str) -> str:
 def encode_path_for_uri(path: str) -> str:
     """Percent-encode non-URL characters in a path."""
     return quote(path.replace(os.sep, "/"))
+
+
+def is_http_url(url: str) -> bool:
+    return urlparse(url).scheme in {"http", "https"}


Test patch:
diff --git a/tests/e2e/bento_new_sdk/test_asgi.py b/tests/e2e/bento_new_sdk/test_asgi.py
index 84ec7b20..a3795d33 100644
--- a/tests/e2e/bento_new_sdk/test_asgi.py
+++ b/tests/e2e/bento_new_sdk/test_asgi.py
@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+from pathlib import Path
+
+import pydantic
 from starlette.testclient import TestClient
 
 import bentoml
@@ -14,6 +17,28 @@ class Model:
         return self.model.get(key, "")
 
 
+@bentoml.service(metrics={"enabled": False})
+class RootInput:
+    @bentoml.api()
+    def get(self, key: str, /) -> str:
+        return f"root-{key}"
+
+    @bentoml.api()
+    def read(self, path: Path, /) -> str:
+        return path.read_text()
+
+
+class InputModel(pydantic.BaseModel):
+    key: str
+
+
+@bentoml.service(metrics={"enabled": False})
+class RootModelInput:
+    @bentoml.api
+    def get(self, input: InputModel, /) -> str:
+        return f"root-{input.key}"
+
+
 @bentoml.service(metrics={"enabled": False})
 class Service:
     model = bentoml.depends(Model)
@@ -37,3 +62,23 @@ def test_composed_service():
 
         assert resp.status_code == 200
         assert resp.json() == {"value": "world"}
+
+
+def test_service_root_input(tmp_path: Path):
+    with TestClient(app=RootInput.to_asgi()) as client:
+        resp = client.post("/get", content="hello")
+
+        assert resp.status_code == 200
+        assert resp.text == "root-hello"
+
+        tmp_path.joinpath("test.txt").write_text("hello")
+        resp = client.post("/read", content=tmp_path.joinpath("test.txt").read_bytes())
+
+        assert resp.status_code == 200
+        assert resp.text == "hello"
+
+    with TestClient(app=RootModelInput.to_asgi()) as client:
+        resp = client.post("/get", json={"key": "hello"})
+
+        assert resp.status_code == 200
+        assert resp.text == "root-hello"
diff --git a/tests/unit/bentoml_io/test_decorators.py b/tests/unit/bentoml_io/test_decorators.py
index f3c3b44f..ef540eb0 100644
--- a/tests/unit/bentoml_io/test_decorators.py
+++ b/tests/unit/bentoml_io/test_decorators.py
@@ -8,6 +8,7 @@ from starlette.testclient import TestClient
 from typing_extensions import Annotated
 
 import bentoml
+from _bentoml_sdk.validators import TensorSchema
 from bentoml.validators import DataframeSchema
 from bentoml.validators import DType
 from bentoml.validators import Shape
@@ -100,10 +101,16 @@ def test_api_decorator_numpy():
     ) -> Annotated[npt.NDArray[np.int64], DType("int64"), Shape((1,))]:
         return arr.astype(np.int64)
 
-    numpy_func.input_spec.model_fields["arr"].annotation is npt.NDArray[np.float64]
-    numpy_func.output_spec.model_fields["root"].annotation is Annotated[
-        npt.NDArray[np.int64], DType("int64"), Shape((1,))
-    ]
+    assert (
+        numpy_func.input_spec.model_fields["arr"].annotation == npt.NDArray[np.float64]
+    )
+    assert (
+        numpy_func.output_spec.model_fields["root"].annotation == npt.NDArray[np.int64]
+    )
+    assert (
+        TensorSchema(format="numpy-array", dtype="int64")
+        in numpy_func.output_spec.model_fields["root"].metadata
+    )
 
     with pytest.raises(
         TypeError,
@@ -130,12 +137,39 @@ def test_api_decorator_pandas():
     ]:
         return pd.concat([df1, df2], axis=1)
 
-    pandas_func.input_spec.model_fields["df1"].annotation is pd.DataFrame
-    pandas_func.input_spec.model_fields["df2"].annotation is Annotated[
-        pd.DataFrame,
-        DataframeSchema(columns=("b",)),
-    ]
-    pandas_func.output_spec.model_fields["root"].annotation is Annotated[
-        pd.DataFrame,
-        DataframeSchema(orient="columns", columns=("a", "b")),
-    ]
+    assert pandas_func.input_spec.model_fields["df1"].annotation is pd.DataFrame
+    assert pandas_func.input_spec.model_fields["df2"].annotation is pd.DataFrame
+    assert (
+        DataframeSchema(columns=("b",))
+        in pandas_func.input_spec.model_fields["df2"].metadata
+    )
+    assert pandas_func.output_spec.model_fields["root"].annotation is pd.DataFrame
+    assert (
+        DataframeSchema(orient="columns", columns=("a", "b"))
+        in pandas_func.output_spec.model_fields["root"].metadata
+    )
+
+
+def test_api_root_input():
+    from _bentoml_sdk.io_models import IORootModel
+
+    @bentoml.api
+    def root_input(_, name: str, /) -> str:
+        return name
+
+    assert issubclass(root_input.input_spec, IORootModel)
+    assert root_input.input_spec.model_fields["root"].annotation is str
+
+
+def test_api_root_input_illegal():
+    with pytest.raises(TypeError):
+
+        @bentoml.api
+        def root_input(_, name: str, age: int, /) -> str:
+            return name
+
+    with pytest.raises(TypeError):
+
+        @bentoml.api
+        def root_input(_, name: str, /, age: int) -> str:
+            return name