Files
AI-VideoAssistant-Engine-V2/models/ws_v1.py
2026-02-17 10:39:23 +08:00

74 lines
2.1 KiB
Python

"""WS v1 protocol message models and helpers."""
from typing import Optional, Dict, Any, Literal
from pydantic import BaseModel, Field
def now_ms() -> int:
"""Current unix timestamp in milliseconds."""
import time
return int(time.time() * 1000)
# Client -> Server messages
class HelloMessage(BaseModel):
type: Literal["hello"]
version: str = Field(..., description="Protocol version, currently v1")
auth: Optional[Dict[str, str]] = Field(default=None, description="Auth payload, e.g. {'apiKey': '...'}")
class SessionStartMessage(BaseModel):
type: Literal["session.start"]
audio: Optional[Dict[str, Any]] = Field(default=None, description="Optional audio format metadata")
metadata: Optional[Dict[str, Any]] = Field(default=None, description="Optional session metadata")
class SessionStopMessage(BaseModel):
type: Literal["session.stop"]
reason: Optional[str] = None
class InputTextMessage(BaseModel):
type: Literal["input.text"]
text: str
class ResponseCancelMessage(BaseModel):
type: Literal["response.cancel"]
graceful: bool = False
class ToolCallResultsMessage(BaseModel):
type: Literal["tool_call.results"]
results: list[Dict[str, Any]] = Field(default_factory=list)
CLIENT_MESSAGE_TYPES = {
"hello": HelloMessage,
"session.start": SessionStartMessage,
"session.stop": SessionStopMessage,
"input.text": InputTextMessage,
"response.cancel": ResponseCancelMessage,
"tool_call.results": ToolCallResultsMessage,
}
def parse_client_message(data: Dict[str, Any]) -> BaseModel:
"""Parse and validate a WS v1 client message."""
msg_type = data.get("type")
if not msg_type:
raise ValueError("Missing 'type' field")
msg_class = CLIENT_MESSAGE_TYPES.get(msg_type)
if not msg_class:
raise ValueError(f"Unknown client message type: {msg_type}")
return msg_class(**data)
# Server -> Client event helpers
def ev(event_type: str, **payload: Any) -> Dict[str, Any]:
"""Create a WS v1 server event payload."""
base = {"type": event_type, "timestamp": now_ms()}
base.update(payload)
return base