From 3b3b386a8651d582d1dec08597781533aa0cc02d Mon Sep 17 00:00:00 2001 From: Mohamed Elkholy Date: Mon, 27 Apr 2026 00:01:57 -0400 Subject: [PATCH 01/26] chore(obs): add OpenTelemetry SDK and Arize Phoenix deps --- requirements.txt | 3 +++ 1 file changed, 3 insertions(+) diff --git a/requirements.txt b/requirements.txt index de2d70de..fcfcb1db 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,3 +17,6 @@ chromadb>=1.0.0 datasets>=2.16,<3 sentence-transformers>=2.7 sqlmodel>=0.0.24 +arize-phoenix>=4.0 +opentelemetry-exporter-otlp-proto-http>=1.25 +opentelemetry-sdk>=1.25 From 2a00c47306be25bb5f5484566bf91df28048f042 Mon Sep 17 00:00:00 2001 From: Mohamed Elkholy Date: Mon, 27 Apr 2026 00:04:30 -0400 Subject: [PATCH 02/26] feat(obs): add OpenTelemetry tracer init and traced_stage decorator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduces src/observability.py with three public symbols: - init_observability(): idempotent, OTLP exporter setup, fails quietly - get_tracer(): always returns a usable tracer (no-op if uninit) - traced_stage(name): (payload, attrs) → span decorator with coercion All 6 tests in tests/test_observability.py pass. --- src/observability.py | 189 ++++++++++++++++++++++++++++++++++++ tests/test_observability.py | 99 +++++++++++++++++++ 2 files changed, 288 insertions(+) create mode 100644 src/observability.py create mode 100644 tests/test_observability.py diff --git a/src/observability.py b/src/observability.py new file mode 100644 index 00000000..c6e3dda2 --- /dev/null +++ b/src/observability.py @@ -0,0 +1,189 @@ +""" +Observability — OpenTelemetry tracer initialization and per-stage +span decorator. + +API Layer Position: + RAGBackend.query → @traced_stage("rag.retrieve") → span exported via OTLP + → Phoenix UI at :6006 + +Design decisions: + - Idempotent init via a module-level flag — multiple imports / lifespan + callbacks won't double-install processors. + - Fail QUIETLY on exporter / endpoint errors. Eval and chat must work + when Phoenix is down; observability is opt-in. + - traced_stage uses the (payload, attrs) return convention so existing + return shapes don't change at call sites; the decorator strips attrs. + - Attribute coercion: OTel restricts attribute types. Non-primitives get + str()-coerced rather than dropped (visibility > strictness). +""" + +from __future__ import annotations + +import logging +import os +from functools import wraps +from typing import Any, Callable, TypeVar + +logger = logging.getLogger(__name__) + +TRACER_NAME = "rag-qa" +_INITIALIZED = False + +# OTel primitive types that span.set_attribute accepts as scalars. +_OTEL_SCALARS = (str, int, float, bool) + + +def init_observability(otlp_endpoint: str | None = None) -> None: + """Initialize global TracerProvider + OTLP exporter. + + Idempotent — safe to call multiple times. + ``otlp_endpoint`` defaults to the ``OTLP_ENDPOINT`` env var, or + ``http://localhost:6006/v1/traces`` when neither is set. + + On import errors or connection failures, fails QUIETLY (logs warning, + spans become no-ops). The system never crashes due to OTel. + + Args: + otlp_endpoint: OTLP HTTP traces endpoint URL. Pass ``None`` to + use the default derived from the environment variable. + """ + # PATTERN: module-level flag makes this idempotent — calling from + # multiple lifespan callbacks or test setups is safe. + global _INITIALIZED + if _INITIALIZED: + return + _INITIALIZED = True + + endpoint = otlp_endpoint or os.getenv( + "OTLP_ENDPOINT", "http://localhost:6006/v1/traces" + ) + + try: + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import BatchSpanProcessor + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter, + ) + import opentelemetry.trace as otel_trace + + # WHY: OTLPSpanExporter is lazy — it won't attempt a connection + # until the first batch is flushed, so construction never raises + # even if the endpoint is unreachable. + exporter = OTLPSpanExporter(endpoint=endpoint) + provider = TracerProvider() + provider.add_span_processor(BatchSpanProcessor(exporter)) + otel_trace.set_tracer_provider(provider) + + except Exception as exc: # noqa: BLE001 + # TRADE-OFF: We catch broadly here because we never want Phoenix + # being unavailable to crash the RAG service. A warning is enough. + logger.warning( + "Observability init failed — spans will be no-ops. Reason: %s", exc + ) + + +def get_tracer(): + """Return the rag-qa tracer (always works; no-op if init never called). + + Returns: + An ``opentelemetry.trace.Tracer`` instance bound to TRACER_NAME. + When no provider has been installed this returns the global no-op + tracer, so call sites never need to guard for None. + """ + import opentelemetry.trace as otel_trace + + # WHY: get_tracer() delegates to whatever TracerProvider is currently + # registered globally. If init_observability was never called that's + # the SDK default (no-op), which is perfectly fine. + return otel_trace.get_tracer(TRACER_NAME) + + +def _coerce_attr(value: Any) -> Any: + """Coerce a value to an OTel-legal span attribute type. + + OTel accepts: str, int, float, bool, or homogeneous Sequence thereof. + Everything else is str()-converted so we retain visibility over dropping. + + Args: + value: Raw attribute value from the decorated function. + + Returns: + A value safe to pass to ``span.set_attribute``. + """ + if isinstance(value, _OTEL_SCALARS): + return value + + # Sequences: pass through only when every element is a scalar of the + # same OTel-legal type. Mixed or complex elements fall back to str(). + if isinstance(value, (list, tuple)): + if value and all(isinstance(el, _OTEL_SCALARS) for el in value): + # OTel SDK coerces list → tuple internally; list is fine here. + return value + # Empty or mixed-type list — convert whole thing. + return str(value) + + # Dicts, None, and anything else: str-coerce for visibility. + return str(value) + + +def traced_stage(name: str): + """Decorator that opens a span around the wrapped function. + + The wrapped function MUST return ``(payload, attrs_dict)``. The + decorator opens a span named ``name``, calls the function, records + each ``attrs_dict`` entry as a span attribute, then returns only + ``payload`` — callers see the same shape as before instrumentation. + + On exception the span is marked ERROR and the exception is re-raised. + + Args: + name: Span name (e.g. ``"rag.retrieve"``). + + Returns: + A decorator that wraps ``f(*args, **kwargs) -> tuple[Any, dict]`` + and exposes only the payload to callers. + + Example:: + + @traced_stage("rag.retrieve") + def retrieve(query: str): + chunks = _do_search(query) + return chunks, {"chunk_count": len(chunks)} + + result = retrieve("what is RAG?") # returns chunks directly + """ + def decorator(f: Callable) -> Callable: + @wraps(f) + def wrapper(*args: Any, **kwargs: Any) -> Any: + tracer = get_tracer() + with tracer.start_as_current_span(name) as span: + try: + payload, attrs = f(*args, **kwargs) + except Exception as exc: + # PATTERN: Record exception details on the span so + # Phoenix shows the stack trace, then propagate. + from opentelemetry.trace import Status, StatusCode + span.record_exception(exc) + span.set_status(Status(StatusCode.ERROR)) + raise + + # Set attributes after successful return. + for key, raw_val in attrs.items(): + coerced = _coerce_attr(raw_val) + try: + span.set_attribute(key, coerced) + except Exception: # noqa: BLE001 + logger.warning( + "traced_stage: could not set attribute %r=%r, skipping", + key, + coerced, + ) + + from opentelemetry.trace import Status, StatusCode + span.set_status(Status(StatusCode.OK)) + + return payload + + return wrapper + + return decorator diff --git a/tests/test_observability.py b/tests/test_observability.py new file mode 100644 index 00000000..8e93f929 --- /dev/null +++ b/tests/test_observability.py @@ -0,0 +1,99 @@ +"""Tests for src.observability.""" + +from __future__ import annotations + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) + +from src.observability import ( + TRACER_NAME, + get_tracer, + init_observability, + traced_stage, +) + + +@pytest.fixture +def in_memory_exporter(): + """Install an InMemorySpanExporter on a fresh TracerProvider for the test.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + # Force-install our test provider, overriding any prior init. + trace._TRACER_PROVIDER = provider # type: ignore[attr-defined] + yield exporter + exporter.clear() + + +class TestInitObservability: + def test_idempotent(self): + """Calling init twice doesn't crash and doesn't double-install.""" + init_observability() + init_observability() + # No assertion — just doesn't raise. + + def test_unreachable_endpoint_does_not_raise(self): + """A bad endpoint is logged but doesn't crash.""" + # Reset the idempotency flag so this call actually attempts init. + import src.observability as obs + obs._INITIALIZED = False # type: ignore[attr-defined] + init_observability(otlp_endpoint="http://127.0.0.1:1/v1/traces") + # Subsequent traced spans must still work (as no-ops or local). + @traced_stage("test.bad-endpoint") + def f(): + return "x", {"k": 1} + assert f() == "x" + + +class TestGetTracer: + def test_returns_tracer(self): + t = get_tracer() + assert t is not None + + +class TestTracedStage: + def test_records_span_with_attrs(self, in_memory_exporter): + @traced_stage("rag.retrieve") + def retrieve(query: str): + return ["chunk1", "chunk2"], {"top_k": 5, "chunk_count": 2} + + result = retrieve("test query") + assert result == ["chunk1", "chunk2"] + + spans = in_memory_exporter.get_finished_spans() + assert len(spans) == 1 + span = spans[0] + assert span.name == "rag.retrieve" + assert span.attributes["top_k"] == 5 + assert span.attributes["chunk_count"] == 2 + + def test_propagates_exception_and_records_error(self, in_memory_exporter): + @traced_stage("rag.fail") + def boom(): + raise RuntimeError("kaboom") + + with pytest.raises(RuntimeError, match="kaboom"): + boom() + + spans = in_memory_exporter.get_finished_spans() + assert len(spans) == 1 + # Span recorded ERROR status + from opentelemetry.trace import StatusCode + assert spans[0].status.status_code == StatusCode.ERROR + + def test_attribute_coercion(self, in_memory_exporter): + """Non-primitive attribute values are coerced to str.""" + @traced_stage("rag.coerce") + def f(): + return None, {"a_dict": {"x": 1}, "a_list_of_str": ["a", "b"]} + f() + spans = in_memory_exporter.get_finished_spans() + # a_dict gets coerced to str representation; a_list_of_str passes through. + attrs = spans[0].attributes + assert "a_dict" in attrs + assert tuple(attrs["a_list_of_str"]) == ("a", "b") From 895bd1d6287675667c1127f39c4f395616cad7b4 Mon Sep 17 00:00:00 2001 From: Mohamed Elkholy Date: Mon, 27 Apr 2026 00:05:11 -0400 Subject: [PATCH 03/26] feat(api): add StageTelemetry DTO for per-stage timings --- src/api/schemas/telemetry.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) create mode 100644 src/api/schemas/telemetry.py diff --git a/src/api/schemas/telemetry.py b/src/api/schemas/telemetry.py new file mode 100644 index 00000000..b7b57ba7 --- /dev/null +++ b/src/api/schemas/telemetry.py @@ -0,0 +1,26 @@ +""" +StageTelemetry — per-stage timing, token, and cost numbers returned to the +chat client alongside the answer. + +API Layer Position: + RAGBackend.query → returns (answer, sources, StageTelemetry) + /api/query response includes StageTelemetry as `telemetry` field + WebSocket emits a final `telemetry` event with this payload after `done` + +Frontend renders these as a muted footer under each assistant chat bubble: + "Retrieve 142ms · Generate 2.1s · 4,217 tok · $0.0083" +""" + +from __future__ import annotations + +from pydantic import BaseModel, Field + + +class StageTelemetry(BaseModel): + """Per-stage observability numbers for one chat turn.""" + + retrieve_ms: float = Field(ge=0.0) + generate_ms: float = Field(ge=0.0) + prompt_tokens: int = Field(ge=0) + completion_tokens: int = Field(ge=0) + cost_usd: float = Field(ge=0.0) From cecdba9446f3a86d1530ccca154558dfede549be Mon Sep 17 00:00:00 2001 From: Mohamed Elkholy Date: Mon, 27 Apr 2026 00:16:03 -0400 Subject: [PATCH 04/26] feat(backend): instrument RAGBackend with traced stages and StageTelemetry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add query_with_telemetry() sibling and stream_query telemetry event: - query_with_telemetry() wraps retrieve + generate each in an OTel span (rag.retrieve / rag.generate via get_tracer()), times both phases with time.perf_counter(), counts prompt and completion tokens via count_tokens(), and computes cost_usd(). Returns (result_dict, StageTelemetry). - query() becomes a thin wrapper that calls query_with_telemetry() and discards the telemetry — existing tests and call sites are unaffected. - stream_query() gains retrieve and generate spans and yields a final ("telemetry", StageTelemetry.model_dump()) event AFTER ("done", ...). All existing event shapes (status, reasoning, token, done) are unchanged. The empty-results early-return branch also emits a zero telemetry event so the route layer always receives one. Option A (sibling) chosen over modifying query() in-place because existing tests assert dict access on query()'s return value; changing to a tuple would break them with no benefit — the route layer (Task 5) calls the new sibling directly. --- src/backend.py | 223 +++++++++++++++++++++++++----- tests/test_backend_telemetry.py | 233 ++++++++++++++++++++++++++++++++ 2 files changed, 423 insertions(+), 33 deletions(-) create mode 100644 tests/test_backend_telemetry.py diff --git a/src/backend.py b/src/backend.py index 604398f1..a6748e4d 100644 --- a/src/backend.py +++ b/src/backend.py @@ -32,6 +32,7 @@ import hashlib import logging import tempfile +import time import uuid from datetime import datetime, timezone from pathlib import Path @@ -40,6 +41,7 @@ from sqlalchemy import Engine from sqlmodel import Session, select +from .api.schemas.telemetry import StageTelemetry from .config import ( DEFAULT_MODEL, EVAL_MODEL, @@ -49,6 +51,8 @@ TOP_K_RESULTS, ) from .document_loader import DocumentLoader, TextChunker +from .eval._telemetry import count_tokens +from .eval.pricing import cost_usd from .evaluation import ( evaluate_answer_relevancy, evaluate_context_precision, @@ -59,6 +63,7 @@ from .models.document import DocumentRecord from .models.evaluation import MessageEvaluation from .models.message import Message, MessageSource +from .observability import get_tracer from .vector_store import ChromaVectorStore logger = logging.getLogger(__name__) @@ -298,6 +303,9 @@ def query( ) -> dict[str, Any]: """Run a full RAG query: retrieve chunks -> build context -> generate answer. + Thin wrapper over query_with_telemetry() that discards the StageTelemetry + so existing callers (tests, route layer) see no behaviour change. + Args: question: Natural language question from the user. top_k: Number of chunks to retrieve (default from config). @@ -306,19 +314,71 @@ def query( Returns: Dict with answer (str), sources (list[dict]), confidence (float). """ - k = top_k or TOP_K_RESULTS + result, _ = self.query_with_telemetry(question, top_k=top_k, model=model) + return result + + def query_with_telemetry( + self, + question: str, + top_k: int | None = None, + model: str | None = None, + ) -> tuple[dict[str, Any], StageTelemetry]: + """Run a full RAG query and return per-stage observability data. + + Identical to query() in output, but also returns a StageTelemetry + object with retrieve_ms, generate_ms, prompt_tokens, completion_tokens, + and cost_usd. The route layer (Task 5) calls this method so the REST + response can include telemetry without changing the query() contract. + + WHY a sibling instead of modifying query(): + Existing tests assert on result["answer"] and result["sources"] from + query(). Changing query() to return a tuple would break them silently + at dict-access time. The sibling keeps the tested contract intact. + + RAG Pipeline Position: + Question -> [RETRIEVE (traced)] -> [GENERATE (traced)] -> (answer + telemetry) - # WHY query_text: ChromaDB auto-embeds the question using the same - # embedding function that was used to embed the chunks. This - # ensures query and document embeddings live in the same space. - results = self.vector_store.query(query_text=question, top_k=k) + Args: + question: Natural language question from the user. + top_k: Number of chunks to retrieve (default from config). + model: LLM model override (creates a new handler if different). + + Returns: + Tuple of (result_dict, StageTelemetry). result_dict has the same + shape as query(): {answer, sources, confidence}. + """ + k = top_k or TOP_K_RESULTS + tracer = get_tracer() + + # ---- PHASE 1: Retrieval (timed + traced) ---------------------------- + # WHY: We open a span here rather than using @traced_stage because the + # retrieval call is a one-liner on self.vector_store — refactoring + # it to return (payload, attrs) would require an indirection wrapper + # that adds more lines than the inline approach. + t_retrieve_start = time.perf_counter() + with tracer.start_as_current_span("rag.retrieve") as retrieve_span: + retrieve_span.set_attribute("top_k", k) + retrieve_span.set_attribute("question_len", len(question)) + results = self.vector_store.query(query_text=question, top_k=k) + retrieve_span.set_attribute("results_count", len(results)) + retrieve_ms = (time.perf_counter() - t_retrieve_start) * 1000 if not results: - return { - "answer": "No documents indexed yet. Please upload documents first.", - "sources": [], - "confidence": 0.0, - } + # PATTERN: Early return with zero telemetry — no LLM call was made. + return ( + { + "answer": "No documents indexed yet. Please upload documents first.", + "sources": [], + "confidence": 0.0, + }, + StageTelemetry( + retrieve_ms=round(retrieve_ms, 2), + generate_ms=0.0, + prompt_tokens=0, + completion_tokens=0, + cost_usd=0.0, + ), + ) # Build context string from retrieved chunks context = "\n\n".join( @@ -331,7 +391,22 @@ def query( if model and model != self.llm.model: handler = LLMHandler(model=model) - answer = handler.generate_with_context(question, context) + # WHY reconstruct prompt strings here: generate_with_context() builds + # these internally before calling self.generate(). We reproduce them to + # count exactly the tokens that were billed, not a proxy string. + answer_system_prompt = ( + "You are a helpful assistant. Answer the user's question based solely on the " + "provided context. If the context does not contain enough information, say so." + ) + answer_user_prompt = f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:" + + # ---- PHASE 2: Generation (timed + traced) --------------------------- + t_generate_start = time.perf_counter() + with tracer.start_as_current_span("rag.generate") as generate_span: + generate_span.set_attribute("model", handler.model) + answer = handler.generate_with_context(question, context) + generate_span.set_attribute("answer_len", len(answer)) + generate_ms = (time.perf_counter() - t_generate_start) * 1000 sources = [ { @@ -346,16 +421,33 @@ def query( ] # PATTERN: Confidence = clamped average of top-3 similarity scores. - # This gives a rough signal of retrieval quality without - # requiring a separate calibration model. top_scores = [r.score for r in results[: min(3, len(results))]] confidence = max(0.0, min(1.0, sum(top_scores) / len(top_scores))) - return { - "answer": answer, - "sources": sources, - "confidence": round(confidence, 4), - } + # ---- Telemetry assembly --------------------------------------------- + # WHY count system_prompt + user_prompt together: together they form + # the full input that the provider billed as prompt tokens. + prompt_text = answer_system_prompt + "\n" + answer_user_prompt + prompt_tokens = count_tokens(prompt_text, handler.model) + completion_tokens = count_tokens(answer, handler.model) + total_cost = cost_usd(handler.model, prompt_tokens, completion_tokens) + + telemetry = StageTelemetry( + retrieve_ms=round(retrieve_ms, 2), + generate_ms=round(generate_ms, 2), + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cost_usd=total_cost, + ) + + return ( + { + "answer": answer, + "sources": sources, + "confidence": round(confidence, 4), + }, + telemetry, + ) def stream_query( self, @@ -367,10 +459,15 @@ def stream_query( """Retrieve context and stream reasoning + answer with chain-of-thought events. Event stream shape (in order): - ("status", str) — programmatic retrieval milestones - ("reasoning", str) — LLM chain-of-thought tokens (brief pre-answer pass) - ("token", str) — final answer tokens - ("done", dict) — sources + persistence metadata + ("status", str) — programmatic retrieval milestones + ("reasoning", str) — LLM chain-of-thought tokens (brief pre-answer pass) + ("token", str) — final answer tokens + ("done", dict) — sources + persistence metadata [UNCHANGED SHAPE] + ("telemetry", dict) — StageTelemetry dict (NEW — additive, after done) + + WHY telemetry is last: Old consumers that handle status/reasoning/token/done + and ignore unknown event types continue to work. The ("telemetry", ...) event + is purely additive — no existing event shape is modified. WHY two LLM calls: The reasoning pass uses a focused prompt that asks the model to think out loud about the retrieved context BEFORE giving @@ -381,19 +478,41 @@ def stream_query( for UX, not a durable artifact of the conversation. Persisting it would double storage and confuse the sliding-window context. + WHY telemetry covers the answer pass only (not reasoning): + The reasoning pass uses a separate model (REASONING_MODEL) with its own + cost. Telemetry here tracks the user-visible answer generation; mixing + two model costs into one StageTelemetry would confuse the "per-query + cost" display. Reasoning cost is a separate concern. + Yields: Tuples as described above. """ k = top_k or TOP_K_RESULTS + tracer = get_tracer() - # ---- PHASE 0: Retrieval ------------------------------------------------ + # ---- PHASE 0: Retrieval (timed + traced) -------------------------------- yield ("status", "Searching indexed documents...") - results = self.vector_store.query(query_text=question, top_k=k) + t_retrieve_start = time.perf_counter() + with tracer.start_as_current_span("rag.retrieve") as retrieve_span: + retrieve_span.set_attribute("top_k", k) + retrieve_span.set_attribute("question_len", len(question)) + results = self.vector_store.query(query_text=question, top_k=k) + retrieve_span.set_attribute("results_count", len(results)) + retrieve_ms = (time.perf_counter() - t_retrieve_start) * 1000 if not results: yield ("status", "No indexed documents — nothing to retrieve.") yield ("token", "No documents indexed yet. Please upload documents first.") yield ("done", {"sources": []}) + # PATTERN: Emit zero telemetry even on early return so the route + # layer always gets a telemetry event it can forward. + yield ("telemetry", StageTelemetry( + retrieve_ms=round(retrieve_ms, 2), + generate_ms=0.0, + prompt_tokens=0, + completion_tokens=0, + cost_usd=0.0, + ).model_dump()) return # WHY: Summarise retrieval in one status line so the user can see which @@ -477,7 +596,7 @@ def stream_query( logger.warning("Reasoning pass failed: %s", exc) yield ("status", "Reasoning unavailable — skipping to answer.") - # ---- PHASE 2: Answer pass ---------------------------------------------- + # ---- PHASE 2: Answer pass (timed + traced) ----------------------------- yield ("status", "Composing answer...") system_prompt = ( @@ -497,8 +616,10 @@ def stream_query( ) user_prompt = f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:" - # Accumulate full response for persistence - full_response = [] + # Accumulate full response for persistence and token counting + full_response: list[str] = [] + + t_generate_start = time.perf_counter() if conversation_id: # PHASE 1: Save user message BEFORE streaming @@ -515,9 +636,15 @@ def stream_query( messages.append({"role": "user", "content": user_prompt}) # PHASE 4: Stream via messages API (multi-turn aware) - for token in handler.stream_messages(messages): - full_response.append(token) - yield ("token", token) + with tracer.start_as_current_span("rag.generate") as generate_span: + generate_span.set_attribute("model", handler.model) + generate_span.set_attribute("has_conversation", True) + for token in handler.stream_messages(messages): + full_response.append(token) + yield ("token", token) + generate_span.set_attribute("answer_len", sum(len(t) for t in full_response)) + + generate_ms = (time.perf_counter() - t_generate_start) * 1000 # PHASE 5: Save assistant message + sources assistant_content = "".join(full_response) @@ -538,14 +665,44 @@ def stream_query( "conversation_id": conversation_id, }) + # WHY count full messages list for prompt: the multi-turn path sends + # the sliding window + system prompt + user turn as one request. + # We join all message content to estimate the billed input tokens. + prompt_text = "\n".join(m["content"] for m in messages) + else: # No conversation — simple single-turn streaming - for token in handler.stream_response(user_prompt, system_prompt=system_prompt): - full_response.append(token) - yield ("token", token) + with tracer.start_as_current_span("rag.generate") as generate_span: + generate_span.set_attribute("model", handler.model) + generate_span.set_attribute("has_conversation", False) + for token in handler.stream_response(user_prompt, system_prompt=system_prompt): + full_response.append(token) + yield ("token", token) + generate_span.set_attribute("answer_len", sum(len(t) for t in full_response)) + + generate_ms = (time.perf_counter() - t_generate_start) * 1000 yield ("done", {"sources": sources}) + prompt_text = system_prompt + "\n" + user_prompt + + # ---- Telemetry assembly (after done, additive) ----------------------- + # WHY after done: the done event is what the client waits for to show + # sources. Telemetry is a secondary signal — emit it last so done's + # latency is not affected by token-counting arithmetic. + answer_text = "".join(full_response) + prompt_tokens = count_tokens(prompt_text, handler.model) + completion_tokens = count_tokens(answer_text, handler.model) + total_cost = cost_usd(handler.model, prompt_tokens, completion_tokens) + + yield ("telemetry", StageTelemetry( + retrieve_ms=round(retrieve_ms, 2), + generate_ms=round(generate_ms, 2), + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cost_usd=total_cost, + ).model_dump()) + # ------------------------------------------------------------------ # # Document management # # ------------------------------------------------------------------ # diff --git a/tests/test_backend_telemetry.py b/tests/test_backend_telemetry.py new file mode 100644 index 00000000..31a90102 --- /dev/null +++ b/tests/test_backend_telemetry.py @@ -0,0 +1,233 @@ +"""Tests that RAGBackend instrumentation produces a StageTelemetry payload. + +RAG Pipeline Position: + Question -> Retrieve -> Generate -> (answer, StageTelemetry) + ^^^ + These tests verify that query_with_telemetry() and stream_query() both emit + well-formed StageTelemetry with non-negative numeric fields. + +What concept it teaches: + Testing the observability layer in isolation from real LLM providers. + LLMHandler falls back to a dummy response when no API keys are configured, + so these tests run without any network access and with zero cost. + +Why this fixture pattern: + Mirrors test_backend.py exactly — EphemeralClient for ChromaDB (no disk I/O, + isolated per test) and sqlite:// for SQLite (in-memory, vanishes on teardown). + Both stores are fully isolated; no shared mutable state between tests. +""" + +from __future__ import annotations + +import uuid +from pathlib import Path + +import chromadb +import pytest + +from src.backend import RAGBackend +from src.api.schemas.telemetry import StageTelemetry +from src.database import create_db_and_tables, get_engine + + +# --------------------------------------------------------------------------- # +# Fixtures (mirrored from test_backend.py) # +# --------------------------------------------------------------------------- # + +@pytest.fixture +def tmp_sqlite_engine(): + """In-memory SQLite engine with all tables created.""" + engine = get_engine("sqlite://") + create_db_and_tables(engine) + return engine + + +@pytest.fixture +def chroma_backend_collection(): + """Ephemeral ChromaDB collection with auto-embedding enabled. + + WHY unique name: EphemeralClient shares an in-process store. + A UUID suffix ensures complete isolation between test runs. + """ + client = chromadb.EphemeralClient() + return client.get_or_create_collection( + name=f"test_backend_telemetry_{uuid.uuid4().hex}", + metadata={"hnsw:space": "cosine"}, + ) + + +@pytest.fixture +def backend(tmp_sqlite_engine, chroma_backend_collection): + """Fully-wired RAGBackend with ephemeral ChromaDB + in-memory SQLite.""" + return RAGBackend(engine=tmp_sqlite_engine, collection=chroma_backend_collection) + + +@pytest.fixture +def ingested_backend(backend: RAGBackend, tmp_path: Path) -> RAGBackend: + """A backend that has already ingested a small text document. + + WHY pre-ingest: query_with_telemetry requires at least one indexed chunk + to take the non-empty retrieval path and call the LLM. Without a document, + both retrieve and generate phases return zeros — not useful to test. + """ + content = ( + "Retrieval-Augmented Generation (RAG) combines document retrieval " + "with large language model generation to produce grounded answers. " + "The retrieval step finds relevant chunks via vector similarity search. " + "The generation step builds a prompt from those chunks and queries the LLM." + ) + doc = tmp_path / "rag_intro.txt" + doc.write_text(content, encoding="utf-8") + backend.ingest_file(doc) + return backend + + +# --------------------------------------------------------------------------- # +# Tests # +# --------------------------------------------------------------------------- # + +class TestQueryWithTelemetry: + """Tests for RAGBackend.query_with_telemetry().""" + + def test_telemetry_fields_are_non_negative_after_ingest( + self, ingested_backend: RAGBackend + ): + """query_with_telemetry returns a StageTelemetry with all non-negative fields. + + PATTERN: With no real LLM configured, LLMHandler falls back to a dummy + response. The dummy response still produces valid answer text, which + means token counting and cost computation run on real strings — the + telemetry shape is exercised even without a real LLM call. + """ + result, telemetry = ingested_backend.query_with_telemetry("What is RAG?") + + # Result dict has the same shape as query() + assert "answer" in result + assert "sources" in result + assert len(result["sources"]) > 0 + + # Telemetry is the right type + assert isinstance(telemetry, StageTelemetry) + + # All five numeric fields are non-negative + assert telemetry.retrieve_ms >= 0.0 + assert telemetry.generate_ms >= 0.0 + assert telemetry.prompt_tokens >= 0 + assert telemetry.completion_tokens >= 0 + assert telemetry.cost_usd >= 0.0 + + def test_telemetry_zeros_when_no_documents(self, backend: RAGBackend): + """When no documents are indexed, StageTelemetry has zero generate/token fields. + + WHY: The early-return branch skips the LLM call entirely. + retrieve_ms records real elapsed time (even for an empty search); + generate_ms, prompt_tokens, completion_tokens, and cost_usd are all 0. + """ + result, telemetry = backend.query_with_telemetry("What is RAG?") + + assert result["answer"].startswith("No documents indexed") + assert telemetry.generate_ms == 0.0 + assert telemetry.prompt_tokens == 0 + assert telemetry.completion_tokens == 0 + assert telemetry.cost_usd == 0.0 + # retrieve_ms is still measured (we did make the call, it just returned empty) + assert telemetry.retrieve_ms >= 0.0 + + def test_query_unchanged_after_sibling_added(self, ingested_backend: RAGBackend): + """query() still returns the plain dict — the new sibling has no side effects. + + WHY: This is the regression guard for the Option A design choice. + Existing tests use result["answer"] / result["sources"] on the dict + returned by query(). If we accidentally broke the dict shape, this + test catches it immediately. + """ + result = ingested_backend.query("What is RAG?") + + assert isinstance(result, dict) + assert "answer" in result + assert "sources" in result + assert "confidence" in result + + +class TestStreamQueryTelemetry: + """Tests for the telemetry event emitted by stream_query().""" + + def test_stream_query_emits_telemetry_event_last( + self, ingested_backend: RAGBackend + ): + """stream_query yields a ("telemetry", dict) as the final event after ("done", ...). + + WHY last: The done event is what the client waits for to display sources. + Telemetry is a secondary signal. Emitting it last ensures done latency + is not delayed by token-counting arithmetic. + """ + events = list(ingested_backend.stream_query("What is RAG?")) + + # The last event must be the telemetry event + last_type, last_data = events[-1] + assert last_type == "telemetry", ( + f"Expected last event type 'telemetry', got {last_type!r}. " + f"All event types: {[e[0] for e in events]}" + ) + + # The telemetry payload must be a dict (model_dump() output) + assert isinstance(last_data, dict) + + # All five fields must be present and non-negative + for field in ("retrieve_ms", "generate_ms", "prompt_tokens", "completion_tokens", "cost_usd"): + assert field in last_data, f"Missing telemetry field: {field}" + assert last_data[field] >= 0, f"Telemetry field {field} is negative: {last_data[field]}" + + def test_stream_query_done_event_unchanged(self, ingested_backend: RAGBackend): + """The ("done", ...) event shape is not modified by the telemetry addition. + + STRICT: The done event must still carry 'sources' (and nothing else + unexpected). This test guards against accidental mutation of the done + event dict. + """ + events = list(ingested_backend.stream_query("What is RAG?")) + + done_events = [(t, d) for t, d in events if t == "done"] + assert len(done_events) == 1, f"Expected exactly one done event, got {len(done_events)}" + + _, done_data = done_events[0] + assert "sources" in done_data + + def test_stream_query_emits_telemetry_on_empty_store(self, backend: RAGBackend): + """The early-return (no documents) path still emits a telemetry event. + + WHY: The route layer always expects a telemetry event. If the early-return + path omitted it, the frontend would never receive telemetry data for + failed queries — a silent gap that's hard to debug. + """ + events = list(backend.stream_query("What is RAG?")) + + event_types = [e[0] for e in events] + assert "telemetry" in event_types + + last_type, last_data = events[-1] + assert last_type == "telemetry" + assert last_data["generate_ms"] == 0.0 + assert last_data["prompt_tokens"] == 0 + + def test_stream_query_existing_events_order_preserved( + self, ingested_backend: RAGBackend + ): + """Existing event types appear in the expected order before telemetry. + + The protocol guarantees: status* → reasoning* → status → token* → done → telemetry + (where * = zero or more). This test checks that all mandatory events + are still present and that telemetry is appended AFTER done. + """ + events = list(ingested_backend.stream_query("What is RAG?")) + event_types = [e[0] for e in events] + + assert "status" in event_types + assert "done" in event_types + assert "telemetry" in event_types + + done_idx = next(i for i, (t, _) in enumerate(events) if t == "done") + telemetry_idx = next(i for i, (t, _) in enumerate(events) if t == "telemetry") + assert telemetry_idx > done_idx, ( + "telemetry event must come after done event" + ) From ab1e65af1d6bb00a421111055a2c6f7f385715a1 Mon Sep 17 00:00:00 2001 From: Mohamed Elkholy Date: Mon, 27 Apr 2026 00:23:43 -0400 Subject: [PATCH 05/26] feat(api): emit StageTelemetry in REST and WebSocket query responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit POST /api/query now calls query_with_telemetry() instead of query() and returns a `telemetry` field (StageTelemetry) alongside the existing answer, sources, confidence, and latency_ms fields. The WebSocket /api/chat handler gains an explicit branch for the ("telemetry", dict) event emitted by stream_query() after done, forwarding it as {"type": "telemetry", "content": {...}} so clients can render per-stage timing and cost without polling. QueryResponse in models.py is extended with an Optional[StageTelemetry] field (defaults to None for backward compatibility). StageTelemetry is imported once at module level via models.py to keep the routes clean. Nine route-layer tests added in tests/test_api_query_telemetry.py using mocked backends — isolating serialization from real LLM / ChromaDB I/O. --- src/api/models.py | 14 +- src/api/routes/query.py | 12 +- tests/test_api_query_telemetry.py | 264 ++++++++++++++++++++++++++++++ 3 files changed, 288 insertions(+), 2 deletions(-) create mode 100644 tests/test_api_query_telemetry.py diff --git a/src/api/models.py b/src/api/models.py index c64abe64..a9a04f14 100644 --- a/src/api/models.py +++ b/src/api/models.py @@ -9,6 +9,7 @@ from pydantic import BaseModel, Field +from src.api.schemas.telemetry import StageTelemetry from src.config import DEFAULT_MODEL @@ -50,7 +51,11 @@ class SourceInfo(BaseModel): class QueryResponse(BaseModel): - """Response body for the query endpoint.""" + """Response body for the query endpoint. + + ADDITIVE: `telemetry` was added in Task 5 (Sub-plan 1D). All other fields + are unchanged. Clients that ignore unknown fields are unaffected. + """ answer: str = Field(description="Generated answer.") sources: List[SourceInfo] = Field(default_factory=list, description="Retrieved source chunks.") @@ -58,6 +63,13 @@ class QueryResponse(BaseModel): ge=0.0, le=1.0, description="Estimated answer confidence (0–1)." ) latency_ms: float = Field(description="Total request latency in milliseconds.") + # WHY: StageTelemetry carries per-stage timing and token-cost numbers. + # Optional with None default so existing callers that construct + # QueryResponse without telemetry (e.g., older tests) still validate. + telemetry: Optional[StageTelemetry] = Field( + default=None, + description="Per-stage timing, token counts, and cost for this request.", + ) class UploadResponse(BaseModel): diff --git a/src/api/routes/query.py b/src/api/routes/query.py index bcf20470..0d61f6ee 100644 --- a/src/api/routes/query.py +++ b/src/api/routes/query.py @@ -37,7 +37,10 @@ async def query(request_body: QueryRequest, request: Request) -> QueryResponse: start = time.perf_counter() backend = request.app.state.backend - result = backend.query( + # WHY query_with_telemetry: replaces the plain query() call so we get + # per-stage timing and token-cost numbers in the response. The + # result_dict has the same shape as before — only telemetry is new. + result, telemetry = backend.query_with_telemetry( request_body.query, top_k=request_body.top_k, model=request_body.model, @@ -61,6 +64,8 @@ async def query(request_body: QueryRequest, request: Request) -> QueryResponse: sources=sources, confidence=result.get("confidence", 0.0), latency_ms=latency_ms, + # ADDITIVE: telemetry field added in Task 5 (Sub-plan 1D). + telemetry=telemetry, ) @@ -177,6 +182,11 @@ async def chat_websocket(websocket: WebSocket) -> None: "message_id": data.get("message_id"), "conversation_id": data.get("conversation_id"), }) + elif event_type == "telemetry": + # WHY: stream_query yields ("telemetry", StageTelemetry.model_dump()) + # after the done event. Forward it verbatim so the frontend + # can render per-stage timing and cost without polling. + await websocket.send_json({"type": "telemetry", "content": data}) except Exception as exc: logger.error("Streaming error: %s", exc) await websocket.send_json( diff --git a/tests/test_api_query_telemetry.py b/tests/test_api_query_telemetry.py new file mode 100644 index 00000000..fe5aef06 --- /dev/null +++ b/tests/test_api_query_telemetry.py @@ -0,0 +1,264 @@ +"""Tests for telemetry surfacing in the REST and WebSocket query routes. + +API Layer Position: + RAGBackend.query_with_telemetry → (result_dict, StageTelemetry) + POST /api/query → QueryResponse with `telemetry` field [REST] + GET /api/chat → WebSocket stream ends with telemetry event [WS] + +What concept it teaches: + Route-layer serialization testing: mock the backend dependency, assert the + route converts the returned StageTelemetry into the right JSON shape. + This isolates the serialization concern from real LLM / ChromaDB I/O. + +Why mock approach instead of full-ingest: + - RAGBackend telemetry behaviour is already tested in test_backend_telemetry.py. + - Task 5 only changes the route layer (query.py + models.py). + - Mocking the backend keeps tests fast (<1 s) and dependency-free. +""" + +from __future__ import annotations + +import json +from typing import Iterator +from unittest.mock import MagicMock, patch + +import pytest +from fastapi.testclient import TestClient + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + +_FAKE_TELEMETRY = { + "retrieve_ms": 42.5, + "generate_ms": 310.0, + "prompt_tokens": 512, + "completion_tokens": 128, + "cost_usd": 0.0034, +} + +_FAKE_RESULT = { + "answer": "RAG combines retrieval with generation.", + "sources": [ + { + "doc_id": "doc-1", + "chunk_id": "chunk-1", + "filename": "rag.txt", + "score": 0.92, + "excerpt": "RAG stands for Retrieval-Augmented Generation.", + } + ], + "confidence": 0.85, +} + + +def _make_telemetry_model(): + """Return a StageTelemetry instance with the fake values.""" + from src.api.schemas.telemetry import StageTelemetry + return StageTelemetry(**_FAKE_TELEMETRY) + + +# --------------------------------------------------------------------------- +# REST endpoint — POST /api/query +# --------------------------------------------------------------------------- + + +class TestRestTelemetry: + """POST /api/query response must include a well-formed `telemetry` field.""" + + @pytest.fixture + def client(self): + """TestClient with a mocked backend attached to app.state. + + WHY: TestClient runs the FastAPI lifespan, which creates a real + RAGBackend on app.state. We replace it with a MagicMock after + startup so query_with_telemetry() returns predictable values + without touching ChromaDB or an LLM. + """ + from src.api.main import app + with TestClient(app) as c: + mock_backend = MagicMock() + mock_backend.query_with_telemetry.return_value = ( + _FAKE_RESULT, _make_telemetry_model() + ) + # evaluate_faithfulness_realtime must not raise during WS teardown + mock_backend.evaluate_faithfulness_realtime.return_value = {} + app.state.backend = mock_backend + yield c + + def test_response_includes_telemetry_field(self, client): + """The telemetry object appears at the top level of the JSON response.""" + r = client.post("/api/query", json={"query": "What is RAG?"}) + assert r.status_code == 200, r.text + body = r.json() + assert "telemetry" in body, f"'telemetry' key missing from response: {body}" + + def test_telemetry_has_all_five_fields(self, client): + """All five StageTelemetry fields are present in the response body.""" + r = client.post("/api/query", json={"query": "What is RAG?"}) + assert r.status_code == 200 + t = r.json()["telemetry"] + for field in ("retrieve_ms", "generate_ms", "prompt_tokens", + "completion_tokens", "cost_usd"): + assert field in t, f"Missing telemetry field: {field}" + assert t[field] >= 0, f"Telemetry field {field} must be >= 0" + + def test_telemetry_values_match_backend(self, client): + """Telemetry values round-trip correctly from backend to JSON response.""" + r = client.post("/api/query", json={"query": "What is RAG?"}) + t = r.json()["telemetry"] + assert t["retrieve_ms"] == _FAKE_TELEMETRY["retrieve_ms"] + assert t["generate_ms"] == _FAKE_TELEMETRY["generate_ms"] + assert t["prompt_tokens"] == _FAKE_TELEMETRY["prompt_tokens"] + assert t["completion_tokens"] == _FAKE_TELEMETRY["completion_tokens"] + assert t["cost_usd"] == pytest.approx(_FAKE_TELEMETRY["cost_usd"]) + + def test_existing_response_fields_unchanged(self, client): + """Adding telemetry must not drop or alter existing response fields.""" + r = client.post("/api/query", json={"query": "What is RAG?"}) + body = r.json() + assert "answer" in body + assert "sources" in body + assert "confidence" in body + assert "latency_ms" in body + + def test_query_with_telemetry_is_called_not_query(self, client): + """The route must call query_with_telemetry, not the plain query(). + + WHY: If someone reverts to backend.query() the telemetry field would + silently become None (the Optional default). This assertion catches it. + """ + from src.api.main import app + client.post("/api/query", json={"query": "test"}) + app.state.backend.query_with_telemetry.assert_called_once() + app.state.backend.query.assert_not_called() + + +# --------------------------------------------------------------------------- +# WebSocket endpoint — /api/chat +# --------------------------------------------------------------------------- + + +class TestWebSocketTelemetry: + """WebSocket /api/chat must forward the telemetry event from stream_query.""" + + def _make_stream(self) -> list[tuple[str, object]]: + """Minimal event stream: status → done → telemetry.""" + return [ + ("status", "Searching indexed documents..."), + ("token", "RAG combines retrieval with generation."), + ("done", { + "sources": [ + { + "doc_id": "doc-1", + "chunk_id": "chunk-1", + "filename": "rag.txt", + "score": 0.92, + "excerpt": "RAG stands for Retrieval-Augmented Generation.", + } + ], + "message_id": "msg-abc", + "conversation_id": "conv-xyz", + }), + ("telemetry", _FAKE_TELEMETRY), + ] + + @pytest.fixture + def ws_client(self): + """TestClient with a mocked streaming backend.""" + from src.api.main import app + with TestClient(app) as c: + mock_backend = MagicMock() + + def _fake_stream(*args, **kwargs) -> Iterator: + yield from self._make_stream() + + mock_backend.stream_query.side_effect = _fake_stream + mock_backend.evaluate_faithfulness_realtime.return_value = {} + app.state.backend = mock_backend + yield c + + def test_stream_emits_telemetry_event(self, ws_client): + """The WebSocket stream must include a telemetry event.""" + from src.api.main import app + with ws_client.websocket_connect("/api/chat") as ws: + ws.send_json({"query": "What is RAG?", "top_k": 3}) + events = [] + # Drain all events until we see telemetry or hit 20 messages. + # WHY cap: if telemetry is never emitted, we stop instead of hanging. + for _ in range(20): + try: + msg = ws.receive_json() + events.append(msg) + if msg.get("type") == "telemetry": + break + except Exception: + break + + assert any(e.get("type") == "telemetry" for e in events), ( + f"No telemetry event received. Got event types: {[e.get('type') for e in events]}" + ) + + def test_telemetry_event_has_content_key(self, ws_client): + """The telemetry event must be shaped: {type: 'telemetry', content: {...}}.""" + with ws_client.websocket_connect("/api/chat") as ws: + ws.send_json({"query": "What is RAG?", "top_k": 3}) + tele_event = None + for _ in range(20): + try: + msg = ws.receive_json() + if msg.get("type") == "telemetry": + tele_event = msg + break + except Exception: + break + + assert tele_event is not None, "No telemetry event received." + assert "content" in tele_event, f"telemetry event missing 'content' key: {tele_event}" + + def test_telemetry_content_has_all_five_fields(self, ws_client): + """All five StageTelemetry fields must appear inside the content dict.""" + with ws_client.websocket_connect("/api/chat") as ws: + ws.send_json({"query": "What is RAG?", "top_k": 3}) + tele_event = None + for _ in range(20): + try: + msg = ws.receive_json() + if msg.get("type") == "telemetry": + tele_event = msg + break + except Exception: + break + + content = tele_event["content"] + for field in ("retrieve_ms", "generate_ms", "prompt_tokens", + "completion_tokens", "cost_usd"): + assert field in content, f"Missing telemetry field: {field}" + assert content[field] >= 0, f"Telemetry field {field} must be >= 0" + + def test_done_event_shape_unchanged(self, ws_client): + """Adding telemetry must not modify the done event shape. + + STRICT: done must still carry sources, message_id, conversation_id. + This guards against accidentally merging telemetry into done. + """ + with ws_client.websocket_connect("/api/chat") as ws: + ws.send_json({"query": "What is RAG?", "top_k": 3}) + done_event = None + for _ in range(20): + try: + msg = ws.receive_json() + if msg.get("type") == "done": + done_event = msg + if msg.get("type") == "telemetry": + break + except Exception: + break + + assert done_event is not None, "No done event received." + assert "sources" in done_event, f"done event missing 'sources': {done_event}" + assert "message_id" in done_event + assert "conversation_id" in done_event + # telemetry must NOT be embedded inside done + assert "telemetry" not in done_event From 4fa63e0ae8aacf69c22915886356f96f0b52226b Mon Sep 17 00:00:00 2001 From: Mohamed Elkholy Date: Mon, 27 Apr 2026 00:26:05 -0400 Subject: [PATCH 06/26] feat(frontend): handle telemetry event and expose on chat message Add TelemetryPayload interface and WsTelemetryMessage to types.ts, extend ChatMessage with an optional telemetry field, and wire up a new "telemetry" case in the useChat WebSocket handler so per-request timing and cost data is attached to the assistant message for the upcoming TelemetryFooter component. --- frontend/src/api/types.ts | 25 ++++++++++++++++++++++++- frontend/src/hooks/use-chat.ts | 23 ++++++++++++++++++++++- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/frontend/src/api/types.ts b/frontend/src/api/types.ts index c7a5a994..8c96be4d 100644 --- a/frontend/src/api/types.ts +++ b/frontend/src/api/types.ts @@ -76,6 +76,26 @@ export interface WsErrorMessage { content: string; } +/** + * Per-request timing and token-cost summary emitted by the backend after + * the "done" event. Mirrors the backend `StageTelemetry` dataclass. + * + * WHY a dedicated interface: keeps the telemetry shape self-documenting and + * lets the UI consume it without casting or runtime duck-typing. + */ +export interface TelemetryPayload { + retrieve_ms: number; + generate_ms: number; + prompt_tokens: number; + completion_tokens: number; + cost_usd: number; +} + +export interface WsTelemetryMessage { + type: "telemetry"; + content: TelemetryPayload; +} + export interface EvaluationScore { metric: string; score: number; @@ -96,7 +116,8 @@ export type WsMessage = | WsStatusMessage | WsDoneMessage | WsErrorMessage - | WsEvaluationMessage; + | WsEvaluationMessage + | WsTelemetryMessage; export interface ChatMessage { id: string; @@ -115,6 +136,8 @@ export interface ChatMessage { streamDone?: boolean; /** LLM-as-judge evaluation scores (faithfulness, relevancy, precision). */ evaluation?: EvaluationScore[]; + /** Per-request timing and cost from the backend telemetry event. */ + telemetry?: TelemetryPayload; } export interface ConversationSummary { diff --git a/frontend/src/hooks/use-chat.ts b/frontend/src/hooks/use-chat.ts index 3cbe9c22..f5e9f3b4 100644 --- a/frontend/src/hooks/use-chat.ts +++ b/frontend/src/hooks/use-chat.ts @@ -1,7 +1,7 @@ import { useCallback, useRef, useState } from "react"; import { useQueryClient } from "@tanstack/react-query"; import { wsUrl } from "@/api/client"; -import type { ChatMessage, EvaluationScore, SourceInfo, WsMessage } from "@/api/types"; +import type { ChatMessage, EvaluationScore, SourceInfo, TelemetryPayload, WsMessage } from "@/api/types"; let msgCounter = 0; function nextId() { @@ -127,6 +127,27 @@ export function useChat() { setTimeout(() => { if (ws.readyState === WebSocket.OPEN) ws.close(); }, 30_000); + } else if (data.type === "telemetry") { + // WHY: The backend emits a "telemetry" event immediately after "done" + // (before evaluation) with per-request timing and token cost. + // We attach it to the assistant message so the UI can render a + // TelemetryFooter without polling or a separate API call. + // + // PATTERN: Same setMessages updater as the other event handlers — + // merge the payload into the target message by ID. + // + // TEST TODO: A React Testing Library test should feed the sequence + // status → reasoning → token → done → telemetry → evaluation + // and assert that the last assistant message has both a non-empty + // `content` string and a `telemetry` object with numeric fields + // (retrieve_ms, generate_ms, prompt_tokens, completion_tokens, cost_usd). + setMessages((prev) => + prev.map((m) => + m.id === assistantIdRef.current + ? { ...m, telemetry: data.content as TelemetryPayload } + : m + ) + ); } else if (data.type === "evaluation") { // WHY: The backend fires a separate WebSocket event after the "done" // event with real-time faithfulness scores. We attach it to the From 1254b1bdf3312631694d30aae4c1569c638a009a Mon Sep 17 00:00:00 2001 From: Mohamed Elkholy Date: Mon, 27 Apr 2026 00:29:21 -0400 Subject: [PATCH 07/26] feat(frontend): add TelemetryFooter under each assistant chat message MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Renders a muted xs one-liner (Retrieve · Generate · Tokens · Cost) below every assistant bubble that carries a telemetry payload. A hover tooltip surfaces the full prompt/completion token breakdown. Uses the existing @base-ui/react Tooltip primitive (already wrapped in ui/tooltip.tsx and provided at App root) with the render-prop pattern to avoid an implicit button role on a purely informational element. --- frontend/src/components/chat/chat-message.tsx | 9 + .../src/components/chat/telemetry-footer.tsx | 189 ++++++++++++++++++ 2 files changed, 198 insertions(+) create mode 100644 frontend/src/components/chat/telemetry-footer.tsx diff --git a/frontend/src/components/chat/chat-message.tsx b/frontend/src/components/chat/chat-message.tsx index 43c68c58..2292a208 100644 --- a/frontend/src/components/chat/chat-message.tsx +++ b/frontend/src/components/chat/chat-message.tsx @@ -5,6 +5,7 @@ import { cn } from "@/lib/utils"; import { ThinkingPanel } from "./thinking-panel"; import { MarkdownRenderer } from "./markdown-renderer"; import { EvaluationBadge } from "./evaluation-badge"; +import { TelemetryFooter } from "./telemetry-footer"; // WHY removed TypingIndicator: The bouncing-dots placeholder used to show // for the brief window between sending a query and the first server event. @@ -134,6 +135,14 @@ export function ChatMessage({ message, onEvaluate }: ChatMessageProps) { )} + {/* TelemetryFooter — muted one-liner: Retrieve · Generate · Tokens · Cost. + Only renders once telemetry arrives (after the "done" WebSocket event) + and only on assistant messages. Purely additive; does not affect layout + of surrounding elements. */} + {!isUser && message.telemetry ? ( + + ) : null} + {/* PATTERN: EvaluationBadge is shown only after streaming completes (streamDone=true) so it doesn't flash up mid-response. The badge triggers the LLM-as-judge evaluation call on first render and diff --git a/frontend/src/components/chat/telemetry-footer.tsx b/frontend/src/components/chat/telemetry-footer.tsx new file mode 100644 index 00000000..022ce7b3 --- /dev/null +++ b/frontend/src/components/chat/telemetry-footer.tsx @@ -0,0 +1,189 @@ +/** + * TelemetryFooter — per-message timing and cost summary bar. + * + * RAG Pipeline Position: + * Query → Retrieval → Generation → Response → [TELEMETRY FOOTER] + * ^^^ + * This component sits at the DISPLAY step, after a full assistant turn + * completes. It surfaces backend timing (retrieve_ms, generate_ms), + * token counts (prompt + completion), and cost so the developer can spot + * slow retrievals, verbose generations, or unexpectedly expensive calls + * without leaving the chat window. + * + * WHY a dedicated component: + * The telemetry data arrives as a separate WebSocket event ("telemetry") + * after the "done" event. Keeping its rendering isolated means the parent + * (ChatMessage) never needs to know about formatting logic — it just + * passes `message.telemetry` through. + * + * PATTERN: Tooltip-wrapped trigger — the one-liner summary is always visible; + * a hover tooltip surfaces the full prompt/completion split without + * cluttering the message layout. + */ + +import type { TelemetryPayload } from "@/api/types"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { cn } from "@/lib/utils"; + +// --------------------------------------------------------------------------- +// Public API +// --------------------------------------------------------------------------- + +export interface TelemetryFooterProps { + telemetry: TelemetryPayload; + className?: string; +} + +// --------------------------------------------------------------------------- +// Formatting helpers +// --------------------------------------------------------------------------- + +/** + * Format retrieval latency as an integer millisecond value. + * Retrieval is fast enough that sub-second display is always the right unit. + */ +function formatRetrieve(ms: number): string { + return `Retrieve ${Math.round(ms)}ms`; +} + +/** + * Format generation latency: seconds for ≥1 s, milliseconds otherwise. + * + * WHY dual-unit: generation times span 100 ms (fast local model) to 30 s + * (slow hosted model). Showing "30,000ms" is less readable than "30.0s". + */ +function formatGenerate(ms: number): string { + if (ms >= 1000) { + return `Generate ${(ms / 1000).toFixed(1)}s`; + } + return `Generate ${Math.round(ms)}ms`; +} + +/** + * Format total token count with locale-aware thousands separators. + * e.g. 4217 → "4,217 tok" + */ +function formatTokens(prompt: number, completion: number): string { + const total = prompt + completion; + return `${total.toLocaleString()} tok`; +} + +/** + * Format cost in USD to 4 decimal places. + * Always renders, even when cost is exactly 0 ($0.0000). + * + * WHY always render zero: hiding a zero cost would make users think the + * telemetry event didn't arrive, rather than that the call was free (e.g. + * when a local Ollama model is used). + */ +function formatCost(cost: number): string { + return `$${cost.toFixed(4)}`; +} + +// --------------------------------------------------------------------------- +// Tooltip content — token breakdown table +// --------------------------------------------------------------------------- + +/** + * TooltipBreakdown — rendered inside the shadcn Tooltip popup. + * + * WHY tabular: the four rows (Prompt, Completion, Total, Cost) need alignment + * to be scannable. A two-column grid achieves this without a element. + */ +function TooltipBreakdown({ telemetry }: { telemetry: TelemetryPayload }) { + const { prompt_tokens, completion_tokens, cost_usd } = telemetry; + const total = prompt_tokens + completion_tokens; + + const rows: [string, string][] = [ + ["Prompt", `${prompt_tokens.toLocaleString()} tokens`], + ["Completion", `${completion_tokens.toLocaleString()} tokens`], + ["Total", `${total.toLocaleString()} tokens`], + ["Cost USD", cost_usd.toFixed(4)], + ]; + + return ( +
+ {rows.map(([label, value]) => ( + <> + + {label}: + + {value} + + ))} +
+ ); +} + +// --------------------------------------------------------------------------- +// TelemetryFooter — public export +// --------------------------------------------------------------------------- + +/** + * TelemetryFooter renders a muted one-liner beneath an assistant message + * showing retrieve latency, generate latency, token count, and cost. + * + * Usage: + * {message.telemetry ? : null} + * + * @param telemetry - Backend timing and token-cost payload from the + * "telemetry" WebSocket event. See TelemetryPayload in api/types.ts. + * @param className - Optional extra Tailwind classes for the container. + */ +export function TelemetryFooter({ + telemetry, + className, +}: TelemetryFooterProps): React.JSX.Element { + const segments = [ + formatRetrieve(telemetry.retrieve_ms), + formatGenerate(telemetry.generate_ms), + formatTokens(telemetry.prompt_tokens, telemetry.completion_tokens), + formatCost(telemetry.cost_usd), + ]; + + // Dot separator — aria-hidden so screen readers don't announce bullet chars. + const Dot = () => ( + + ); + + return ( + // PATTERN: TooltipProvider is already mounted at the App root, so we only + // need Tooltip > TooltipTrigger > TooltipContent here. + + {/* + * WHY render prop: @base-ui/react TooltipTrigger renders a