#!/usr/bin/env python3 import sys import os import json import time import argparse import numpy as np import httpx DEFAULT_BASE_URL = "http://localhost:8000" PROMPTS = [ "Hello! Who are you and what do you do?", "Can you explain the CoALA framework in simple terms?", "What are the four memory subsystems defined in CoALA?", "How does Working Memory differ from Episodic Memory?", "What role does Semantic Memory play in grounding knowledge?", "How does Procedural Memory evolve through interaction?", "Give an example of an action cycle: Planning -> Execution -> Evaluation.", "Summarize our conversation so far in 3 bullet points.", "Thanks! How can external tools or MCP be integrated?", "Thank you for the detailed explanations, goodbye!" ] STAGES = [ "guard", "episodic_search", "semantic_search", "llm_summarize", "llm_agent", "llm_reflect", "llm_procedural", "snapshot_capture", "diff", "audit_save" ] def parse_sse_stream(response): current_event = None for line in response.iter_lines(): if not line: continue if line.startswith("event:"): current_event = line.split("event:", 1)[1].strip() elif line.startswith("data:"): data_str = line.split("data:", 1)[1].strip() yield current_event, data_str def run_turn_stream(client, base_url, session_id, prompt): t_start = time.perf_counter() ttft_ms = None first_token_time = None done_payload = None token_count = 0 url = f"{base_url}/api/v1/chat/stream" payload = {"session_id": session_id, "message": prompt} with client.stream("POST", url, json=payload, timeout=120.0) as response: response.raise_for_status() for event, data_str in parse_sse_stream(response): if event == "token": token_count += 1 if ttft_ms is None: ttft_ms = (time.perf_counter() - t_start) * 1000 elif event == "done": try: done_payload = json.loads(data_str) except Exception as e: print(f"Error parsing done payload: {e}") total_turn_ms = (time.perf_counter() - t_start) * 1000 if ttft_ms is None: ttft_ms = total_turn_ms return { "ttft_ms": ttft_ms, "total_turn_ms": total_turn_ms, "tokens_streamed": token_count, "done": done_payload or {} } def verify_server_timing_header(client, base_url): print("Verifying Server-Timing header on /api/v1/chat...") res = client.post(f"{base_url}/api/v1/chat", json={"session_id": "test-header-check", "message": "Ping check for header"}, timeout=60.0) server_timing = res.headers.get("server-timing") or res.headers.get("Server-Timing") print(f"Server-Timing Header: {server_timing}") assert server_timing is not None, "Server-Timing header was missing from /api/v1/chat!" return server_timing def run_session(base_url=DEFAULT_BASE_URL, session_id="session-perf-baseline", output_file=None, label="baseline"): client = httpx.Client() # Check Server-Timing header verify_server_timing_header(client, base_url) # Reset session client.post(f"{base_url}/api/v1/chat/reset", json={"session_id": session_id}, timeout=30.0) print(f"\nStarting 10-turn performance benchmark session: '{session_id}' against {base_url}\n" + "="*80) turn_results = [] for i, prompt in enumerate(PROMPTS, 1): print(f"\n[Turn {i}/10] Prompt: '{prompt}'") res = run_turn_stream(client, base_url, session_id, prompt) timing_info = res["done"].get("timing", {}) stages_info = timing_info.get("stages", {}) llm_calls_info = timing_info.get("llm_calls", {}) turn_data = { "turn": i, "prompt": prompt, "ttft_ms": res["ttft_ms"], "total_turn_ms": res["total_turn_ms"], "stages": {s: stages_info.get(s, 0.0) for s in STAGES}, "llm_calls": { "summarize": llm_calls_info.get("summarize", 0), "agent": llm_calls_info.get("agent", 0), "reflect": llm_calls_info.get("reflect", 0), "procedural": llm_calls_info.get("procedural", 0), "total": llm_calls_info.get("total", sum(llm_calls_info.values()) if llm_calls_info else 0) } } turn_results.append(turn_data) print(f" -> TTFT: {res['ttft_ms']:.1f} ms | Total Turn: {res['total_turn_ms']:.1f} ms | LLM Calls: {turn_data['llm_calls']['total']}") print(f" Stages: guard={turn_data['stages']['guard']:.1f}ms, episodic={turn_data['stages']['episodic_search']:.1f}ms, semantic={turn_data['stages']['semantic_search']:.1f}ms, agent={turn_data['stages']['llm_agent']:.1f}ms, reflect={turn_data['stages']['llm_reflect']:.1f}ms, procedural={turn_data['stages']['llm_procedural']:.1f}ms, snap={turn_data['stages']['snapshot_capture']:.1f}ms, diff={turn_data['stages']['diff']:.1f}ms, audit={turn_data['stages']['audit_save']:.1f}ms") # Aggregate p50 / p95 p50_stages = {} p95_stages = {} for s in STAGES: vals = [t["stages"][s] for t in turn_results] p50_stages[s] = float(np.percentile(vals, 50)) p95_stages[s] = float(np.percentile(vals, 95)) ttft_vals = [t["ttft_ms"] for t in turn_results] total_vals = [t["total_turn_ms"] for t in turn_results] llm_calls_vals = [t["llm_calls"]["total"] for t in turn_results] p50_ttft = float(np.percentile(ttft_vals, 50)) p95_ttft = float(np.percentile(ttft_vals, 95)) p50_total = float(np.percentile(total_vals, 50)) p95_total = float(np.percentile(total_vals, 95)) avg_llm_calls = float(np.mean(llm_calls_vals)) max_llm_calls = int(np.max(llm_calls_vals)) print("\n" + "="*80) print(f"BENCHMARK RESULTS SUMMARY ({label.upper()} - 10 TURNS)") print("="*80) print(f"Time to First Token (TTFT): p50={p50_ttft:.1f}ms | p95={p95_ttft:.1f}ms") print(f"Total Turn Time: p50={p50_total:.1f}ms | p95={p95_total:.1f}ms") print(f"LLM Calls Per Turn: avg={avg_llm_calls:.1f} | max={max_llm_calls}") print("-" * 80) for s in STAGES: print(f" Stage {s:<20}: p50={p50_stages[s]:8.1f}ms | p95={p95_stages[s]:8.1f}ms") print("="*80) # Format Markdown Report title = "Baseline Performance Report" if label == "baseline" else "Post-Optimization Performance Report" md_content = f"""# {title} **Environment & Setup:** - Date: {time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime())} - Total Turns Evaluated: 10 scripted turns (sequential multi-turn conversation) - Session ID: `{session_id}` - Model: `qwen/qwen3.8-27b` (ChatGroq) - Vector DB: Weaviate Cloud (with in-memory fallback adapter) - Script: `scripts/perf_session.py` ## 1. Summary Metrics (p50 / p95) | Metric / Stage | p50 (Median) | p95 (95th Percentile) | Unit | |---|---|---|---| | **Time to First Token (TTFT)** | **{p50_ttft:.1f}** | **{p95_ttft:.1f}** | ms | | **Total Turn Time** | **{p50_total:.1f}** | **{p95_total:.1f}** | ms | | **LLM Calls Per Turn** | **{avg_llm_calls:.1f} (avg)** | **{max_llm_calls} (max)** | calls | | Stage: `guard` | {p50_stages['guard']:.1f} | {p95_stages['guard']:.1f} | ms | | Stage: `episodic_search` | {p50_stages['episodic_search']:.1f} | {p95_stages['episodic_search']:.1f} | ms | | Stage: `semantic_search` | {p50_stages['semantic_search']:.1f} | {p95_stages['semantic_search']:.1f} | ms | | Stage: `llm_summarize` | {p50_stages['llm_summarize']:.1f} | {p95_stages['llm_summarize']:.1f} | ms | | Stage: `llm_agent` | {p50_stages['llm_agent']:.1f} | {p95_stages['llm_agent']:.1f} | ms | | Stage: `llm_reflect` | {p50_stages['llm_reflect']:.1f} | {p95_stages['llm_reflect']:.1f} | ms | | Stage: `llm_procedural` | {p50_stages['llm_procedural']:.1f} | {p95_stages['llm_procedural']:.1f} | ms | | Stage: `snapshot_capture` | {p50_stages['snapshot_capture']:.1f} | {p95_stages['snapshot_capture']:.1f} | ms | | Stage: `diff` | {p50_stages['diff']:.1f} | {p95_stages['diff']:.1f} | ms | | Stage: `audit_save` | {p50_stages['audit_save']:.1f} | {p95_stages['audit_save']:.1f} | ms | ## 2. Turn-by-Turn Breakdown | Turn | Prompt | TTFT (ms) | Total (ms) | LLM Calls (Sum/Agent/Refl/Proc) | Total LLM Calls | |---|---|---|---|---|---| """ for t in turn_results: calls = t["llm_calls"] calls_str = f"{calls['summarize']}/{calls['agent']}/{calls['reflect']}/{calls['procedural']}" clean_prompt = t["prompt"].replace("|", "\\|") md_content += f"| {t['turn']} | {clean_prompt} | {t['ttft_ms']:.1f} | {t['total_turn_ms']:.1f} | {calls_str} | {calls['total']} |\n" md_content += f""" ## 3. Analysis of Baseline Bottlenecks Identified 1. **Faked Streaming & Late TTFT:** - In `/chat/stream`, TTFT is {p50_ttft:.1f} ms because the entire turn (including searches, agent invocation, and synchronous reflections) is executed *before* streaming synthetic word tokens with `asyncio.sleep(0.02)`. 2. **Duplicate & Blocking Reflection In Request Path:** - On multiple turns, reflection (`llm_reflect`) and procedural rewrite (`llm_procedural`) run synchronously in the request path (in both the LangGraph `update_memory` node and in `SnapshotCoordinatorService.execute_turn` step 6). - This causes turns to execute multiple sequential LLM calls (up to {max_llm_calls} LLM calls in a single turn). 3. **Sequential Vector Retrieval & Collection Checking:** - `episodic_search` and `semantic_search` run sequentially and each issue `collections.exists()` HTTP roundtrips over Weaviate Cloud. 4. **History Summarization Escalation:** - Once conversation history exceeds threshold, summarization fires synchronously before agent execution, increasing turn latency and LLM call counts. """ if output_file: os.makedirs(os.path.dirname(output_file), exist_ok=True) with open(output_file, "w") as f: f.write(md_content) print(f"\nSaved report to: {output_file}") return { "p50_ttft": p50_ttft, "p95_ttft": p95_ttft, "p50_total": p50_total, "p95_total": p95_total, "avg_llm_calls": avg_llm_calls, "max_llm_calls": max_llm_calls, "p50_stages": p50_stages, "p95_stages": p95_stages, "turns": turn_results } if __name__ == "__main__": parser = argparse.ArgumentParser(description="Mnemos Chat Latency Benchmark") parser.add_argument("--url", default=DEFAULT_BASE_URL, help="Base API URL") parser.add_argument("--session", default="session-perf-baseline", help="Session ID") parser.add_argument("--output", default="docs/perf/baseline.md", help="Output Markdown path") parser.add_argument("--label", default="baseline", help="Label (baseline or after)") args = parser.parse_args() run_session(base_url=args.url, session_id=args.session, output_file=args.output, label=args.label)