refactor: Collapse four inference backends into one MLX path

The project targets Apple Silicon. It carried CUDA/PyTorch, GGUF/llama.cpp,
Ollama and a Hugging Face Spaces ZeroGPU path anyway, and maintaining four
backends for one machine was most of the complexity in the codebase.
orchestrator.py (1,327 lines) and tools.py (427) go; engine.py, agent.py and
toolcall.py replace them.

engine.py keeps one persistent KV cache for the process. _align_cache finds the
longest prefix of the incoming prompt the cache already holds, trims to it, and
feeds only the remainder, so the fixed system-prompt-plus-tool-schema prefix is
prefilled once at startup by prewarm() rather than once per tool step. Generated
tokens are tracked in the cache too, which is what makes continuing a
tool-calling turn nearly free. First-token latency per step: ~6.4s -> ~0.7s.

Sampling uses presence_penalty, not a flat repetition_penalty. A blanket
repetition penalty punishes the repeated structural tokens that matrices and
JSON are made of ("[", "0", ",") exactly when the model is emitting a tool call.
think_budget caps tokens spent inside a <think> block and closes it by hand on
overrun, bounding worst-case latency on a reasoning model.

agent.py streams throughout. _StreamGate releases text as soon as it cannot be
the start of <tool_call>, so prose appears immediately while a tool call never
leaks into the chat. Retrieved passages attach to the user turn, never spliced
into the system prompt, which keeps the cached prefix byte-identical across
questions.

The parameter-provenance check is gone. It refused any matrix it could not trace
back to the user's message, which blocked the single most useful thing the
assistant does -- working an example the user asked for. Schema validation and
sandboxed execution in registry.execute are the real guarantees, and
toolcall.degenerate_reason catches genuine decoding loops with thresholds set
well above any hand-written matrix. prompts.py drops 2,587 tokens to 506 by
stating what to do instead of what not to do; the old "never invent a parameter"
clause was what taught the model to refuse worked examples.

Two model failures are fixed in tool results rather than in prompt wording,
which is the pattern to reach for first:
- continuous_lqr, discrete_lqr and place_state_feedback return closed_loop_A.
Given a correct gain, the model wrote A-BK for a double integrator as
[[-6,1],[-5,-6]] instead of [[0,1],[-6,-5]] -- right gain, wrong write-up.
- bode_analysis returns gain and phase margins. Asked for a phase margin the
model reached for it instead of stability_margins, got back sampled curves,
and correctly but uselessly concluded the margin "cannot be determined".

app.py routes every inference call through one dedicated worker thread for the
process lifetime, since MLX keeps its compute stream in thread-local state, and
serialises turns with a lock because they all mutate the same KV cache.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

CLAUDE.md

@@ -0,0 +1,182 @@
+# CLAUDE.md
+
+Guidance for Claude Code (claude.ai/code) when working in this repository.
+
+## What this is
+
+ControlAI is an offline control-systems engineering assistant. A locally held model
+(`mlx-community/Qwen3-14B-4bit` by default, via MLX) answers control questions, and every number it
+states comes from a deterministic solver — SciPy/LAPACK/CVXPY behind a validated tool registry —
+never from the model's own arithmetic. A BM25 + dense hybrid retriever over a local control-theory
+corpus grounds conceptual answers. A FastAPI + vanilla-JS console and a terminal CLI are the two
+front ends. Nothing leaves the machine.
+
+Apple Silicon only. There is no CUDA, GGUF, Ollama, or Hugging Face Spaces path — an earlier
+version carried all four, and maintaining four backends for one machine was most of the complexity
+in the codebase.
+
+## Commands
+
+```bash
+./run.sh                      # web console at http://127.0.0.1:8000
+./run.sh --cli                # interactive terminal chat
+./run.sh --cli "design an LQR for A=[[0,1],[-2,-3]], B=[[0],[1]], Q=eye(2), R=1"
+./run.sh --fetch-index        # download the prebuilt index from the private Hub dataset repo
+./run.sh --build-index        # rebuild the dense retrieval index (~45 min, checkpointed/resumable)
+./run.sh --ingest-corpus      # merge data/processed/ chunks into the index
+./run.sh --calibrate          # re-measure MIN_COSINE against the current corpus
+
+python -m unittest discover tests -v          # unittest, not pytest -- pytest is not installed
+python -m unittest tests.test_agent_tools     # deterministic numerics
+python -m unittest tests.test_agent_core      # parsing, stream gating, truncation
+
+pip install -r requirements.txt               # runtime
+pip install -r requirements-training.txt      # dataset generation / ground truth
+pip install -r requirements-corpus.txt        # corpus crawling and extraction
+```
+
+Environment: `CONTROLAI_MODEL`, `CONTROLAI_ADAPTER`, `CONTROLAI_THINKING` (`off|auto|on`),
+`CONTROLAI_THINK_BUDGET`, `CONTROLAI_EMBED_MODEL`.
+
+There is no lint or format command configured — don't invent one.
+
+## Architecture
+
+### Request flow
+`web/` (SSE console) or `cli.py` → `app.py` (`/api/chat`, `/api/chat/stream`) →
+`ControlAgent` (`controlai_agent/agent.py`) → `LocalEngine` (`controlai_agent/engine.py`, MLX) and
+`registry.execute` (`controlai_agent/registry.py`) → tools in `controlai_agent/tools/*.py`.
+
+### `engine.py` — inference
+One persistent KV cache lives for the process. `_align_cache` finds the longest prefix of the
+incoming prompt that the cache already holds, trims to that point, and feeds only the remainder, so
+the fixed system-prompt-plus-tool-schema prefix (~8k tokens) is prefilled once at startup by
+`prewarm()` rather than once per tool step. Generated tokens are tracked in the cache too, which is
+what makes continuing a tool-calling turn nearly free. Measured effect: first-token latency fell
+from ~6.4 s per step to ~0.7 s.
+
+`stream()` yields tokens as they are produced. `think_budget` caps tokens spent inside a `<think>`
+block and closes it by hand on overrun, bounding worst-case latency on a reasoning model.
+
+Sampling uses `presence_penalty`, not a flat `repetition_penalty`. A blanket repetition penalty
+punishes the repeated structural tokens that matrices and JSON are made of (`[`, `0`, `,`) exactly
+when the model is emitting a tool call.
+
+### `agent.py` — the loop
+A short tool-calling loop (`MAX_TOOL_STEPS = 2`, `MAX_CALLS_PER_TOOL = 2`) that streams throughout.
+`_StreamGate` releases text as soon as it cannot be the start of `<tool_call>`, so prose appears
+immediately while a tool call never leaks into the chat. Retrieved passages are attached to the
+*user* turn, never spliced into the system prompt, because that keeps the cached prefix
+byte-identical across questions.
+
+**There is deliberately no parameter-provenance check.** An earlier version refused any matrix it
+could not trace back to the user's message. It blocked the single most useful thing the assistant
+does — working an example that the user asked for — so it was removed. Schema validation and
+sandboxed execution in `registry.execute` are the real guarantees; `toolcall.degenerate_reason`
+catches genuine decoding loops with thresholds set well above any hand-written matrix.
+
+### `prompts.py`
+506 tokens, down from 2,587. States what to do rather than what not to do. The previous prompt's
+"never invent a parameter" clause taught the model to refuse worked examples; the current one
+explicitly authorises choosing an illustrative system when the user asks for a demonstration.
+
+### Tool registry (`registry.py` + `controlai_agent/tools/`)
+Tools are plain functions decorated with `@registry.register(name, description, parameters_schema)`
+across `linear.py`, `frequency.py`, `synthesis.py`, `estimation.py`, `nonlinear.py`, `robust.py`,
+`allocation.py`, `simulation.py`, `matrix_ops.py`, `plotting.py`, `python_executor.py`, `rag.py`.
+`registry.execute()` coerces stringified arrays, validates against JSON Schema, runs the function,
+and rounds every float to 6 significant figures. `controlai_agent/verifier.py` cross-checks results
+with independent invariants — Riccati residual, pole-placement error, CBF forward invariance.
+
+**Prefer putting a fact in a tool result over instructing the model to derive it.** Two observed
+failures were fixed that way rather than by prompt wording, and it is the pattern to reach for
+first:
+- `continuous_lqr`, `discrete_lqr` and `place_state_feedback` return `closed_loop_A`. Given a
+  correct gain, the model wrote $A-BK$ for a double integrator as `[[-6,1],[-5,-6]]` instead of
+  `[[0,1],[-6,-5]]` — right gain, wrong write-up.
+- `bode_analysis` returns the gain/phase margins as well. Asked for a phase margin the model
+  reached for it instead of `stability_margins`, got back sampled curves, and correctly but
+  uselessly concluded the margin "cannot be determined".
+
+### Retrieval (`controlai_rag/`)
+
+Three data bugs were found here and are worth knowing about, because each one silently degraded
+retrieval rather than failing:
+- **Chunk ids were not unique.** `chunk_document` is called once per page and restarted its counter
+  each time, so every page's first chunk was `<file>_c0000` — 154 distinct ids across 9,976 chunks.
+  Anything keyed on chunk_id resolved to the wrong row. Fixed in `chunker.py` (ids are now
+  page-qualified); `scripts/repair_chunk_ids.py` migrated the existing index in place.
+- **Embeddings were pooled without an EOS token.** Qwen3-Embedding pools the last position and was
+  trained with `<|endoftext|>` there. Omitting it dropped the margin between relevant and irrelevant
+  passages from +0.32 to +0.15. Note the checkpoint's own `eos_token_id` is `<|im_end|>`, which is
+  the wrong token here — `controlai_rag/embeddings.py` pins the right one.
+- **71.5% of the Nise textbook was mojibake.** Its PDF has a broken symbol-font ToUnicode map, so
+  every extractor returns `L½ f ðtÞ/C138 ¼FðsÞ` for `L[f(t)] = F(s)`. `controlai_rag/textfix.py`
+  reverses the substitution (it is deterministic); `scripts/repair_corpus_text.py` migrated the
+  index. `document_loader` now applies it on ingest.
+
+`document_loader.py` → `chunker.py` → `index.py` (BM25, persisted to `data/rag_index/`) and
+`embeddings.py` (Qwen3-Embedding-0.6B under MLX, last-token pooling, instruction-prefixed queries).
+
+**The index holds 80,370 chunks.** 9,976 come from `data/user_docs/` (course notes, Nise, Ogata);
+the other 70,394 were bridged in from `data/processed/*_chunks/` by
+`scripts/ingest_processed_corpus.py`. That bridge did not exist before: the `scripts/` pipeline
+(`raw → extracted → processed`) fed *training dataset generation* only, while `ControlRAGIndex` read
+`user_docs` alone. Retrieval was therefore seeing 12% of the corpus, and none of the canonical
+texts — Doyle/Francis/Tannenbaum, Åström & Murray, Rawlings/Mayne/Diehl, Sontag, Liberzon, Boyd,
+Söderström & Stoica — which were downloaded, extracted and chunked on disk, unread. Re-run the
+bridge after adding anything to `data/processed/`.
+
+**The index is not in git.** `chunks.json` (187MB), `embeddings.npz` (144MB) and `bm25.pkl`
+(124MB) are each past GitHub's 100MB per-file limit, and `chunks.json` holds the full extracted text
+of 666 documents including commercial textbooks (Nise, Ogata) — fine to keep locally, wrong to
+redistribute. They live in the **private** Hub dataset repo `atakankahya/controlai-rag-index`;
+`controlai_rag/fetch_index.py` (`./run.sh --fetch-index`) pulls them with your HF token. It copies
+out of the Hub cache rather than symlinking, because the index is mutated in place by uploads and by
+`scripts/repair_*.py`. A public clone has no index and builds its own from `data/user_docs/`.
+
+`build()` is checkpointed every 4,000 chunks to `embeddings.partial.npz` and resumes from it. An
+80k-chunk build takes ~45 minutes, and writing only at the end meant one interruption discarded all
+of it (observed at 72,008 of 80,370).
+`retriever.py` fuses the two rankings with reciprocal rank fusion and gates on cosine similarity
+(`MIN_COSINE = 0.62`). That number is a property of the model *and* the corpus, so re-measure it
+with `./run.sh --calibrate` whenever the corpus changes size materially — it moved when the index
+grew 8x. Current margins: in-domain worst best-match 0.678, off-domain best 0.571. Lexical candidates are scored densely too — without that, a chunk BM25 ranked first but
+that fell outside the dense top-K was dropped for having no score rather than for being irrelevant.
+`_is_low_value` discards back-of-book index pages, which match almost any control query because they
+contain every term in the field, while explaining nothing.
+
+The gate matters. BM25 scores are unbounded and corpus-relative, so the old threshold of 2.5 passed
+essentially everything: a question about the Bode sensitivity integral retrieved Routh-Hurwitz
+tables at score 19.7 and injected them as authoritative context. Returning nothing is a valid and
+frequent outcome — the model then answers from its own knowledge.
+
+`display_source_name()` maps raw indexed filenames (which carry owner initials, course codes, scan
+artifacts) to clean citable labels. Never let a raw filename reach an answer.
+
+Uploads through `/api/upload` are embedded immediately via `HybridRetriever.add_chunks`, or they
+would have no vector and the cosine gate would make them permanently unreachable.
+
+### `app.py` — serving
+Local-only FastAPI. MLX keeps its compute stream in thread-local state, so every inference call
+goes through one dedicated worker thread for the process lifetime; a lock serialises turns because
+they all mutate the same KV cache. `_to_wire_events` translates the agent's event vocabulary
+(`text`/`thinking`/…) into what `web/app.js` consumes (`token`/`thought`/…).
+
+### Model artifacts
+`adapters/` and `models/` hold earlier fine-tuning output. **They are gitignored — not backed up by
+git.** The LoRA adapters are not loaded by default: `behavior_v1` emits a spurious empty
+`<tool_call></tool_call>` as its first output on essentially every prompt, so no tool ever runs, and
+it refuses fully-specified problems; `sft_v2` produces empty output when no tools are exposed and
+string-typed numbers when they are. Both were measured against the base model, which routes and
+formats correctly. `CONTROLAI_ADAPTER=<path>` loads one anyway for A/B work.
+
+`configs/*.yaml` are MLX-LoRA training configs and `scripts/` holds the offline pipeline (corpus
+discovery → extraction → dataset generation → training → evaluation). That pipeline is independent
+of the serving path and uses `requirements-training.txt`/`requirements-corpus.txt`.
+
+### Benchmark (`benchmarks/`)
+`controlbench_v1.jsonl` is the eval set; `SCOPE.md` defines the taxonomy and `README.md` the
+data-hygiene rules (never let benchmark prompts leak into training data; split by `family`, not by
+individual question). `scripts/eval_answer_quality.py` is the fast qualitative check across 18
+domain cases.

README.md

@@ -1,12 +1,3 @@
----
-title: ControlAI Agent
-colorFrom: blue
-colorTo: indigo
-sdk: gradio
-app_file: app.py
-pinned: false
----
-
 # ControlAI: Open-Source Safety-Critical AI Agent for Control Systems Engineering
 
 [![Hugging Face Space](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Space-blue)](https://huggingface.co/spaces/atakankahya/ControlAI-Agent)
@@ -19,23 +10,29 @@ pinned: false
 
 ---
 
-## Quickstart: Run Locally in 2 Steps
+## Quickstart
 
-ControlAI runs locally on your workstation without cloud dependencies:
+ControlAI runs entirely on your own machine. Requires **Apple Silicon** (inference is MLX) with
+at least 16 GB of unified memory; the default model needs about 8 GB on disk and is downloaded on
+first run.
 
 ```bash
-# 1. Clone the repository
 git clone https://github.com/atakankahya/controlai-agent.git
 cd controlai-agent
+pip install -r requirements.txt
 
-# 2. Launch the application (Automatically opens browser at http://127.0.0.1:8000)
-./run.sh
+./run.sh                  # web console, opens http://127.0.0.1:8000
+./run.sh --cli            # interactive terminal chat
+./run.sh --cli "design an LQR for A=[[0,1],[-2,-3]], B=[[0],[1]], Q=eye(2), R=1"
+./run.sh --build-index    # build the dense retrieval index from data/user_docs/ (~45 min)
 ```
 
-To run in interactive terminal CLI mode instead:
-```bash
-./run.sh --cli
-```
+| Variable | Default | Purpose |
+| :-- | :-- | :-- |
+| `CONTROLAI_MODEL` | `mlx-community/Qwen3-14B-4bit` | any MLX model id or local path |
+| `CONTROLAI_ADAPTER` | *(none)* | optional LoRA adapter |
+| `CONTROLAI_THINKING` | `auto` | `off`, `auto` (conceptual questions only), or `on` |
+| `CONTROLAI_THINK_BUDGET` | `512` | ceiling on tokens spent reasoning |
 
 ---
 
@@ -50,7 +47,13 @@ Standard large language models (LLMs) operate probabilistically without determin
 1. **Deterministic Scientific Sandbox:** Computes continuous/discrete algebraic Riccati equations (CARE/DARE), matrix exponentials, and Bode diagrams using LAPACK, SciPy, and CVXPY.
 2. **4-Stage Mathematical Proof Standard:** Formulates system class, analytical theorems, closed-form derivations, and engineering breakdown limits.
 3. **Dynamic Simulation & Plotting:** Solves nonlinear differential equations and renders verified trajectories directly in the interface.
-4. **Offline RAG Knowledge Engine:** Grounded with 9,900+ chunks indexed across classical and modern control engineering literature.
+4. **Offline Retrieval:** 80,000+ chunks of classical and modern control literature, indexed locally.
+   The prebuilt index is not distributed -- it carries the full text of copyrighted textbooks --
+   so `--build-index` builds one from whatever you put in `data/user_docs/`.
+   Lexical and dense rankings are fused and gated on cosine similarity, so an unrelated passage is
+   dropped rather than cited — the retriever returning nothing is a normal outcome.
+5. **Bounded latency:** the fixed system-prompt and tool-schema prefix is prefilled once at startup
+   and reused across turns, so answers begin streaming in well under a second.
 
 ---
 
@@ -58,28 +61,27 @@ Standard large language models (LLMs) operate probabilistically without determin
 
 ```mermaid
 graph TD
-    User([Engineering Query]) --> WebApp[Web Console & CLI]
-    WebApp --> Orchestrator[ControlAI Agent Orchestrator]
-    
-    subgraph Core Engine [Hybrid Verification Engine]
-        Brain[Foundation Model - Qwen3-4B]
-        LoRA[Theory & Tool Adapter - SFT]
-        RAG[Offline Knowledge Engine - 9,900+ Chunks]
-        Tools[Deterministic Numerical Tool Suite]
+    User([Engineering Query]) --> Frontend[Web Console / CLI]
+    Frontend --> Agent[ControlAgent streaming tool loop]
+
+    subgraph Local[Runs entirely on-device]
+        Engine[LocalEngine - MLX, cached KV prefix]
+        Model[Qwen3-14B-4bit]
+        Retriever[Hybrid Retriever - BM25 + dense]
+        Tools[Deterministic Tool Registry]
     end
-    
-    Orchestrator --> Brain
-    Orchestrator --> LoRA
-    Orchestrator --> RAG
-    Orchestrator --> Tools
-    
-    Tools --> SciPy[SciPy / LAPACK / BLAS Matrix Solvers]
-    Tools --> PythonControl[Python Control ct.tf / step_response]
-    Tools --> PyExecutor[Python Simulation Sandbox]
-    Tools --> Verifier[Numerical Residual Verifier]
-    
-    Verifier --> FinalResponse[Verified Technical Synthesis & Plots]
-    FinalResponse --> WebApp
+
+    Agent --> Engine --> Model
+    Agent --> Retriever
+    Agent --> Tools
+
+    Tools --> SciPy[SciPy / LAPACK / BLAS]
+    Tools --> PythonControl[python-control]
+    Tools --> PyExecutor[Sandboxed Python]
+    Tools --> Verifier[Residual Verifier]
+
+    Verifier --> Answer[Verified answer, plots, citations]
+    Answer --> Frontend
 ```
 
 ---
@@ -107,6 +109,11 @@ graph TD
 
 ## Benchmark Results (ControlBench-v1)
 
+> **These figures were measured against the previous architecture** — the fine-tuned Qwen3-4B LoRA
+> served through the old orchestrator. That serving path has been replaced (base Qwen3-14B, no
+> adapter, rewritten agent loop and retriever), and the suite has not yet been re-run against it, so
+> treat the table as historical rather than as a description of the current build.
+
 Evaluated across **50 multi-pillar benchmark problems**:
 
 | Benchmark Pillar | Qwen3-4B Base (Text-Only) | ControlAI Agent (Ours) | Verification Mechanism |

app.py

@@ -1,10 +1,21 @@
-"""ControlAI Web Application Backend (FastAPI)."""
+"""ControlAI web server.
+
+Local-only by design. The previous version carried a Hugging Face Spaces
+deployment inside it -- Gradio blocks that existed purely to satisfy a ZeroGPU
+SDK check, `@spaces.GPU` decorators, a CUDA/GGUF branch, and a threading split
+whose two halves each existed to work around the other environment. None of it
+ran on this machine, and all of it had to be reasoned about on every change.
+What is left is a FastAPI app serving one MLX-backed agent.
+
+MLX keeps its compute stream in thread-local state: the model must be used from
+the same OS thread that loaded it, or it raises "There is no Stream(gpu, 0) in
+current thread." Every call therefore goes through one dedicated worker thread.
+"""
 
 from __future__ import annotations
 
 import asyncio
 import json
-import os
 import queue
 import shutil
 import sys
@@ -15,11 +26,9 @@ from contextlib import asynccontextmanager
 from pathlib import Path
 from typing import Any
 
-import gradio as gr
-import spaces
 from fastapi import FastAPI, File, HTTPException, UploadFile
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse, HTMLResponse, JSONResponse
+from fastapi.responses import HTMLResponse, StreamingResponse
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel
 
@@ -27,56 +36,47 @@ PROJECT_ROOT = Path(__file__).resolve().parent
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
 
-from controlai_agent.orchestrator import HAS_MLX, ControlAIAgent
+from controlai_agent.agent import ControlAgent
 from controlai_rag.chunker import chunk_document
 from controlai_rag.document_loader import load_single_file
 from controlai_rag.index import get_shared_index
 
-# MLX keeps its compute stream in thread-local state: the model must be
-# loaded on the exact same OS thread that later runs generation, or MLX
-# raises "There is no Stream(cpu, 0) in current thread." Routing every agent
-# call through this single dedicated worker thread keeps MLX on one
-# consistent thread for the whole process lifetime.
-#
-# This must NEVER be used for the CUDA/ZeroGPU path: ZeroGPU's `spaces`
-# library only intercepts CUDA calls within the exact context it manages
-# (the main thread during startup; the request-handling context for
-# @spaces.GPU calls). Moving a CUDA-touching call into this manually-created
-# thread let a real `torch._C._cuda_init()` leak through outside any
-# @spaces.GPU context and crashed the Space on startup:
-#   "Low-level CUDA init reached. ZeroGPU's PyTorch CUDA emulation mode
-#    did not intercept a CUDA operation in your code."
-# HAS_MLX is only ever True on the machine that actually has mlx_lm
-# installed (Apple Silicon) -- never on a HF Spaces Linux/CUDA container --
-# so gating on it keeps MLX's fix local while restoring the CUDA/GGUF path
-# to calling directly on whatever thread FastAPI/spaces already controls,
-# exactly as it worked before.
-inference_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="controlai-inference") if HAS_MLX else None
-# llama.cpp's Llama object is not safe for concurrent generation calls from
-# multiple threads; on the non-MLX path (no dedicated executor serializing
-# things for us) a plain lock does that job instead.
+STATIC_DIR = PROJECT_ROOT / "web"
+PLOTS_DIR = PROJECT_ROOT / "outputs" / "plots"
+UPLOADS_DIR = PROJECT_ROOT / "data" / "user_docs" / "user_uploaded"
+for directory in (STATIC_DIR, PLOTS_DIR, UPLOADS_DIR):
+    directory.mkdir(parents=True, exist_ok=True)
+
+# One thread, for the lifetime of the process: see the module docstring.
+inference_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="controlai")
+# The agent holds a single KV cache that every turn mutates, so turns must not
+# interleave even though they all land on the same thread.
 inference_lock = threading.Lock()
 
+_agent: ControlAgent | None = None
+
+
+def get_agent() -> ControlAgent:
+    global _agent
+    if _agent is None:
+        _agent = ControlAgent()
+    return _agent
+
 
-async def _run_inference(fn, *args):
-    """Run an inference call on the MLX-safe dedicated thread if MLX is in
-    play, otherwise directly (correct for CUDA/ZeroGPU and GGUF)."""
-    if inference_executor is not None:
-        return await asyncio.get_event_loop().run_in_executor(inference_executor, fn, *args)
-    with inference_lock:
-        return fn(*args)
+async def _on_inference_thread(fn, *args):
+    return await asyncio.get_running_loop().run_in_executor(inference_executor, fn, *args)
 
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    print("Pre-loading ControlAI Core Engine on startup...")
-    await _run_inference(get_agent)
-    print("ControlAI Core Engine is online and ready for traffic.")
+    started = time.time()
+    print("Loading ControlAI…")
+    await _on_inference_thread(get_agent)
+    print(f"ControlAI ready in {time.time() - started:.1f}s -> http://127.0.0.1:8000")
     yield
 
 
-app = FastAPI(title="ControlAI", version="1.0.0", lifespan=lifespan)
-
+app = FastAPI(title="ControlAI", version="2.0.0", lifespan=lifespan)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -84,56 +84,8 @@ app.add_middleware(
     allow_methods=["*"],
     allow_headers=["*"],
 )
-
-# Static directories
-STATIC_DIR = PROJECT_ROOT / "web"
-PLOTS_DIR = PROJECT_ROOT / "outputs" / "plots"
-USER_DOCS_DIR = PROJECT_ROOT / "data" / "user_docs"
-UPLOADED_DOCS_DIR = USER_DOCS_DIR / "user_uploaded"
-
-STATIC_DIR.mkdir(parents=True, exist_ok=True)
-PLOTS_DIR.mkdir(parents=True, exist_ok=True)
-UPLOADED_DOCS_DIR.mkdir(parents=True, exist_ok=True)
-
-# Mount plots for web rendering
 app.mount("/plots", StaticFiles(directory=str(PLOTS_DIR)), name="plots")
-if STATIC_DIR.exists():
-    app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
-
-# Initialize Agent
-agent_instance: ControlAIAgent | None = None
-
-
-def get_agent() -> ControlAIAgent:
-    global agent_instance
-    if agent_instance is None:
-        print("Initializing ControlAI Core Engine...")
-        agent_instance = ControlAIAgent()
-        print("ControlAI Core Engine initialized.")
-    return agent_instance
-
-
-# On ZeroGPU Spaces, actual CUDA work may only happen inside a function
-# decorated with @spaces.GPU (it requests physical GPU time for the call and
-# releases it afterward). Outside of a ZeroGPU Space this decorator is a
-# harmless no-op, so it's safe to always wrap these.
-@spaces.GPU(duration=300)
-def _run_stream_on_gpu(message: str, history: list[dict[str, str]]):
-    yield from get_agent().run_stream(message, history=history)
-
-
-@spaces.GPU(duration=300)
-def _run_on_gpu(message: str, history: list[dict[str, str]]):
-    return get_agent().run(message, history=history, verbose=False)
-
-
-# ZeroGPU's startup check statically looks for a @spaces.GPU function wired
-# to a Gradio event handler -- must be a module-level function referenced by
-# name, not one defined inline inside a `with gr.Blocks():` block.
-@spaces.GPU(duration=300)
-def _zerogpu_registration_probe(message: str) -> str:
-    result = get_agent().run(message or "hi", history=[], verbose=False)
-    return result.final_response
+app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
 
 
 class ChatRequest(BaseModel):
@@ -148,239 +100,211 @@ class ChatResponse(BaseModel):
     elapsed_seconds: float
 
 
+def _to_wire_events(message: str, history: list[dict[str, str]]):
+    """Translate agent events into the shape the browser client consumes.
+
+    The client speaks `thought`/`token`/`tool_end`/`plot`/`done`; the agent
+    speaks `thinking`/`text`/... Doing the mapping here keeps the wire format
+    stable for the existing UI while the agent's own vocabulary stays clean.
+    """
+    thoughts: list[str] = []
+    for event in get_agent().stream(message, history):
+        kind = event["type"]
+        if kind == "text":
+            yield {"type": "token", "content": event["text"]}
+        elif kind == "thinking":
+            yield {"type": "thought", "content": event["text"]}
+        elif kind == "tool_start":
+            note = f"Calling {event['tool']} with {json.dumps(event['arguments'], ensure_ascii=False, default=str)[:300]}"
+            thoughts.append(note)
+            yield {"type": "thought", "content": note}
+            yield {"type": "tool_start", "tool": event["tool"], "args": event["arguments"]}
+        elif kind == "tool_end":
+            trace = {"tool": event["tool"], "status": event["status"]}
+            note = f"{event['tool']} -> {event['status']}"
+            thoughts.append(note)
+            yield {"type": "thought", "content": note}
+            yield {"type": "tool_end", "trace": trace}
+        elif kind == "plot":
+            yield {"type": "plot", "url": event["url"]}
+        elif kind == "done":
+            yield {
+                "type": "done",
+                "response": event["answer"],
+                "traces": event["traces"],
+                "plots": event["plots"],
+                "sources": event["sources"],
+                "thoughts": thoughts,
+                "stats": event["stats"],
+            }
+
+
 @app.get("/", response_class=HTMLResponse)
 async def serve_index() -> HTMLResponse:
     index_file = STATIC_DIR / "index.html"
-    if index_file.exists():
-        return HTMLResponse(
-            content=index_file.read_text(encoding="utf-8"),
-            headers={"Cache-Control": "no-cache, no-store, must-revalidate", "Pragma": "no-cache", "Expires": "0"},
-        )
-    return HTMLResponse("<h1>ControlAI Web UI initializing...</h1>")
+    if not index_file.exists():
+        return HTMLResponse("<h1>ControlAI: web/index.html is missing</h1>", status_code=500)
+    return HTMLResponse(
+        index_file.read_text(encoding="utf-8"),
+        headers={"Cache-Control": "no-cache, no-store, must-revalidate"},
+    )
 
 
 @app.get("/api/status")
-async def get_status() -> dict[str, Any]:
-    rag_index_path = PROJECT_ROOT / "data" / "rag_index"
-    chunk_count = 0
-    if rag_index_path.exists():
-        chunks_file = rag_index_path / "chunks.json"
-        if chunks_file.exists():
-            try:
-                data = json.loads(chunks_file.read_text(encoding="utf-8"))
-                chunk_count = len(data)
-            except Exception:
-                pass
-
-    return {
+async def status() -> dict[str, Any]:
+    index = get_shared_index()
+    agent = _agent
+    payload: dict[str, Any] = {
         "system": "ControlAI",
-        "status": "ready",
-        "indexed_chunks": chunk_count,
+        "status": "ready" if agent else "loading",
+        "indexed_chunks": len(index.chunks),
     }
+    if agent:
+        payload.update(
+            model=agent.engine.model_id,
+            adapter=agent.engine.adapter_path,
+            thinking=agent.thinking,
+            dense_retrieval=bool(agent.retriever and agent.retriever.has_dense),
+        )
+    return payload
 
 
 @app.get("/api/documents")
 async def list_documents() -> dict[str, Any]:
     categories: dict[str, list[str]] = {}
-    if USER_DOCS_DIR.exists():
-        for item in USER_DOCS_DIR.iterdir():
+    root = UPLOADS_DIR.parent
+    if root.exists():
+        for item in sorted(root.iterdir()):
             if item.is_dir() and not item.name.startswith("."):
-                cat_name = item.name.replace("_", " ").title()
                 files = [f.name for f in item.rglob("*.pdf") if not f.name.startswith(".")]
                 if files:
-                    categories[cat_name] = files[:10]  # Show up to 10 per category
+                    categories[item.name.replace("_", " ").title()] = files[:10]
     return {"categories": categories}
 
 
 @app.post("/api/upload")
 async def upload_document(file: UploadFile = File(...)) -> dict[str, Any]:
-    filename = file.filename
-    if not filename:
+    if not file.filename:
         raise HTTPException(status_code=400, detail="Invalid filename")
-
-    ext = Path(filename).suffix.lower()
-    if ext not in (".pdf", ".txt", ".md"):
+    if Path(file.filename).suffix.lower() not in (".pdf", ".txt", ".md"):
         raise HTTPException(status_code=400, detail="Only PDF, TXT, and Markdown files are supported.")
 
-    target_path = UPLOADED_DOCS_DIR / filename
-    with target_path.open("wb") as f:
-        shutil.copyfileobj(file.file, f)
+    target = UPLOADS_DIR / Path(file.filename).name
+    with target.open("wb") as handle:
+        shutil.copyfileobj(file.file, handle)
 
-    # Ingest directly into RAG index
-    try:
-        pages = load_single_file(target_path)
-        new_chunks = []
-        for p in pages:
-            chunks = chunk_document(p)
-            new_chunks.extend(chunks)
-
-        # Mutating the shared index makes the upload live for the running
-        # agent immediately -- a fresh instance would only persist to disk and
-        # stay invisible until restart.
+    def _ingest() -> dict[str, Any]:
+        pages = load_single_file(target)
+        new_chunks = [c for page in pages for c in chunk_document(page)]
         index = get_shared_index()
         index.add_chunks(new_chunks)
-
+        # Embedding happens on the inference thread because it uses MLX too.
+        embedded = 0
+        agent = get_agent()
+        if agent.retriever is not None:
+            embedded = agent.retriever.add_chunks(
+                [c for c in index.chunks if c["chunk_id"] in {n["chunk_id"] for n in new_chunks}]
+            )
         return {
             "status": "success",
-            "filename": filename,
+            "filename": target.name,
             "pages_parsed": len(pages),
             "chunks_added": len(new_chunks),
+            "chunks_embedded": embedded,
             "total_indexed_chunks": len(index.chunks),
         }
-    except Exception as exc:
-        raise HTTPException(status_code=500, detail=f"Indexing failed: {str(exc)}")
-
 
-from fastapi.responses import StreamingResponse
+    try:
+        with inference_lock:
+            return await _on_inference_thread(_ingest)
+    except Exception as exc:
+        raise HTTPException(status_code=500, detail=f"Indexing failed: {exc}") from exc
 
 
 @app.post("/api/chat/stream")
-async def chat_stream_endpoint(req: ChatRequest):
-    if not req.message.strip():
+async def chat_stream(req: ChatRequest) -> StreamingResponse:
+    message = req.message.strip()
+    if not message:
         raise HTTPException(status_code=400, detail="Message cannot be empty")
 
-    message = req.message.strip()
-    history = req.history
-
-    if inference_executor is not None:
-        # MLX path: a persistent dedicated thread is required (see the
-        # inference_executor comment above), so bridge it into the async
-        # response via a queue -- Starlette's own rotating threadpool would
-        # violate MLX's single-thread requirement.
-        event_queue: queue.Queue = queue.Queue()
-        _DONE = object()
-
-        def _produce() -> None:
-            try:
-                with inference_lock:
-                    for event in _run_stream_on_gpu(message, history):
-                        event_queue.put(event)
-            except Exception as exc:
-                event_queue.put({"type": "error", "error": str(exc)})
-            finally:
-                event_queue.put(_DONE)
-
-        async def event_generator():
-            loop = asyncio.get_event_loop()
-            loop.run_in_executor(inference_executor, _produce)
-            while True:
-                # Draining the queue never touches MLX, so this can safely
-                # run on the default threadpool.
-                event = await loop.run_in_executor(None, event_queue.get)
-                if event is _DONE:
-                    break
-                yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
-    else:
-        # CUDA/ZeroGPU/GGUF path: a plain sync generator handed directly to
-        # StreamingResponse, exactly as this ran before the MLX fix existed.
-        # Starlette wraps this in its own threadpool (iterate_in_threadpool),
-        # which -- unlike a manually created ThreadPoolExecutor -- is a
-        # context ZeroGPU's CUDA interception correctly recognizes.
-        def event_generator():
-            try:
-                with inference_lock:
-                    for event in _run_stream_on_gpu(message, history):
-                        yield f"data: {json.dumps(event, ensure_ascii=False)}\n\n"
-            except Exception as exc:
-                yield f"data: {json.dumps({'type': 'error', 'error': str(exc)}, ensure_ascii=False)}\n\n"
+    events: queue.Queue = queue.Queue()
+    sentinel = object()
+
+    def produce() -> None:
+        try:
+            with inference_lock:
+                for event in _to_wire_events(message, req.history):
+                    events.put(event)
+        except Exception as exc:  # surface the failure in the chat, don't hang
+            print(f"[chat] {type(exc).__name__}: {exc}")
+            events.put({"type": "error", "error": f"{type(exc).__name__}: {exc}"})
+        finally:
+            events.put(sentinel)
+
+    async def relay():
+        loop = asyncio.get_running_loop()
+        loop.run_in_executor(inference_executor, produce)
+        while True:
+            # Draining the queue never touches MLX, so the default threadpool
+            # is fine here and keeps the single inference thread free.
+            event = await loop.run_in_executor(None, events.get)
+            if event is sentinel:
+                break
+            yield f"data: {json.dumps(event, ensure_ascii=False, default=str)}\n\n"
 
     return StreamingResponse(
-        event_generator(),
+        relay(),
         media_type="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-            "X-Accel-Buffering": "no",
-        },
+        headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
     )
 
 
 @app.post("/api/chat", response_model=ChatResponse)
-async def chat_endpoint(req: ChatRequest) -> ChatResponse:
-    if not req.message.strip():
+async def chat(req: ChatRequest) -> ChatResponse:
+    message = req.message.strip()
+    if not message:
         raise HTTPException(status_code=400, detail="Message cannot be empty")
 
-    t0 = time.time()
-    try:
-        result = await _run_inference(_run_on_gpu, req.message.strip(), req.history)
-        elapsed = time.time() - t0
-
-        # Collect tool traces
-        traces = []
-        plots = []
-        for t in result.tool_traces:
-            trace_data = {
-                "tool": t.tool_name,
-                "args": t.arguments,
-                "status": t.result.get("status", "success"),
-                "residual": t.result.get("residual"),
-            }
-            traces.append(trace_data)
+    started = time.time()
 
-            # Check if plot was generated by any tool
-            if "plot_path" in t.result:
-                p_path = Path(t.result["plot_path"])
-                if p_path.exists():
-                    plots.append(f"/plots/{p_path.name}")
+    def _run():
+        with inference_lock:
+            return get_agent().run(message, req.history)
 
-        return ChatResponse(
-            response=result.final_response,
-            tool_traces=traces,
-            plots=plots,
-            elapsed_seconds=round(elapsed, 2),
-        )
+    try:
+        result = await _on_inference_thread(_run)
     except Exception as exc:
-        elapsed = time.time() - t0
-        print(f"Chat execution error: {exc}")
+        print(f"[chat] {type(exc).__name__}: {exc}")
         return ChatResponse(
-            response=f"Computation error during inference: {str(exc)}",
-            tool_traces=[],
-            plots=[],
-            elapsed_seconds=round(elapsed, 2),
+            response=f"Inference failed: {exc}",
+            elapsed_seconds=round(time.time() - started, 2),
         )
 
-
-# A minimal Gradio Blocks app exists purely so this Space is recognized as a
-# Gradio SDK app -- required for ZeroGPU hardware -- and so the @spaces.GPU
-# registration probe below is wired to a real Gradio event handler ("No
-# @spaces.GPU function detected" otherwise). It is launched on its own
-# internal port rather than mounted into our FastAPI app: mounting it and
-# then binding our own uvicorn to the same public port collided with
-# Gradio's own server startup ("address already in use"). The real UI and
-# API are still served entirely by our own FastAPI routes above, on the
-# public port.
-_gpu_demo = gr.Blocks()
-with _gpu_demo:
-    gr.Markdown("ControlAI is running. Visit the Space's root URL for the full app.")
-    _probe_in = gr.Textbox(visible=False)
-    _probe_out = gr.Textbox(visible=False)
-    _probe_btn = gr.Button(visible=False)
-    _probe_btn.click(fn=_zerogpu_registration_probe, inputs=_probe_in, outputs=_probe_out)
+    return ChatResponse(
+        response=result.answer,
+        tool_traces=[{"tool": t.name, "status": t.status} for t in result.traces],
+        plots=result.plots,
+        elapsed_seconds=round(time.time() - started, 2),
+    )
 
 
 def main() -> None:
-    import uvicorn
-
-    if os.environ.get("SPACE_ID"):
-        port = int(os.environ.get("PORT", 7860))
-        _gpu_demo.launch(server_name="0.0.0.0", server_port=port + 1, prevent_thread_lock=True, show_error=False, quiet=True)
-        print(f"ControlAI Web UI running on Hugging Face Spaces (port {port})...")
-        uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")
-        return
-
-    import threading
+    import threading as _threading
     import webbrowser
 
-    def _open_browser() -> None:
-        time.sleep(1.2)
+    import uvicorn
+
+    def open_browser() -> None:
+        time.sleep(1.5)
         try:
             webbrowser.open("http://127.0.0.1:8000")
         except Exception:
             pass
 
-    threading.Thread(target=_open_browser, daemon=True).start()
-    print("ControlAI Web UI running at: http://127.0.0.1:8000 (Opening browser...)")
-    uvicorn.run("app:app", host="127.0.0.1", port=8000, reload=True, log_level="info")
+    _threading.Thread(target=open_browser, daemon=True).start()
+    # No reload: the model load is far too expensive to repeat on every file
+    # save, and reload would fork a second copy of it.
+    uvicorn.run(app, host="127.0.0.1", port=8000, log_level="info")
 
 
 if __name__ == "__main__":

cli.py

@@ -1,85 +1,114 @@
 #!/usr/bin/env python3
-"""Interactive CLI for ControlAI Agent with deterministic tool calling."""
+"""Interactive terminal client for ControlAI."""
 
 from __future__ import annotations
 
 import argparse
 import sys
+import time
 from pathlib import Path
 
 PROJECT_ROOT = Path(__file__).resolve().parent
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
 
-from controlai_agent.orchestrator import ControlAIAgent
+from controlai_agent.agent import ControlAgent
+from controlai_agent.engine import DEFAULT_ADAPTER, DEFAULT_MODEL, LocalEngine
+
+DIM, RESET, BOLD = "\033[2m", "\033[0m", "\033[1m"
+
+
+def converse(agent: ControlAgent, question: str, history: list[dict], verbose: bool) -> str:
+    started = time.time()
+    first_text: float | None = None
+    answer = ""
+    in_thought = False
+
+    for event in agent.stream(question, history):
+        kind = event["type"]
+        if kind == "thinking":
+            if not in_thought:
+                print(f"{DIM}thinking… ", end="", flush=True)
+                in_thought = True
+            if verbose:
+                print(f"{DIM}{event['text']}{RESET}", end="", flush=True)
+        elif kind == "text":
+            if in_thought:
+                print(RESET, flush=True)
+                in_thought = False
+            if first_text is None:
+                first_text = time.time() - started
+            print(event["text"], end="", flush=True)
+        elif kind == "tool_start":
+            args = str(event["arguments"])
+            print(f"\n{DIM}→ {event['tool']}({args[:100]}{'…' if len(args) > 100 else ''}){RESET}", flush=True)
+        elif kind == "tool_end":
+            print(f"{DIM}  {event['status']}{RESET}", flush=True)
+        elif kind == "plot":
+            print(f"\n{DIM}[plot saved: outputs/plots/{Path(event['url']).name}]{RESET}", flush=True)
+        elif kind == "done":
+            answer = event["answer"]
+            if event["sources"]:
+                print(f"\n\n{DIM}Sources: {'; '.join(dict.fromkeys(event['sources']))}{RESET}")
+            stats = event["stats"]
+            print(
+                f"\n{DIM}{time.time() - started:.1f}s total · first token {first_text or 0:.2f}s · "
+                f"{stats['cached_tokens']}/{stats['prompt_tokens']} prompt tokens cached · "
+                f"{stats['decode_tps']} tok/s{RESET}"
+            )
+    return answer
 
 
 def main() -> int:
-    parser = argparse.ArgumentParser(description="ControlAI Offline Engineering Agent CLI")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="mlx-community/Qwen3-4B-Instruct-2507-4bit",
-        help="Base model path or HuggingFace repo",
-    )
+    parser = argparse.ArgumentParser(description="ControlAI -- offline control engineering agent")
+    parser.add_argument("prompt", nargs="*", help="ask one question and exit")
+    parser.add_argument("--model", default=DEFAULT_MODEL, help="MLX model id or local path")
+    parser.add_argument("--adapter", default=DEFAULT_ADAPTER, help="optional LoRA adapter path")
     parser.add_argument(
-        "--adapter-path",
-        type=str,
+        "--thinking",
+        choices=("off", "auto", "on"),
         default=None,
-        help="Optional fine-tuned LoRA adapter path",
-    )
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        default=None,
-        help="Single prompt to execute in non-interactive mode",
-    )
-    parser.add_argument(
-        "--verbose",
-        action="store_true",
-        help="Print raw tool calls and intermediate steps",
+        help="reasoning before answering (default: auto -- on for conceptual questions only)",
     )
+    parser.add_argument("--think-budget", type=int, default=None, help="max tokens spent reasoning")
+    parser.add_argument("--verbose", action="store_true", help="print reasoning and full tool output")
     args = parser.parse_args()
 
-    print("Initializing ControlAI Agent...")
-    agent = ControlAIAgent(model_path=args.model, adapter_path=args.adapter_path)
-    print("ControlAI Agent initialized with deterministic control tools.")
+    print(f"Loading {args.model}…")
+    engine = LocalEngine(model_id=args.model, adapter_path=args.adapter)
+    kwargs = {}
+    if args.thinking:
+        kwargs["thinking"] = args.thinking
+    if args.think_budget:
+        kwargs["think_budget"] = args.think_budget
+    agent = ControlAgent(engine=engine, **kwargs)
+    print(f"Ready in {engine.load_seconds:.1f}s.\n")
 
     if args.prompt:
-        result = agent.run(args.prompt, verbose=args.verbose)
-        print("\n" + "=" * 50)
-        print("AGENT RESPONSE:")
-        print("=" * 50)
-        print(result.final_response)
-        if result.tool_traces:
-            print("\n" + "-" * 50)
-            print(f"EXECUTED TOOLS ({len(result.tool_traces)}):")
-            for trace in result.tool_traces:
-                print(f"  * {trace.tool_name}({trace.arguments}) -> {trace.result.get('status', 'done')}")
+        converse(agent, " ".join(args.prompt), [], args.verbose)
         return 0
 
-    print("\n" + "=" * 60)
-    print("🤖 Welcome to Control-LLM (Offline Control Engineering Agent)")
-    print("=" * 60)
-    print("Type your control engineering question or design problem.")
-    print("Commands: 'exit' or 'quit' to end, 'verbose' to toggle tool details.\n")
+    print(f"{BOLD}ControlAI{RESET} — control engineering assistant. Ctrl-D or 'exit' to quit.")
+    history: list[dict] = []
     while True:
         try:
-            user_input = input("\n[Engineer] > ").strip()
-            if not user_input:
-                continue
-            if user_input.lower() in ("exit", "quit", "q"):
-                break
-            result = agent.run(user_input, verbose=args.verbose)
-            print(f"\n[ControlAI]\n{result.final_response}")
-            if result.tool_traces and not args.verbose:
-                tools_used = ", ".join(t.tool_name for t in result.tool_traces)
-                print(f"\n(Verified using tools: {tools_used})")
+            question = input(f"\n{BOLD}you ›{RESET} ").strip()
         except (KeyboardInterrupt, EOFError):
-            print("\nExiting ControlAI.")
-            break
-
-    return 0
+            print("\nBye.")
+            return 0
+        if not question:
+            continue
+        if question.lower() in ("exit", "quit", "q"):
+            return 0
+        if question.lower() in ("clear", "reset"):
+            history.clear()
+            agent.engine.reset_cache()
+            print(f"{DIM}history cleared{RESET}")
+            continue
+        print()
+        answer = converse(agent, question, history, args.verbose)
+        history.append({"role": "user", "content": question})
+        history.append({"role": "assistant", "content": answer})
 
 
 if __name__ == "__main__":

controlai_agent/agent.py

@@ -0,0 +1,418 @@
+"""The ControlAI agent loop.
+
+Replaces the previous `orchestrator.py`. The differences that matter:
+
+*   **It streams.** Text reaches the caller as the model produces it. The old
+    loop blocked for a full generation and then re-emitted the finished string
+    word by word, which looked like streaming but meant the user waited for the
+    entire answer before seeing anything.
+*   **It reuses the KV cache** across tool steps via `LocalEngine`, so the
+    multi-thousand-token tool-schema prefix is prefilled once per process
+    rather than once per step.
+*   **It trusts the model with parameters.** The old loop ran a "provenance"
+    check that refused any matrix it could not trace back to the user's text.
+    That blocked the most useful thing the agent does -- working an example the
+    user asked for -- so it is gone. What remains is schema validation and
+    execution in `registry.execute`, which are real guarantees, plus a
+    repetition guard for genuinely degenerate output.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Generator, Iterable, Sequence
+
+from controlai_agent import toolcall
+from controlai_agent.engine import LocalEngine, SamplingConfig
+from controlai_agent.prompts import RETRIEVAL_PREAMBLE, SYNTHESIS_NUDGE, SYSTEM_PROMPT
+from controlai_agent.registry import registry
+from controlai_agent.toolcall import ToolCall
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+
+MAX_HISTORY_TOKENS = 8000
+# Two rounds of tools answers essentially every real question (design, then
+# simulate). The old default of four mostly bought extra latency and gave a
+# stuck model more room to loop.
+MAX_TOOL_STEPS = 2
+MAX_CALLS_PER_TOOL = 2
+
+THINKING_MODE = os.environ.get("CONTROLAI_THINKING", "auto").lower()
+THINK_BUDGET = int(os.environ.get("CONTROLAI_THINK_BUDGET", "512"))
+
+# Questions that are about a concept rather than a specific system. Used only
+# to decide whether to spend thinking tokens -- never to block a tool call.
+_CONCEPTUAL_RE = re.compile(
+    r"\b(why|explain|what is|what are|difference between|compare|derive|"
+    r"derivation|prove|proof|intuition|when should|trade-?off|meaning of)\b",
+    re.IGNORECASE,
+)
+
+
+@dataclass
+class ToolTrace:
+    name: str
+    arguments: dict[str, Any]
+    result: dict[str, Any]
+
+    @property
+    def status(self) -> str:
+        return str(self.result.get("status", "success"))
+
+
+@dataclass
+class AgentResult:
+    answer: str
+    traces: list[ToolTrace] = field(default_factory=list)
+    plots: list[str] = field(default_factory=list)
+    sources: list[str] = field(default_factory=list)
+    stats: dict[str, Any] = field(default_factory=dict)
+
+
+class _StreamGate:
+    """Emits streamed text while withholding anything from `marker` onward.
+
+    The model decides between answering and calling a tool by what it emits
+    first, and that decision is only visible partway through a token. This
+    releases text as soon as it cannot be the start of `marker`, so prose
+    streams with no perceptible delay while a tool call never leaks into the
+    chat.
+    """
+
+    def __init__(self, marker: str = "<tool_call>") -> None:
+        self.marker = marker
+        self._pending = ""
+        self.suppressed = False
+
+    def feed(self, text: str) -> str:
+        if self.suppressed:
+            return ""
+        self._pending += text
+        idx = self._pending.find(self.marker)
+        if idx != -1:
+            out, self._pending, self.suppressed = self._pending[:idx], "", True
+            return out
+        # Hold back only a possible partial marker at the very end.
+        hold = 0
+        for n in range(min(len(self.marker) - 1, len(self._pending)), 0, -1):
+            if self._pending.endswith(self.marker[:n]):
+                hold = n
+                break
+        out, self._pending = (self._pending[:-hold], self._pending[-hold:]) if hold else (self._pending, "")
+        return out
+
+    def flush(self) -> str:
+        if self.suppressed:
+            return ""
+        out, self._pending = self._pending, ""
+        return out
+
+
+class ControlAgent:
+    """Control-engineering agent over a local model and deterministic tools."""
+
+    def __init__(
+        self,
+        engine: LocalEngine | None = None,
+        tool_registry=registry,
+        retriever: Any | None = None,
+        max_tool_steps: int = MAX_TOOL_STEPS,
+        thinking: str = THINKING_MODE,
+        think_budget: int = THINK_BUDGET,
+    ) -> None:
+        import controlai_agent.tools  # noqa: F401  -- registers every tool
+
+        self.engine = engine or LocalEngine()
+        self.registry = tool_registry
+        self.max_tool_steps = max_tool_steps
+        self.thinking = thinking
+        self.think_budget = think_budget
+        self.tool_schemas = self.registry.get_tool_schemas()
+
+        if retriever is None:
+            try:
+                from controlai_rag.retriever import get_retriever
+
+                retriever = get_retriever()
+            except Exception as exc:  # retrieval is an enhancement, not a dependency
+                print(f"[agent] retrieval unavailable ({type(exc).__name__}: {exc}); continuing without it")
+        self.retriever = retriever
+
+        self._prewarm()
+
+    # ------------------------------------------------------------- setup
+
+    def _prewarm(self) -> None:
+        """Prefill the fixed system-prompt-plus-tool-schema prefix.
+
+        Everything after it in a real prompt is conversation, so this is the
+        one part of every request that is byte-identical every time. Paying for
+        it at startup is what makes the first question feel instant.
+        """
+        prefix = self.engine.render(
+            [{"role": "system", "content": SYSTEM_PROMPT}], tools=self.tool_schemas
+        )
+        # Cut at the end of the system block: the generation prompt that
+        # `render` appends belongs to the user's turn, not to the prefix.
+        anchor = prefix.rfind("<|im_end|>")
+        if anchor != -1:
+            prefix = prefix[: anchor + len("<|im_end|>\n")]
+        n = self.engine.prewarm(prefix)
+        print(f"[agent] prewarmed {n} prefix tokens ({self.engine.model_id})")
+
+    # --------------------------------------------------------- prompting
+
+    def _retrieve(self, question: str) -> tuple[str, list[str]]:
+        if self.retriever is None:
+            return "", []
+        try:
+            hits = self.retriever.search(question, top_k=4)
+        except Exception as exc:
+            print(f"[agent] retrieval failed ({type(exc).__name__}: {exc})")
+            return "", []
+        if not hits:
+            return "", []
+        blocks, labels = [], []
+        for hit in hits:
+            label = hit.get("label") or hit.get("source_name") or "Reference"
+            page = hit.get("page")
+            label = f"{label}, p. {page}" if page else label
+            text = " ".join(str(hit.get("text", "")).split())[:800]
+            blocks.append(f"[{label}]\n{text}")
+            labels.append(label)
+        return RETRIEVAL_PREAMBLE + "\n\n" + "\n\n".join(blocks), labels
+
+    def _build_messages(
+        self, question: str, history: Sequence[dict[str, Any]] | None
+    ) -> tuple[list[dict[str, Any]], list[str]]:
+        messages: list[dict[str, Any]] = [{"role": "system", "content": SYSTEM_PROMPT}]
+        messages += self._truncate(history or [])
+
+        context, sources = self._retrieve(question)
+        # Retrieved passages ride along with the user's turn rather than being
+        # spliced into the system prompt. That keeps the cached prefix
+        # byte-stable across questions, which is worth more than the tidier
+        # placement.
+        content = f"{context}\n\n---\n\n{question}" if context else question
+        messages.append({"role": "user", "content": content})
+        return messages, sources
+
+    def _truncate(self, history: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Drop the oldest turns until the history fits the token budget.
+
+        The web client resends the whole conversation every request and trims
+        nothing, so the server has to.
+        """
+        kept: list[dict[str, Any]] = []
+        used = 0
+        for item in reversed(list(history)):
+            role, content = item.get("role"), (item.get("content") or "").strip()
+            if role not in ("user", "assistant") or not content:
+                continue
+            cost = self.engine.count_tokens(content) + 8
+            if used + cost > MAX_HISTORY_TOKENS:
+                break
+            kept.append({"role": role, "content": content})
+            used += cost
+        return list(reversed(kept))
+
+    def _wants_thinking(self, question: str) -> bool:
+        if self.thinking == "on":
+            return True
+        if self.thinking == "off":
+            return False
+        # "auto": reasoning earns its latency on conceptual questions, which is
+        # where this model is weakest and where no solver can help it.
+        return bool(_CONCEPTUAL_RE.search(question))
+
+    # ------------------------------------------------------------- tools
+
+    def _execute(self, call: ToolCall) -> dict[str, Any]:
+        for key, value in call.arguments.items():
+            reason = toolcall.degenerate_reason(value)
+            if reason:
+                return {
+                    "status": "error",
+                    "error_type": "DegenerateArgument",
+                    "error": (
+                        f"The value passed as '{key}' {reason}, which is not a real system. "
+                        f"Re-read the question and pass the actual values, or say what is missing."
+                    ),
+                }
+        return self.registry.execute(call.name, call.arguments)
+
+    # ----------------------------------------------------------- running
+
+    def stream(
+        self,
+        question: str,
+        history: Sequence[dict[str, Any]] | None = None,
+        max_tokens: int = 1536,
+    ) -> Generator[dict[str, Any], None, None]:
+        """Run one turn, yielding events as they happen.
+
+        Event types: `thinking`, `text`, `tool_start`, `tool_end`, `plot`,
+        `done`.
+        """
+        messages, sources = self._build_messages(question, history)
+        traces: list[ToolTrace] = []
+        plots: list[str] = []
+        answer_parts: list[str] = []
+        call_counts: dict[str, int] = {}
+        seen: set[str] = set()
+        # Accumulated across every generation in the turn: the engine's own
+        # stats only describe its most recent call, which for a tool-using
+        # question is the short synthesis pass and badly understates the work.
+        totals = {"prompt_tokens": 0, "cached_tokens": 0, "generated_tokens": 0, "prefill_seconds": 0.0, "decode_seconds": 0.0}
+
+        # Decided once, before the loop. Deciding it per-pass looked equivalent
+        # but was not: a conceptual question is answered on the *first* pass,
+        # which is never the final pass, so reasoning was silently never
+        # enabled for exactly the questions "auto" exists to help.
+        think_turn = self._wants_thinking(question)
+
+        for step in range(self.max_tool_steps + 1):
+            final_pass = step == self.max_tool_steps
+            tools = None if final_pass else self.tool_schemas
+            # Once a solver has produced the number, the number is the answer;
+            # reasoning over it only adds latency.
+            think = think_turn and not traces
+
+            if final_pass and traces:
+                messages.append({"role": "user", "content": SYNTHESIS_NUDGE})
+
+            prompt = self.engine.render(messages, tools=tools, enable_thinking=think)
+            gate = _StreamGate()
+            raw: list[str] = []
+
+            for chunk in self.engine.stream(
+                prompt,
+                sampling=self.engine.sampling.with_(max_tokens=max_tokens),
+                stop=("</tool_call>",) if tools else (),
+                think_budget=self.think_budget if think else None,
+            ):
+                if chunk.thinking:
+                    yield {"type": "thinking", "text": chunk.text}
+                    continue
+                raw.append(chunk.text)
+                visible = gate.feed(chunk.text)
+                if visible:
+                    answer_parts.append(visible)
+                    yield {"type": "text", "text": visible}
+
+            tail = gate.flush()
+            if tail:
+                answer_parts.append(tail)
+                yield {"type": "text", "text": tail}
+
+            stats = self.engine.last_stats
+            totals["prompt_tokens"] += stats.prompt_tokens
+            totals["cached_tokens"] += stats.cached_tokens
+            totals["generated_tokens"] += stats.generated_tokens
+            totals["prefill_seconds"] += stats.prefill_seconds
+            totals["decode_seconds"] += stats.decode_seconds
+
+            output = "".join(raw)
+            calls, _ = toolcall.parse(output)
+
+            if not calls:
+                break
+
+            # A tool call is arriving, so whatever prose preceded it was
+            # narration ("Let me compute that"), not the answer. Drop it from
+            # the answer text; the user already saw it stream past.
+            answer_parts.clear()
+            messages.append({"role": "assistant", "content": output})
+
+            for call in calls:
+                if call_counts.get(call.name, 0) >= MAX_CALLS_PER_TOOL:
+                    continue
+                signature = f"{call.name}:{json.dumps(call.arguments, sort_keys=True, default=str)}"
+                if signature in seen:
+                    continue
+                seen.add(signature)
+                call_counts[call.name] = call_counts.get(call.name, 0) + 1
+
+                yield {"type": "tool_start", "tool": call.name, "arguments": call.arguments}
+                result = self._execute(call)
+                traces.append(ToolTrace(call.name, call.arguments, result))
+                yield {
+                    "type": "tool_end",
+                    "tool": call.name,
+                    "status": result.get("status", "success"),
+                    "result": result,
+                }
+
+                plot_path = result.get("plot_path")
+                if plot_path and Path(plot_path).exists():
+                    url = f"/plots/{Path(plot_path).name}"
+                    plots.append(url)
+                    yield {"type": "plot", "url": url}
+
+                messages.append(
+                    {
+                        "role": "tool",
+                        "name": call.name,
+                        "content": json.dumps(result, ensure_ascii=False, default=str),
+                    }
+                )
+
+        answer = "".join(answer_parts).strip()
+        if not answer:
+            answer = self._recover(question, messages)
+            if answer:
+                yield {"type": "text", "text": answer}
+
+        yield {
+            "type": "done",
+            "answer": answer,
+            "traces": [{"tool": t.name, "arguments": t.arguments, "status": t.status} for t in traces],
+            "plots": plots,
+            "sources": sources,
+            "stats": {
+                "prompt_tokens": totals["prompt_tokens"],
+                "cached_tokens": totals["cached_tokens"],
+                "generated_tokens": totals["generated_tokens"],
+                "prefill_seconds": round(totals["prefill_seconds"], 3),
+                "decode_tps": round(
+                    totals["generated_tokens"] / totals["decode_seconds"], 1
+                ) if totals["decode_seconds"] else 0.0,
+            },
+        }
+
+    def _recover(self, question: str, messages: list[dict[str, Any]]) -> str:
+        """Last resort when the loop produced no prose.
+
+        Re-asks with the tool results kept but the tool schemas withdrawn, so
+        the model has nothing to answer with except words.
+        """
+        messages = messages + [{"role": "user", "content": SYNTHESIS_NUDGE}]
+        prompt = self.engine.render(messages, tools=None, enable_thinking=False)
+        _, prose = toolcall.parse(self.engine.generate(prompt))
+        return prose.strip()
+
+    def run(
+        self,
+        question: str,
+        history: Sequence[dict[str, Any]] | None = None,
+        max_tokens: int = 1536,
+    ) -> AgentResult:
+        """Blocking variant of `stream`."""
+        result = AgentResult(answer="")
+        traces: list[ToolTrace] = []
+        for event in self.stream(question, history, max_tokens):
+            if event["type"] == "tool_end":
+                traces.append(ToolTrace(event["tool"], {}, event["result"]))
+            elif event["type"] == "done":
+                result = AgentResult(
+                    answer=event["answer"],
+                    traces=traces,
+                    plots=event["plots"],
+                    sources=event["sources"],
+                    stats=event["stats"],
+                )
+        return result

controlai_agent/engine.py

@@ -0,0 +1,393 @@
+"""Local inference engine: MLX on Apple Silicon, with a prefix-reusing prompt cache.
+
+Design notes that matter for anyone changing this file:
+
+*   **One persistent KV cache per engine, reused across every generation.**
+    The agent's prompt is dominated by a fixed prefix -- the system prompt plus
+    the JSON schemas of ~28 tools, which together are several thousand tokens.
+    Re-prefilling that on every tool step is what made the previous
+    implementation feel slow (measured: 6.4 s of prefill per step, five to six
+    steps per question). Here the cache is kept between calls and only the
+    tokens that actually differ from what the cache already holds are fed to
+    the model, so the fixed prefix is prefilled exactly once per process.
+
+*   **Generated tokens stay in the cache too.** A tool-calling turn is
+    append-only: prompt, then the assistant's tool call, then the tool result,
+    then more assistant text. Tracking generated tokens alongside prompt tokens
+    means continuing that turn costs only the tool-result tokens.
+
+*   **Streaming is real.** `stream()` yields text as the model produces it.
+    Nothing here buffers a whole response and re-emits it word by word.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Generator, Iterable, Sequence
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+
+DEFAULT_MODEL = os.environ.get("CONTROLAI_MODEL", "mlx-community/Qwen3-14B-4bit")
+# The project's own LoRA adapters are deliberately NOT loaded by default. Both
+# of them regressed the behaviour they were meant to improve: `behavior_v1`
+# emits a spurious empty `<tool_call></tool_call>` as its first output on
+# essentially every prompt (so no tool ever runs), and `sft_v2` generates empty
+# output when no tools are exposed and string-typed numbers when they are.
+# Set CONTROLAI_ADAPTER=<path> to load one anyway for A/B work.
+DEFAULT_ADAPTER = os.environ.get("CONTROLAI_ADAPTER") or None
+
+
+@dataclass
+class SamplingConfig:
+    """Decoding parameters. Defaults follow Qwen3's own non-thinking recipe."""
+
+    temperature: float = 0.7
+    top_p: float = 0.8
+    top_k: int = 20
+    # Qwen3 recommends presence_penalty over the blunt repetition_penalty that
+    # the previous implementation applied at 1.15 across the board. A flat
+    # repetition penalty is actively harmful for this workload: it penalises
+    # the repeated structural tokens that matrices and JSON are made of
+    # (`[`, `0`, `,`) exactly when the model is emitting a tool call.
+    presence_penalty: float = 0.5
+    max_tokens: int = 1024
+
+    def with_(self, **kw: Any) -> "SamplingConfig":
+        merged = {**self.__dict__, **{k: v for k, v in kw.items() if v is not None}}
+        return SamplingConfig(**merged)
+
+
+@dataclass
+class Chunk:
+    """One streamed piece of model output."""
+
+    text: str
+    token: int
+    thinking: bool = False
+    tool_call: bool = False
+
+
+@dataclass
+class Stats:
+    prompt_tokens: int = 0
+    cached_tokens: int = 0
+    generated_tokens: int = 0
+    prefill_seconds: float = 0.0
+    decode_seconds: float = 0.0
+
+    @property
+    def decode_tps(self) -> float:
+        return self.generated_tokens / self.decode_seconds if self.decode_seconds else 0.0
+
+
+class LocalEngine:
+    """Streaming text generation against a locally held MLX model."""
+
+    def __init__(
+        self,
+        model_id: str = DEFAULT_MODEL,
+        adapter_path: str | None = DEFAULT_ADAPTER,
+        sampling: SamplingConfig | None = None,
+        max_cache_tokens: int = 32768,
+    ) -> None:
+        from mlx_lm import load
+
+        self.model_id = model_id
+        self.adapter_path = adapter_path
+        self.sampling = sampling or SamplingConfig()
+        self.max_cache_tokens = max_cache_tokens
+
+        t0 = time.time()
+        if adapter_path:
+            self.model, self.tokenizer = load(model_id, adapter_path=adapter_path)
+        else:
+            self.model, self.tokenizer = load(model_id)
+        self.load_seconds = time.time() - t0
+
+        self._cache: list[Any] | None = None
+        self._cache_tokens: list[int] = []
+        self.last_stats = Stats()
+
+        # Resolved once: whether this checkpoint's chat template understands
+        # Qwen3-style `enable_thinking`, and the ids of the think delimiters.
+        self.supports_thinking = self._probe_thinking_support()
+        # `<think>`, `</think>`, `<tool_call>` and `</tool_call>` are each a
+        # single special token in the Qwen3 vocabulary. Watching for the token
+        # id rather than matching the rendered string is exact: it cannot be
+        # defeated by a tag split across two streamed chunks, and it costs one
+        # integer comparison per token instead of a substring scan.
+        self._think_open = self._single_token("<think>")
+        self._think_close = self._single_token("</think>")
+        self._tool_open = self._single_token("<tool_call>")
+        self._tool_close = self._single_token("</tool_call>")
+
+    # ------------------------------------------------------------------ setup
+
+    def _token_ids(self, text: str) -> list[int]:
+        try:
+            return self.tokenizer.encode(text, add_special_tokens=False)
+        except TypeError:
+            return self.tokenizer.encode(text)
+
+    def _single_token(self, text: str) -> int | None:
+        """The id of `text` if the tokenizer represents it as one token."""
+        ids = self._token_ids(text)
+        return ids[0] if len(ids) == 1 else None
+
+    def _probe_thinking_support(self) -> bool:
+        probe = [{"role": "user", "content": "hi"}]
+        try:
+            on = self.tokenizer.apply_chat_template(
+                probe, tokenize=False, add_generation_prompt=True, enable_thinking=True
+            )
+            off = self.tokenizer.apply_chat_template(
+                probe, tokenize=False, add_generation_prompt=True, enable_thinking=False
+            )
+        except Exception:
+            return False
+        return on != off
+
+    # -------------------------------------------------------------- rendering
+
+    def render(
+        self,
+        messages: Sequence[dict[str, Any]],
+        tools: Sequence[dict[str, Any]] | None = None,
+        enable_thinking: bool = False,
+    ) -> str:
+        kwargs: dict[str, Any] = {
+            "tokenize": False,
+            "add_generation_prompt": True,
+        }
+        if tools:
+            kwargs["tools"] = list(tools)
+        if self.supports_thinking:
+            kwargs["enable_thinking"] = enable_thinking
+        return self.tokenizer.apply_chat_template(list(messages), **kwargs)
+
+    def encode(self, text: str) -> list[int]:
+        return self.tokenizer.encode(text)
+
+    def count_tokens(self, text: str) -> int:
+        return len(self.tokenizer.encode(text))
+
+    # ------------------------------------------------------------------ cache
+
+    def _align_cache(self, tokens: list[int]) -> list[int]:
+        """Point the persistent cache at the longest prefix of `tokens` it
+        already holds, and return the tokens that still need prefilling."""
+        from mlx_lm.models.cache import (
+            can_trim_prompt_cache,
+            make_prompt_cache,
+            trim_prompt_cache,
+        )
+
+        reusable = 0
+        if self._cache is not None:
+            limit = min(len(self._cache_tokens), len(tokens))
+            while reusable < limit and self._cache_tokens[reusable] == tokens[reusable]:
+                reusable += 1
+
+        # A cache that cannot be trimmed back to the divergence point is worse
+        # than no cache: it would silently condition generation on stale
+        # tokens. Rebuild instead.
+        if self._cache is not None and reusable < len(self._cache_tokens):
+            if can_trim_prompt_cache(self._cache):
+                trim_prompt_cache(self._cache, len(self._cache_tokens) - reusable)
+            else:
+                self._cache, reusable = None, 0
+
+        if self._cache is None or reusable == 0:
+            self._cache = make_prompt_cache(self.model)
+            self._cache_tokens = []
+            reusable = 0
+
+        # MLX must be fed at least one token; an exact cache hit therefore
+        # rewinds by one and replays the final token.
+        if reusable == len(tokens) and reusable > 0:
+            from mlx_lm.models.cache import trim_prompt_cache as _trim
+
+            _trim(self._cache, 1)
+            reusable -= 1
+
+        self._cache_tokens = list(tokens[:reusable])
+        self.last_stats.cached_tokens = reusable
+        return list(tokens[reusable:])
+
+    def reset_cache(self) -> None:
+        self._cache = None
+        self._cache_tokens = []
+
+    def prewarm(self, text: str) -> int:
+        """Prefill a prompt prefix so the first real question doesn't pay for it.
+
+        Returns the number of tokens now resident in the cache. Called at
+        startup with the system-prompt-plus-tool-schemas prefix, which turns
+        first-question latency from a multi-second prefill into a cache hit.
+        """
+        import mlx.core as mx
+        from mlx_lm import stream_generate
+
+        tokens = self.encode(text)
+        to_feed = self._align_cache(tokens)
+        if to_feed:
+            self.model(mx.array(to_feed)[None], cache=self._cache)
+            mx.eval([c.state for c in self._cache])
+        self._cache_tokens = list(tokens)
+
+        # Prefilling alone leaves the single-token decode kernels uncompiled,
+        # so the first real question still paid several seconds of Metal
+        # warm-up. Generate and discard one token against a throwaway cache to
+        # force that compilation now, without disturbing the prefix cache.
+        from mlx_lm.models.cache import make_prompt_cache
+
+        scratch = make_prompt_cache(self.model)
+        for _ in stream_generate(
+            self.model, self.tokenizer, [tokens[-1]], max_tokens=1, prompt_cache=scratch
+        ):
+            break
+        return len(tokens)
+
+    # ------------------------------------------------------------- generation
+
+    def stream(
+        self,
+        prompt: str | list[int],
+        sampling: SamplingConfig | None = None,
+        stop: Iterable[str] = (),
+        think_budget: int | None = None,
+    ) -> Generator[Chunk, None, None]:
+        """Yield output chunks as they are generated.
+
+        `stop` sequences end generation as soon as they appear (the sequence
+        itself is emitted, since tool-call parsing wants the closing tag).
+        `think_budget` caps how many tokens may be spent inside a `<think>`
+        block: on overrun the block is closed by hand and the model is made to
+        answer, which bounds worst-case latency on a reasoning model.
+        """
+        from mlx_lm import stream_generate
+        from mlx_lm.sample_utils import make_logits_processors, make_sampler
+
+        cfg = sampling or self.sampling
+        tokens = self.encode(prompt) if isinstance(prompt, str) else list(prompt)
+
+        t0 = time.time()
+        to_feed = self._align_cache(tokens)
+        self.last_stats = Stats(
+            prompt_tokens=len(tokens),
+            cached_tokens=len(tokens) - len(to_feed),
+        )
+
+        sampler = make_sampler(temp=cfg.temperature, top_p=cfg.top_p, top_k=cfg.top_k)
+        logits_processors = (
+            make_logits_processors(presence_penalty=cfg.presence_penalty)
+            if cfg.presence_penalty
+            else None
+        )
+
+        stop = tuple(s for s in stop if s)
+        stop_ids = {self._tool_close} if "</tool_call>" in stop and self._tool_close else set()
+        text_stops = tuple(s for s in stop if not (s == "</tool_call>" and self._tool_close))
+        # Only the tail can contain a partial stop sequence, so matching over a
+        # bounded window keeps this O(1) per token instead of rescanning the
+        # whole response.
+        window = max((len(s) for s in text_stops), default=0) + 8
+
+        emitted: list[int] = []
+        tail = ""
+        budget_left = think_budget
+        in_think = False
+        in_tool_call = False
+        first_token_at: float | None = None
+        remaining = cfg.max_tokens
+
+        while remaining > 0:
+            forced_close = False
+            for resp in stream_generate(
+                self.model,
+                self.tokenizer,
+                to_feed,
+                max_tokens=remaining,
+                sampler=sampler,
+                logits_processors=logits_processors,
+                prompt_cache=self._cache,
+            ):
+                if first_token_at is None:
+                    first_token_at = time.time()
+                    self.last_stats.prefill_seconds = first_token_at - t0
+                emitted.append(resp.token)
+                self._cache_tokens.append(resp.token)
+                remaining -= 1
+                text = resp.text
+
+                # Exact, token-id state transitions. The tags themselves are
+                # never emitted -- the caller gets the content and the flags.
+                if resp.token == self._think_open:
+                    in_think = True
+                    continue
+                if resp.token == self._think_close:
+                    in_think = False
+                    continue
+                if resp.token == self._tool_open:
+                    # Unlike the think tags, this one is emitted: the caller
+                    # parses the `<tool_call>...</tool_call>` block out of the
+                    # raw text. The flag lets it suppress the same text from
+                    # the user-visible stream.
+                    in_tool_call = True
+                if text and not in_think:
+                    tail = (tail + text)[-window:] if window else ""
+                    # Fallback for a checkpoint whose think tags are not single
+                    # tokens; harmless when the ids above already matched.
+                    if self._think_open is None and "<think>" in tail:
+                        in_think = True
+                    if self._think_close is None and "</think>" in tail:
+                        in_think = False
+
+                yield Chunk(text=text, token=resp.token, thinking=in_think, tool_call=in_tool_call)
+
+                if in_think and budget_left is not None:
+                    budget_left -= 1
+                    if budget_left <= 0:
+                        forced_close = True
+                        break
+
+                if resp.token in stop_ids or (text_stops and any(s in tail for s in text_stops)):
+                    remaining = 0
+                    break
+            else:
+                remaining = 0
+
+            if not forced_close:
+                break
+
+            # Overran the thinking budget: close the block ourselves and let
+            # the same cache continue straight into the answer.
+            closer = "\n</think>\n\n"
+            closer_ids = self._token_ids(closer)
+            self._cache_tokens.extend(closer_ids)
+            to_feed = closer_ids
+            # Deliberately not yielded: this is a control action on the model,
+            # not model output. Emitting it put a bare "</think>" at the top of
+            # the answer whenever the budget was reached.
+            in_think = False
+            budget_left = None
+
+        now = time.time()
+        self.last_stats.generated_tokens = len(emitted)
+        self.last_stats.decode_seconds = now - (first_token_at or now)
+        if len(self._cache_tokens) > self.max_cache_tokens:
+            self.reset_cache()
+
+    def generate(
+        self,
+        prompt: str | list[int],
+        sampling: SamplingConfig | None = None,
+        stop: Iterable[str] = (),
+        think_budget: int | None = None,
+    ) -> str:
+        return "".join(
+            c.text for c in self.stream(prompt, sampling, stop, think_budget)
+        )

controlai_agent/orchestrator.py

@@ -1,1327 +0,0 @@
-"""Control-LLM Agent Orchestrator: Multi-step tool calling, execution, grounding, and streaming synthesis loop."""
-
-from __future__ import annotations
-
-import json
-import re
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any, Generator
-
-import os
-import sys
-
-try:
-    from mlx_lm import generate as mlx_generate, load as mlx_load
-    from mlx_lm.sample_utils import make_logits_processors
-    HAS_MLX = True
-except ImportError:
-    HAS_MLX = False
-
-try:
-    import llama_cpp
-    HAS_LLAMA_CPP = True
-except ImportError:
-    HAS_LLAMA_CPP = False
-
-try:
-    import ollama
-    HAS_OLLAMA = True
-except ImportError:
-    HAS_OLLAMA = False
-
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-from controlai_agent.prompts import CONTROLAI_SYSTEM_PROMPT
-from controlai_agent.registry import ToolRegistry, registry
-import controlai_agent.tools  # noqa: F401 (ensure all tools are registered)
-from controlai_rag.index import get_shared_index
-
-
-@dataclass
-class ToolExecutionTrace:
-    tool_name: str
-    arguments: dict[str, Any]
-    result: dict[str, Any]
-
-
-@dataclass
-class AgentResult:
-    final_response: str
-    tool_traces: list[ToolExecutionTrace] = field(default_factory=list)
-    total_steps: int = 0
-    raw_messages: list[dict[str, Any]] = field(default_factory=list)
-    is_grounded: bool = True
-    plots: list[str] = field(default_factory=list)
-
-
-def sanitize_json_escapes(raw: str) -> str:
-    """Escape unescaped backslashes in JSON strings (e.g. \\dot, \\omega, \\mu in math/code)."""
-    return re.sub(r"\\(?![/\"\\bfnrtu])", r"\\\\", raw)
-
-
-def fix_space_separated_arrays(json_str: str) -> str:
-    """Convert MATLAB/Numpy space-separated lists inside brackets [1 2 3] to [1, 2, 3]."""
-    def _fix_brackets(match: re.Match) -> str:
-        content = match.group(1).strip()
-        # Replace spaces between numbers with commas
-        fixed = re.sub(r"([0-9eE\.\-+]+)\s+([0-9eE\.\-+]+)", r"\1, \2", content)
-        fixed = re.sub(r"([0-9eE\.\-+]+)\s+([0-9eE\.\-+]+)", r"\1, \2", fixed)
-        fixed = re.sub(r"([0-9eE\.\-+]+)\s+([0-9eE\.\-+]+)", r"\1, \2", fixed)
-        return f"[{fixed}]"
-    return re.sub(r"\[\s*([0-9eE\.\-+\s]+?)\s*\]", _fix_brackets, json_str)
-
-
-def close_unbalanced_json(raw: str) -> str:
-    """Append the closing braces/brackets a truncated JSON object is missing.
-
-    Small quantized models routinely drop the final `}` of a tool call (emitting
-    `{"name": ..., "arguments": {...}` with the outer object left open). Without
-    repair the call fails to parse, the tool never runs, and the turn collapses
-    into an empty answer -- so balance the delimiters and let it parse.
-    """
-    stack: list[str] = []
-    in_string = False
-    escaped = False
-    for ch in raw:
-        if in_string:
-            if escaped:
-                escaped = False
-            elif ch == "\\":
-                escaped = True
-            elif ch == '"':
-                in_string = False
-            continue
-        if ch == '"':
-            in_string = True
-        elif ch in "{[":
-            stack.append(ch)
-        elif ch in "}]":
-            if stack and ((ch == "}" and stack[-1] == "{") or (ch == "]" and stack[-1] == "[")):
-                stack.pop()
-
-    repaired = raw
-    if in_string:
-        repaired += '"'
-    # A generation cut off by the token budget mid-array/object usually stops
-    # right after a comma (about to write the next element) -- a dangling
-    # trailing comma is invalid JSON even once the brackets below are
-    # balanced ("[1, 0, 0,]" still fails to parse), so drop it first.
-    if not in_string:
-        repaired = re.sub(r",\s*$", "", repaired)
-    for opener in reversed(stack):
-        repaired += "}" if opener == "{" else "]"
-    return repaired
-
-
-def parse_flexible_json(raw_str: str) -> dict[str, Any] | None:
-    """Parse JSON with fallback to escape sanitization, array fixing, and non-strict control characters."""
-    raw_str = raw_str.strip()
-    try:
-        obj = json.loads(raw_str, strict=False)
-        if isinstance(obj, dict):
-            return obj
-    except Exception:
-        pass
-
-    try:
-        sanitized = sanitize_json_escapes(raw_str)
-        obj = json.loads(sanitized, strict=False)
-        if isinstance(obj, dict):
-            return obj
-    except Exception:
-        pass
-
-    try:
-        fixed_arrays = fix_space_separated_arrays(sanitize_json_escapes(raw_str))
-        obj = json.loads(fixed_arrays, strict=False)
-        if isinstance(obj, dict):
-            return obj
-    except Exception:
-        pass
-
-    # Last resort: the object is well-formed but truncated (a dropped closing
-    # brace), so balance the delimiters and retry each variant.
-    for candidate in (
-        raw_str,
-        sanitize_json_escapes(raw_str),
-        fix_space_separated_arrays(sanitize_json_escapes(raw_str)),
-    ):
-        try:
-            obj = json.loads(close_unbalanced_json(candidate), strict=False)
-            if isinstance(obj, dict):
-                return obj
-        except Exception:
-            continue
-
-    return None
-
-
-# A doubled backslash is an accidental JSON-style escape when it precedes:
-#   - a command name       \\zeta  \\frac  \\sum
-#   - a delimiter escape   \\|  \\{  \\}  \\(  \\)  \\]
-#   - \\[ that does NOT begin a spacing argument such as a genuine "\\[2pt]"
-# A real LaTeX row break inside matrix/aligned is followed by whitespace, a
-# newline, or a digit-led spacing option -- never by any of the above -- so
-# these rewrites leave legitimate line breaks untouched.
-_LATEX_DOUBLE_BACKSLASH_RE = re.compile(r"\\\\(?=[a-zA-Z|{}()\]])|\\\\(?=\[\s*[^\d\s])")
-
-
-def _fix_doubled_latex_backslashes(text: str) -> str:
-    """Collapse an accidental JSON-escape-style double backslash before a LaTeX
-    command or delimiter (\\\\zeta, \\\\|) down to the single backslash KaTeX
-    expects (\\zeta, \\|).
-
-    The base model was heavily trained on JSON-argument function calling, where
-    a literal backslash must be written as `\\\\` inside a JSON string. That
-    habit leaks into plain-text math even outside of a tool call. Left alone,
-    `\\\\|x - x_f\\\\|^2` renders as a line break followed by a stray pipe,
-    tearing a norm across two lines instead of drawing it.
-    """
-    return _LATEX_DOUBLE_BACKSLASH_RE.sub(lambda m: "\\", text)
-
-
-def _extract_tool_calls(text: str) -> tuple[list[dict[str, Any]], str]:
-    """Parse tool calls from <tool_call>, markdown code blocks, or raw JSON robustly."""
-    calls: list[dict[str, Any]] = []
-
-    # 1. Standard <tool_call> tags
-    for match in re.finditer(r"<tool_call>\s*([\s\S]*?)\s*</tool_call>", text):
-        obj = parse_flexible_json(match.group(1))
-        if obj and "name" in obj:
-            calls.append(obj)
-
-    # 2. Markdown json blocks with tool call schema
-    if not calls:
-        for match in re.finditer(r"```(?:json)?\s*(\{\s*[\"']name[\"']\s*:[\s\S]*?\})\s*```", text):
-            obj = parse_flexible_json(match.group(1))
-            if obj and "name" in obj:
-                calls.append(obj)
-
-    # 3. Raw JSON object containing "name" and "arguments" / "parameters"
-    if not calls and ('"name"' in text or "'name'" in text):
-        match = re.search(r"(\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*\})", text)
-        if not match:
-            # Generation was truncated mid-JSON (e.g. a repetition loop hit the
-            # token budget before the array closed), so there's no closing "}"
-            # to anchor on -- fall back to matching to the end of the string so
-            # close_unbalanced_json can still repair it and _degenerate_array_reason
-            # can refuse it, instead of the raw truncated JSON leaking into the chat.
-            match = re.search(r"(\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*)", text)
-        if match:
-            obj = parse_flexible_json(match.group(1))
-            if obj and "name" in obj:
-                calls.append(obj)
-
-    # Clean pre-tool thought / raw JSON artifacts
-    cleaned = re.sub(r"<tool_call>[\s\S]*?</tool_call>", "", text, flags=re.DOTALL)
-    cleaned = re.sub(r"```(?:json)?\s*\{\s*[\"']name[\"']\s*:[\s\S]*?\}\s*```", "", cleaned)
-    cleaned = re.sub(r"\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*\}", "", cleaned)
-    # Same truncated-JSON fallback as above, so the leftover raw JSON is
-    # stripped from the visible text even when it never closed.
-    cleaned = re.sub(r"\{\s*[\"']name[\"']\s*:\s*[\"'][a-zA-Z0-9_]+[\"'][\s\S]*", "", cleaned)
-    # An orphaned <tool_call> tag with no matching close (the model started a
-    # call, then abandoned it mid-generation for plain text) survives the
-    # paired regex above -- strip any leftover tag so it never reaches the UI.
-    cleaned = re.sub(r"</?tool_call>", "", cleaned)
-    cleaned = cleaned.strip()
-    cleaned = _fix_doubled_latex_backslashes(cleaned)
-    return calls, cleaned
-
-
-# Applied across every inference backend below. Without it, a small quantized
-# model under low-temperature decoding has no defense against falling into a
-# token-repetition loop once it starts (observed in production as
-# routh_hurwitz_analysis coefficient arrays of 300+ repeated zeros): it burns
-# the entire max_tokens budget on garbage, which is both the direct cause of
-# the degenerate-array tool-call failures and a major source of latency.
-REPETITION_PENALTY = 1.15
-
-PROJECT_ROOT = Path(__file__).resolve().parent.parent
-
-# Our own fine-tuned model's repo -- every GGUF/Ollama/CUDA backend must load
-# ITS tokenizer and chat template (never a generic base model's), since the
-# tool-call format, special tokens, and vocab are specific to this fine-tune.
-CONTROLAI_HF_REPO = "atakankahya/ControlAI-Agent"
-
-# How many times a single tool may be invoked within one user turn. Two allows
-# a legitimate retry with corrected arguments after an error, while stopping
-# the model from spending its whole step budget re-running the same lookup.
-MAX_CALLS_PER_TOOL = 2
-
-
-# ---------------------------------------------------------------------------
-# Parameter provenance enforcement
-#
-# Prompt rules against inventing parameters demonstrably do not hold on a 4B
-# model: given "A times A" after an earlier LQR conversation, it fabricated a
-# brand-new B, resized Q, and ran a confident LQR synthesis -- three separate
-# prompt formulations failed to stop it. So faithfulness is enforced in code:
-# every 2D matrix handed to a tool must be traceable to the user's messages or
-# a prior tool result in this conversation, or the call is refused before it
-# executes and the model is told to ask the user for the missing value.
-# ---------------------------------------------------------------------------
-
-# Parameters exempt from the guard:
-#  - C/D: conventional output-selector / feedthrough matrices (entries 0/1),
-#    routinely and legitimately chosen by the designer, and harmless.
-#  - desired_poles: pole locations are design choices derived from specs
-#    ("settling time under 2s"), not data the user must dictate literally.
-_PROVENANCE_EXEMPT_PARAMS = {"C", "D", "desired_poles"}
-
-_NUMBER_RE = re.compile(r"-?\d+(?:\.\d+)?(?:[eE][+-]?\d+)?")
-
-
-def _extract_arrays_from_text(text: str) -> list[list]:
-    """Find every parseable bracketed numeric array (1D or 2D) in free text."""
-    arrays: list[list] = []
-    n = len(text)
-    i = 0
-    while i < n:
-        if text[i] != "[":
-            i += 1
-            continue
-        depth = 0
-        j = i
-        while j < n:
-            if text[j] == "[":
-                depth += 1
-            elif text[j] == "]":
-                depth -= 1
-                if depth == 0:
-                    break
-            j += 1
-        if depth != 0:
-            i += 1
-            continue
-        candidate = text[i : j + 1]
-        # Third variant: users mix separators freely ("[3 0, 1]"), which the
-        # comma-only fixer can't handle -- insert a comma between any two
-        # adjacent number tokens regardless of other separators present.
-        mixed_fixed = re.sub(r"(?<=[\d.])\s+(?=[-\d.])", ", ", candidate)
-        for variant in (candidate, fix_space_separated_arrays(candidate), mixed_fixed):
-            try:
-                parsed = json.loads(variant)
-            except json.JSONDecodeError:
-                continue
-            if isinstance(parsed, list) and parsed:
-                arrays.append(parsed)
-                break
-        # Only skip past this bracket char, not the whole span: inner arrays
-        # of a 2D matrix should also be collected individually (rows are
-        # legitimate 1D vectors in their own right).
-        i += 1
-    return arrays
-
-
-def _collect_numeric_leaves(obj: Any, arrays: list[list], numbers: set[float]) -> None:
-    """Harvest nested numeric lists and scalars from a parsed JSON object."""
-    if isinstance(obj, bool):
-        return
-    if isinstance(obj, (int, float)):
-        numbers.add(float(obj))
-        return
-    if isinstance(obj, list):
-        if obj and all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in obj):
-            arrays.append(list(obj))
-        elif obj and all(isinstance(v, list) for v in obj):
-            arrays.append(obj)
-        for v in obj:
-            _collect_numeric_leaves(v, arrays, numbers)
-        return
-    if isinstance(obj, dict):
-        for v in obj.values():
-            _collect_numeric_leaves(v, arrays, numbers)
-
-
-def _as_2d_float_array(value: Any):
-    """Return value as a 2D float ndarray, or None if it isn't one."""
-    import numpy as _np
-
-    try:
-        arr = _np.array(value, dtype=float)
-    except (TypeError, ValueError):
-        return None
-    if arr.ndim == 1 and arr.size > 0:
-        arr = arr.reshape(1, -1)
-    if arr.ndim != 2 or arr.size == 0:
-        return None
-    return arr
-
-
-class ParameterProvenance:
-    """Everything numeric the conversation has actually provided so far."""
-
-    def __init__(self, messages: list[dict[str, Any]]) -> None:
-        import numpy as _np
-
-        self._np = _np
-        self.arrays: list[Any] = []
-        self.current_numbers: set[float] = set()
-        all_user_texts: list[str] = []
-        current_user_text = ""
-
-        for msg in messages:
-            role = msg.get("role")
-            content = msg.get("content", "")
-            if role == "user" and isinstance(content, str):
-                all_user_texts.append(content)
-                current_user_text = content
-            elif role == "tool" and isinstance(content, str):
-                try:
-                    parsed = json.loads(content)
-                except json.JSONDecodeError:
-                    continue
-                found_arrays: list[list] = []
-                found_numbers: set[float] = set()
-                _collect_numeric_leaves(parsed, found_arrays, found_numbers)
-                for a in found_arrays:
-                    arr = _as_2d_float_array(a)
-                    if arr is not None:
-                        self.arrays.append(arr)
-                # Scalars computed by tools (gains, margins) are legitimate
-                # inputs for later steps.
-                self.current_numbers |= found_numbers
-
-        for text in all_user_texts:
-            for a in _extract_arrays_from_text(text):
-                arr = _as_2d_float_array(a)
-                if arr is not None:
-                    self.arrays.append(arr)
-        # Loose scalars only count from the CURRENT user message: numbers from
-        # an earlier, different problem are exactly what must not silently
-        # seed a new Q or R.
-        for m in _NUMBER_RE.finditer(current_user_text):
-            try:
-                self.current_numbers.add(float(m.group()))
-            except ValueError:
-                pass
-
-    def _matches_known_array(self, arr) -> bool:
-        np_ = self._np
-        for known in self.arrays:
-            for cand in (known, known.T):
-                if cand.shape == arr.shape and np_.allclose(cand, arr, rtol=1e-6, atol=1e-9):
-                    return True
-        return False
-
-    def _number_provided(self, x: float) -> bool:
-        if x in (0.0, 1.0, -1.0):
-            return True
-        return any(abs(x - n) <= 1e-9 * max(1.0, abs(n)) for n in self.current_numbers)
-
-    def verify(self, value: Any) -> bool:
-        """True if this matrix is traceable to the conversation."""
-        np_ = self._np
-        arr = _as_2d_float_array(value)
-        if arr is None:
-            return True  # not a matrix -- out of scope for this guard
-
-        if self._matches_known_array(arr):
-            return True
-
-        # 1x1 "matrix" wrapping a scalar the user stated (R=1 -> [[1]]).
-        if arr.shape == (1, 1):
-            return self._number_provided(float(arr[0, 0]))
-
-        # Identity / zero matrices are structural, not data.
-        if arr.shape[0] == arr.shape[1]:
-            if np_.allclose(arr, np_.eye(arr.shape[0])) or np_.allclose(arr, 0.0):
-                return True
-            # diag(...) built from numbers in the current message, e.g. the
-            # user wrote "Q=diag([10, 1])" in prose rather than as an array.
-            if np_.allclose(arr, np_.diag(np_.diag(arr))):
-                return all(self._number_provided(float(d)) for d in np_.diag(arr))
-
-        return False
-
-
-# Real control-engineering coefficient/numerator/denominator arrays are
-# essentially always short -- a 10th-order polynomial (already an extreme
-# hand-solved case) has 11 coefficients. A flat numeric array cannot be
-# provenance-traced the way a matrix can (a legitimately *expanded* factored
-# polynomial, e.g. (s+1)(s+2)(s+3) -> [1,6,11,6], never appears verbatim in
-# the user's text, so a strict "must match the conversation" rule would
-# wrongly block correct derivations). Instead this catches the actual
-# observed failure mode directly: a runaway repetition loop, where a small
-# model gets stuck emitting the same value and the token budget cuts it off
-# mid-array -- e.g. 300+ elements of mostly zeros for routh_hurwitz_analysis
-# on a question that never gave it a polynomial at all.
-_MAX_SANE_1D_ARRAY_LEN = 25
-_MAX_IDENTICAL_RUN = 6
-
-
-def _degenerate_array_reason(values: list) -> str | None:
-    if not all(isinstance(v, (int, float)) and not isinstance(v, bool) for v in values):
-        return None
-    if len(values) > _MAX_SANE_1D_ARRAY_LEN:
-        return f"has {len(values)} elements, far beyond any real control-engineering array of this kind"
-    run = 1
-    for i in range(1, len(values)):
-        run = run + 1 if values[i] == values[i - 1] else 1
-        if run >= _MAX_IDENTICAL_RUN:
-            return f"repeats the value {values[i]!r} {run}+ times in a row -- a runaway generation loop, not real data"
-    return None
-
-
-def _check_parameter_provenance(
-    tool_args: dict[str, Any], messages: list[dict[str, Any]]
-) -> tuple[str, Any, str] | None:
-    """Return (param_name, value, reason) for the first bad matrix/array, else None."""
-    provenance = ParameterProvenance(messages)
-    for param, value in tool_args.items():
-        if param in _PROVENANCE_EXEMPT_PARAMS:
-            continue
-        if not (isinstance(value, list) and value):
-            continue
-        if isinstance(value[0], list):
-            if not provenance.verify(value):
-                return param, value, "was not provided by the user in this conversation and did not come from any prior tool result"
-        else:
-            reason = _degenerate_array_reason(value)
-            if reason:
-                return param, value, reason
-    return None
-
-
-# ---------------------------------------------------------------------------
-# RAG fast path
-#
-# The system prompt already tells the model "for definitional questions,
-# answer from the retrieved reference passages, do not run a numeric solver"
-# -- but that instruction demonstrably does not hold on a 4B model either: in
-# production, "what is Routh-Hurwitz, explain with an example" (fully
-# answerable from the Nise/Ogata passages already injected into the system
-# prompt by _get_grounded_instruction) still drove the model through 4 tool
-# calls (a redundant re-search plus three unrelated numeric tools) and 5
-# sequential full-length generations before it produced an answer. Since
-# prompting alone can't be trusted here any more than it can for parameter
-# provenance above, this is enforced the same way: in code.
-#
-# A real computational request in this domain essentially always names a
-# concrete number (a coefficient, a gain, a frequency) or an explicit
-# computational verb ("plot", "design", "simulate"); a pure "what is X" /
-# "explain X" / "how does X work" question has neither. This heuristic only
-# ever widens back to the full tool loop on a false negative (a computational
-# question with no digits and none of these verbs) -- it never narrows
-# correctness, since the loop it skips still runs whenever this returns True.
-_COMPUTATION_KEYWORDS = (
-    "plot", "simulate", "simulation", "compute", "calculate", "design",
-    "solve", "gain", "matrix", "matrices", "pole", "place", "locus", "bode",
-    "nyquist", "margin", "response", "transfer function", "eigen",
-    "controllab", "observab", "lyapunov", "kalman", "mpc", "invert",
-    "determinant", "transpose", "multiply", "rank", "code", "script",
-)
-
-
-def _needs_tools(user_prompt: str) -> bool:
-    """True if the question plausibly needs a numeric tool rather than being
-    answerable straight from grounded reference text."""
-    if any(ch.isdigit() for ch in user_prompt):
-        return True
-    lower = user_prompt.lower()
-    return any(kw in lower for kw in _COMPUTATION_KEYWORDS)
-
-
-def _fabrication_refusal_note(refusals: list[str]) -> str:
-    """Extra synthesis instruction appended when a tool call was refused for
-    inventing a parameter.
-
-    Without this, the model's forced final-answer turn happily hallucinates a
-    plausible-looking numeric replacement anyway: observed directly with
-    sft_v2 on "design an LQR controller for A = [[0, 1], [-2, -3]]" (no B, Q,
-    R given) -- the tool call was correctly REFUSED by the provenance guard,
-    but the final prose still confidently printed a fully invented gain
-    $K = [1.83, 1.83]$, silently routing around its own refusal. The system
-    prompt already forbids this (rule 6), but that alone doesn't hold on a 4B
-    model any more than the provenance rules did in prose form -- so the
-    refusal is restated directly in the synthesis turn itself, where it can't
-    be missed.
-    """
-    if not refusals:
-        return ""
-    return (
-        "\n\nIMPORTANT: at least one tool call above was REFUSED for inventing a parameter you were never "
-        "given (see the REFUSED error message(s) in the tool results above for exactly which one and why). "
-        "This means you do NOT have enough information to compute a numeric answer for that part of the "
-        "request. Do NOT invent a substitute number, gain, matrix, or result to answer anyway -- state "
-        "plainly what is missing and ask the user for it instead of guessing."
-    )
-
-
-# The web frontend resends the entire growing conversation on every turn
-# (web/app.js) with no client-side trimming, and nothing here capped it
-# either: a long chat would eventually exceed the GGUF backend's
-# n_ctx=16384 and error outright, or just keep getting slower with no
-# ceiling at all on the PyTorch/MLX backends. Budget is deliberately well
-# under 16384 to leave headroom for the system prompt, injected RAG
-# passages, the full tool-schema catalog, and the current turn's own
-# multi-step tool loop -- all of which share the same context window.
-MAX_HISTORY_TOKENS = 6000
-
-
-def _truncate_history(
-    history: list[dict[str, Any]] | None, tokenizer: Any, max_tokens: int = MAX_HISTORY_TOKENS
-) -> list[dict[str, Any]] | None:
-    """Keep the most recent turns of `history` that fit within max_tokens,
-    dropping the oldest first so a long conversation degrades (forgets early
-    turns) instead of failing outright."""
-    if not history:
-        return history
-
-    kept: list[dict[str, Any]] = []
-    total = 0
-    for item in reversed(history):
-        content = item.get("content", "")
-        if not isinstance(content, str) or not content:
-            continue
-        try:
-            n = len(tokenizer.encode(content))
-        except Exception:
-            n = len(content) // 4  # rough fallback if the tokenizer call itself fails
-        if kept and total + n > max_tokens:
-            break
-        kept.append(item)
-        total += n
-
-    kept.reverse()
-    # A leading assistant turn with no preceding user turn would read as a
-    # reply to nothing -- drop it.
-    while kept and kept[0].get("role") != "user":
-        kept.pop(0)
-    return kept
-
-
-class ControlAIAgent:
-    """Universal Control Engineering Agent supporting GGUF, Ollama C++, Apple MLX, and PyTorch."""
-
-    def __init__(
-        self,
-        model_path: str = "mlx-community/Qwen3-4B-Instruct-2507-4bit",
-        adapter_path: str | None = None,
-        tool_registry: ToolRegistry = registry,
-        max_tool_steps: int = 4,
-    ) -> None:
-        self.model_path = model_path
-        
-        # Auto-detect trained LoRA adapter if none explicitly provided
-        default_adapter = PROJECT_ROOT / "adapters" / "controlai_qwen3_4b_sft_v2"
-        if adapter_path is None and default_adapter.exists():
-            adapter_path = str(default_adapter)
-
-        self.adapter_path = adapter_path
-        self.registry = tool_registry
-        self.max_tool_steps = max_tool_steps
-
-        # Detect platform & backend
-        self.is_ollama = str(model_path).startswith("ollama")
-        self.is_gguf = str(model_path).endswith(".gguf") or "gguf" in str(model_path).lower()
-        self.is_mlx = HAS_MLX and not self.is_ollama and not self.is_gguf and not str(model_path).startswith("Qwen/") and not os.environ.get("FORCE_TRANSFORMERS")
-
-        if self.is_ollama:
-            self.ollama_model = model_path.split(":", 1)[1] if ":" in str(model_path) else "controlai"
-            self.hf_tokenizer = AutoTokenizer.from_pretrained(CONTROLAI_HF_REPO, trust_remote_code=True)
-        elif self.is_gguf:
-            if not HAS_LLAMA_CPP:
-                raise ImportError("llama-cpp-python is required to run GGUF models. Install it with: pip install llama-cpp-python")
-            self.llama_model = llama_cpp.Llama(
-                model_path=str(model_path),
-                n_gpu_layers=-1,  # Offload all layers to Metal / CUDA GPU
-                n_ctx=16384,
-                verbose=False,
-            )
-            # A local .gguf path carries no tokenizer/chat-template of its own --
-            # always load ours (never a generic base model's), same as the
-            # auto-downloaded deployment path below.
-            local_fused = PROJECT_ROOT / "models" / "controlai_fused"
-            tokenizer_source = str(local_fused) if local_fused.exists() else CONTROLAI_HF_REPO
-            self.hf_tokenizer = AutoTokenizer.from_pretrained(tokenizer_source, trust_remote_code=True)
-        elif self.is_mlx:
-            if adapter_path:
-                self.model, self.mlx_tokenizer = mlx_load(model_path, adapter_path=adapter_path)
-            else:
-                self.model, self.mlx_tokenizer = mlx_load(model_path)
-            self.hf_tokenizer = AutoTokenizer.from_pretrained(model_path)
-        else:
-            # Universal Linux / Cloud / HuggingFace Spaces backend: fast C++ GGUF
-            # of OUR OWN fine-tuned ControlAI model (never a generic base model).
-            #
-            # Q8_0, not Q4_K_M: measured directly against this exact model on the
-            # "simulate a step response" prompt (3 runs each). Q4_K_M mis-routed
-            # tool calls in 3/3 runs (reached for continuous_lqr/bode_analysis
-            # instead of simulate_step_response, then hallucinated inconsistent
-            # overshoot values -- 8.6%, 43.1%, 25.4%, none near the true ~9.5%).
-            # Q6_K, Q8_0, and full f16 all called the correct tool with
-            # consistent, tool-verified numbers in 3/3 runs apiece, matching
-            # local MLX behavior -- Q8_0 was picked among those three because it
-            # was both the fastest of the three in this test (8-19s vs Q6_K's
-            # 11-25s and f16's 11-22s) and, at 8.5 bits/weight, close enough to
-            # f16 to be considered near-lossless, at roughly half f16's size.
-            gguf_repo = os.environ.get("CONTROLAI_GGUF_REPO", CONTROLAI_HF_REPO)
-            gguf_filename = os.environ.get("CONTROLAI_GGUF_FILENAME", "*controlai-q8_0.gguf")
-            # Escape hatch for measuring the PyTorch/CUDA path against the GGUF
-            # path on real deployment hardware -- set to "pytorch" to skip the
-            # GGUF attempt entirely and go straight to the branch below.
-            force_backend = os.environ.get("CONTROLAI_BACKEND", "").lower()
-
-            self.llama_model = None
-            if HAS_LLAMA_CPP and force_backend != "pytorch":
-                try:
-                    threads = min(os.cpu_count() or 2, 4)
-                    print(f"Loading GGUF ControlAI model ({gguf_filename}) from {gguf_repo} (threads: {threads})...")
-                    self.llama_model = llama_cpp.Llama.from_pretrained(
-                        repo_id=gguf_repo,
-                        filename=gguf_filename,
-                        n_ctx=16384,
-                        n_threads=threads,
-                        verbose=False,
-                    )
-                    self.is_gguf = True
-                    self.hf_tokenizer = AutoTokenizer.from_pretrained(CONTROLAI_HF_REPO, trust_remote_code=True)
-                    print("GGUF ControlAI model loaded successfully via llama_cpp.")
-                except Exception as exc:
-                    print(f"Notice: llama_cpp GGUF auto-load failed, falling back to PyTorch: {exc}")
-                    self.is_gguf = False
-
-            if not self.is_gguf:
-                # PyTorch fallback on CUDA or if the GGUF engine is unavailable.
-                # Always our own fine-tuned model, never a generic base model.
-                import torch
-                num_threads = min(os.cpu_count() or 2, 4)
-                try:
-                    torch.set_num_threads(num_threads)
-                except Exception:
-                    pass
-
-                hf_id = CONTROLAI_HF_REPO if "mlx" in str(model_path) or str(model_path).startswith("Qwen/") else model_path
-                print(f"Loading PyTorch model: {hf_id} (threads: {num_threads})...")
-                self.hf_tokenizer = AutoTokenizer.from_pretrained(hf_id, trust_remote_code=True)
-                has_cuda = torch.cuda.is_available()
-                # float16, not bfloat16: live-tested bf16 against this exact
-                # model on the zeta/wn step-response prompt and it reproducibly
-                # derived the wrong closed-loop coefficients (2.56/2.1952-ish
-                # instead of the correct 1.68/1.96) across both greedy and
-                # temperature=0.2 sampling -- the one thing GGUF/MLX (which got
-                # this right from the same weights) don't share with bf16 is
-                # its coarser 7-bit mantissa. Testing fp16's extra precision as
-                # the fix for that specific divergence.
-                dtype = torch.float16 if has_cuda else torch.float32
-                # SDPA is a large, free speedup over the "eager" attention
-                # default -- meaningful for the long tool-schema prompt prefix.
-                attn_impl = "sdpa" if has_cuda else None
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    hf_id,
-                    torch_dtype=dtype,
-                    low_cpu_mem_usage=True,
-                    trust_remote_code=True,
-                    attn_implementation=attn_impl,
-                )
-                if has_cuda:
-                    # device_map="auto" infers available VRAM at load time; on
-                    # ZeroGPU no physical GPU is attached to the process yet at
-                    # this point (that only happens inside a @spaces.GPU call),
-                    # so it can silently offload layers to CPU. An explicit
-                    # move avoids that and any ambiguity about what actually
-                    # ran where.
-                    self.model = self.model.to("cuda")
-                # The fused ControlAI model already has the LoRA weights merged in;
-                # only apply a separate adapter when loading a plain base model.
-                if hf_id != CONTROLAI_HF_REPO and adapter_path and Path(adapter_path).exists():
-                    try:
-                        from peft import PeftModel
-                        self.model = PeftModel.from_pretrained(self.model, adapter_path)
-                        print(f"Loaded PEFT LoRA adapter from: {adapter_path}")
-                    except Exception as exc:
-                        print(f"Warning: Could not load LoRA adapter in PyTorch: {exc}")
-                print(f"PyTorch model loaded successfully. Model device: {next(self.model.parameters()).device}")
-
-        # Initialize local offline RAG index
-        try:
-            self.rag_index = get_shared_index()
-        except Exception:
-            self.rag_index = None
-
-    def _generate(self, prompt: str, max_tokens: int = 2000) -> str:
-        """Universal text generation handling Ollama C++, GGUF llama_cpp, MLX, and PyTorch."""
-        if self.is_ollama:
-            res = ollama.generate(
-                model=self.ollama_model,
-                prompt=prompt,
-                options={"temperature": 0.2, "num_predict": max_tokens, "repeat_penalty": REPETITION_PENALTY},
-            )
-            return res.get("response", "").strip()
-        elif self.is_gguf:
-            # llama_cpp defaults repeat_penalty to 1.0 (fully off) when unset --
-            # it does NOT inherit any sane default, so this must be passed explicitly.
-            output = self.llama_model(
-                prompt,
-                max_tokens=max_tokens,
-                stop=["<|im_end|>", "<|endoftext|>"],
-                temperature=0.2,
-                repeat_penalty=REPETITION_PENALTY,
-            )
-            return output["choices"][0]["text"].strip()
-        elif self.is_mlx:
-            return mlx_generate(
-                self.model,
-                self.mlx_tokenizer,
-                prompt=prompt,
-                max_tokens=max_tokens,
-                logits_processors=make_logits_processors(repetition_penalty=REPETITION_PENALTY),
-                verbose=False,
-            ).strip()
-        else:
-            import time as _time
-            import torch
-            inputs = self.hf_tokenizer(prompt, return_tensors="pt").to(self.model.device)
-            prompt_tokens = inputs["input_ids"].shape[1]
-            eos_ids = [
-                tid
-                for tid in (self.hf_tokenizer.eos_token_id, self.hf_tokenizer.convert_tokens_to_ids("<|im_end|>"))
-                if tid is not None and tid >= 0
-            ] or None
-            t0 = _time.time()
-            print(f"[_generate] starting: prompt_tokens={prompt_tokens} max_new_tokens={max_tokens} eos_ids={eos_ids}")
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_new_tokens=max_tokens,
-                    # Pure greedy (do_sample=False) reproducibly derived the
-                    # WRONG closed-loop coefficients for a plain zeta/wn step
-                    # response on this exact model (2.56/2.1952 instead of the
-                    # correct 1.68/1.96, in 3/3 identical runs on deployed
-                    # bf16) -- low-temperature sampling, matching the GGUF/MLX
-                    # backends, is the fix being tested for that failure.
-                    do_sample=True,
-                    temperature=0.2,
-                    top_p=None,
-                    top_k=None,
-                    repetition_penalty=REPETITION_PENALTY,
-                    eos_token_id=eos_ids,
-                    pad_token_id=self.hf_tokenizer.pad_token_id or self.hf_tokenizer.eos_token_id,
-                )
-            new_tokens = outputs[0][prompt_tokens:]
-            elapsed = _time.time() - t0
-            print(f"[_generate] done: generated_tokens={len(new_tokens)} elapsed={elapsed:.1f}s ({len(new_tokens)/max(elapsed,0.001):.1f} tok/s)")
-            return self.hf_tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
-
-    def _get_grounded_instruction(self, user_prompt: str, base_instruction: str) -> str:
-        """Retrieve relevant textbook theorems and inject grounding into system instructions."""
-        if not self.rag_index or not self.rag_index.chunks:
-            return base_instruction
-
-        try:
-            rag_hits = self.rag_index.search(user_prompt, top_k=6)
-            high_rel = [h for h in rag_hits if h.get("score", 0) > 2.5]
-            if not high_rel:
-                return base_instruction
-
-            ref_texts = []
-            for h in high_rel[:4]:
-                # Cite the cleaned display name, never the raw indexed filename.
-                label = h.get("source_name") or h.get("filename", "Reference")
-                page = h.get("page")
-                page_str = f", p. {page}" if page else ""
-                clean_chunk = " ".join(h.get("text", "").split())[:900].strip()
-                ref_texts.append(f"[{label}{page_str}]\n{clean_chunk}")
-
-            return (
-                base_instruction
-                + "\n\n### Grounded Reference Context from Canonical Control Literature:\n"
-                + "\n\n".join(ref_texts)
-                + "\n\nThese passages were retrieved from the local library for THIS question. When the "
-                "question asks what a specific author or textbook says, answer directly from these "
-                "passages -- that is already the reference lookup, so do not call a numeric solver tool "
-                "to answer a conceptual question. Where a passage conflicts with your own recollection, "
-                "TRUST THE PASSAGE. Cite using exactly the bracketed label shown above (for example "
-                "[Nise, Control Systems Engineering, p. 583]); never invent a citation and never print a "
-                "raw filename, file extension, or course code."
-            )
-        except Exception as exc:
-            print(f"[_get_grounded_instruction] FAILED, continuing ungrounded: {type(exc).__name__}: {exc}")
-            return base_instruction
-
-    def _direct_answer(
-        self,
-        user_prompt: str,
-        system_instruction: str,
-        history: list[dict[str, Any]] | None = None,
-        max_tokens: int = 1400,
-    ) -> str:
-        """Answer the question with no tools exposed at all.
-
-        This is the recovery path for when the tool loop fails to produce usable
-        prose -- the model spent its steps on tool calls, or its synthesis turn
-        came back empty or as yet another tool call. Re-asking with tools=None
-        and without the tool-result transcript removes whatever was derailing it
-        (and shortens the prompt considerably), which reliably yields a real
-        answer instead of a canned "analysis complete" placeholder. The grounded
-        system instruction is kept, so retrieved reference passages are still
-        available to answer from.
-        """
-        messages: list[dict[str, Any]] = []
-        if system_instruction:
-            messages.append({"role": "system", "content": system_instruction})
-        if history:
-            for item in history:
-                r = item.get("role")
-                c = item.get("content")
-                if r in ("user", "assistant") and c:
-                    messages.append({"role": r, "content": c})
-        messages.append({"role": "user", "content": user_prompt})
-
-        try:
-            rendered = self.hf_tokenizer.apply_chat_template(
-                messages, tools=None, tokenize=False, add_generation_prompt=True
-            )
-            _, cleaned = _extract_tool_calls(self._generate(rendered, max_tokens=max_tokens))
-            return cleaned
-        except Exception as exc:
-            # This is the last line of defense before the user sees the canned
-            # placeholder -- if it fails, that failure must leave a trace
-            # instead of vanishing, or a real bug here is undiagnosable.
-            print(f"[_direct_answer] FAILED for prompt {user_prompt[:80]!r}: {type(exc).__name__}: {exc}")
-            return ""
-
-    def _execute_with_provenance(
-        self,
-        tool_name: str,
-        tool_args: dict[str, Any],
-        messages: list[dict[str, Any]],
-    ) -> dict[str, Any]:
-        """Run a tool only after every matrix argument is traced to the
-        conversation; refuse fabricated inputs before they execute."""
-        try:
-            fabricated = _check_parameter_provenance(tool_args, messages)
-        except Exception:
-            fabricated = None  # the guard must never take down a legitimate call
-        if fabricated is not None:
-            param, value, reason = fabricated
-            # Never echo a runaway array (possibly hundreds of elements) back
-            # into the context -- it wastes the token budget and risks
-            # priming the exact same repetition pathology again.
-            shown = value if len(json.dumps(value)) < 200 else f"[{len(value)}-element array, truncated]"
-            return {
-                "status": "error",
-                "error_type": "FabricatedParameter",
-                "error": (
-                    f"REFUSED: the value passed as '{param}' ({shown}) {reason}. Inventing or "
-                    f"malformed parameter values is forbidden. Do NOT retry this tool with another "
-                    f"guessed or partially-repeated '{param}' -- if you don't actually have a concrete "
-                    f"value for it, this tool cannot be used for this question. Answer from your own "
-                    f"knowledge or the retrieved reference passages instead, or tell the user what's "
-                    f"missing."
-                ),
-            }
-        return self.registry.execute(tool_name, tool_args)
-
-    def run(
-        self,
-        user_prompt: str,
-        system_instruction: str = CONTROLAI_SYSTEM_PROMPT,
-        history: list[dict[str, Any]] | None = None,
-        max_tokens_per_step: int = 2500,
-        verbose: bool = False,
-    ) -> AgentResult:
-        """Execute a complete agent interaction loop synchronously."""
-        history = _truncate_history(history, self.hf_tokenizer)
-        messages: list[dict[str, Any]] = []
-        effective_sys = self._get_grounded_instruction(user_prompt, system_instruction)
-
-        # RAG fast path: a conceptual/definitional question that already has
-        # strong grounded passages needs one generation, not a multi-step tool
-        # loop. See the _needs_tools docstring for why this is safe.
-        if effective_sys != system_instruction and not _needs_tools(user_prompt):
-            fast_answer = self._direct_answer(user_prompt, effective_sys, history)
-            if fast_answer:
-                return AgentResult(
-                    final_response=fast_answer,
-                    tool_traces=[],
-                    total_steps=1,
-                    raw_messages=[{"role": "system", "content": effective_sys}, {"role": "user", "content": user_prompt}],
-                    is_grounded=True,
-                    plots=[],
-                )
-
-        if effective_sys:
-            messages.append({"role": "system", "content": effective_sys})
-
-        if history:
-            for item in history:
-                r = item.get("role")
-                c = item.get("content")
-                if r in ("user", "assistant") and c:
-                    messages.append({"role": r, "content": c})
-
-        messages.append({"role": "user", "content": user_prompt})
-
-        tools_schema = self.registry.get_tool_schemas()
-        traces: list[ToolExecutionTrace] = []
-        plots: list[str] = []
-        called_signatures: set[str] = set()
-        tool_call_counts: dict[str, int] = {}
-        fabrication_refusals: list[str] = []
-
-        for step in range(1, self.max_tool_steps + 1):
-            rendered_prompt = self.hf_tokenizer.apply_chat_template(
-                messages,
-                tools=tools_schema,
-                tokenize=False,
-                add_generation_prompt=True,
-            )
-
-            model_output = self._generate(rendered_prompt, max_tokens=max_tokens_per_step)
-
-            tool_calls, pre_text = _extract_tool_calls(model_output)
-
-            if not tool_calls and not fabrication_refusals:
-                # The model chose to stop calling tools but produced no usable
-                # text either (typically right after a tool error it has no
-                # good way to recover from in-context). Retry once with no
-                # tools and no error-laden transcript rather than returning
-                # nothing.
-                final_response = pre_text or self._direct_answer(user_prompt, effective_sys, history)
-                if not final_response:
-                    final_response = "The computational analysis has been completed as detailed above."
-                return AgentResult(
-                    final_response=final_response,
-                    tool_traces=traces,
-                    total_steps=step,
-                    raw_messages=messages,
-                    is_grounded=True,
-                    plots=plots,
-                )
-
-            if not tool_calls:
-                # A fabrication refusal happened earlier this turn -- pre_text
-                # here is exactly the kind of confident, ungrounded prose that
-                # refusal was meant to prevent (observed directly: a blocked
-                # LQR call still produced a fully invented gain in this same
-                # spot). Don't trust it at face value; fall through to the
-                # shared forced-synthesis turn below instead of returning,
-                # since that turn explicitly instructs against inventing a
-                # substitute value.
-                break
-
-            # Drop repeats. An exact-signature check alone is not enough: the
-            # model will re-search with a lightly reworded query ("... MPC",
-            # "... MPC algorithm", "... MPC definition"), exhausting the step
-            # budget on near-identical lookups and leaving nothing for the
-            # answer. Cap how many times any single tool may run per turn.
-            new_calls = []
-            for call in tool_calls:
-                name = call.get("name")
-                sig = f"{name}:{json.dumps(call.get('arguments', {}), sort_keys=True)}"
-                if sig in called_signatures:
-                    continue
-                if tool_call_counts.get(name, 0) >= MAX_CALLS_PER_TOOL:
-                    continue
-                called_signatures.add(sig)
-                tool_call_counts[name] = tool_call_counts.get(name, 0) + 1
-                new_calls.append(call)
-
-            if not new_calls:
-                break
-
-            messages.append({"role": "assistant", "content": model_output})
-
-            for call_data in new_calls:
-                tool_name = call_data.get("name")
-                tool_args = call_data.get("arguments", {})
-
-                tool_result = self._execute_with_provenance(tool_name, tool_args, messages)
-                traces.append(ToolExecutionTrace(tool_name=tool_name, arguments=tool_args, result=tool_result))
-                if tool_result.get("error_type") == "FabricatedParameter":
-                    fabrication_refusals.append(tool_result.get("error", ""))
-
-                if "plot_path" in tool_result:
-                    p_path = Path(tool_result["plot_path"])
-                    if p_path.exists():
-                        plots.append(f"/plots/{p_path.name}")
-
-                messages.append({
-                    "role": "tool",
-                    "name": tool_name,
-                    "content": json.dumps(tool_result, ensure_ascii=False),
-                })
-
-        # Final synthesis after tool execution
-        messages.append({
-            "role": "user",
-            "content": (
-                "Now answer the user's most recent question directly and completely, in LaTeX-formatted "
-                "prose. If the tool results above are relevant to that question, incorporate them; if the "
-                "question is conceptual, definitional, or about what a source says and the tool results "
-                "above don't actually address it, answer the question from your own knowledge instead of "
-                "describing the tool results. Do not call any more tools and do not output JSON or "
-                "tool-call tags -- write the final answer now."
-            ) + _fabrication_refusal_note(fabrication_refusals),
-        })
-        forced_prompt = self.hf_tokenizer.apply_chat_template(
-            messages,
-            tools=None,
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-        final_output = self._generate(forced_prompt, max_tokens=max_tokens_per_step)
-
-        _, clean_final = _extract_tool_calls(final_output)
-        # Never fall back to the raw final_output: if the model's closing turn
-        # degenerates into another (unparseable) tool-call attempt, clean_final
-        # is stripped down to "" and showing final_output would leak a raw
-        # <tool_call>{...}</tool_call> JSON blob straight into the chat.
-        final_response = clean_final or self._direct_answer(user_prompt, effective_sys, history)
-        if not final_response:
-            # _direct_answer reliably produces real text when tested in
-            # isolation, so an empty result here is most likely a transient
-            # failure (resource contention, a stray exception) rather than a
-            # repeatable one -- one retry clears most of those before giving
-            # up and showing the placeholder.
-            final_response = self._direct_answer(user_prompt, effective_sys, history)
-        if not final_response:
-            final_response = "The computational analysis has been completed as detailed above."
-        return AgentResult(
-            final_response=final_response,
-            tool_traces=traces,
-            total_steps=len(traces) + 1,
-            raw_messages=messages,
-            is_grounded=True,
-            plots=plots,
-        )
-
-    def run_stream(
-        self,
-        user_prompt: str,
-        system_instruction: str = CONTROLAI_SYSTEM_PROMPT,
-        history: list[dict[str, Any]] | None = None,
-        max_tokens_per_step: int = 2500,
-    ) -> Generator[dict[str, Any], None, None]:
-        """Stream token-by-token generation and tool execution events with zero JSON leakage."""
-        history = _truncate_history(history, self.hf_tokenizer)
-        messages: list[dict[str, Any]] = []
-        effective_sys = self._get_grounded_instruction(user_prompt, system_instruction)
-
-        # RAG fast path: a conceptual/definitional question that already has
-        # strong grounded passages needs one generation, not a multi-step tool
-        # loop. See the _needs_tools docstring for why this is safe.
-        if effective_sys != system_instruction and not _needs_tools(user_prompt):
-            fast_answer = self._direct_answer(user_prompt, effective_sys, history)
-            if fast_answer:
-                words = re.split(r"(\s+)", fast_answer)
-                for w in words:
-                    if w:
-                        yield {"type": "token", "content": w}
-                yield {
-                    "type": "done",
-                    "response": fast_answer,
-                    "traces": [],
-                    "plots": [],
-                    "thoughts": [],
-                }
-                return
-
-        if effective_sys:
-            messages.append({"role": "system", "content": effective_sys})
-
-        if history:
-            for item in history:
-                r = item.get("role")
-                c = item.get("content", "")
-                if r in ("user", "assistant") and c:
-                    # Strip any legacy JSON tool call artifacts from previous chat sessions
-                    c_clean = re.sub(r"\{\s*[\"']name[\"']\s*:[\s\S]*?\}\s*\}", "", c).strip()
-                    c_clean = re.sub(r"<tool_call>[\s\S]*?</tool_call>", "", c_clean).strip()
-                    if c_clean:
-                        messages.append({"role": r, "content": c_clean})
-
-        messages.append({"role": "user", "content": user_prompt})
-
-        tools_schema = self.registry.get_tool_schemas()
-        traces: list[dict[str, Any]] = []
-        plots: list[str] = []
-        thoughts: list[str] = []
-        called_signatures: set[str] = set()
-        tool_call_counts: dict[str, int] = {}
-        fabrication_refusals: list[str] = []
-
-        # Dynamic thought generation - only show thoughts when tools or derivations occur
-        for step in range(1, self.max_tool_steps + 1):
-            rendered_prompt = self.hf_tokenizer.apply_chat_template(
-                messages,
-                tools=tools_schema,
-                tokenize=False,
-                add_generation_prompt=True,
-            )
-
-            model_output = self._generate(rendered_prompt, max_tokens=max_tokens_per_step)
-
-            tool_calls, pre_text = _extract_tool_calls(model_output)
-
-            if not tool_calls and not fabrication_refusals:
-                # Direct final answer without tools -> stream tokens directly
-                clean_output = pre_text or self._direct_answer(user_prompt, effective_sys, history)
-                if not clean_output:
-                    clean_output = "The computational analysis has been completed as detailed above."
-                words = re.split(r"(\s+)", clean_output)
-                for w in words:
-                    if w:
-                        yield {"type": "token", "content": w}
-
-                yield {
-                    "type": "done",
-                    "response": clean_output,
-                    "traces": traces,
-                    "plots": plots,
-                    "thoughts": thoughts,
-                }
-                return
-
-            if not tool_calls:
-                # A fabrication refusal happened earlier this turn -- pre_text
-                # here is exactly the kind of confident, ungrounded prose that
-                # refusal was meant to prevent. Don't trust it at face value;
-                # fall through to the shared forced-synthesis turn below,
-                # which explicitly instructs against inventing a substitute
-                # value, instead of streaming it straight to the user.
-                break
-
-            if pre_text:
-                thoughts.append(pre_text)
-                yield {"type": "thought", "content": pre_text}
-
-            # Drop repeats. An exact-signature check alone is not enough: the
-            # model will re-search with a lightly reworded query ("... MPC",
-            # "... MPC algorithm", "... MPC definition"), exhausting the step
-            # budget on near-identical lookups and leaving nothing for the
-            # answer. Cap how many times any single tool may run per turn.
-            new_calls = []
-            for call in tool_calls:
-                name = call.get("name")
-                sig = f"{name}:{json.dumps(call.get('arguments', {}), sort_keys=True)}"
-                if sig in called_signatures:
-                    continue
-                if tool_call_counts.get(name, 0) >= MAX_CALLS_PER_TOOL:
-                    continue
-                called_signatures.add(sig)
-                tool_call_counts[name] = tool_call_counts.get(name, 0) + 1
-                new_calls.append(call)
-
-            if not new_calls:
-                break
-
-            messages.append({"role": "assistant", "content": model_output})
-
-            for call_data in new_calls:
-                tool_name = call_data.get("name")
-                tool_args = call_data.get("arguments", {})
-
-                t_start_msg = f"Executing tool: {tool_name} with parameters: {json.dumps(tool_args, ensure_ascii=False)}"
-                thoughts.append(t_start_msg)
-                yield {"type": "thought", "content": t_start_msg}
-                yield {"type": "tool_start", "tool": tool_name, "args": tool_args}
-
-                tool_result = self._execute_with_provenance(tool_name, tool_args, messages)
-                trace_item = {
-                    "tool": tool_name,
-                    "args": tool_args,
-                    "status": tool_result.get("status", "success"),
-                    "residual": tool_result.get("residual"),
-                }
-                traces.append(trace_item)
-                if tool_result.get("error_type") == "FabricatedParameter":
-                    fabrication_refusals.append(tool_result.get("error", ""))
-
-                if "plot_path" in tool_result:
-                    p_path = Path(tool_result["plot_path"])
-                    if p_path.exists():
-                        plot_url = f"/plots/{p_path.name}"
-                        plots.append(plot_url)
-                        yield {"type": "plot", "url": plot_url}
-
-                yield {"type": "tool_end", "trace": trace_item}
-                t_end_msg = f"Tool {tool_name} returned status: {trace_item['status']}"
-                thoughts.append(t_end_msg)
-                yield {"type": "thought", "content": t_end_msg}
-
-                messages.append({
-                    "role": "tool",
-                    "name": tool_name,
-                    "content": json.dumps(tool_result, ensure_ascii=False),
-                })
-
-        # Final Synthesis phase
-        synth_thought = "Synthesizing verified engineering response and LaTeX formulations..."
-        thoughts.append(synth_thought)
-        yield {"type": "thought", "content": synth_thought}
-
-        messages.append({
-            "role": "user",
-            "content": (
-                "Now answer the user's most recent question directly and completely, in LaTeX-formatted "
-                "prose. If the tool results above are relevant to that question, incorporate them; if the "
-                "question is conceptual, definitional, or about what a source says and the tool results "
-                "above don't actually address it, answer the question from your own knowledge instead of "
-                "describing the tool results. Do not call any more tools and do not output JSON or "
-                "tool-call tags -- write the final answer now."
-            ) + _fabrication_refusal_note(fabrication_refusals),
-        })
-
-        forced_prompt = self.hf_tokenizer.apply_chat_template(
-            messages,
-            tools=None,
-            tokenize=False,
-            add_generation_prompt=True,
-        )
-
-        final_output = self._generate(forced_prompt, max_tokens=max_tokens_per_step)
-
-        # If model generated another tool call during final synthesis, execute it!
-        synth_tool_calls, clean_synth = _extract_tool_calls(final_output)
-        if synth_tool_calls:
-            for call_data in synth_tool_calls:
-                tool_name = call_data.get("name")
-                tool_args = call_data.get("arguments", {})
-                t_res = self._execute_with_provenance(tool_name, tool_args, messages)
-                if "plot_path" in t_res:
-                    p_path = Path(t_res["plot_path"])
-                    if p_path.exists():
-                        plot_url = f"/plots/{p_path.name}"
-                        plots.append(plot_url)
-                        yield {"type": "plot", "url": plot_url}
-                messages.append({"role": "tool", "name": tool_name, "content": json.dumps(t_res, ensure_ascii=False)})
-
-            # Re-generate synthesis after executing the tool
-            re_prompt = self.hf_tokenizer.apply_chat_template(messages, tools=None, tokenize=False, add_generation_prompt=True)
-            final_output = self._generate(re_prompt, max_tokens=max_tokens_per_step)
-            _, clean_synth = _extract_tool_calls(final_output)
-
-        # Never fall back to the raw final_output here: if the model's closing
-        # turn degenerates into another (unparseable) tool-call attempt,
-        # clean_synth is stripped down to "" and showing final_output would
-        # leak a raw <tool_call>{...}</tool_call> JSON blob into the chat.
-        final_text = clean_synth
-        # Strip any lingering raw json
-        final_text = re.sub(r"\{\s*[\"']name[\"']\s*:[\s\S]*?\}\s*\}", "", final_text).strip()
-        if not final_text:
-            # A fresh, tool-free, minimal-context generation reliably produces
-            # real text even when the tool-heavy synthesis turn degenerated --
-            # try it (twice, for a transient failure) before the placeholder.
-            final_text = self._direct_answer(user_prompt, effective_sys, history)
-            if not final_text:
-                final_text = self._direct_answer(user_prompt, effective_sys, history)
-        if not final_text:
-            final_text = "The computational analysis and simulation have been executed successfully as detailed above."
-
-        # Stream words smoothly
-        words = re.split(r"(\s+)", final_text)
-        for w in words:
-            if w:
-                yield {"type": "token", "content": w}
-
-        yield {
-            "type": "done",
-            "response": final_text,
-            "traces": traces,
-            "plots": plots,
-            "thoughts": thoughts,
-        }

controlai_agent/prompts.py

@@ -1,42 +1,67 @@
-CONTROLAI_SYSTEM_PROMPT = r"""You are ControlAI, a premier AI research scientist and expert engineering agent for control systems engineering in ALL of its application domains -- aerospace and flight control, automotive and vehicle dynamics, robotics and motion control, industrial automation and PLC/process control, mechatronics, power and energy systems, marine and space vehicles -- together with the underlying applied mathematics, system identification, estimation, and dynamical systems theory.
-
-Treat every control question as in scope regardless of the industry it comes from, and answer it in the vocabulary of that domain: an autopilot question deserves flight-dynamics language (short-period mode, actuator rate limits, gain scheduling on airspeed), a vehicle question deserves automotive language (yaw rate, tire slip, ESP/ABS intervention), a robot question deserves manipulator/mobile-robot language (Jacobians, impedance control, odometry drift), and an automation question deserves plant language (loop tuning, valve saturation, cascade and ratio control, sensor filtering).
-
-### Matching the Answer to the Question:
-- **Design / applied / "how does X work" questions**: lead with the engineering answer -- the architecture, the loop structure, the tuning trade-off, the sensors and actuators involved, the failure modes and safety interlocks. Use mathematics where it clarifies, not as ceremony.
-- **Numerical questions**: call the deterministic tools and report exact computed values.
-- **Theory, derivation, proof, comparison, and limitation questions**: use the 4-stage standard below.
-- **Definitional or "what does <author/textbook> say" questions**: answer from the retrieved reference passages and cite them; do not run a numeric solver.
-
-### 4-Stage Mathematical Reasoning & Proof Standard:
-When answering theoretical principles, derivations, proofs, comparisons, or limitation questions:
-1. **Mathematical Context & System Class**: State the state space equations (e.g. $\dot{x} = Ax + Bu, y = Cx + Du$ or $\dot{x} = f(x) + g(x)u$), signal spaces ($\mathcal{L}_2, \mathcal{H}_\infty$), domain definitions, and underlying assumptions.
-2. **Canonical Theorem / Analytical Principle**: State the governing theorem with exact mathematical rigor (e.g. PBH rank test, Doyle 1978 LQG robustness counterexample, Poisson Integral for RHP zeros, Small Gain vs Passivity, Lyapunov Invariance).
-3. **Exact Derivation & Closed-Form Formulas**: Provide the full, exact mathematical relationship in pure LaTeX (e.g. both the standard linear approximation $\zeta \approx PM/100$ and the exact non-linear relationship $PM = \arctan\left(\frac{2\zeta}{\sqrt{\sqrt{1+4\zeta^4}-2\zeta^2}}\right)$).
-4. **Engineering Caveats & Breakdown Conditions**: Explicitly state where approximations break down (e.g. $PM > 60^\circ$ or high-frequency non-dominant dynamics), numerical ill-conditioning (e.g. high-order Kalman matrices vs PBH), and physical conservation trade-offs (e.g. Bode sensitivity integral / waterbed effect).
-
-### Core Deterministic Capabilities & Tool Calling:
-- **State-feedback simulation (`simulate_state_feedback_response`)**: THE tool for "design a controller and simulate it". Pass the original `A`, `B` and the gain `K` returned by `continuous_lqr` / `discrete_lqr` / `place_state_feedback`; it forms the closed loop $A - BK$ internally and returns poles, damping, the exact closed-loop transfer function, transient metrics, and a plot. NEVER expand closed-loop polynomial coefficients by hand and feed them to `simulate_step_response` -- that algebra is the single most common source of silently wrong answers. Add `normalize_dc_gain: true` when the response should track a unit step.
-- **Transfer-function step response (`simulate_step_response`)**: only when the system is genuinely *given* as $G(s) = num/den$.
-- **Frequency domain**: `stability_margins` (GM/PM/crossovers), `bode_analysis`, `nyquist_analysis` (encirclements and the $Z = N + P$ criterion), `root_locus_analysis` (asymptotes, breakaway points, critical gain at instability).
-- Python Code Execution (`execute_python_code`): Write valid Python scripts using `scipy.signal`, `scipy.linalg`, `control` (`import control as ct`), `numpy`, `matplotlib.pyplot`. (Note: `T` in `signal.step` must be a 1D array like `np.linspace(...)`). Prefer a dedicated tool above when one fits -- it is verified and cannot be mis-transcribed. Only reach for this when the question actually requires a number, a simulation, or a plot for a SPECIFIC system. A "design considerations for X", "how does X compare to Y", or "explain X" question with no concrete numbers in it does not need code -- answer in prose. Writing code for a question that doesn't need it only risks a wasted tool step on a shape/dimension mistake with nothing to show for it.
-- **Plain matrix math (`matrix_arithmetic`)**: "A times B", inverse, determinant, rank, transpose, eigenvalues of a given matrix. A request to multiply or invert matrices is a LINEAR ALGEBRA question -- it is not an invitation to design a controller, and it needs no B, Q, or R.
-- Deterministic Math Tools: `continuous_lqr`, `discrete_lqr`, `place_state_feedback`, `exact_zoh`, `eigen_analysis`, `controllability_analysis`, `observability_analysis`, `solve_lyapunov`, `mpc_solve_qp`, `kalman_*`.
-- `plot_math_expression` plots a REAL function of one real variable ($t$ or $x$) only. Never pass a transfer function, a Laplace-domain expression, or anything containing the imaginary unit to it.
-- Reference Lookup (`search_control_references`): Use this -- not a numeric tool -- for conceptual, definitional, or "how does textbook/author X explain Y" questions. A question about what a concept means or how a source presents it is never a reason to (re-)run stability_margins, bode_analysis, routh_hurwitz_analysis, or any other numeric solver. If retrieved reference passages are already provided in this system prompt, answer straight from them and cite `[filename, p. N]`.
-- **One lookup is enough**: after a reference search returns passages, ANSWER from them. Do not re-run the same search with a reworded query ("X", then "X algorithm", then "X definition") -- repeated lookups consume the step budget and leave nothing for the answer itself. If the retrieved passages are thin, answer from your own knowledge and say what is uncertain.
-- **Don't recompute what's already in the conversation**: if a transfer function's margins, poles, or response were already computed earlier in this conversation, reuse those results instead of calling the same numeric tool again for a follow-up question about something else. Answer the question that was actually asked.
-- **Coefficient care**: when you must expand a factored transfer function like $s(s+1)(s+5)$ into polynomial form, expand one factor at a time and re-check each coefficient -- the deterministic tools verify their own arithmetic, not the coefficients you hand them.
-- **NEVER invent a missing parameter, ever -- not a reused one, not a new one, not a "reasonable-looking" one.** Every number you pass to a tool -- every matrix, gain, coefficient -- must come from the user's CURRENT message, from a tool result already in this conversation, or from a standard formula you can name. If a computation needs a parameter (e.g. B, Q, R for LQR) that is not present anywhere and not derivable, STOP. Do not guess a plausible value. Do not carry one over, resized or not, from a different problem earlier in the conversation. Do not make one up because a shape needs to match. Tell the user exactly which parameter is missing and ask for it, or answer only the part of the question you actually can with what was given. A short, honest "I don't have B for this system -- what is it?" is the correct answer. A confident computation built on an invented number is not a partial answer, it is a fabricated one, and it is the single worst thing you can do in this domain. If you catch yourself about to write a matrix or number that did not come from the message, the history, or a named formula, that is the signal to stop and ask instead of proceeding.
-- **A confusing or malformed message is a request for clarification, not a license to substitute a different, cleaner-looking problem.** If the user's wording is ambiguous (e.g. it could mean matrix multiplication, or could mean a controller design, and you cannot tell which), say what you think they might mean and ask, rather than silently picking one interpretation and inventing whatever inputs that interpretation requires.
-
-### Strict Negative & Formatting Constraints:
-1. **ZERO EMOJIS**: NEVER use any emojis anywhere in your response (absolutely NO checkmarks, NO pins, NO graphs, NO rockets, NO lightbulbs).
-2. **ALWAYS USE DOLLAR DELIMITERS FOR MATH**: ALWAYS enclose EVERY mathematical formula, transfer function, variable, fraction, and Greek letter in dollar signs: `$ ... $` for inline math or `$$ ... $$` for centered equations. Example: `$$G(s) = \frac{1}{s(s+1)(s+2)}$$` and `$\omega \to \infty$`. NEVER write raw LaTeX keywords (`\frac`, `\omega`, `\to`) without `$` or `$$` delimiters!
-2b. **SINGLE BACKSLASH ONLY**: Every LaTeX command starts with exactly one backslash character, as in `\zeta`, `\left(`, `\sqrt{}`, `\sin`. Count the backslashes before you write a command and stop at one.
-3. **NO RAW LATEX DOCUMENT TAGS**: NEVER output `\begin{figure}`, `\includegraphics`, `\caption`, `\centering`, `\section*`, or `\end{figure}`. Use standard markdown headers (e.g. `### Section Title`).
-4. **MANDATORY TOOL CALL FOR PLOTS**: When asked to plot, visualize, or simulate (Nyquist plot, Bode diagram, Root Locus, Step Response, Phase Portrait), NEVER hallucinate a fake image filename. You MUST explicitly call `execute_python_code` (using `control as ct` or `matplotlib.pyplot`) or `simulate_step_response` to generate and display the real plot!
-5. **Deterministic Grounding**: When a tool executes successfully, present the exact numerical results and generated plot.
-6. **No Infinite Retries, but READ the error first**: If a tool returns an error, the error text itself is often the entire answer to what's actually being asked -- "uncontrollable", "singular", "poles can't be placed", "not positive semi-definite" name the exact concept the question is testing. Before doing anything else: quote or restate that reason in your final answer as the headline finding, in plain engineering language (e.g. "this system is uncontrollable because..."). Only after stating it should you correct parameters and retry, or move on. NEVER silently produce a numeric-looking final answer (a bare $$K = ...$$, a computed matrix) that routes around an error without ever mentioning what it said -- an ignored error is a worse answer than no answer.
-7. **No Forced Citations**: Do NOT cite arbitrary random file paths unless the user explicitly requests literature references.
+"""System prompts for ControlAI.
+
+The previous prompt was 2,587 tokens of mostly prohibitions written in shouted
+capitals. It cost real latency on every single turn, and on a small model the
+prohibitions backfired: the "never invent a parameter" clause taught the model
+to refuse worked examples, which is the single most useful thing a teaching
+assistant does. This one is short, states what to do rather than what not to
+do, and leaves the genuinely structural guarantees (schema validation, result
+rounding, independent verification) to code where they belong.
 """
+
+SYSTEM_PROMPT = """You are ControlAI, an expert control systems engineer. You cover classical and \
+modern control, state-space methods, optimal and robust control, estimation, nonlinear and \
+adaptive control, and system identification, across aerospace, automotive, robotics, process \
+automation, power systems, and mechatronics. Answer in the vocabulary of whichever domain the \
+question comes from.
+
+## Using tools
+
+You have deterministic solvers (SciPy/LAPACK/CVXPY). Every number you state must come from one \
+of them or from the user -- never compute a Riccati solution, a set of closed-loop poles, a gain \
+margin, or a polynomial expansion in your head.
+
+- Call a tool when the question involves a concrete system and needs a number, a matrix, or a plot.
+- Answer directly, with no tool call, when the question is conceptual, definitional, comparative, \
+or a derivation. Explaining what a phase margin *is* needs prose, not a solver.
+- When a system is given in factored form such as $G(s) = K/(s(s+1)(s+5))$, use \
+`expand_polynomial_from_roots` to get the coefficients rather than multiplying the factors out \
+yourself. Expand the denominator roots with `gain: 1` and pass $K$ as the numerator.
+- For gain and phase margins call `stability_margins`, which returns them along with both \
+crossover frequencies. `bode_analysis` returns a magnitude and phase curve; do not read margins \
+off it by eye.
+- To design a controller and then see its behaviour, pass the original $A$, $B$ and the returned \
+gain $K$ to `simulate_state_feedback_response`; it closes the loop internally.
+- Do not form $A - BK$, multiply matrices, or expand a characteristic polynomial by hand when \
+writing up a result. Hand algebra in the write-up is where a correct solver output turns into a \
+wrong answer. The LQR and pole-placement tools already return `closed_loop_A` -- quote that. \
+Otherwise use `matrix_arithmetic`, or state what the solver returned and stop there.
+- If a tool returns an error, that error is usually the answer: "uncontrollable", "singular", \
+"not stabilizable" name the exact property the question is about. Lead with it in plain language.
+
+## Worked examples
+
+When the user asks for an example, a demonstration, or "show me how this works" without giving a \
+system, choose a clean illustrative one yourself, say plainly that you are choosing it, and run \
+the real tools on it. A concrete worked example is the correct answer to that request.
+
+When the user asks about *their* system but a parameter you need is genuinely missing, ask for \
+that one parameter. Do not silently substitute a value and present the result as theirs.
+
+## Style
+
+Lead with the engineering answer: the loop structure, the trade-off, the number that was asked \
+for. Add derivations when they clarify. State where an approximation breaks down.
+
+Write all mathematics in LaTeX delimited by `$...$` inline or `$$...$$` displayed, using single \
+backslashes. Use markdown headings, never LaTeX document commands. Do not use emoji."""
+
+
+RETRIEVAL_PREAMBLE = """The following passages were retrieved from the local control-engineering \
+library for this question. Where a passage covers the question, answer from it and cite it with \
+the bracketed label exactly as shown. Where it does not, ignore it and answer from your own \
+knowledge -- do not force a citation, and never print a raw filename."""
+
+
+SYNTHESIS_NUDGE = """Write the final answer now, using the tool results above. State the computed \
+values explicitly. Do not call any more tools."""

controlai_agent/toolcall.py

@@ -0,0 +1,119 @@
+"""Parsing of Qwen-style `<tool_call>` blocks out of a model response.
+
+The previous implementation carried several hundred lines of regex repair --
+unbalanced-brace closing, space-separated MATLAB array rewriting, LaTeX
+backslash unescaping. Almost all of it existed to compensate for a fine-tuned
+adapter that emitted malformed JSON; the base model emits well-formed calls, so
+what remains here is a small, readable tolerance margin rather than a repair
+pipeline. Argument-level coercion (stringified arrays and the like) already
+happens in `registry.execute`, so it is deliberately not duplicated here.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass, field
+from typing import Any
+
+TOOL_CALL_RE = re.compile(r"<tool_call>\s*(.*?)\s*(?:</tool_call>|$)", re.DOTALL)
+# Some checkpoints emit a bare JSON object with no wrapping tags at all.
+BARE_CALL_RE = re.compile(r'\{\s*"name"\s*:\s*"[^"]+"\s*,\s*"arguments"\s*:\s*\{.*?\}\s*\}', re.DOTALL)
+
+
+@dataclass
+class ToolCall:
+    name: str
+    arguments: dict[str, Any] = field(default_factory=dict)
+
+
+def _close_json(raw: str) -> str:
+    """Append the brackets needed to balance a truncated JSON object.
+
+    Generation can hit the token limit mid-call. Closing the structure recovers
+    a usable call instead of discarding one that was merely cut short.
+    """
+    in_string = escaped = False
+    stack: list[str] = []
+    for ch in raw:
+        if escaped:
+            escaped = False
+            continue
+        if ch == "\\":
+            escaped = True
+        elif ch == '"':
+            in_string = not in_string
+        elif not in_string:
+            if ch in "{[":
+                stack.append(ch)
+            elif ch in "}]" and stack:
+                stack.pop()
+    return raw + ('"' if in_string else "") + "".join("}" if c == "{" else "]" for c in reversed(stack))
+
+
+def _loads(raw: str) -> dict[str, Any] | None:
+    for candidate in (raw, _close_json(raw)):
+        try:
+            parsed = json.loads(candidate)
+        except (json.JSONDecodeError, TypeError):
+            continue
+        if isinstance(parsed, dict) and parsed.get("name"):
+            return parsed
+    return None
+
+
+def parse(text: str) -> tuple[list[ToolCall], str]:
+    """Split a response into its tool calls and its prose.
+
+    Returns `(calls, prose)`. An empty `calls` list means the response was a
+    direct answer.
+    """
+    calls: list[ToolCall] = []
+    for block in TOOL_CALL_RE.findall(text):
+        parsed = _loads(block.strip())
+        if parsed:
+            args = parsed.get("arguments") or parsed.get("parameters") or {}
+            if isinstance(args, str):
+                args = _loads(args) or {}
+            calls.append(ToolCall(name=str(parsed["name"]), arguments=args if isinstance(args, dict) else {}))
+
+    prose = TOOL_CALL_RE.sub("", text)
+    if not calls:
+        for block in BARE_CALL_RE.findall(prose):
+            parsed = _loads(block)
+            if parsed:
+                args = parsed.get("arguments") or {}
+                calls.append(ToolCall(name=str(parsed["name"]), arguments=args if isinstance(args, dict) else {}))
+                prose = prose.replace(block, "")
+
+    prose = re.sub(r"<think>.*?</think>", "", prose, flags=re.DOTALL)
+    prose = re.sub(r"</?(?:think|tool_call|tool_response)>", "", prose)
+    return calls, prose.strip()
+
+
+_MAX_ARRAY_LEN = 400
+_MAX_IDENTICAL_RUN = 12
+
+
+def degenerate_reason(value: Any, depth: int = 0) -> str | None:
+    """Detect runaway repetition in a tool argument.
+
+    A decoding loop can produce a several-hundred-element array of the same
+    number, which is never a real system and wastes a step plus a large slice
+    of context if it is executed. The thresholds are set well above any
+    plausible hand-written matrix so ordinary inputs are never touched.
+    """
+    if depth > 4 or not isinstance(value, list):
+        return None
+    if len(value) > _MAX_ARRAY_LEN:
+        return f"has {len(value)} elements"
+    run = 1
+    for prev, cur in zip(value, value[1:]):
+        run = run + 1 if prev == cur and isinstance(cur, (int, float)) else 1
+        if run >= _MAX_IDENTICAL_RUN:
+            return f"repeats the value {cur} {run} times in a row"
+    for item in value:
+        reason = degenerate_reason(item, depth + 1)
+        if reason:
+            return reason
+    return None

controlai_agent/tools.py

@@ -1,427 +0,0 @@
-"""Deterministic control engineering tools backed by NumPy, SciPy, and CVXPY."""
-
-from __future__ import annotations
-
-import math
-from typing import Any
-
-import numpy as np
-from scipy import linalg, signal
-
-from controlai_agent.registry import registry
-
-
-@registry.register(
-    name="exact_zoh",
-    description="Discretize continuous-time state-space matrices (A, B) under exact Zero-Order Hold (ZOH) at sample time Ts.",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "A": {
-                "type": "array",
-                "items": {"type": "array", "items": {"type": "number"}},
-                "description": "Continuous system matrix A (n x n)",
-            },
-            "B": {
-                "type": "array",
-                "items": {"type": "array", "items": {"type": "number"}},
-                "description": "Continuous input matrix B (n x m)",
-            },
-            "Ts": {
-                "type": "number",
-                "description": "Sample time in seconds (Ts > 0)",
-            },
-        },
-        "required": ["A", "B", "Ts"],
-    },
-)
-def exact_zoh(A: list[list[float]], B: list[list[float]], Ts: float) -> dict[str, Any]:
-    A_mat = np.array(A, dtype=float)
-    B_mat = np.array(B, dtype=float)
-    n = A_mat.shape[0]
-    m = B_mat.shape[1] if B_mat.ndim > 1 else 1
-    C_dummy = np.eye(n)
-    D_dummy = np.zeros((n, m))
-    Ad, Bd, _, _, _ = signal.cont2discrete((A_mat, B_mat, C_dummy, D_dummy), Ts, method="zoh")
-    return {
-        "Ad": Ad.tolist(),
-        "Bd": Bd.tolist(),
-        "sample_time": Ts,
-        "method": "exact_zoh",
-    }
-
-
-@registry.register(
-    name="eigen_analysis",
-    description="Compute eigenvalues, eigenvectors, and asymptotic stability of a continuous or discrete system matrix A.",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "A": {
-                "type": "array",
-                "items": {"type": "array", "items": {"type": "number"}},
-                "description": "System matrix A (n x n)",
-            },
-            "discrete": {
-                "type": "boolean",
-                "description": "True if discrete-time (criterion |lambda| < 1), False if continuous (Re(lambda) < 0)",
-            },
-        },
-        "required": ["A"],
-    },
-)
-def eigen_analysis(A: list[list[float]], discrete: bool = False) -> dict[str, Any]:
-    A_mat = np.array(A, dtype=float)
-    eigenvalues = np.linalg.eigvals(A_mat)
-    if discrete:
-        stable = bool(np.all(np.abs(eigenvalues) < 1.0))
-        criterion = "|lambda_i| < 1"
-    else:
-        stable = bool(np.all(np.real(eigenvalues) < 0.0))
-        criterion = "Re(lambda_i) < 0"
-
-    formatted_eigs = []
-    for eig in eigenvalues:
-        if abs(eig.imag) < 1e-9:
-            formatted_eigs.append(f"{eig.real:.6g}")
-        else:
-            sign = "+" if eig.imag >= 0 else "-"
-            formatted_eigs.append(f"{eig.real:.6g} {sign} {abs(eig.imag):.6g}j")
-
-    return {
-        "eigenvalues": [[float(e.real), float(e.imag)] for e in eigenvalues],
-        "eigenvalues_formatted": formatted_eigs,
-        "is_stable": stable,
-        "criterion": criterion,
-        "time_domain": "discrete" if discrete else "continuous",
-    }
-
-
-@registry.register(
-    name="controllability_analysis",
-    description="Compute controllability matrix, rank, and PBH modal controllability test for pair (A, B).",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "A": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "B": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-        },
-        "required": ["A", "B"],
-    },
-)
-def controllability_analysis(A: list[list[float]], B: list[list[float]]) -> dict[str, Any]:
-    A_mat = np.array(A, dtype=float)
-    B_mat = np.array(B, dtype=float)
-    n = A_mat.shape[0]
-    blocks = [B_mat]
-    current = B_mat
-    for _ in range(1, n):
-        current = A_mat @ current
-        blocks.append(current)
-    C_mat = np.hstack(blocks)
-    rank = int(np.linalg.matrix_rank(C_mat))
-    is_controllable = rank == n
-    return {
-        "controllability_matrix": C_mat.tolist(),
-        "rank": rank,
-        "state_dimension": n,
-        "is_controllable": is_controllable,
-    }
-
-
-@registry.register(
-    name="observability_analysis",
-    description="Compute observability matrix, rank, and PBH modal observability test for pair (A, C).",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "A": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "C": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-        },
-        "required": ["A", "C"],
-    },
-)
-def observability_analysis(A: list[list[float]], C: list[list[float]]) -> dict[str, Any]:
-    A_mat = np.array(A, dtype=float)
-    C_mat = np.array(C, dtype=float)
-    n = A_mat.shape[0]
-    blocks = [C_mat]
-    current = C_mat
-    for _ in range(1, n):
-        current = current @ A_mat
-        blocks.append(current)
-    O_mat = np.vstack(blocks)
-    rank = int(np.linalg.matrix_rank(O_mat))
-    is_observable = rank == n
-    return {
-        "observability_matrix": O_mat.tolist(),
-        "rank": rank,
-        "state_dimension": n,
-        "is_observable": is_observable,
-    }
-
-
-@registry.register(
-    name="continuous_lqr",
-    description="Solve Continuous-time Linear Quadratic Regulator (CARE) problem: min integral (x^T Q x + u^T R u) dt.",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "A": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "B": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "Q": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "R": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-        },
-        "required": ["A", "B", "Q", "R"],
-    },
-)
-def continuous_lqr(
-    A: list[list[float]], B: list[list[float]], Q: list[list[float]], R: list[list[float]]
-) -> dict[str, Any]:
-    A_mat = np.array(A, dtype=float)
-    B_mat = np.array(B, dtype=float)
-    Q_mat = np.array(Q, dtype=float)
-    R_mat = np.array(R, dtype=float)
-    P = linalg.solve_continuous_are(A_mat, B_mat, Q_mat, R_mat)
-    K = np.linalg.solve(R_mat, B_mat.T @ P)
-    A_cl = A_mat - B_mat @ K
-    poles = np.linalg.eigvals(A_cl)
-    residual = float(np.max(np.abs(A_mat.T @ P + P @ A_mat - P @ B_mat @ np.linalg.inv(R_mat) @ B_mat.T @ P + Q_mat)))
-    return {
-        "P": P.tolist(),
-        "K": K.tolist(),
-        "closed_loop_poles": [[float(p.real), float(p.imag)] for p in poles],
-        "is_stable": bool(np.all(np.real(poles) < 0)),
-        "riccati_residual": residual,
-    }
-
-
-@registry.register(
-    name="discrete_lqr",
-    description="Solve Discrete-time Linear Quadratic Regulator (DARE) problem: min sum (x_k^T Q x_k + u_k^T R u_k).",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "A": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "B": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "Q": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "R": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-        },
-        "required": ["A", "B", "Q", "R"],
-    },
-)
-def discrete_lqr(
-    A: list[list[float]], B: list[list[float]], Q: list[list[float]], R: list[list[float]]
-) -> dict[str, Any]:
-    A_mat = np.array(A, dtype=float)
-    B_mat = np.array(B, dtype=float)
-    Q_mat = np.array(Q, dtype=float)
-    R_mat = np.array(R, dtype=float)
-    P = linalg.solve_discrete_are(A_mat, B_mat, Q_mat, R_mat)
-    K = np.linalg.solve(R_mat + B_mat.T @ P @ B_mat, B_mat.T @ P @ A_mat)
-    A_cl = A_mat - B_mat @ K
-    poles = np.linalg.eigvals(A_cl)
-    return {
-        "P": P.tolist(),
-        "K": K.tolist(),
-        "closed_loop_poles": [[float(p.real), float(p.imag)] for p in poles],
-        "is_stable": bool(np.all(np.abs(poles) < 1.0)),
-    }
-
-
-@registry.register(
-    name="cbf_safety_filter",
-    description="Scalar Control Barrier Function (CBF) quadratic program safety filter for single integrator x_dot=u with constraint x >= x_min.",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "x": {"type": "number", "description": "Current state"},
-            "u_nom": {"type": "number", "description": "Nominal desired control input"},
-            "alpha": {"type": "number", "description": "CBF class-K gain parameter"},
-            "x_min": {"type": "number", "description": "Safety lower bound x >= x_min"},
-        },
-        "required": ["x", "u_nom", "alpha", "x_min"],
-    },
-)
-def cbf_safety_filter(x: float, u_nom: float, alpha: float, x_min: float) -> dict[str, Any]:
-    h = x - x_min
-    lower_bound = -alpha * h
-    u_safe = max(u_nom, lower_bound)
-    active = bool(u_nom < lower_bound)
-    return {
-        "h": h,
-        "cbf_lower_bound": lower_bound,
-        "u_safe": u_safe,
-        "is_constraint_active": active,
-        "h_dot_plus_alpha_h": u_safe + alpha * h,
-    }
-
-
-@registry.register(
-    name="dynamic_inversion",
-    description="Compute exact Nonlinear Dynamic Inversion (NDI) input u for scalar affine system x_dot = a*x + b*u to track target error dynamics x_dot = -gain*(x - r).",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "a": {"type": "number", "description": "Open-loop plant coefficient a"},
-            "b": {"type": "number", "description": "Control effectiveness coefficient b (b != 0)"},
-            "x": {"type": "number", "description": "Current state x"},
-            "reference": {"type": "number", "description": "Setpoint reference r"},
-            "gain": {"type": "number", "description": "Tracking convergence gain k > 0"},
-        },
-        "required": ["a", "b", "x", "reference", "gain"],
-    },
-)
-def dynamic_inversion(a: float, b: float, x: float, reference: float, gain: float) -> dict[str, Any]:
-    virtual_control = -gain * (x - reference)
-    u = (virtual_control - a * x) / b
-    x_dot = a * x + b * u
-    return {
-        "virtual_control_v": virtual_control,
-        "control_input_u": u,
-        "achieved_x_dot": x_dot,
-        "residual": abs(x_dot - virtual_control),
-    }
-
-
-@registry.register(
-    name="kharitonov_stability_test",
-    description="Evaluate robust Hurwitz stability of a real interval polynomial using Kharitonov's Theorem (tests the 4 extreme polynomials).",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "lower_bounds": {
-                "type": "array",
-                "items": {"type": "number"},
-                "description": "Lower coefficient bounds [a0^-, a1^-, a2^-, a3^-] in ascending powers",
-            },
-            "upper_bounds": {
-                "type": "array",
-                "items": {"type": "number"},
-                "description": "Upper coefficient bounds [a0^+, a1^+, a2^+, a3^+] in ascending powers",
-            },
-        },
-        "required": ["lower_bounds", "upper_bounds"],
-    },
-)
-def kharitonov_stability_test(lower_bounds: list[float], upper_bounds: list[float]) -> dict[str, Any]:
-    lows = np.array(lower_bounds, dtype=float)
-    highs = np.array(upper_bounds, dtype=float)
-    polys = [
-        [lows[0], lows[1], highs[2], highs[3]],
-        [highs[0], highs[1], lows[2], lows[3]],
-        [highs[0], lows[1], lows[2], highs[3]],
-        [lows[0], highs[1], highs[2], lows[3]],
-    ]
-    pole_sets = []
-    stable_flags = []
-    for poly in polys:
-        desc_poly = poly[::-1]  # descending for np.roots
-        roots = np.roots(desc_poly)
-        pole_sets.append([[float(r.real), float(r.imag)] for r in roots])
-        stable_flags.append(bool(np.all(np.real(roots) < 0)))
-
-    robustly_stable = all(stable_flags)
-    return {
-        "kharitonov_polynomials_ascending": polys,
-        "stable_flags": stable_flags,
-        "is_robustly_hurwitz": robustly_stable,
-        "poles": pole_sets,
-    }
-
-
-@registry.register(
-    name="minimum_norm_control_allocation",
-    description="Compute minimum 2-norm control allocation for redundant actuators: min ||u||_2 subject to B*u = tau.",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "B": {
-                "type": "array",
-                "items": {"type": "number"},
-                "description": "Actuator effectiveness row vector B (1 x m)",
-            },
-            "desired_tau": {
-                "type": "number",
-                "description": "Desired virtual control torque/force tau",
-            },
-        },
-        "required": ["B", "desired_tau"],
-    },
-)
-def minimum_norm_control_allocation(B: list[float], desired_tau: float) -> dict[str, Any]:
-    B_vec = np.array(B, dtype=float)
-    b_norm_sq = float(np.dot(B_vec, B_vec))
-    u = (desired_tau / b_norm_sq) * B_vec
-    achieved = float(np.dot(B_vec, u))
-    return {
-        "u": u.tolist(),
-        "achieved_tau": achieved,
-        "residual": abs(achieved - desired_tau),
-        "norm_u": float(np.linalg.norm(u)),
-    }
-
-
-@registry.register(
-    name="jump_system_stability",
-    description="Compute second-moment contraction factor E[a^2] = sum(p_i * a_i^2) for i.i.d. scalar jump linear systems.",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "probabilities": {
-                "type": "array",
-                "items": {"type": "number"},
-                "description": "Probabilities p_i summing to 1",
-            },
-            "multipliers": {
-                "type": "array",
-                "items": {"type": "number"},
-                "description": "Multipliers a_i for each jump mode",
-            },
-        },
-        "required": ["probabilities", "multipliers"],
-    },
-)
-def jump_system_stability(probabilities: list[float], multipliers: list[float]) -> dict[str, Any]:
-    p = np.array(probabilities, dtype=float)
-    a = np.array(multipliers, dtype=float)
-    expected_square = float(np.sum(p * (a**2)))
-    is_stable = expected_square < 1.0
-    return {
-        "expected_squared_multiplier": expected_square,
-        "is_mean_square_stable": is_stable,
-        "contraction_margin": 1.0 - expected_square,
-    }
-
-
-@registry.register(
-    name="place_state_feedback",
-    description="Compute state feedback gain matrix K such that eig(A - B*K) matches target desired poles.",
-    parameters_schema={
-        "type": "object",
-        "properties": {
-            "A": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "B": {"type": "array", "items": {"type": "array", "items": {"type": "number"}}},
-            "desired_poles": {
-                "type": "array",
-                "items": {"type": "number"},
-                "description": "Target closed-loop pole locations",
-            },
-        },
-        "required": ["A", "B", "desired_poles"],
-    },
-)
-def place_state_feedback(
-    A: list[list[float]], B: list[list[float]], desired_poles: list[float]
-) -> dict[str, Any]:
-    A_mat = np.array(A, dtype=float)
-    B_mat = np.array(B, dtype=float)
-    des = np.array(desired_poles, dtype=float)
-    placed = signal.place_poles(A_mat, B_mat, des)
-    K = placed.gain_matrix
-    closed_poles = np.linalg.eigvals(A_mat - B_mat @ K)
-    return {
-        "K": K.tolist(),
-        "closed_loop_poles": [[float(p.real), float(p.imag)] for p in closed_poles],
-        "target_poles": desired_poles,
-    }

controlai_agent/tools/frequency.py

@@ -17,9 +17,58 @@ from controlai_agent.registry import registry
 ARTIFACT_DIR = Path("outputs/plots")
 
 
+@registry.register(
+    name="expand_polynomial_from_roots",
+    description=(
+        "Expand a transfer function or characteristic polynomial given in factored/root form -- "
+        "e.g. G(s) = K / (s*(s+1)*(s+5)), roots at s = 0, -1, -5 -- into exact polynomial "
+        "coefficients in descending powers. ALWAYS call this instead of multiplying the factors "
+        "out by hand before calling bode_analysis, stability_margins, simulate_step_response, "
+        "routh_hurwitz_analysis, root_locus, or any other tool that takes numerator/denominator "
+        "coefficients: hand-expanding factors is the single most common source of a silently "
+        "wrong tool call, since nothing downstream can verify an argument that was already wrong "
+        "going in. A factor '(s + a)' contributes the root -a; '(s - a)' contributes root a."
+    ),
+    parameters_schema={
+        "type": "object",
+        "properties": {
+            "roots": {
+                "type": "array",
+                "items": {"type": "number"},
+                "description": "The roots of the polynomial, one per linear factor -- e.g. [0, -1, -5] for s*(s+1)*(s+5)",
+            },
+            "gain": {
+                "type": "number",
+                "default": 1.0,
+                "description": (
+                    "Multiplies every returned coefficient. Leave this at 1 when expanding a "
+                    "DENOMINATOR -- the overall constant K belongs in the numerator, not scaled "
+                    "into the denominator. For G(s) = 10/(s(s+1)(s+5)), expand roots [0,-1,-5] "
+                    "with gain 1 to get the denominator [1,6,5,0] and pass numerator [10] "
+                    "separately. Passing gain=10 here instead yields [10,60,50,0], which is the "
+                    "same transfer function scaled by 1/10 and silently wrong."
+                ),
+            },
+        },
+        "required": ["roots"],
+    },
+)
+def expand_polynomial_from_roots(roots: list[float], gain: float = 1.0) -> dict[str, Any]:
+    coefficients = (gain * np.poly(roots)).tolist()
+    return {
+        "status": "success",
+        "coefficients_descending": [float(c) for c in coefficients],
+        "degree": len(roots),
+    }
+
+
 @registry.register(
     name="bode_analysis",
-    description="Compute frequency response magnitude (linear & dB), phase (degrees), and bandwidth for transfer function G(s) = num(s) / den(s).",
+    description=(
+        "Compute the frequency response of G(s) = num(s)/den(s): magnitude in dB, phase in "
+        "degrees, resonant peak, DC gain, and the gain/phase margins with both crossover "
+        "frequencies."
+    ),
     parameters_schema={
         "type": "object",
         "properties": {
@@ -63,6 +112,13 @@ def bode_analysis(
         "frequencies_sample": w_out[::10].tolist(),
         "magnitudes_db_sample": mag_db[::10].tolist(),
         "phases_deg_sample": phase_deg[::10].tolist(),
+        # Margins are included even though `stability_margins` is the tool that
+        # advertises them. Asked for a phase margin, the model was observed
+        # reaching for bode_analysis, getting back only sampled curves, and
+        # concluding the margin "cannot be determined" -- a correct reading of
+        # a sampled plot, and a useless answer. Computing them here makes that
+        # routing choice harmless rather than fatal.
+        "margins": stability_margins(numerator, denominator),
     }
 
 

controlai_agent/tools/synthesis.py

@@ -43,6 +43,12 @@ def continuous_lqr(
     return {
         "P": P.tolist(),
         "K": K.tolist(),
+        # Returned so the answer never has to derive it. The model was
+        # observed writing A - BK as [[-6, 1], [-5, -6]] for a double
+        # integrator whose true closed loop is [[0, 1], [-6, -5]] -- correct
+        # gain, wrong write-up. Handing it the computed matrix removes the
+        # arithmetic from the answer entirely.
+        "closed_loop_A": (A - B @ K).tolist(),
         "closed_loop_poles": [[float(p.real), float(p.imag)] for p in poles],
         "is_stable": bool(np.all(np.real(poles) < 0)),
         "verification": v_report,
@@ -79,6 +85,7 @@ def discrete_lqr(
     return {
         "P": P.tolist(),
         "K": K.tolist(),
+        "closed_loop_A": (A - B @ K).tolist(),
         "closed_loop_poles": [[float(p.real), float(p.imag)] for p in poles],
         "is_stable": bool(np.all(np.abs(poles) < 1.0)),
         "verification": v_report,
@@ -138,6 +145,7 @@ def place_state_feedback(
     v_report = verifier.verify_pole_placement(A_mat, B_mat, K, list(des))
     return {
         "K": K.tolist(),
+        "closed_loop_A": (A - B @ K).tolist(),
         "closed_loop_poles": [[float(p.real), float(p.imag)] for p in closed_poles],
         "target_poles": [[float(p.real), float(p.imag)] for p in des],
         "verification": v_report,

requirements.txt

@@ -1,32 +1,29 @@
-gradio>=5.9.0
-spaces>=0.30.0
+# ControlAI runtime dependencies -- local, Apple Silicon.
+#
+# Inference is MLX only. The previous file also carried torch, transformers,
+# accelerate, peft, llama-cpp-python, gradio and spaces to support a Hugging
+# Face Spaces deployment and three alternative backends; none of that is used
+# any more. transformers is still pulled in indirectly by mlx-lm for tokenizer
+# loading, so it is not listed here.
+
+mlx-lm>=0.31.0
 huggingface-hub>=0.23.0
-torch>=2.0.0
-transformers>=4.40.0
-accelerate>=0.28.0
-peft>=0.10.0
-jsonschema>=4.20.0
-scipy>=1.11.0
+
+# Deterministic numerics -- every number in an answer comes from these.
 numpy>=1.24.0
-matplotlib>=3.7.0
+scipy>=1.11.0
 control>=0.9.4
 cvxpy>=1.4.0
-rich>=13.7.0
+matplotlib>=3.7.0
+jsonschema>=4.20.0
+
+# Retrieval over the local control-engineering library.
 rank-bm25>=0.2.2
-# Pulls a prebuilt CPU wheel instead of compiling llama.cpp's C++ core from
-# source -- a source build routinely exceeds HF Spaces' build-job timeout
-# ("Job timeout" BUILD_ERROR, observed directly). This backend is CPU-only by
-# design (see the GGUF branch in orchestrator.py): llama.cpp's raw CUDA calls
-# aren't visible to ZeroGPU's torch-based interception, so a GPU wheel
-# wouldn't get real GPU time here anyway.
---extra-index-url https://abetlen.github.io/llama-cpp-python/whl/cpu
-llama-cpp-python==0.3.35
 pypdf>=3.17.0
 pymupdf>=1.23.0
-scikit-learn>=1.3.0
-requests>=2.31.0
-pyyaml>=6.0.0
-pydantic>=2.0.0
+
+# Web console.
 fastapi>=0.110.0
 uvicorn>=0.28.0
 python-multipart>=0.0.9
+pydantic>=2.0.0

scripts/eval_answer_quality.py

@@ -31,7 +31,7 @@ PROJECT_ROOT = Path(__file__).resolve().parent.parent
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
 
-from controlai_agent.orchestrator import ControlAIAgent
+from controlai_agent.agent import ControlAgent
 
 PLACEHOLDER = "computational analysis has been completed"
 
@@ -67,7 +67,7 @@ DOUBLE_BS_PAT = re.compile(r"\\\\(?=[a-zA-Z|{}()])")
 
 def grade(question: str, result) -> tuple[list[str], dict]:
     """Return (list of problems, metrics) for one answer."""
-    text = result.final_response or ""
+    text = result.answer or ""
     problems: list[str] = []
 
     if not text.strip():
@@ -86,8 +86,8 @@ def grade(question: str, result) -> tuple[list[str], dict]:
     if RAW_FILENAME_PAT.search(text):
         problems.append("RAW filename in citation")
 
-    names = [t.tool_name for t in result.tool_traces]
-    failed = [t.tool_name for t in result.tool_traces if t.result.get("status") == "error"]
+    names = [t.name for t in result.traces]
+    failed = [t.name for t in result.traces if t.result.get("status") == "error"]
     if failed:
         problems.append(f"TOOL ERROR: {', '.join(sorted(set(failed)))}")
     for n in set(names):
@@ -108,7 +108,7 @@ def main() -> int:
         cases = cases[: args.limit]
 
     print(f"Loading agent...\n")
-    agent = ControlAIAgent()
+    agent = ControlAgent()
 
     failures: list[tuple[str, str, list[str]]] = []
     for i, (domain, q) in enumerate(cases, 1):

scripts/run_agent_benchmark.py

@@ -13,7 +13,8 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1]
 if str(PROJECT_ROOT) not in sys.path:
     sys.path.insert(0, str(PROJECT_ROOT))
 
-from controlai_agent.orchestrator import ControlAIAgent
+from controlai_agent.agent import ControlAgent
+from controlai_agent.engine import LocalEngine
 
 
 def main() -> int:
@@ -59,7 +60,7 @@ def main() -> int:
     print(f"Total benchmark items: {len(items)}")
 
     print(f"Initializing ControlAI Agent ({args.model})...")
-    agent = ControlAIAgent(model_path=args.model, adapter_path=args.adapter_path)
+    agent = ControlAgent(engine=LocalEngine(model_id=args.model, adapter_path=args.adapter_path))
     print("Agent ready.")
 
     responses = []
@@ -72,27 +73,27 @@ def main() -> int:
         print(f"[{idx}/{len(items)}] Evaluating {item_id}...", end=" ", flush=True)
 
         item_start = time.time()
-        result = agent.run(prompt, max_tokens_per_step=args.max_tokens)
+        result = agent.run(prompt, max_tokens=args.max_tokens)
         elapsed = time.time() - item_start
 
-        tool_names = [t.tool_name for t in result.tool_traces]
-        print(f"done in {elapsed:.2f}s | Steps: {result.total_steps} | Tools: {tool_names}")
+        tool_names = [t.name for t in result.traces]
+        print(f"done in {elapsed:.2f}s | Steps: {len(result.traces)} | Tools: {tool_names}")
 
         record = {
             "benchmark_id": item_id,
             "id": item_id,
             "family": item.get("family", ""),
             "prompt": prompt,
-            "response": result.final_response,
+            "response": result.answer,
             "tool_calls": [
                 {
-                    "name": t.tool_name,
+                    "name": t.name,
                     "arguments": t.arguments,
                     "result": t.result,
                 }
-                for t in result.tool_traces
+                for t in result.traces
             ],
-            "total_steps": result.total_steps,
+            "total_steps": len(result.traces),
             "finish_reason": "stop",
         }
         responses.append(record)

tests/test_agent_core.py

@@ -0,0 +1,140 @@
+"""Tests for the model-independent parts of the agent: tool-call parsing,
+streamed tool-call suppression, and history truncation."""
+
+from __future__ import annotations
+
+import sys
+import unittest
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from controlai_agent.agent import _StreamGate
+from controlai_agent.toolcall import degenerate_reason, parse
+
+
+class TestToolCallParsing(unittest.TestCase):
+    def test_plain_call(self):
+        calls, prose = parse(
+            '<tool_call>\n{"name": "continuous_lqr", "arguments": {"A": [[0, 1], [-2, -3]]}}\n</tool_call>'
+        )
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(calls[0].name, "continuous_lqr")
+        self.assertEqual(calls[0].arguments["A"], [[0, 1], [-2, -3]])
+        self.assertEqual(prose, "")
+
+    def test_narration_before_call_is_kept_separate(self):
+        calls, prose = parse(
+            'Let me compute that.\n<tool_call>\n{"name": "eigen_analysis", "arguments": {}}\n</tool_call>'
+        )
+        self.assertEqual([c.name for c in calls], ["eigen_analysis"])
+        self.assertEqual(prose, "Let me compute that.")
+
+    def test_truncated_call_is_recovered(self):
+        """Hitting the token limit mid-call should not discard the call."""
+        calls, _ = parse('<tool_call>\n{"name": "bode_analysis", "arguments": {"numerator": [10]')
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(calls[0].arguments["numerator"], [10])
+
+    def test_empty_tool_call_yields_no_call(self):
+        """The failure mode of the old fine-tuned adapter: an empty call plus
+        real prose. The prose must survive and no tool must run."""
+        calls, prose = parse("<tool_call>\n\n</tool_call>\n\nThe phase margin is the answer.")
+        self.assertEqual(calls, [])
+        self.assertEqual(prose, "The phase margin is the answer.")
+
+    def test_math_is_not_mangled(self):
+        _, prose = parse(r"Gain margin is $6$ dB with $\zeta = 0.5$.")
+        self.assertEqual(prose, r"Gain margin is $6$ dB with $\zeta = 0.5$.")
+
+    def test_multiple_calls(self):
+        calls, _ = parse(
+            '<tool_call>\n{"name": "a", "arguments": {}}\n</tool_call>'
+            '<tool_call>\n{"name": "b", "arguments": {}}\n</tool_call>'
+        )
+        self.assertEqual([c.name for c in calls], ["a", "b"])
+
+    def test_thinking_block_stripped(self):
+        _, prose = parse("<think>weighing options</think>\n\nThe answer is 3 dB.")
+        self.assertEqual(prose, "The answer is 3 dB.")
+
+
+class TestDegenerateGuard(unittest.TestCase):
+    def test_repetition_loop_detected(self):
+        self.assertIsNotNone(degenerate_reason([[0.5] * 40]))
+
+    def test_runaway_length_detected(self):
+        self.assertIsNotNone(degenerate_reason(list(range(500))))
+
+    def test_ordinary_matrices_pass(self):
+        for value in ([[0, 1], [-2, -3]], [[1.0]], [1, 6, 5, 0], [[0], [1]]):
+            self.assertIsNone(degenerate_reason(value), value)
+
+    def test_identity_matrix_passes(self):
+        """A legitimate matrix full of repeated values must not be rejected."""
+        self.assertIsNone(degenerate_reason([[1 if i == j else 0 for j in range(8)] for i in range(8)]))
+
+
+class TestStreamGate(unittest.TestCase):
+    def test_prose_passes_through(self):
+        gate = _StreamGate()
+        self.assertEqual(gate.feed("The gain margin "), "The gain margin ")
+        self.assertEqual(gate.feed("is 6 dB."), "is 6 dB.")
+
+    def test_tool_call_never_leaks(self):
+        gate = _StreamGate()
+        emitted = "".join(gate.feed(part) for part in ("Computing.", "<tool", "_call>", '{"name"', "}"))
+        self.assertEqual(emitted, "Computing.")
+        self.assertTrue(gate.suppressed)
+        self.assertEqual(gate.flush(), "")
+
+    def test_marker_split_across_chunks_is_held_back(self):
+        gate = _StreamGate()
+        self.assertEqual(gate.feed("done<too"), "done")
+        self.assertEqual(gate.feed("l_call>x"), "")
+
+    def test_partial_lookalike_is_released(self):
+        """`<t` that turns out to be something else must not be swallowed."""
+        gate = _StreamGate()
+        gate.feed("value <t")
+        self.assertEqual(gate.feed("hreshold>"), "<threshold>")
+
+
+class TestHistoryTruncation(unittest.TestCase):
+    def test_oldest_turns_dropped_first(self):
+        from controlai_agent import agent as agent_module
+
+        class FakeEngine:
+            def count_tokens(self, text):
+                return len(text.split())
+
+        holder = object.__new__(agent_module.ControlAgent)
+        holder.engine = FakeEngine()
+        history = [
+            {"role": "user", "content": "old " * 100},
+            {"role": "assistant", "content": "reply " * 100},
+            {"role": "user", "content": "recent question"},
+        ]
+        kept = holder._truncate(history)
+        self.assertEqual(kept[-1]["content"], "recent question")
+        self.assertLessEqual(len(kept), 3)
+
+    def test_blank_and_tool_turns_dropped(self):
+        from controlai_agent import agent as agent_module
+
+        class FakeEngine:
+            def count_tokens(self, text):
+                return len(text.split())
+
+        holder = object.__new__(agent_module.ControlAgent)
+        holder.engine = FakeEngine()
+        kept = holder._truncate(
+            [{"role": "tool", "content": "x"}, {"role": "user", "content": "   "}, {"role": "user", "content": "hi"}]
+        )
+        self.assertEqual(kept, [{"role": "user", "content": "hi"}])
+
+
+if __name__ == "__main__":
+    unittest.main()