定义 WebSocket 端点

BentoML 允许您在服务中设置 WebSocket 端点,用于实时通信。WebSocket 支持非常适合实时消息传递、数据流式传输或语音 AI 应用程序等用例。

用法

BentoML 服务可以通过 FastAPI 的 WebSocket 支持处理 WebSocket 连接。要设置 WebSocket 服务器

  1. 创建 FastAPI 应用。

  2. 将其挂载到您的 BentoML 服务。.

一旦 BentoML 服务器启动,WebSocket 服务器也将被初始化并准备好接受连接。

这是一个简单的示例

service.py
import bentoml
from fastapi import FastAPI, WebSocket

# Initialize FastAPI app
app = FastAPI()

# Define BentoML Service and mount the app
@bentoml.service(
    traffic={"timeout": 30}
)
@bentoml.asgi_app(app, path="/chat")
class WebSocketService:
    def __init__(self):
        # Initialize your resources here (e.g., models, configurations)
        print("Service initialized")

    @app.websocket("/ws")
    async def websocket_endpoint(self, websocket: WebSocket):
        await websocket.accept()
        # Define your custom logic here
        print("WebSocket connection accepted")
        try:
            while True:
                data = await websocket.receive_text()
                print(f"Received: {data}")
                await websocket.send_text(f"Echo: {data}")
        except Exception as e:
            print(f"Connection closed: {e}")

发送请求

BentoML Python 客户端尚不支持 WebSocket 端点。您需要自己实现客户端。以下是使用 websockets 库调用上面定义的端点的示例:

import asyncio
import websockets

async def test_websocket():
    # /chat comes from the asgi_app path, /ws from the endpoint
    # Adjust URL as needed
    uri = "ws://:3000/chat/ws"
    async with websockets.connect(uri) as websocket:
        # Send a test message
        await websocket.send("Hello BentoML")
        response = await websocket.recv()
        print(f"Response: {response}")

# Run the test
asyncio.run(test_websocket())

了解更多

有关更实际的示例,请参阅如何使用 Twilio ConversationRelay 和 BentoML 通过开源模型构建语音 AI 应用