Sync from GitHub f6782fa

README.md

@@ -105,3 +105,34 @@ is running show up without a restart. See `pipelines/mock_ask.py`.
   `flash_attention_2` if flash-attn is installed).
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## Tracing (Langfuse)
+
+Each find turn is traced to [Langfuse](https://langfuse.com) when keys are
+present, otherwise tracing is a complete no-op and the Space runs unchanged.
+Set these as **Space secrets** (or shell env / a `.env` for local runs):
+
+```bash
+LANGFUSE_PUBLIC_KEY=pk-lf-...
+LANGFUSE_SECRET_KEY=sk-lf-...
+LANGFUSE_HOST=https://cloud.langfuse.com   # or https://us.cloud.langfuse.com, or self-hosted
+```
+
+(`LANGFUSE_BASE_URL` is accepted as an alias for `LANGFUSE_HOST`.) Keys come from
+the Langfuse project's **Settings → API Keys**.
+
+One turn is one trace — an `agent` observation named `agent-find` (input = the
+mechanic's request, output = the terminal action), with a child per step:
+
+- `agent-decide` (`generation`) — the 1B brain picking a tool, with the resident
+  model id and input/output token counts.
+- `search` / `find_answer` (`retriever`) — the ColEmbed / dense lookups and their
+  page hits.
+- `ground-circle` (`generation`) — MiniCPM-V placing the circle box.
+
+The frontend sends a per-page-load `session_id`, so a visit's turns group into
+one **session** in the Langfuse UI. The wiring lives in `core/tracing.py` (a
+no-op-safe wrapper that builds and flushes the trace inside the ZeroGPU worker,
+attaching children to the root explicitly rather than via context — the turn is
+a streamed generator). The Langfuse working skill is vendored under
+`.claude/skills/langfuse/`.

app.py

@@ -300,6 +300,7 @@ def api_find(
     think: bool = GROUND_ENABLE_THINKING,
     agent_model: str = DEFAULT_AGENT_MODEL,
     vram_log: bool = False,
+    session_id: str = "",
 ) -> dict:  # the per-yield type: Server.api infers outputs from this annotation
     """One agent turn (one ZeroGPU call), streamed as events (see
     pipelines/agent_ask.py for the protocol). page/section are what the viewer
@@ -345,6 +346,7 @@ def api_find(
         events = PIPELINE.run_find(
             VISUAL_STORE, PARSED_STORE, request, [manual], int(k), options,
             viewer, history, bool(think), agent_model, bool(vram_log),
+            str(session_id or ""),
         )
         for ev in events:
             if ev.get("type") == "tool_result" and "gallery" in ev:

core/tracing.py

@@ -0,0 +1,188 @@
+"""Langfuse tracing for the find-and-point agent turn — optional, no-op safe.
+
+One agent turn (pipelines/agent_ask.agent_events) is the unit of observability:
+the root is an `agent` observation whose children are the brain's decisions
+(`generation`), the searches (`retriever`) and the circle grounding
+(`generation`). The mapping to Langfuse:
+
+    agent-find (agent)            input=the request, output=the terminal action
+    ├─ agent-decide (generation)  the 1B "brain" picks a tool      [models/minicpm_agent]
+    ├─ search / find_answer (retriever)  ColEmbed / dense lookup    [pipelines/agent_ask]
+    └─ ground-circle (generation) MiniCPM-V places the box          [models/minicpm]
+
+Two constraints shape the design, and are why this does NOT use the SDK's
+context-manager "current span" nesting:
+
+  1. The turn runs inside a @spaces.GPU worker (ZeroGPU forks a worker that
+     shares the resident CUDA models). The whole trace must therefore be built
+     and flushed *inside* that worker — so the root span is opened in
+     agent_events, not in the FastAPI endpoint that streams it.
+  2. agent_events is a generator. OTel's contextvar-based "current span" is
+     fragile across generator yields, so instead we open a root span object,
+     stash it on a threadlocal (one worker thread per turn), and attach every
+     child to it EXPLICITLY via parent.start_observation(...). No reliance on
+     the active context survives a yield.
+
+Tracing is fully disabled — every call a cheap no-op — unless
+LANGFUSE_PUBLIC_KEY and LANGFUSE_SECRET_KEY are set, so the Space (and local
+MOCK_MODELS runs) work unchanged without any Langfuse config.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import os
+import threading
+
+log = logging.getLogger("repairguy.tracing")
+
+# Keys present → tracing on. Read once at import; the Space sets these as Space
+# secrets, local dev via the shell / .env.
+_ENABLED = bool(
+    os.environ.get("LANGFUSE_PUBLIC_KEY") and os.environ.get("LANGFUSE_SECRET_KEY")
+)
+# The skill/CLI use LANGFUSE_BASE_URL; the SDK reads LANGFUSE_HOST. Bridge them
+# so a single var configures both.
+if os.environ.get("LANGFUSE_BASE_URL") and not os.environ.get("LANGFUSE_HOST"):
+    os.environ["LANGFUSE_HOST"] = os.environ["LANGFUSE_BASE_URL"]
+
+_client = None
+_client_failed = False
+# The root span of the turn currently running on this (worker) thread. Children
+# attach to it explicitly; cleared by finish_turn so a later untraced turn on
+# the same worker thread can't graft onto a stale parent.
+_state = threading.local()
+
+
+def _get_client():
+    """The shared Langfuse client, or None when disabled / unconfigurable.
+    Lazy so importing this module never touches the network or hard-requires
+    the SDK; cached, including the failure case."""
+    global _client, _client_failed
+    if not _ENABLED or _client_failed:
+        return None
+    if _client is None:
+        try:
+            from langfuse import get_client
+
+            _client = get_client()
+        except Exception as e:  # SDK missing or misconfigured — degrade to no-op
+            log.warning("Langfuse tracing disabled: %s", e)
+            _client_failed = True
+            return None
+    return _client
+
+
+def enabled() -> bool:
+    return _get_client() is not None
+
+
+def _parent():
+    return getattr(_state, "parent", None)
+
+
+def start_turn(
+    *,
+    name: str,
+    input,
+    session_id: str | None = None,
+    user_id: str | None = None,
+    tags: list[str] | None = None,
+    metadata: dict | None = None,
+):
+    """Open the root `agent` span for one turn and register it as the active
+    parent for this thread. Returns the span (or None when tracing is off).
+    Always pair with finish_turn in a finally."""
+    client = _get_client()
+    if client is None:
+        _state.parent = None
+        return None
+    try:
+        from langfuse import propagate_attributes
+
+        # session_id/user_id/tags are trace-level; propagate_attributes stamps
+        # them onto the span we create inside the context (and thus its trace).
+        # Coerce empty strings to None so we never set blank attributes.
+        with propagate_attributes(
+            session_id=session_id or None,
+            user_id=user_id or None,
+            tags=tags or None,
+            trace_name=name,
+        ):
+            span = client.start_observation(
+                name=name, as_type="agent", input=input, metadata=metadata
+            )
+        _state.parent = span
+        return span
+    except Exception as e:
+        log.warning("tracing start_turn failed: %s", e)
+        _state.parent = None
+        return None
+
+
+def finish_turn(span, *, output=None, level: str | None = None, status_message: str | None = None):
+    """Close the turn: record its output, end the root span, clear the parent,
+    and flush (the worker is short-lived, so unflushed events would be lost)."""
+    _state.parent = None
+    if span is None:
+        return
+    try:
+        if output is not None or level is not None or status_message is not None:
+            span.update(output=output, level=level, status_message=status_message)
+        span.end()
+    except Exception as e:
+        log.warning("tracing finish_turn failed: %s", e)
+    finally:
+        client = _get_client()
+        if client is not None:
+            with contextlib.suppress(Exception):
+                client.flush()
+
+
+@contextlib.contextmanager
+def generation(name: str, *, model: str | None = None, input=None, metadata: dict | None = None):
+    """A child `generation` under the active turn — one model decode. Yields the
+    span (or None when off / no active turn) so the caller can .update(output=,
+    usage_details=). Synchronous body only (no yields inside), so the parent
+    link is explicit and contextvar-safe."""
+    parent = _parent()
+    if parent is None:
+        yield None
+        return
+    gen = None
+    try:
+        gen = parent.start_observation(
+            name=name, as_type="generation", model=model, input=input, metadata=metadata
+        )
+    except Exception as e:
+        log.warning("tracing generation(%s) failed: %s", name, e)
+    try:
+        yield gen
+    finally:
+        if gen is not None:
+            with contextlib.suppress(Exception):
+                gen.end()
+
+
+@contextlib.contextmanager
+def retriever(name: str, *, input=None, metadata: dict | None = None):
+    """A child `retriever` under the active turn — one search/lookup. Yields the
+    span (or None) for the caller to .update(output=the hits)."""
+    parent = _parent()
+    if parent is None:
+        yield None
+        return
+    span = None
+    try:
+        span = parent.start_observation(
+            name=name, as_type="retriever", input=input, metadata=metadata
+        )
+    except Exception as e:
+        log.warning("tracing retriever(%s) failed: %s", name, e)
+    try:
+        yield span
+    finally:
+        if span is not None:
+            with contextlib.suppress(Exception):
+                span.end()

frontend/index.html

@@ -433,6 +433,9 @@ const ICONS = {
 function repairGuy() {
   return {
     client:null, ready:false,
+    // One id per page load — groups this visit's find turns into a single
+    // Langfuse session so a whole conversation can be replayed together.
+    sessionId:(self.crypto&&crypto.randomUUID)?crypto.randomUUID():'sess-'+Date.now()+'-'+Math.random().toString(36).slice(2),
     manuals:[], manual:'',
     k:__DEFAULT_K__, maxK:__MAX_K__,
     // When on, MiniCPM-V reasons before committing to a circle box (slower,
@@ -731,6 +734,7 @@ function repairGuy() {
       try{
         const job = this.client.submit('/find', {
           request:q, manual:this.manual, k:this.k, think:this.think, agent_model:this.agentModel,
+          session_id:this.sessionId,
           page: this.viewDoc ? this.viewPage : 0, section: this.sectionTitle(this.viewPage),
           pages: this.spreadPages,  // every page currently on screen — the agent may circle on any
           history: this.history,

models/minicpm.py

@@ -25,6 +25,7 @@ import torch
 from PIL import Image
 from transformers import AutoModel, AutoProcessor, AutoTokenizer
 
+from core import tracing
 from core.constants import (
     ANSWER_MAX_NEW_TOKENS,
     DESCRIBE_MAX_NEW_TOKENS,
@@ -171,24 +172,33 @@ def ground_box(
     GROUND_ENABLE_THINKING default; the UI settings panel passes an explicit
     bool per request."""
     think = GROUND_ENABLE_THINKING if enable_thinking is None else enable_thinking
-    with torch.no_grad():
-        out = _MODEL.chat(
-            msgs=[
-                {
-                    "role": "user",
-                    "content": [image.convert("RGB"), GROUND_PROMPT.format(query=query)],
-                }
-            ],
-            tokenizer=_TOKENIZER,
-            # The think trace needs room (a bare box fits in 64 tokens; a think
-            # trace does not) or it gets cut off before emitting the box — so
-            # the token budget tracks the flag.
-            enable_thinking=think,
-            max_new_tokens=(
-                GROUND_THINK_MAX_NEW_TOKENS if think else GROUND_BOX_MAX_NEW_TOKENS
-            ),
-        )
-    raw = str(out).strip()
+    # One `generation` (the VLM "eyes" placing the box): query in, raw reply out.
+    with tracing.generation(
+        "ground-circle",
+        model=MINICPM_MODEL_ID,
+        input=query,
+        metadata={"thinking": bool(think)},
+    ) as gen:
+        with torch.no_grad():
+            out = _MODEL.chat(
+                msgs=[
+                    {
+                        "role": "user",
+                        "content": [image.convert("RGB"), GROUND_PROMPT.format(query=query)],
+                    }
+                ],
+                tokenizer=_TOKENIZER,
+                # The think trace needs room (a bare box fits in 64 tokens; a think
+                # trace does not) or it gets cut off before emitting the box — so
+                # the token budget tracks the flag.
+                enable_thinking=think,
+                max_new_tokens=(
+                    GROUND_THINK_MAX_NEW_TOKENS if think else GROUND_BOX_MAX_NEW_TOKENS
+                ),
+            )
+        raw = str(out).strip()
+        if gen is not None:
+            gen.update(output=raw)
     # The thinking trace can mention "not found" mid-reasoning ("at first this
     # looked not found, but…"), so test only the FINAL answer after </think>,
     # not the whole reply.

models/minicpm_agent.py

@@ -42,6 +42,7 @@ import re
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+from core import tracing
 from core.constants import (
     AGENT_MAX_NEW_TOKENS,
     AGENT_MODELS,
@@ -254,6 +255,11 @@ def _spec(key: str | None) -> dict:
     return _REGISTRY.get(key or "", _REGISTRY[DEFAULT_AGENT_MODEL])
 
 
+def _active_model_id() -> str:
+    """The HF id of the resident brain — the `model` for its generation spans."""
+    return _spec(_active_key).get("model_id", _active_key or "unknown")
+
+
 def use_model(key: str | None = None) -> str:
     """Make `key` the resident agent brain, loading it (and evicting the
     previous one) when it isn't already active — one model in VRAM at a time.
@@ -319,25 +325,40 @@ def _template_kwargs() -> dict:
     return {"enable_thinking": False} if _THINKING else {}
 
 
-def _generate(messages: list[dict], max_new_tokens: int) -> str:
-    """Greedy decode the assistant's next message."""
-    inputs = _TOKENIZER.apply_chat_template(
-        messages,
-        tokenize=True,
-        add_generation_prompt=True,
-        return_dict=True,
-        return_tensors="pt",
-        **_template_kwargs(),
-    ).to(_MODEL.device)
-    # apply_chat_template emits token_type_ids, which this LlamaForCausalLM's
-    # generate() rejects as an unused kwarg.
-    inputs.pop("token_type_ids", None)
-    with torch.no_grad():
-        out = _MODEL.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False)
-    text = _TOKENIZER.decode(
-        out[0, inputs["input_ids"].shape[1] :], skip_special_tokens=True
-    )
-    return text.strip()
+def _generate(
+    messages: list[dict], max_new_tokens: int, trace_name: str = "agent-generate"
+) -> str:
+    """Greedy decode the assistant's next message. Traced as one `generation`
+    (the resident brain as the model, the messages as input, the reply and the
+    in/out token counts attached) when Langfuse is configured."""
+    with tracing.generation(
+        trace_name, model=_active_model_id(), input=messages
+    ) as gen:
+        inputs = _TOKENIZER.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_dict=True,
+            return_tensors="pt",
+            **_template_kwargs(),
+        ).to(_MODEL.device)
+        # apply_chat_template emits token_type_ids, which this LlamaForCausalLM's
+        # generate() rejects as an unused kwarg.
+        inputs.pop("token_type_ids", None)
+        n_in = int(inputs["input_ids"].shape[1])
+        with torch.no_grad():
+            out = _MODEL.generate(
+                **inputs, max_new_tokens=max_new_tokens, do_sample=False
+            )
+        new_ids = out[0, n_in:]
+        text = _TOKENIZER.decode(new_ids, skip_special_tokens=True)
+        if gen is not None:
+            n_out = int(new_ids.shape[0])
+            gen.update(
+                output=text.strip(),
+                usage_details={"input": n_in, "output": n_out, "total": n_in + n_out},
+            )
+        return text.strip()
 
 
 def render_prompt(messages: list[dict]) -> str:
@@ -419,7 +440,7 @@ def decide(messages: list[dict]) -> tuple[dict | None, str]:
     maintains (system + past turns + this turn's state and any tool results).
     Returns (parsed tool call, raw reply); the tool is None when the reply isn't
     a usable JSON tool call. Must run on GPU."""
-    raw = _generate(messages, AGENT_MAX_NEW_TOKENS)
+    raw = _generate(messages, AGENT_MAX_NEW_TOKENS, trace_name="agent-decide")
     return _parse_tool(raw), raw
 
 
@@ -443,7 +464,9 @@ def rerank(query: str, candidates: list[tuple[int, str]]) -> tuple[int, str]:
         f"PAGE {page}:\n{text or '(no text)'}" for page, text in candidates
     )
     prompt = RERANK_PROMPT.format(query=query, n=len(candidates), candidates=listing)
-    raw = _generate([{"role": "user", "content": prompt}], max_new_tokens=8)
+    raw = _generate(
+        [{"role": "user", "content": prompt}], max_new_tokens=8, trace_name="agent-rerank"
+    )
     m = re.search(r"\d+", raw)
     if m:
         picked = int(m.group())

pipelines/agent_ask.py

@@ -47,6 +47,7 @@ import logging
 
 import spaces
 
+from core import tracing
 from core.constants import (
     AGENT_HISTORY_TURNS,
     AGENT_MAX_STEPS,
@@ -91,6 +92,7 @@ def agent_events(
     ground_thinking: bool | None = None,
     agent_model: str | None = None,
     vram_log: bool = False,
+    session_id: str | None = None,
 ):
     """Yield the events of one agent turn (see module docstring). sections is the
     numbered table of contents shown to the agent ([{title, page}]); the agent's
@@ -182,167 +184,205 @@ def agent_events(
             )
         )
 
-    for step in range(AGENT_MAX_STEPS):
-        # Render the exact prompt BEFORE deciding so the trace can show what the
-        # brain was asked, not just what it answered.
-        prompt = minicpm_agent.render_prompt(messages)
-        tool, raw = minicpm_agent.decide(messages)
-        log.info("step %d: tool=%s | raw=%r", step, tool, raw[:200])
-        # Diagnostic event: the prompt fed in, the raw 1B reply, and the parsed
-        # tool for this step, so the UI's trace view shows exactly what the brain
-        # was asked and decided (and why a reply was rejected). Not used by the
-        # normal chip flow.
-        yield {"type": "trace", "step": step, "tool": tool, "raw": raw,
-               "prompt": prompt}
-        if tool is None:
-            # Unusable reply (bad JSON, or an echoed placeholder target). Correct
-            # it and let the agent try again rather than abandon the turn.
-            messages.append(
-                minicpm_agent.tool_result_message(
-                    "Your last reply was not one complete JSON object. Reply with "
-                    "ONE complete JSON object and nothing else, e.g. "
-                    '{"tool": "search", "query": "fuel filter"}. If you circle, the '
-                    "target MUST be copied from the page text above — never invent "
-                    "a part that is not printed there."
-                )
-            )
-            continue
-        messages.append(minicpm_agent.assistant_action_message(tool))
-
-        if tool["tool"] == "go_to_section":
-            idx = tool["section"] - 1
-            if not 0 <= idx < len(sections):
+    # Open the trace for this turn (no-op when Langfuse is unconfigured). The
+    # whole loop runs under try/finally so the root span is always ended and
+    # flushed — on a terminal return, a mid-turn error, OR an early client
+    # disconnect (GeneratorExit raised at a yield). `turn_output` is the terminal
+    # `done` event, recorded as the trace's output. Children (the brain's
+    # decisions, searches, the grounding) attach to this span explicitly.
+    span = tracing.start_turn(
+        name="agent-find",
+        input=request,
+        session_id=session_id,
+        tags=[t for t in (manual, active) if t],
+        metadata={
+            "manual": manual,
+            "agent_model": active,
+            "k": int(top_k),
+            "thinking": bool(ground_thinking),
+            "viewer_pages": shown_pages,
+        },
+    )
+    turn_output = None
+    try:
+        for step in range(AGENT_MAX_STEPS):
+            # Render the exact prompt BEFORE deciding so the trace can show what the
+            # brain was asked, not just what it answered.
+            prompt = minicpm_agent.render_prompt(messages)
+            tool, raw = minicpm_agent.decide(messages)
+            log.info("step %d: tool=%s | raw=%r", step, tool, raw[:200])
+            # Diagnostic event: the prompt fed in, the raw 1B reply, and the parsed
+            # tool for this step, so the UI's trace view shows exactly what the brain
+            # was asked and decided (and why a reply was rejected). Not used by the
+            # normal chip flow.
+            yield {"type": "trace", "step": step, "tool": tool, "raw": raw,
+                   "prompt": prompt}
+            if tool is None:
+                # Unusable reply (bad JSON, or an echoed placeholder target). Correct
+                # it and let the agent try again rather than abandon the turn.
                 messages.append(
                     minicpm_agent.tool_result_message(
-                        f"There is no section {tool['section']}. Pick a number from "
-                        "the table of contents, or use search."
+                        "Your last reply was not one complete JSON object. Reply with "
+                        "ONE complete JSON object and nothing else, e.g. "
+                        '{"tool": "search", "query": "fuel filter"}. If you circle, the '
+                        "target MUST be copied from the page text above — never invent "
+                        "a part that is not printed there."
                     )
                 )
                 continue
-            opt = sections[idx]
-            yield {"type": "step", "tool": "go_to_section",
-                   "title": opt["title"], "page": int(opt["page"])}
-            yield {"type": "done", "kind": "navigate", "nav": "section",
-                   "page": int(opt["page"]), "title": opt["title"]}
-            return
+            messages.append(minicpm_agent.assistant_action_message(tool))
 
-        if tool["tool"] == "go_to_page":
-            page = tool["page"]
-            n = page_count(visual_store.pdf_path(doc_id))
-            if not 1 <= page <= n:
-                messages.append(
-                    minicpm_agent.tool_result_message(
-                        f"There is no page {page}; this manual has pages 1–{n}. "
-                        "Pick a page in range, search, or go to a section."
+            if tool["tool"] == "go_to_section":
+                idx = tool["section"] - 1
+                if not 0 <= idx < len(sections):
+                    messages.append(
+                        minicpm_agent.tool_result_message(
+                            f"There is no section {tool['section']}. Pick a number from "
+                            "the table of contents, or use search."
+                        )
                     )
-                )
-                continue
-            yield {"type": "step", "tool": "go_to_page", "page": page}
-            yield {"type": "done", "kind": "navigate", "nav": "page",
-                   "page": page, "title": f"Page {page}"}
-            return
+                    continue
+                opt = sections[idx]
+                yield {"type": "step", "tool": "go_to_section",
+                       "title": opt["title"], "page": int(opt["page"])}
+                turn_output = {"type": "done", "kind": "navigate", "nav": "section",
+                               "page": int(opt["page"]), "title": opt["title"]}
+                yield turn_output
+                return
 
-        if tool["tool"] == "search":
-            query = tool["query"]
-            yield {"type": "step", "tool": "search", "query": query}
-            yield {"type": "status", "text": f"Searching for “{query}”…"}
-            # k (the viewer's slider) is the shortlist size; ColEmbed's top page
-            # is the one shown. A 1B text rerank measured WORSE than raw ColEmbed
-            # top-1 (0.68 vs 0.84 hit@1) — visual late interaction already ranks
-            # these (figure-heavy) pages better than re-judging from page text.
-            hits = maxsim_search(query, visual_store, doc_ids, top_k)
-            log.info("search(%r) → %s", query, [(p, round(s, 3)) for _, p, s in hits])
-            if not hits:
-                messages.append(
-                    minicpm_agent.tool_result_message(f"Search for {query!r} found nothing.")
-                )
-                continue
-            yield from present_hits(hits, "search:" + " ".join(query.lower().split()))
-            continue
+            if tool["tool"] == "go_to_page":
+                page = tool["page"]
+                n = page_count(visual_store.pdf_path(doc_id))
+                if not 1 <= page <= n:
+                    messages.append(
+                        minicpm_agent.tool_result_message(
+                            f"There is no page {page}; this manual has pages 1–{n}. "
+                            "Pick a page in range, search, or go to a section."
+                        )
+                    )
+                    continue
+                yield {"type": "step", "tool": "go_to_page", "page": page}
+                turn_output = {"type": "done", "kind": "navigate", "nav": "page",
+                               "page": page, "title": f"Page {page}"}
+                yield turn_output
+                return
 
-        if tool["tool"] == "find_answer":
-            query = tool["query"]
-            yield {"type": "step", "tool": "find_answer", "query": query}
-            yield {"type": "status", "text": f"Looking up “{query}”…"}
-            # Dense retrieval over the PARSED chunks (text/semantic) — the index
-            # the parsed store was built for. A fact lookup ("what fuel does it
-            # take") is a TEXT match: ColEmbed ranks pages by VISUAL similarity
-            # and misses the plain specs page, so fact questions route here. Same
-            # (doc_id, page, score) shape as maxsim_search; the agent then circles
-            # the answering line on the page shown.
-            hits = retrieve_pages(query, parsed_store, doc_ids, top_k)
-            log.info("find_answer(%r) → %s", query, [(p, round(s, 3)) for _, p, s in hits])
-            if not hits:
-                messages.append(
-                    minicpm_agent.tool_result_message(f"Looking up {query!r} found nothing.")
-                )
+            if tool["tool"] == "search":
+                query = tool["query"]
+                yield {"type": "step", "tool": "search", "query": query}
+                yield {"type": "status", "text": f"Searching for “{query}”…"}
+                # k (the viewer's slider) is the shortlist size; ColEmbed's top page
+                # is the one shown. A 1B text rerank measured WORSE than raw ColEmbed
+                # top-1 (0.68 vs 0.84 hit@1) — visual late interaction already ranks
+                # these (figure-heavy) pages better than re-judging from page text.
+                with tracing.retriever("search", input=query,
+                                       metadata={"retriever": "colembed", "k": int(top_k)}) as rsp:
+                    hits = maxsim_search(query, visual_store, doc_ids, top_k)
+                    if rsp is not None:
+                        rsp.update(output=[{"page": p, "score": round(float(s), 4)}
+                                           for _, p, s in hits])
+                log.info("search(%r) → %s", query, [(p, round(s, 3)) for _, p, s in hits])
+                if not hits:
+                    messages.append(
+                        minicpm_agent.tool_result_message(f"Search for {query!r} found nothing.")
+                    )
+                    continue
+                yield from present_hits(hits, "search:" + " ".join(query.lower().split()))
                 continue
-            yield from present_hits(hits, "answer:" + " ".join(query.lower().split()))
-            continue
 
-        if tool["tool"] == "circle":
-            target = tool["target"]
-            # The agent says which shown page the target is on (it has both pages'
-            # text). Default to the active page when it's unspecified or not one of
-            # the pages on screen — so the box is grounded on, and drawn over, the
-            # RIGHT page.
-            page = tool.get("page")
-            if page not in circleable:
-                page = current_page
-            yield {"type": "step", "tool": "circle", "target": target, "page": page}
-            yield {"type": "status", "text": "Pinning it down…"}
-            img = render_page(visual_store.pdf_path(doc_id), page)
-            box, braw = minicpm.ground_box(img, target, enable_thinking=ground_thinking)
-            # Heaviest op of the turn — the VLM vision encoder runs on a full-res
-            # page. peak here is the turn's activation high-water (since
-            # turn-start), the number that decides whether a big brain still fits.
-            log_vram("after-ground")
-            log.info("ground_box(%r) on p.%d → %s | raw=%r",
-                     target, page, box, braw[:200])
-            # The VLM couldn't find the target on this page — almost always
-            # because it's on a DIFFERENT page (the agent circled too early).
-            # Don't end the turn with an empty pin: push it to relocate and try
-            # again. Only fall through to showing the page un-pinned once we've
-            # already missed this exact (page, target) — a repeat means retrying
-            # here won't help, same guard as the no-op search.
-            tkey = (page, " ".join(target.lower().split()))
-            if box is None and tkey not in ground_failed:
-                ground_failed.add(tkey)
-                messages.append(
-                    minicpm_agent.tool_result_message(
-                        minicpm_agent.ground_failed_message(request, target, page)
+            if tool["tool"] == "find_answer":
+                query = tool["query"]
+                yield {"type": "step", "tool": "find_answer", "query": query}
+                yield {"type": "status", "text": f"Looking up “{query}”…"}
+                # Dense retrieval over the PARSED chunks (text/semantic) — the index
+                # the parsed store was built for. A fact lookup ("what fuel does it
+                # take") is a TEXT match: ColEmbed ranks pages by VISUAL similarity
+                # and misses the plain specs page, so fact questions route here. Same
+                # (doc_id, page, score) shape as maxsim_search; the agent then circles
+                # the answering line on the page shown.
+                with tracing.retriever("find_answer", input=query,
+                                       metadata={"retriever": "parsed-dense", "k": int(top_k)}) as rsp:
+                    hits = retrieve_pages(query, parsed_store, doc_ids, top_k)
+                    if rsp is not None:
+                        rsp.update(output=[{"page": p, "score": round(float(s), 4)}
+                                           for _, p, s in hits])
+                log.info("find_answer(%r) → %s", query, [(p, round(s, 3)) for _, p, s in hits])
+                if not hits:
+                    messages.append(
+                        minicpm_agent.tool_result_message(f"Looking up {query!r} found nothing.")
                     )
-                )
+                    continue
+                yield from present_hits(hits, "answer:" + " ".join(query.lower().split()))
                 continue
-            yield {
-                "type": "done",
-                "kind": "point",
-                "found": True,
-                "target": target,
-                "page": page,
-                "bbox": [round(v) for v in box] if box is not None else None,
-                # The pixel size of the image the box was GROUNDED on — the bbox
-                # is in this coordinate space. The frontend sizes its SVG viewBox
-                # from this (not the browser-loaded <img>), so the circle lands
-                # correctly even if the displayed page PNG is served at a
-                # different/stale resolution than this grounding render.
-                "dims": [img.width, img.height],
-                # the VLM's raw grounding reply — diagnostic only (helps explain
-                # where/why a box landed); shown in the trace view.
-                "ground_raw": braw[:300],
-            }
-            return
 
-        if tool["tool"] == "done":
-            yield {"type": "done", "kind": "reply",
-                   "message": tool.get("message") or "Done."}
-            return
+            if tool["tool"] == "circle":
+                target = tool["target"]
+                # The agent says which shown page the target is on (it has both pages'
+                # text). Default to the active page when it's unspecified or not one of
+                # the pages on screen — so the box is grounded on, and drawn over, the
+                # RIGHT page.
+                page = tool.get("page")
+                if page not in circleable:
+                    page = current_page
+                yield {"type": "step", "tool": "circle", "target": target, "page": page}
+                yield {"type": "status", "text": "Pinning it down…"}
+                img = render_page(visual_store.pdf_path(doc_id), page)
+                box, braw = minicpm.ground_box(img, target, enable_thinking=ground_thinking)
+                # Heaviest op of the turn — the VLM vision encoder runs on a full-res
+                # page. peak here is the turn's activation high-water (since
+                # turn-start), the number that decides whether a big brain still fits.
+                log_vram("after-ground")
+                log.info("ground_box(%r) on p.%d → %s | raw=%r",
+                         target, page, box, braw[:200])
+                # The VLM couldn't find the target on this page — almost always
+                # because it's on a DIFFERENT page (the agent circled too early).
+                # Don't end the turn with an empty pin: push it to relocate and try
+                # again. Only fall through to showing the page un-pinned once we've
+                # already missed this exact (page, target) — a repeat means retrying
+                # here won't help, same guard as the no-op search.
+                tkey = (page, " ".join(target.lower().split()))
+                if box is None and tkey not in ground_failed:
+                    ground_failed.add(tkey)
+                    messages.append(
+                        minicpm_agent.tool_result_message(
+                            minicpm_agent.ground_failed_message(request, target, page)
+                        )
+                    )
+                    continue
+                turn_output = {
+                    "type": "done",
+                    "kind": "point",
+                    "found": True,
+                    "target": target,
+                    "page": page,
+                    "bbox": [round(v) for v in box] if box is not None else None,
+                    # The pixel size of the image the box was GROUNDED on — the bbox
+                    # is in this coordinate space. The frontend sizes its SVG viewBox
+                    # from this (not the browser-loaded <img>), so the circle lands
+                    # correctly even if the displayed page PNG is served at a
+                    # different/stale resolution than this grounding render.
+                    "dims": [img.width, img.height],
+                    # the VLM's raw grounding reply — diagnostic only (helps explain
+                    # where/why a box landed); shown in the trace view.
+                    "ground_raw": braw[:300],
+                }
+                yield turn_output
+                return
 
-    yield {
-        "type": "done",
-        "kind": "reply",
-        "message": "I went in circles on that one — try rephrasing?",
-    }
+            if tool["tool"] == "done":
+                turn_output = {"type": "done", "kind": "reply",
+                               "message": tool.get("message") or "Done."}
+                yield turn_output
+                return
+
+        turn_output = {
+            "type": "done",
+            "kind": "reply",
+            "message": "I went in circles on that one — try rephrasing?",
+        }
+        yield turn_output
+    finally:
+        tracing.finish_turn(span, output=turn_output)
 
 
 class AgentPipeline:
@@ -362,9 +402,11 @@ class AgentPipeline:
         ground_thinking: bool | None = None,
         agent_model: str | None = None,
         vram_log: bool = False,
+        session_id: str | None = None,
     ):
         """One streamed agent turn (the event generator of agent_events).
-        vram_log forwards the UI's VRAM-logging toggle to the probe."""
+        vram_log forwards the UI's VRAM-logging toggle to the probe; session_id
+        groups a page session's turns together in Langfuse."""
         request = (request or "").strip()
         if not request:
             raise ValueError("Tell me what to find.")
@@ -375,5 +417,5 @@ class AgentPipeline:
         return agent_events(
             request, visual_store, parsed_store, doc_ids or list(names),
             int(top_k), names, sections, viewer, history, ground_thinking,
-            agent_model, vram_log,
+            agent_model, vram_log, session_id,
         )

pipelines/mock_ask.py

@@ -145,11 +145,13 @@ class MockAskPipeline:
         history: list | None = None,
         ground_thinking: bool | None = None,
         agent_model: str | None = None,
+        vram_log: bool = False,
+        session_id: str | None = None,
     ):
         """Yield the same event sequence as pipelines/agent_ask.py, with a
         keyword-driven stand-in for the agent's tool choice. Both stores are the
-        one MockStore; history, ground_thinking and agent_model are ignored.
-        MOCK_DELAY (seconds) paces events."""
+        one MockStore; history, ground_thinking, agent_model, vram_log and
+        session_id are ignored. MOCK_DELAY (seconds) paces events."""
         request = (request or "").strip()
         if not request:
             raise ValueError("Tell me what to find.")

requirements.txt

@@ -6,3 +6,4 @@ torchvision
 pymupdf
 pillow
 numpy
+langfuse>=3,<5