From 32b895da9eb741a3e804fd4e10577a355d3e5577 Mon Sep 17 00:00:00 2001 From: makeitdouble Date: Thu, 28 May 2026 15:12:04 +0300 Subject: [PATCH] chore: runtime console - add flow telemetry - separate brain context from provider usage --- agents/translation_node.py | 46 ++-- contracts/context_contract.py | 30 +-- runtime/runtime_context.py | 4 + runtime/runtime_stream.py | 142 +++++++++++-- static/logger.js | 5 + tests/test_brain_runtime_actions.py | 23 +- tests/test_runtime_stream_tokens.py | 268 +++++++++++++++++++++++ tests/test_token_usage.py | 50 +++++ tests/test_websocket_pending_usage.py | 126 +++++++++++ utils/token_usage.py | 295 ++++++++++++++++++++++++++ utils/tokens.py | 69 +++--- websocket.py | 100 +++++++++ 12 files changed, 1050 insertions(+), 108 deletions(-) create mode 100644 tests/test_runtime_stream_tokens.py create mode 100644 tests/test_token_usage.py create mode 100644 tests/test_websocket_pending_usage.py create mode 100644 utils/token_usage.py diff --git a/agents/translation_node.py b/agents/translation_node.py index 6ad93edc..2867066f 100644 --- a/agents/translation_node.py +++ b/agents/translation_node.py @@ -8,8 +8,8 @@ settings, ) -from utils.runtime_state_sync import ( - refresh_runtime_state, +from utils.token_usage import ( + record_token_usage, ) @@ -45,20 +45,32 @@ async def run( translated_text ) - await refresh_runtime_state( - context, - runtime_id=( - settings.TRANSLATOR_MODEL_UID - ), - add_tokens=usage.get( - "total_tokens", - 0, - ), - max_tokens=( - settings.TRANSLATOR_CONTEXT_WINDOW - ), - last_error=None, - status="online", - ) + if usage: + record_token_usage( + context, + runtime_id=( + settings.TRANSLATOR_MODEL_UID + ), + role="translator", + kind="service", + prompt_tokens=( + usage.get( + "prompt_tokens", + 0, + ) + ), + completion_tokens=( + usage.get( + "completion_tokens", + 0, + ) + ), + total_tokens=( + usage.get( + "total_tokens", + 0, + ) + ), + ) state.translated_input = translated_text diff --git a/contracts/context_contract.py b/contracts/context_contract.py index aed4beb4..d263cc11 100644 --- a/contracts/context_contract.py +++ b/contracts/context_contract.py @@ -13,33 +13,14 @@ ) -def cdata( - value: str, -) -> str: - - return ( - "', ']]]]>')}" - "]]>" - ) - - def format_xml_field( tag: str, value, ) -> str: - if tag.endswith( - "_ACTION" - ): - rendered_value = cdata( - str(value) - ) - - else: - rendered_value = escape( - str(value) - ) + rendered_value = escape( + str(value) + ) return f"<{tag}>{rendered_value}" @@ -55,7 +36,7 @@ def format_available_actions( ( " " f"" - f"{cdata(template)}" + f"{template}" "" ) for name, template @@ -96,13 +77,12 @@ def build_runtime_fields(self) -> str: fields = { "TIMESTAMP": self.timestamp, "WEEKDAY": self.weekday, - "YEAR": self.year, } available_actions = [] if self.can_deep_thought: - fields["DEEP_THOUGHT_COUNTER"] = self.deep_thought_count + fields["DEEP_THOUGHT_COUNTER"] = str(self.deep_thought_count) available_actions.append( ( RUNTIME_ACTION_DEEP_THOUGHT, diff --git a/runtime/runtime_context.py b/runtime/runtime_context.py index b466a017..bf9c2c42 100644 --- a/runtime/runtime_context.py +++ b/runtime/runtime_context.py @@ -39,3 +39,7 @@ class RuntimeContext: runtime_action_events: list[dict] = field( default_factory=list ) + + runtime_usage_events: list[dict] = field( + default_factory=list + ) diff --git a/runtime/runtime_stream.py b/runtime/runtime_stream.py index 172f980e..78c14fb9 100644 --- a/runtime/runtime_stream.py +++ b/runtime/runtime_stream.py @@ -12,8 +12,12 @@ StreamHandler, ) +from utils.token_usage import ( + record_stream_token_usage, +) + from utils.tokens import ( - estimate_stream_tokens, + estimate_stream_live_tokens, ) @@ -45,6 +49,7 @@ def __init__( self.log_method = log_method self.emit_to_chat = emit_to_chat + self.context_snapshot = context_snapshot or {} self.stream = StreamHandler( self.websocket, @@ -58,6 +63,113 @@ def __init__( ), ) + def build_input_prompt_text(self) -> str: + + if not isinstance( + self.context_snapshot, + dict, + ): + return "" + + parts = [] + + for key in ( + "system_prompt", + "user_prompt", + "context_payload", + ): + + value = self.context_snapshot.get( + key, + "", + ) + + if value: + parts.append( + str(value) + ) + + return "\n".join( + parts + ) + + def is_brain_context(self) -> bool: + + if not isinstance( + self.context_snapshot, + dict, + ): + return False + + return ( + self.context_snapshot.get( + "context_role" + ) + == "brain" + ) + + def estimate_live_tokens(self) -> int: + + return estimate_stream_live_tokens( + self.stream, + prompt_text=( + self.build_input_prompt_text() + ), + ) + + async def refresh_token_usage(self): + + if not self.is_brain_context(): + return + + used_tokens = ( + self.estimate_live_tokens() + ) + + if not used_tokens: + return + + await refresh_runtime_state( + self.context, + runtime_id=( + self.runtime_id + ), + used_tokens=( + used_tokens + ), + max_tokens=( + self.context_window + ), + last_error=None, + status="online", + ) + + def record_token_usage(self): + is_brain_context = self.is_brain_context() + + record_stream_token_usage( + self.context, + runtime_id=( + self.runtime_id + ), + role=( + "brain" + if is_brain_context + else self.role + ), + kind=( + "brain" + if is_brain_context + else "service" + ), + stream=( + self.stream + ), + prompt_text=( + self.build_input_prompt_text() + ), + ) + def build_action_log( self, action_event_offset: int, @@ -133,6 +245,8 @@ async def run( emit=self.emit_to_chat ) + await self.refresh_token_usage() + await self.logger.log_runtime( f"[STREAM START] role={self.role}" ) @@ -171,6 +285,8 @@ async def run( emit=self.emit_to_chat, ) + await self.refresh_token_usage() + continue # ------------------------------------------------- @@ -196,30 +312,14 @@ async def run( return None + await self.refresh_token_usage() + await self.stream.finish( emit=self.emit_to_chat ) - used_tokens = ( - estimate_stream_tokens( - self.stream - ) - ) - - await refresh_runtime_state( - self.context, - runtime_id=( - self.runtime_id - ), - used_tokens=( - used_tokens - ), - max_tokens=( - self.context_window - ), - last_error=None, - status="online", - ) + await self.refresh_token_usage() + self.record_token_usage() log_response = self.stream.response diff --git a/static/logger.js b/static/logger.js index 4ea66116..1e6cd790 100644 --- a/static/logger.js +++ b/static/logger.js @@ -226,6 +226,11 @@ function appendLog( "text-emerald-500"; } + if (tag.includes("FLOW TELEMETRY")) { + tagClass = + "text-purple-400"; + } + if (tag.includes("ERROR")) { tagClass = "text-red-500 font-bold"; diff --git a/tests/test_brain_runtime_actions.py b/tests/test_brain_runtime_actions.py index 6625ab02..e64aa33c 100644 --- a/tests/test_brain_runtime_actions.py +++ b/tests/test_brain_runtime_actions.py @@ -72,13 +72,18 @@ def test_prompt_uses_passed_agent_runtime_actions(self): self.assertIn( ( - '' f"{SEARCH_ACTION_TEMPLATE}" - "]]>" + "" ), prompt, ) + self.assertNotIn( + "", + prompt, + ) + self.assertNotIn( "RUNTIME_STATE", prompt, @@ -135,13 +145,18 @@ def test_prompt_can_flip_agent_actions_dynamically(self): self.assertIn( ( - '' f"{DEEP_THOUGHT_ACTION}" - "]]>" + "" ), prompt, ) + self.assertNotIn( + " int: + + try: + return int( + value or 0 + ) + + except ( + TypeError, + ValueError, + ): + return 0 + + +def record_token_usage( + context, + *, + runtime_id: str, + role: str, + kind: str = "service", + prompt_tokens: int = 0, + completion_tokens: int = 0, + total_tokens: int = 0, + context_tokens: int = 0, +): + + usage_events = getattr( + context, + "runtime_usage_events", + None, + ) + + if usage_events is None: + usage_events = [] + context.runtime_usage_events = usage_events + + usage_events.append({ + "runtime_id": runtime_id, + "role": role, + "kind": kind, + "prompt_tokens": _as_int( + prompt_tokens + ), + "completion_tokens": _as_int( + completion_tokens + ), + "total_tokens": _as_int( + total_tokens + ), + "context_tokens": _as_int( + context_tokens + ), + }) + + +def record_stream_token_usage( + context, + *, + runtime_id: str, + role: str, + kind: str = "service", + stream, + prompt_text: str = "", +): + + prompt_tokens = ( + _as_int( + getattr( + stream, + "prompt_tokens", + 0, + ) + ) + or estimate_stream_input_tokens( + stream, + prompt_text=prompt_text, + ) + ) + + completion_tokens = ( + _as_int( + getattr( + stream, + "completion_tokens", + 0, + ) + ) + or ( + estimate_optional_tokens( + getattr( + stream, + "response", + "", + ) + ) + + estimate_optional_tokens( + getattr( + stream, + "reasoning", + "", + ) + ) + ) + ) + + total_tokens = ( + _as_int( + getattr( + stream, + "total_tokens", + 0, + ) + ) + or prompt_tokens + + completion_tokens + ) + context_tokens = estimate_stream_live_tokens( + stream, + prompt_text=prompt_text, + ) + + record_token_usage( + context, + runtime_id=runtime_id, + role=role, + kind=kind, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens, + context_tokens=context_tokens, + ) + + +def summarize_token_usage( + context, + *, + kind: str | None = None, +) -> dict: + + summary = { + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + } + + for event in getattr( + context, + "runtime_usage_events", + [], + ): + if ( + kind is not None + and event.get( + "kind" + ) + != kind + ): + continue + + summary["prompt_tokens"] += _as_int( + event.get( + "prompt_tokens", + 0, + ) + ) + summary["completion_tokens"] += _as_int( + event.get( + "completion_tokens", + 0, + ) + ) + summary["total_tokens"] += _as_int( + event.get( + "total_tokens", + 0, + ) + ) + + return summary + + +def summarize_token_usage_by_role( + context, + *, + kind: str | None = None, +) -> list[dict]: + + grouped = {} + + for event in getattr( + context, + "runtime_usage_events", + [], + ): + if ( + kind is not None + and event.get( + "kind" + ) + != kind + ): + continue + + key = ( + event.get( + "role", + "unknown", + ), + event.get( + "runtime_id", + "unknown", + ), + ) + + if key not in grouped: + grouped[key] = { + "role": key[0], + "runtime_id": key[1], + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "context_tokens": 0, + } + + grouped[key]["prompt_tokens"] += _as_int( + event.get( + "prompt_tokens", + 0, + ) + ) + grouped[key]["completion_tokens"] += _as_int( + event.get( + "completion_tokens", + 0, + ) + ) + grouped[key]["total_tokens"] += _as_int( + event.get( + "total_tokens", + 0, + ) + ) + grouped[key]["context_tokens"] += _as_int( + event.get( + "context_tokens", + 0, + ) + ) + + return list( + grouped.values() + ) + + +def format_token_usage_summary( + context, +) -> str: + + summary = summarize_token_usage( + context + ) + breakdown = summarize_token_usage_by_role( + context + ) + + lines = [ + "PROVIDER USAGE", + ] + + for item in breakdown: + lines.append( + ( + f"{item['role']}: " + f"{item['total_tokens']}" + f" (prompt={item['prompt_tokens']}, " + f"completion={item['completion_tokens']})" + ) + ) + + lines.append( + f"total: {summary['total_tokens']}" + ) + + return "\n".join( + lines + ) diff --git a/utils/tokens.py b/utils/tokens.py index e37494c5..38c1d44d 100644 --- a/utils/tokens.py +++ b/utils/tokens.py @@ -1,5 +1,6 @@ from settings.app_settings import settings + def estimate_tokens( text: str, ) -> int: @@ -24,55 +25,37 @@ def estimate_optional_tokens( ) -def estimate_stream_tokens( +def estimate_stream_input_tokens( stream, *, prompt_text: str = "", ) -> int: - visible_tokens = ( - estimate_optional_tokens( - getattr( - stream, - "response", - "", - ) + return estimate_optional_tokens( + prompt_text + ) + + +def estimate_stream_live_tokens( + stream, + *, + prompt_text: str = "", +) -> int: + + return ( + estimate_stream_input_tokens( + stream, + prompt_text=prompt_text, ) + estimate_optional_tokens( getattr( stream, - "reasoning", + "response", "", ) ) ) - prompt_tokens = ( - getattr( - stream, - "prompt_tokens", - 0, - ) - or estimate_optional_tokens( - prompt_text - ) - ) - - provider_total = ( - getattr( - stream, - "total_tokens", - 0, - ) - or 0 - ) - - return max( - provider_total, - prompt_tokens + visible_tokens, - visible_tokens, - ) - def translation_token_limit( text: str, @@ -98,12 +81,16 @@ def estimate_runtime_tokens( reasoning: str = "", ) -> int: - total_text = ( - user_input - + system_prompt - + context_payload - + response - + reasoning + total_text = "\n".join( + value + for value in ( + user_input, + system_prompt, + context_payload, + response, + reasoning, + ) + if value ) return estimate_tokens( diff --git a/websocket.py b/websocket.py index 7699f817..cb43932e 100644 --- a/websocket.py +++ b/websocket.py @@ -19,14 +19,39 @@ AgentRuntime, ) +from clients.brain_client import ( + build_brain_payload, + build_brain_system_prompt, +) + from agents.agent_state import ( AgentState, ) +from utils.brain import ( + get_brain_runtime_config, +) + +from utils.language import ( + contains_cyrillic, +) + +from utils.runtime_state_sync import ( + refresh_runtime_state, +) + +from utils.token_usage import ( + format_token_usage_summary, +) + from utils.telemetry import ( send_telemetry, ) +from utils.tokens import ( + estimate_runtime_tokens, +) + from utils.ws_errors import ( handle_fatal_runtime_error, handle_websocket_error, @@ -103,6 +128,60 @@ async def receive_message( # PROCESS MESSAGE # --------------------------------------------------------- +async def refresh_pending_brain_usage( + context, + user_text: str, +): + + if contains_cyrillic( + user_text + ): + return + + brain_runtime = ( + get_brain_runtime_config() + ) + + runtime_actions = ( + brain_runtime.get( + "runtime_actions", + {}, + ) + ) + + system_prompt = ( + build_brain_system_prompt( + context, + runtime_actions=runtime_actions, + ) + ) + + brain_payload = ( + build_brain_payload( + user_text, + context=context, + ) + ) + + await refresh_runtime_state( + context, + runtime_id=( + brain_runtime["runtime_id"] + ), + used_tokens=( + estimate_runtime_tokens( + system_prompt=system_prompt, + user_input=brain_payload, + ) + ), + max_tokens=( + brain_runtime["context_window"] + ), + last_error=None, + status="online", + ) + + async def process_message( context, message_data: dict, @@ -123,6 +202,15 @@ async def process_message( user_input=user_text ) + if hasattr( + context, + "runtime_usage_events", + ): + context.runtime_usage_events.clear() + + else: + context.runtime_usage_events = [] + runtime = AgentRuntime() await logger.log_system( @@ -142,6 +230,13 @@ async def process_message( context, ) + await logger.log( + "[FLOW TELEMETRY]", + format_token_usage_summary( + context + ), + ) + await logger.log_system( "[WS] agent runtime end" ) @@ -336,6 +431,11 @@ async def websocket_endpoint( # START BACKGROUND TASK # ------------------------------------------------- + await refresh_pending_brain_usage( + context, + user_text, + ) + current_task = ( asyncio.create_task( process_message(