Skip to content

Commit 03288e8

Browse files
committed
feat(realtime): add generation_tick billing event support
Surface periodic billing updates (total billed seconds) during authenticated realtime sessions. generation_ended is parsed but intentionally not exposed publicly.
1 parent 3fcf5a4 commit 03288e8

File tree

5 files changed

+58
-2
lines changed

5 files changed

+58
-2
lines changed

decart/realtime/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
encode_subscribe_token,
66
decode_subscribe_token,
77
)
8+
from .messages import GenerationTickMessage
89
from .types import RealtimeConnectOptions, ConnectionState, AvatarOptions
910

1011
__all__ = [
@@ -14,6 +15,7 @@
1415
"SubscribeOptions",
1516
"encode_subscribe_token",
1617
"decode_subscribe_token",
18+
"GenerationTickMessage",
1719
"RealtimeConnectOptions",
1820
"ConnectionState",
1921
"AvatarOptions",

decart/realtime/client.py

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from pydantic import BaseModel
1010

1111
from .webrtc_manager import WebRTCManager, WebRTCConfiguration
12-
from .messages import PromptMessage, SessionIdMessage
12+
from .messages import PromptMessage, SessionIdMessage, GenerationTickMessage
1313
from .subscribe import (
1414
SubscribeClient,
1515
SubscribeOptions,
@@ -73,6 +73,7 @@ def __init__(
7373
self._is_avatar_live = is_avatar_live
7474
self._connection_callbacks: list[Callable[[ConnectionState], None]] = []
7575
self._error_callbacks: list[Callable[[DecartSDKError], None]] = []
76+
self._generation_tick_callbacks: list[Callable[[GenerationTickMessage], None]] = []
7677
self._session_id: Optional[str] = None
7778
self._subscribe_token: Optional[str] = None
7879
self._buffering = True
@@ -134,6 +135,7 @@ async def connect(
134135
config.on_connection_state_change = client._emit_connection_change
135136
config.on_error = lambda error: client._emit_error(WebRTCError(str(error), cause=error))
136137
config.on_session_id = client._handle_session_id
138+
config.on_generation_tick = client._emit_generation_tick
137139

138140
try:
139141
# For avatar-live, convert and send avatar image before WebRTC connection
@@ -226,6 +228,8 @@ def _do_flush(self) -> None:
226228
self._dispatch_connection_change(data) # type: ignore[arg-type]
227229
elif event == "error":
228230
self._dispatch_error(data) # type: ignore[arg-type]
231+
elif event == "generation_tick":
232+
self._dispatch_generation_tick(data) # type: ignore[arg-type]
229233
self._buffer.clear()
230234

231235
def _dispatch_connection_change(self, state: ConnectionState) -> None:
@@ -254,6 +258,19 @@ def _emit_error(self, error: DecartSDKError) -> None:
254258
else:
255259
self._dispatch_error(error)
256260

261+
def _dispatch_generation_tick(self, message: GenerationTickMessage) -> None:
262+
for callback in list(self._generation_tick_callbacks):
263+
try:
264+
callback(message)
265+
except Exception as e:
266+
logger.exception(f"Error in generation_tick callback: {e}")
267+
268+
def _emit_generation_tick(self, message: GenerationTickMessage) -> None:
269+
if self._buffering:
270+
self._buffer.append(("generation_tick", message))
271+
else:
272+
self._dispatch_generation_tick(message)
273+
257274
async def set(self, input: SetInput) -> None:
258275
if input.prompt is None and input.image is None:
259276
raise InvalidInputError("At least one of 'prompt' or 'image' must be provided")
@@ -350,6 +367,8 @@ def on(self, event: str, callback: Callable) -> None:
350367
self._connection_callbacks.append(callback)
351368
elif event == "error":
352369
self._error_callbacks.append(callback)
370+
elif event == "generation_tick":
371+
self._generation_tick_callbacks.append(callback)
353372

354373
def off(self, event: str, callback: Callable) -> None:
355374
if event == "connection_change":
@@ -362,3 +381,8 @@ def off(self, event: str, callback: Callable) -> None:
362381
self._error_callbacks.remove(callback)
363382
except ValueError:
364383
pass
384+
elif event == "generation_tick":
385+
try:
386+
self._generation_tick_callbacks.remove(callback)
387+
except ValueError:
388+
pass

decart/realtime/messages.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,21 @@ class GenerationStartedMessage(BaseModel):
9393
type: Literal["generation_started"]
9494

9595

96+
class GenerationTickMessage(BaseModel):
97+
"""Periodic billing update during generation."""
98+
99+
type: Literal["generation_tick"]
100+
seconds: int
101+
102+
103+
class GenerationEndedMessage(BaseModel):
104+
"""Server signals that generation has ended. Not exposed publicly."""
105+
106+
type: Literal["generation_ended"]
107+
seconds: int
108+
reason: str
109+
110+
96111
# Discriminated union for incoming messages
97112
IncomingMessage = Annotated[
98113
Union[
@@ -105,6 +120,8 @@ class GenerationStartedMessage(BaseModel):
105120
ReadyMessage,
106121
IceRestartMessage,
107122
GenerationStartedMessage,
123+
GenerationTickMessage,
124+
GenerationEndedMessage,
108125
],
109126
Field(discriminator="type"),
110127
]

decart/realtime/webrtc_connection.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
ErrorMessage,
2929
IceRestartMessage,
3030
SessionIdMessage,
31+
GenerationTickMessage,
3132
OutgoingMessage,
3233
)
3334
from .types import ConnectionState
@@ -42,6 +43,7 @@ def __init__(
4243
on_state_change: Optional[Callable[[ConnectionState], None]] = None,
4344
on_error: Optional[Callable[[Exception], None]] = None,
4445
on_session_id: Optional[Callable[[SessionIdMessage], None]] = None,
46+
on_generation_tick: Optional[Callable[[GenerationTickMessage], None]] = None,
4547
customize_offer: Optional[Callable] = None,
4648
):
4749
self._pc: Optional[RTCPeerConnection] = None
@@ -52,6 +54,7 @@ def __init__(
5254
self._on_state_change = on_state_change
5355
self._on_error = on_error
5456
self._on_session_id = on_session_id
57+
self._on_generation_tick = on_generation_tick
5558
self._customize_offer = customize_offer
5659
self._ws_task: Optional[asyncio.Task] = None
5760
self._ice_candidates_queue: list[RTCIceCandidate] = []
@@ -264,6 +267,14 @@ async def _handle_message(self, data: dict) -> None:
264267
self._handle_set_image_ack(message)
265268
elif message.type == "generation_started":
266269
await self._set_state("generating")
270+
elif message.type == "generation_tick":
271+
if self._on_generation_tick:
272+
self._on_generation_tick(message)
273+
elif message.type == "generation_ended":
274+
# Parsed but intentionally not exposed — unreliable (won't arrive on
275+
# client disconnect/network drop), overlaps with connection_change
276+
# "disconnected", and insufficient_credits is already covered by error event.
277+
logger.debug(f"Generation ended: reason={message.reason}, seconds={message.seconds}")
267278
elif message.type == "error":
268279
self._handle_error(message)
269280
elif message.type == "ready":

decart/realtime/webrtc_manager.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212
)
1313

1414
from .webrtc_connection import WebRTCConnection
15-
from .messages import OutgoingMessage, SessionIdMessage
15+
from .messages import OutgoingMessage, SessionIdMessage, GenerationTickMessage
1616
from .types import ConnectionState
1717
from ..types import ModelState
1818

@@ -44,6 +44,7 @@ class WebRTCConfiguration:
4444
on_connection_state_change: Optional[Callable[[ConnectionState], None]] = None
4545
on_error: Optional[Callable[[Exception], None]] = None
4646
on_session_id: Optional[Callable[[SessionIdMessage], None]] = None
47+
on_generation_tick: Optional[Callable[[GenerationTickMessage], None]] = None
4748
initial_state: Optional[ModelState] = None
4849
customize_offer: Optional[Callable] = None
4950
integration: Optional[str] = None
@@ -208,6 +209,7 @@ def _create_connection(self) -> WebRTCConnection:
208209
on_state_change=self._handle_connection_state_change,
209210
on_error=self._config.on_error,
210211
on_session_id=self._config.on_session_id,
212+
on_generation_tick=self._config.on_generation_tick,
211213
customize_offer=self._config.customize_offer,
212214
)
213215

0 commit comments

Comments
 (0)