import os import sys import code import time import litert_lm import numpy as np try: import readline except ImportError: readline = None MODEL_PATH = os.path.expanduser("~/.litert-lm/models/gemma-4-E2B-it.litertlm/model.litertlm") SYSTEM_CONSTRAINT = ( "Act as an embedded Python co-interpreter. " "Answer the user query in exactly ONE brief sentence or direct assignment. " "No greetings, no preambles, no conversational fluff. " "If asked about variables, inspect the provided state and answer concisely." ) # Variables that should never appear in the namespace snapshot _INTERNAL_VARS = frozenset({ "np", "litert_lm", "shared_globals", "ai", "readline", "engine", "conversation", "reset_ai", "SYSTEM_CONSTRAINT", "_INTERNAL_VARS", "_build_compact_snapshot", "time", }) print("[Initialization] Loading LiteRT-LM Engine on CPU...") engine = litert_lm.Engine( MODEL_PATH, backend=litert_lm.Backend.CPU, # ═══════════════════════════════════════════════════════════════ # FIX #1: Set max_num_tokens explicitly. # # When this is None (default), the engine uses the model's baked-in # default, which for Gemma 4 can be 8192+. A large KV cache makes # EVERY decode step slower — even for short responses — because the # CPU must traverse more memory per attention computation. # # See: https://github.com/google-ai-edge/LiteRT-LM/issues/2568 # max_num_tokens=2048 → 6.28 tok/s decode # max_num_tokens=8192 → 3.52 tok/s decode (1.8x slower!) # max_num_tokens=32768 → 1.44 tok/s decode (4.4x slower!) # # 2048 is plenty for short co-interpreter Q&A (system prompt + # a few turns + short responses). # ═══════════════════════════════════════════════════════════════ max_num_tokens=2048, cache_dir="/tmp/litert-lm-cache", ) # Create ONE persistent conversation with the system message set once. # The system_message is prefilled into the KV cache at creation time and # never re-processed on subsequent turns. conversation = engine.create_conversation(system_message=SYSTEM_CONSTRAINT) shared_globals = { "np": np, "current_coordinates": np.array([5.0, -3.5]), "matrix_z": np.array([[1.0, 2.0], [3.0, np.nan]]) } def _build_compact_snapshot(): """ Build a MINIMAL state string. Keep this as short as possible — every token here must be prefilled on every ai() call. We use repr() instead of str() for arrays to keep it compact, and skip variables that haven't changed if we track a dirty flag. """ parts = [] for var_name, var_val in list(shared_globals.items()): if var_name in _INTERNAL_VARS: continue if isinstance(var_val, np.ndarray): # Compact array repr — avoids multiline formatting parts.append(f"{var_name}={np.array2string(var_val, separator=',')}") else: parts.append(f"{var_name}={repr(eval)}") return "State: " + "; ".join(parts) def ai(query_string): """ Fast co-interpreter call. Key performance choices: 1. Uses send_message (blocking) — NOT send_message_async. The async/streaming path routes every token through a queue.Queue + ctypes callback + json.loads, adding ~2-5ms of Python overhead per token. For short responses (1 sentence), the blocking path does ONE C++ call and is much faster. 2. Sends a compact one-line state snapshot + query, not a multi-line dump. Fewer prefill tokens = faster TTFT. 3. Caps output tokens at 128 to prevent runaway generation. 4. Reuses the persistent conversation (warm KV cache). """ user_message = f"{_build_compact_snapshot()}\nQuery: {query_string}" print("\n[AI Co-Interpreter]: ", end="") sys.stdout.flush() start = time.perf_counter() # ═══════════════════════════════════════════════════════════════ # FIX #2: Use send_message (blocking) instead of send_message_async. # # send_message_async does this per token chunk: # 1. C++ fires ctypes callback # 2. Callback puts item in queue.Queue (acquires GIL) # 3. Main thread gets item from queue.Queue (acquires GIL) # 4. json.loads() the chunk # 5. yield to caller # # For a 30-token response, that's 30 GIL round-trips + 30 JSON parses. # send_message does 1 C++ call, 1 JSON parse, returns the full response. # ═══════════════════════════════════════════════════════════════ try: response = conversation.send_message( user_message, max_output_tokens=128, # FIX #3: cap response length ) text = "" for item in response.get("content", []): if isinstance(item, dict) and item.get("type") == "text": text += item.get("text", "") print(text) except Exception as e: print(f"[Error: {e}]") elapsed = time.perf_counter() - start print(f" [{elapsed:.2f}s]") def reset_ai(): """Discard conversation history and start fresh (frees KV cache space).""" global conversation conversation.close() conversation = engine.create_conversation(system_message=SYSTEM_CONSTRAINT) print("[AI] Conversation reset — KV cache cleared.") shared_globals["ai"] = ai shared_globals["reset_ai"] = reset_ai if readline: readline.parse_and_bind("tab: complete") readline.set_history_length(1000) banner_msg = """ ==================================================================== XYFLOW HYPER-SPEED DUAL INTERPRETER (v7.0-PERF) ==================================================================== System status: OPERATIONAL KV cache: 2048 tokens (small = fast decode) Message mode: Blocking send_message (no streaming overhead) Brevity lock: ACTIVE (system_message + max_output_tokens=128) Performance fixes vs original: • max_num_tokens=2048 (was: model default, likely 8192+) → up to 4x faster decode (Issue #2568) • send_message blocking (was: send_message_async streaming) → eliminates per-token Python/GIL/JSON overhead • max_output_tokens=128 (was: unlimited) → prevents runaway generation • Compact state snapshot (was: multiline array dumps) → fewer prefill tokens per call • Persistent conversation with warm KV cache • System message set once at creation Try: >>> ai("Look at matrix_z. Is it stable or is there an issue?") >>> ai("What are the current_coordinates?") >>> reset_ai() # optional: clear conversation history ==================================================================== """ console = code.InteractiveConsole(locals=shared_globals) console.interact(banner=banner_msg, exitmsg="\n[Shutdown] Safely clearing engine allocations.")