Annotated

from typing import Generator
import bentoml
from pydantic import BaseModel, Field
import pandas as pd


class Input(BaseModel):
    df1: pd.DataFrame = Field(description="df1")


class Output(BaseModel):
    df1: pd.DataFrame = Field(description="df1")
    df2: pd.DataFrame = Field(description="df2")
    df3: pd.DataFrame = Field(description="df3")


@bentoml.service
class Service:
    @bentoml.api
    def echo_df(self, df1: pd.DataFrame) -> Output:
        return Output(df1=df1, df2=df1, df3=df1)

    @bentoml.api
    def stream_df(self, df1: pd.DataFrame) -> Generator[Output, None, None]:
        for i in range(3):
            yield Output(df1=df1, df2=df1, df3=df1)


if __name__ == "__main__":
    from bentoml import SyncHTTPClient

    data = pd.DataFrame({
        "a": [1, 2, 3],
        "b": ["1", "2", "3"],
    })

    with SyncHTTPClient("localhost:3000") as client:
        value = client.echo_df(df1=data)
        print(value)

        for item in client.stream_df(df1=data):
            print(item)