25 lines
1.1 KiB
Python
25 lines
1.1 KiB
Python
from .base_agent import BaseAgent
|
|
from pydantic import BaseModel
|
|
import typing
|
|
from fastapi import APIRouter, WebSocket
|
|
from ..models.agent import AgentStartMessage, AgentReadyMessage, AgentTextMessage, WebSocketAgentMessage, WebSocketAgentMessageType
|
|
|
|
class WebSocketAgent(BaseAgent):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.app.websocket("/respond")(self.respond_websocket)
|
|
|
|
async def respond_websocket(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
AgentStartMessage.parse_obj(await websocket.receive_json())
|
|
await websocket.send_text(AgentReadyMessage().json())
|
|
while True:
|
|
message = WebSocketAgentMessage.parse_obj(await websocket.receive_json())
|
|
if message.type == WebSocketAgentMessageType.AGENT_STOP:
|
|
break
|
|
text_message = typing.cast(AgentTextMessage, message)
|
|
response = await self.respond(text_message.data.text)
|
|
await websocket.send_text(AgentTextMessage.from_text(response).json())
|
|
await websocket.close()
|
|
|