import fal
from pydantic import BaseModel
from fastapi import WebSocket
class Input(BaseModel):
    prompt: str = Field()
class Output(BaseModel):
    output: str = Field()
class RealtimeApp(fal.App):
    @fal.endpoint("/")
    def generate(self, input: Input) -> Output:
        return Output(output=input.prompt)
    @fal.realtime("/realtime")
    def generate_rt(self, input: Input) -> Output:
        return Output(output=input.prompt)
    @fal.endpoint("/ws", is_websocket=True)
    async def generate_ws(self, websocket: WebSocket) -> None:
        await websocket.accept()
        for _ in range(3):
            await websocket.send_json({"message": "Hello world!"})
        await websocket.close()