Spaces:

apiarium
/

labs

Sleeping

labs / src /services /billing /usage_watcher.py

deploy: unified router + dreamy website (2026-06-16T09:46:52Z)

c1a683f 4 months ago

3.08 kB

"""SSE usage watcher — meters streaming responses by capturing the usage chunk.
OpenAI-compatible streams (incl. OpenRouter :free) emit a final chunk carrying a
`usage` object when the client sets `stream_options.include_usage: true`. This
scanner watches the filtered byte stream as it passes through and extracts
(prompt_tokens, completion_tokens) from that final chunk.
If the upstream doesn't emit usage (e.g. raw subnet nodes), tok_in/tok_out stay
0 and the request is effectively unmetered — acceptable for MVP since upstreams
are free and we control the margin.
"""
from __future__ import annotations

import json


class UsageWatcher:
    """Scan SSE bytes for the OpenAI usage object. Buffers across chunk boundaries."""

    def __init__(self):
        self.tok_in = 0   # captured prompt_tokens (filled when the usage chunk arrives)
        self.tok_out = 0   # captured completion_tokens (filled when the usage chunk arrives)
        self._buf = b""   # leftover bytes that didn't end on a newline

    def feed(self, chunk: bytes) -> None:
        """Scan one chunk of SSE bytes for a usage-bearing data line."""
        self._buf += chunk   # append to the carryover buffer
        while b"\n" in self._buf:   # process complete lines only
            line, self._buf = self._buf.split(b"\n", 1)   # peel one line
            self._scan_line(line)   # check it for usage

    def flush(self) -> None:
        """Process any remaining buffered bytes after the stream ends."""
        if self._buf.strip():   # there's a trailing line without a newline
            self._scan_line(self._buf)   # check it too
        self._buf = b""   # clear the buffer

    def _scan_line(self, line: bytes) -> None:
        """Inspect one SSE line; if it's a data: line with usage, capture it."""
        line = line.strip()   # drop trailing \r and whitespace
        if not line.startswith(b"data:"):   # only data: lines carry JSON payloads
            return   # skip comments, event:, id:, etc.
        if b'"usage"' not in line:   # cheap pre-check before JSON parsing
            return   # most chunks have no usage — avoid the parse cost
        payload = line[5:].strip()   # strip the "data:" prefix
        if payload == b"[DONE]":   # stream terminator, not JSON
            return   # ignore
        try:
            obj = json.loads(payload)   # parse the SSE event JSON
        except Exception:   # malformed JSON (shouldn't happen post-filter, but be safe)
            return   # skip
        if not isinstance(obj, dict):   # not an object
            return   # skip
        u = obj.get("usage")   # the usage sub-object (OpenAI shape)
        if isinstance(u, dict):   # this chunk carries usage
            self.tok_in = int(u.get("prompt_tokens", self.tok_in))   # capture input tokens
            self.tok_out = int(u.get("completion_tokens", self.tok_out))   # capture output tokens

    @property
    def has_usage(self) -> bool:
        """True if we captured a usage object during the stream."""
        return self.tok_in > 0 or self.tok_out > 0

	"""SSE usage watcher — meters streaming responses by capturing the usage chunk.

	OpenAI-compatible streams (incl. OpenRouter :free) emit a final chunk carrying a
	`usage` object when the client sets `stream_options.include_usage: true`. This
	scanner watches the filtered byte stream as it passes through and extracts
	(prompt_tokens, completion_tokens) from that final chunk.

	If the upstream doesn't emit usage (e.g. raw subnet nodes), tok_in/tok_out stay
	0 and the request is effectively unmetered — acceptable for MVP since upstreams
	are free and we control the margin.
	"""
	from __future__ import annotations

	import json


	class UsageWatcher:
	"""Scan SSE bytes for the OpenAI usage object. Buffers across chunk boundaries."""

	def __init__(self):
	self.tok_in = 0 # captured prompt_tokens (filled when the usage chunk arrives)
	self.tok_out = 0 # captured completion_tokens (filled when the usage chunk arrives)
	self._buf = b"" # leftover bytes that didn't end on a newline

	def feed(self, chunk: bytes) -> None:
	"""Scan one chunk of SSE bytes for a usage-bearing data line."""
	self._buf += chunk # append to the carryover buffer
	while b"\n" in self._buf: # process complete lines only
	line, self._buf = self._buf.split(b"\n", 1) # peel one line
	self._scan_line(line) # check it for usage

	def flush(self) -> None:
	"""Process any remaining buffered bytes after the stream ends."""
	if self._buf.strip(): # there's a trailing line without a newline
	self._scan_line(self._buf) # check it too
	self._buf = b"" # clear the buffer

	def _scan_line(self, line: bytes) -> None:
	"""Inspect one SSE line; if it's a data: line with usage, capture it."""
	line = line.strip() # drop trailing \r and whitespace
	if not line.startswith(b"data:"): # only data: lines carry JSON payloads
	return # skip comments, event:, id:, etc.
	if b'"usage"' not in line: # cheap pre-check before JSON parsing
	return # most chunks have no usage — avoid the parse cost
	payload = line[5:].strip() # strip the "data:" prefix
	if payload == b"[DONE]": # stream terminator, not JSON
	return # ignore
	try:
	obj = json.loads(payload) # parse the SSE event JSON
	except Exception: # malformed JSON (shouldn't happen post-filter, but be safe)
	return # skip
	if not isinstance(obj, dict): # not an object
	return # skip
	u = obj.get("usage") # the usage sub-object (OpenAI shape)
	if isinstance(u, dict): # this chunk carries usage
	self.tok_in = int(u.get("prompt_tokens", self.tok_in)) # capture input tokens
	self.tok_out = int(u.get("completion_tokens", self.tok_out)) # capture output tokens

	@property
	def has_usage(self) -> bool:
	"""True if we captured a usage object during the stream."""
	return self.tok_in > 0 or self.tok_out > 0