ragflow/rag/advanced_rag/harness/stats.py

"""LLM-call instrumentation for the agentic RAG harness.

Every phase of the agentic pipeline (route / planner / orchestrator / agent /
sufficiency / grounded / finalize, ...) drives the LLM through ``tools.chat_mdl``.
``RAGTools`` wraps that bundle in a :class:`CountingChatModel` proxy which records,
per phase, the number of LLM calls and the token usage reported by the provider.
Phase wall-clock time is measured by :func:`phase` itself — unlike summed LLM
latency, which is meaningless when calls run in parallel. The aggregate is
emitted to the log when each ``rag`` call finishes.
"""

import logging
import time
from collections import defaultdict
from contextlib import contextmanager
from contextvars import ContextVar
from functools import wraps
from inspect import iscoroutinefunction

_LOG = logging.getLogger("rag.advanced_rag.harness.stats")

_CURRENT_PHASE: ContextVar[str] = ContextVar("agentic_rag_llm_phase", default="unknown")
_ACTIVE_PHASES: ContextVar[tuple[str, ...]] = ContextVar("agentic_rag_active_phases", default=())

# Canonical pipeline order for the per-phase usage table. Phases are listed in
# the order they execute across the (high/ultra) agentic pipeline so the log
# reads top-to-bottom like the actual flow, regardless of which phase first
# touched the stats counters. Any phase not in this list (e.g. a future
# addition) is appended afterwards, alphabetically, as a safe fallback.
_PHASE_ORDER = [
    "formalize",
    "route",
    "pre_search",
    "planner",
    "decompose",
    "direct",
    "orchestrator",
    "claim_research",
    "sufficiency",
    "grounded",
    "finalize",
]
_PHASE_RANK = {name: i for i, name in enumerate(_PHASE_ORDER)}


@contextmanager
def phase(name: str):
    """Run the enclosed block with the LLM-call phase set to ``name``.

    Also accrues the block's wall-clock time into the active stats (the
    per-call one bound by :func:`using_stats`, falling back to nothing), so
    each phase reports how long it actually took to execute. Nested phases
    (e.g. ``orchestrator`` containing ``agent``) overlap by design: the outer
    wall-clock includes the inner one, so totals are not summed.

    Re-entrancy: the same phase name may be wrapped several times along one
    call path (graph node + implementation fn + orchestrator wrapper). Only
    the outermost interval is timed — inner re-entries of the same name are
    shadowed and contribute no wall-clock, preventing the time from being
    counted 2-3 times.
    """
    token = _CURRENT_PHASE.set(name)
    active = _ACTIVE_PHASES.get()
    active_token = _ACTIVE_PHASES.set(active + (name,))
    shadowed = name in active
    stats = _CURRENT_STATS.get()
    if stats is not None:
        stats.note_start(name)
        entry_round = stats.current_round
        if not shadowed:
            stats.note_phase_enter(name, entry_round)
    else:
        entry_round = 0
        shadowed = True  # no stats object -> timing guarded by `if stats is not None` below
    try:
        yield
    finally:
        _CURRENT_PHASE.reset(token)
        _ACTIVE_PHASES.reset(active_token)
        stats = _CURRENT_STATS.get()
        if stats is not None and not shadowed:
            stats.note_phase_exit(name, entry_round)


def in_phase(name: str):
    """Decorator: run the (async) function body inside :func:`phase`."""

    def decorate(fn):
        if iscoroutinefunction(fn):

            @wraps(fn)
            async def wrapper(*args, **kwargs):
                with phase(name):
                    return await fn(*args, **kwargs)

            return wrapper

        @wraps(fn)
        def wrapper(*args, **kwargs):
            with phase(name):
                return fn(*args, **kwargs)

        return wrapper

    return decorate


class LLMUsageStats:
    """Per-phase LLM call, wall-clock & token counters for one agentic ``rag`` run."""

    def __init__(self) -> None:
        self.calls: dict[str, int] = defaultdict(int)
        self.failed: dict[str, int] = defaultdict(int)
        self.phase_time_ms: dict[str, float] = defaultdict(float)
        self.prompt_tokens: dict[str, int] = defaultdict(int)
        self.completion_tokens: dict[str, int] = defaultdict(int)
        self.total_tokens: dict[str, int] = defaultdict(int)
        self.rounds: dict[str, int] = defaultdict(int)
        self.round_times: dict[str, list[float]] = defaultdict(list)
        # Phase wall-clock split per orchestrator round (index 0 = round 1).
        # A phase that runs multiple times inside one round (e.g. the agent
        # researching several claims in parallel) accumulates into that round.
        self.round_phase_times_ms: dict[str, list[float]] = defaultdict(list)
        self.round_claim_counts: dict[str, list[int]] = defaultdict(list)
        self._round_starts: dict[str, float] = {}
        self._current_round: int = 0
        self._phase_active_counts: dict[str, int] = defaultdict(int)
        self._phase_starts: dict[str, float] = {}
        self._round_phase_active_counts: dict[tuple[str, int], int] = defaultdict(int)
        self._round_phase_starts: dict[tuple[str, int], float] = {}

    @property
    def current_round(self) -> int:
        """Index (1-based) of the orchestrator round currently executing, 0 outside."""
        return self._current_round

    def note_start(self, phase_name: str) -> None:
        """Hook called when a phase is entered (phase wall-clock is recorded
        separately in :meth:`record_phase_time`). The log row order is fixed by
        the canonical ``_PHASE_ORDER`` list, so nothing needs to be tracked here
        anymore."""
        return

    def record_call(self, phase_name: str) -> None:
        self.calls[phase_name] += 1

    def record_failed(self, phase_name: str) -> None:
        self.failed[phase_name] += 1

    def _accumulate_phase_time(self, phase_name: str, elapsed_ms: float, entry_round: int = 0) -> None:
        self.phase_time_ms[phase_name] += elapsed_ms
        if entry_round > 0:
            times = self.round_phase_times_ms[phase_name]
            while len(times) < entry_round:
                times.append(0.0)
            times[entry_round - 1] += elapsed_ms
        pending = self.rounds[phase_name] - len(self.round_times[phase_name])
        if pending > 0:
            settled = sum(self.round_times[phase_name])
            self.round_times[phase_name].append(max(0.0, self.phase_time_ms[phase_name] - settled))
            self._round_starts.pop(phase_name, None)
        if phase_name == "orchestrator":
            self._current_round = 0

    def note_phase_enter(self, phase_name: str, entry_round: int = 0) -> None:
        now = time.perf_counter()
        self._phase_active_counts[phase_name] += 1
        if self._phase_active_counts[phase_name] == 1:
            self._phase_starts[phase_name] = now
        if entry_round > 0:
            key = (phase_name, entry_round)
            self._round_phase_active_counts[key] += 1
            if self._round_phase_active_counts[key] == 1:
                self._round_phase_starts[key] = now

    def note_phase_exit(self, phase_name: str, entry_round: int = 0) -> None:
        now = time.perf_counter()
        active = self._phase_active_counts.get(phase_name, 0)
        if active > 0:
            active -= 1
            if active == 0:
                start = self._phase_starts.pop(phase_name, now)
                self._phase_active_counts.pop(phase_name, None)
                self._accumulate_phase_time(phase_name, (now - start) * 1000.0, entry_round=0)
            else:
                self._phase_active_counts[phase_name] = active
        if entry_round > 0:
            key = (phase_name, entry_round)
            active = self._round_phase_active_counts.get(key, 0)
            if active > 0:
                active -= 1
                if active == 0:
                    start = self._round_phase_starts.pop(key, now)
                    self._round_phase_active_counts.pop(key, None)
                    times = self.round_phase_times_ms[phase_name]
                    while len(times) < entry_round:
                        times.append(0.0)
                    times[entry_round - 1] += (now - start) * 1000.0
                else:
                    self._round_phase_active_counts[key] = active

    def record_usage(self, phase_name: str, usage: dict | None) -> None:
        if not usage:
            return
        self.prompt_tokens[phase_name] += int(usage.get("prompt_tokens") or 0)
        self.completion_tokens[phase_name] += int(usage.get("completion_tokens") or 0)
        self.total_tokens[phase_name] += int(usage.get("total_tokens") or 0)

    def record_round(self, phase_name: str) -> None:
        """Count one iteration of a looping phase (e.g. an orchestrator cycle)
        and close the previous iteration's wall-clock."""
        self.rounds[phase_name] += 1
        self._current_round = self.rounds[phase_name]
        now = time.perf_counter()
        prev = self._round_starts.pop(phase_name, None)
        if prev is not None:
            self.round_times[phase_name].append((now - prev) * 1000.0)
        self._round_starts[phase_name] = now

    def record_round_claims(self, phase_name: str, count: int) -> None:
        if self._current_round <= 0:
            return
        counts = self.round_claim_counts[phase_name]
        while len(counts) < self._current_round:
            counts.append(0)
        counts[self._current_round - 1] += int(count or 0)

    def snapshot(self) -> dict[str, dict]:
        known = set(self.calls) | set(self.failed) | set(self.total_tokens) | set(self.phase_time_ms) | set(self.rounds)
        # Rows follow the canonical pipeline order (see _PHASE_ORDER) so the
        # table reads like the actual execution flow. Phases not in the known
        # list are appended afterwards, sorted, as a safe fallback.
        phases = [p for p in _PHASE_ORDER if p in known]
        phases += sorted(known - set(phases))
        rows = {}
        for p in phases:
            per_round = self.round_phase_times_ms.get(p) or []
            if per_round:
                rounds, round_times = len(per_round), list(per_round)
            else:
                rounds, round_times = self.rounds[p], list(self.round_times[p])
            rows[p] = {
                "calls": self.calls[p],
                "failed": self.failed[p],
                "phase_time_ms": self.phase_time_ms[p],
                "prompt_tokens": self.prompt_tokens[p],
                "completion_tokens": self.completion_tokens[p],
                "total_tokens": self.total_tokens[p],
                "rounds": rounds,
                "round_times": round_times,
                "round_claim_counts": list(self.round_claim_counts.get(p) or []),
            }
        return rows

    def log(self, logger: logging.Logger | None = None) -> None:
        log = logger or _LOG
        rows = self.snapshot()
        if not rows:
            # No LLM activity recorded (e.g. a cache hit that returned before
            # any call). Still emit a line so every completed ``rag`` call is
            # accounted for in the logs, instead of silently disappearing.
            log.info("[Agentic RAG] LLM usage by phase: (cached / no LLM calls)")
            return
        total_calls = sum(r["calls"] for r in rows.values())
        total_tokens = sum(r["total_tokens"] for r in rows.values())
        header = f"  {'phase':<16} {'llm_calls':>7} {'prompt_tok':>10} {'output_tok':>12} {'total_tok':>10} {'time(s)':>10}"
        lines = ["[Agentic RAG] LLM usage by phase:", header]

        # With orchestrator-round data, expand the table hierarchically: each
        # round repeats its ``orchestrator`` row with the nested sub-phases
        # (agent/sufficiency/grounded) indented underneath. Time(s) is the
        # per-round wall-clock; the token columns are the phase totals, repeated
        # per round for readability. Phases outside the loop (route/planner/
        # finalize) are listed flat afterwards.
        per_round = self.round_phase_times_ms
        orch_rt = self.round_times.get("orchestrator") or []
        n_rounds = max([len(orch_rt)] + [len(v) for v in per_round.values()]) if (per_round or orch_rt) else 0

        def phase_label(p: str, r: dict, round_idx: int | None = None) -> str:
            label = p
            if p == "claim_research":
                counts = r.get("round_claim_counts") or []
                if round_idx is not None and round_idx < len(counts) and counts[round_idx] > 0:
                    label = f"{p} ({counts[round_idx]})"
            return label

        def row(indent: str, p: str, t_ms: float, round_idx: int | None = None) -> str:
            r = rows[p]
            label = phase_label(p, r, round_idx)
            return f"{indent}{label:<16} {r['calls']:>7} {r['prompt_tokens']:>10} {r['completion_tokens']:>12} {r['total_tokens']:>10} {t_ms / 1000.0:>10.1f}"

        in_rounds = set(per_round)

        def custom_row(indent: str, label: str, p: str, t_ms: float) -> str:
            r = rows[p]
            return f"{indent}{label:<16} {r['calls']:>7} {r['prompt_tokens']:>10} {r['completion_tokens']:>12} {r['total_tokens']:>10} {t_ms / 1000.0:>10.1f}"

        # Walk ``rows`` in its existing (canonical) order. When we reach the
        # orchestrator phase, print one block per round with its nested
        # sub-phases underneath, then suppress those sub-phases from appearing
        # again as duplicate top-level rows later in the table.
        for p in rows:
            if p == "orchestrator" and n_rounds:
                for i in range(n_rounds):
                    # ``orch_rt[i]`` is already in ms (matches ``phase_time_ms``);
                    # ``row()`` converts to seconds, so do NOT pre-divide here.
                    orch_t = orch_rt[i] if i < len(orch_rt) else rows[p]["phase_time_ms"]
                    lines.append(custom_row("  ", f"orchestrator round {i + 1}", p, orch_t))
                    for sub in rows:
                        if sub == "orchestrator":
                            continue
                        v = per_round.get(sub)
                        if not v or i >= len(v):
                            continue
                        lines.append(row("    ", sub, v[i], round_idx=i))
            elif p in in_rounds and n_rounds:
                # Already printed inside the per-round orchestrator block.
                continue
            else:
                lines.append(row("  ", p, rows[p]["phase_time_ms"]))
        lines.append(f"  total: {total_calls} LLM calls, {total_tokens} tokens")
        log.info("\n".join(lines))


_CURRENT_STATS: ContextVar["LLMUsageStats | None"] = ContextVar("agentic_rag_llm_stats", default=None)


@contextmanager
def using_stats(stats: LLMUsageStats):
    """Count the enclosed block's LLM calls into ``stats`` instead of the shared sink.

    ``CountingChatModel`` records into the innermost active stats (via
    ``_CURRENT_STATS``, falling back to the bundle-wide one), so parallel ``rag``
    calls — which run in separate asyncio tasks — each get an independent
    accounting while the shared sink keeps counting the whole request.
    """
    token = _CURRENT_STATS.set(stats)
    try:
        yield stats
    finally:
        _CURRENT_STATS.reset(token)


def record_round(name: str) -> None:
    """Count one iteration of ``name`` into the active stats (no-op when no
    stats are bound, mirroring :func:`phase`)."""
    stats = _CURRENT_STATS.get()
    if stats is not None:
        stats.record_round(name)


def record_round_claims(name: str, count: int) -> None:
    """Record how many claim-level tasks ran in the current round for ``name``."""
    stats = _CURRENT_STATS.get()
    if stats is not None:
        stats.record_round_claims(name, count)


def _last_usage(chat_mdl) -> dict | None:
    mdl = getattr(chat_mdl, "mdl", None)
    usage = getattr(mdl, "last_usage", None)
    if isinstance(usage, dict) and usage.get("total_tokens"):
        return usage
    return None


class CountingChatModel:
    """Proxy over an ``LLMBundle`` recording calls/tokens per phase.

    All other attributes are forwarded to the wrapped bundle, so the rest of
    the harness (``max_length``, ``bind_tools``, ``clone``, ...) keeps working
    unchanged while every ``async_chat*`` entry point is counted.
    """

    def __init__(self, chat_mdl, stats: LLMUsageStats):
        self._chat_mdl = chat_mdl
        self._stats = stats

    def _stats_for(self) -> LLMUsageStats:
        return _CURRENT_STATS.get() or self._stats

    def clone(self):
        return CountingChatModel(self._chat_mdl.clone(), self._stats)

    def __getattr__(self, name: str):
        return getattr(self._chat_mdl, name)

    async def async_chat(self, system: str, history: list, gen_conf: dict | None = None, **kwargs):
        stats = self._stats_for()
        phase_name = _CURRENT_PHASE.get()
        stats.record_call(phase_name)
        try:
            txt = await self._chat_mdl.async_chat(system, history, gen_conf or {}, **kwargs)
        except Exception:
            stats.record_failed(phase_name)
            raise
        stats.record_usage(phase_name, _last_usage(self._chat_mdl))
        return txt

    async def async_chat_streamly(self, system: str, history: list, gen_conf: dict | None = None, **kwargs):
        stats = self._stats_for()
        phase_name = _CURRENT_PHASE.get()
        stats.record_call(phase_name)
        try:
            async for txt in self._chat_mdl.async_chat_streamly(system, history, gen_conf or {}, **kwargs):
                yield txt
        except Exception:
            stats.record_failed(phase_name)
            raise
        finally:
            stats.record_usage(phase_name, _last_usage(self._chat_mdl))

    async def async_chat_streamly_delta(self, system: str, history: list, gen_conf: dict | None = None, **kwargs):
        stats = self._stats_for()
        phase_name = _CURRENT_PHASE.get()
        stats.record_call(phase_name)
        try:
            async for txt in self._chat_mdl.async_chat_streamly_delta(system, history, gen_conf or {}, **kwargs):
                yield txt
        except Exception:
            stats.record_failed(phase_name)
            raise
        finally:
            stats.record_usage(phase_name, _last_usage(self._chat_mdl))