diff --git a/src/ai_client.py b/src/ai_client.py index efc85193..12171924 100644 --- a/src/ai_client.py +++ b/src/ai_client.py @@ -3258,8 +3258,10 @@ def send( if chunks: context_block = "## Retrieved Context\n\n" for i, chunk in enumerate(chunks): - path = chunk.get("metadata", {}).get("path", "unknown") - context_block += f"### Chunk {i+1} (Source: {path})\n{chunk.get('document', '')}\n\n" + chunk_meta = chunk["metadata"] if "metadata" in chunk else {} + path = chunk_meta["path"] if "path" in chunk_meta else "unknown" + doc = chunk["document"] if "document" in chunk else "" + context_block += f"### Chunk {i+1} (Source: {path})\n{doc}\n\n" user_message = context_block + user_message _append_comms("OUT", "request", {"message": user_message, "system": _get_combined_system_prompt(_active_tool_preset, _active_bias_profile)}) diff --git a/src/app_controller.py b/src/app_controller.py index 6e8915e0..20081fe4 100644 --- a/src/app_controller.py +++ b/src/app_controller.py @@ -247,8 +247,10 @@ def _api_generate(controller: 'AppController', req: GenerateRequest) -> Metadata if rag_result.ok and rag_result.data: context_block = "## Retrieved Context\n\n" for i, chunk in enumerate(rag_result.data): - path = chunk.get("metadata", {}).get("path", "unknown") - context_block += f"### Chunk {i+1} (Source: {path})\n{chunk.get('document', '')}\n\n" + chunk_meta = chunk["metadata"] if "metadata" in chunk else {} + path = chunk_meta["path"] if "path" in chunk_meta else "unknown" + doc = chunk["document"] if "document" in chunk else "" + context_block += f"### Chunk {i+1} (Source: {path})\n{doc}\n\n" user_msg = context_block + user_msg elif not rag_result.ok: controller._last_request_errors.append(("rag_search", rag_result.errors[0])) @@ -2269,13 +2271,14 @@ class AppController: kind = entry.get("kind", entry.get("type", "")) payload = entry.get("payload", {}) ts = entry.get("ts", "") + comms_entry = CommsLogEntry.from_dict(entry) if kind == 'tool_call': tid = payload.get('id') or payload.get('call_id') script = payload.get('script') or json.dumps(payload.get('args', {}), indent=1) script = _resolve_log_ref(script, session_dir) entry_obj = { - 'source_tier': entry['source_tier'] if 'source_tier' in entry else 'main', + 'source_tier': comms_entry.source_tier, 'script': script, 'result': '', # Waiting for result 'ts': ts @@ -2298,17 +2301,23 @@ class AppController: if kind == 'response' and 'usage' in payload: u = payload['usage'] + u_stats = models.UsageStats( + input_tokens=u.get('input_tokens', 0) or 0, + output_tokens=u.get('output_tokens', 0) or 0, + cache_read_tokens=u.get('cache_read_input_tokens', 0) or 0, + cache_creation_tokens=u.get('cache_creation_input_tokens', 0) or 0, + ) for k in ['input_tokens', 'output_tokens', 'cache_read_input_tokens', 'cache_creation_input_tokens', 'total_tokens']: if k in new_usage: new_usage[k] += u.get(k, 0) or 0 - tier = entry['source_tier'] if 'source_tier' in entry else 'main' + tier = comms_entry.source_tier if tier in new_mma_usage: - new_mma_usage[tier]['input'] += u.get('input_tokens', 0) or 0 - new_mma_usage[tier]['output'] += u.get('output_tokens', 0) or 0 + new_mma_usage[tier]['input'] += u_stats.input_tokens + new_mma_usage[tier]['output'] += u_stats.output_tokens new_token_history.append({ 'time': ts, - 'input': u.get('input_tokens', 0) or 0, - 'output': u.get('output_tokens', 0) or 0, - 'model': entry['model'] if 'model' in entry else 'unknown' + 'input': u_stats.input_tokens, + 'output': u_stats.output_tokens, + 'model': comms_entry.model }) if kind == "history_add": @@ -4160,8 +4169,10 @@ class AppController: if rag_result.ok and rag_result.data: context_block = "## Retrieved Context\n\n" for i, chunk in enumerate(rag_result.data): - path = chunk.get("metadata", {}).get("path", "unknown") - context_block += f"### Chunk {i+1} (Source: {path})\n{chunk.get('document', '')}\n\n" + chunk_meta = chunk["metadata"] if "metadata" in chunk else {} + path = chunk_meta["path"] if "path" in chunk_meta else "unknown" + doc = chunk["document"] if "document" in chunk else "" + context_block += f"### Chunk {i+1} (Source: {path})\n{doc}\n\n" user_msg = context_block + user_msg elif not rag_result.ok: self._last_request_errors.append(("rag_search", rag_result.errors[0])) diff --git a/src/gui_2.py b/src/gui_2.py index 28586cfd..351c343b 100644 --- a/src/gui_2.py +++ b/src/gui_2.py @@ -120,6 +120,7 @@ from src import theme_2 as theme from src import thinking_parser from src import workspace_manager from src.hot_reloader import HotReloader +from src.type_aliases import HistoryMessage, SessionInsights win32gui: Any = None win32con: Any = None @@ -4922,15 +4923,13 @@ def render_session_insights_panel(app: App) -> None: if app.perf_profiling_enabled: app.perf_monitor.start_component("_render_session_insights_panel") imgui.text_colored(C_LBL(), 'Session Insights') imgui.separator() - insights = app.controller.get_session_insights() - imgui.text(f"Total Tokens: {insights.get('total_tokens', 0):,}") - imgui.text(f"API Calls: {insights.get('call_count', 0)}") - imgui.text(f"Burn Rate: {insights.get('burn_rate', 0):.0f} tokens/min") - imgui.text(f"Session Cost: ${insights.get('session_cost', 0):.4f}") - completed = insights.get('completed_tickets', 0) - efficiency = insights.get('efficiency', 0) - imgui.text(f"Completed: {completed}") - imgui.text(f"Tokens/Ticket: {efficiency:.0f}" if efficiency > 0 else "Tokens/Ticket: N/A") + insights = SessionInsights.from_dict(app.controller.get_session_insights()) + imgui.text(f"Total Tokens: {insights.total_tokens:,}") + imgui.text(f"API Calls: {insights.call_count}") + imgui.text(f"Burn Rate: {insights.burn_rate:.0f} tokens/min") + imgui.text(f"Session Cost: ${insights.session_cost:.4f}") + imgui.text(f"Completed: {insights.completed_tickets}") + imgui.text(f"Tokens/Ticket: {insights.efficiency:.0f}" if insights.efficiency > 0 else "Tokens/Ticket: N/A") if app.perf_profiling_enabled: app.perf_monitor.end_component("_render_session_insights_panel") def render_prior_session_view(app: App) -> None: @@ -7791,7 +7790,9 @@ def _handle_history_logic_result(app: "App") -> Result[bool]: ) if not changed and len(current.disc_entries) > 0: - if current.disc_entries[-1].get('content') != app._last_ui_snapshot.disc_entries[-1].get('content'): + curr_msg = HistoryMessage.from_dict(current.disc_entries[-1]) + prev_msg = HistoryMessage.from_dict(app._last_ui_snapshot.disc_entries[-1]) + if curr_msg.content != prev_msg.content: changed = True if changed: diff --git a/src/synthesis_formatter.py b/src/synthesis_formatter.py index 9155791e..51d4f381 100644 --- a/src/synthesis_formatter.py +++ b/src/synthesis_formatter.py @@ -1,10 +1,13 @@ +from src.type_aliases import HistoryMessage + + def format_takes_diff(takes: dict[str, list[dict]]) -> str: """ [C: tests/test_synthesis_formatter.py:test_format_takes_diff_common_prefix, tests/test_synthesis_formatter.py:test_format_takes_diff_empty, tests/test_synthesis_formatter.py:test_format_takes_diff_no_common_prefix, tests/test_synthesis_formatter.py:test_format_takes_diff_single_take] """ if not takes: return "" - + histories = list(takes.values()) if not histories: return "" @@ -20,9 +23,9 @@ def format_takes_diff(takes: dict[str, list[dict]]) -> str: shared_lines = [] for i in range(common_prefix_len): - msg = histories[0][i] - shared_lines.append(f"{msg.get('role', 'unknown')}: {msg.get('content', '')}") - + msg = HistoryMessage.from_dict(histories[0][i]) + shared_lines.append(f"{msg.role}: {msg.content}") + shared_text = "=== Shared History ===" if shared_lines: shared_text += "\n" + "\n".join(shared_lines) @@ -33,8 +36,8 @@ def format_takes_diff(takes: dict[str, list[dict]]) -> str: if len(history) > common_prefix_len: variation_lines.append(f"[{take_name}]") for i in range(common_prefix_len, len(history)): - msg = history[i] - variation_lines.append(f"{msg.get('role', 'unknown')}: {msg.get('content', '')}") + msg = HistoryMessage.from_dict(history[i]) + variation_lines.append(f"{msg.role}: {msg.content}") variation_lines.append("") else: # Single take case