diff --git a/rag/advanced_rag/agentic_rag.py b/rag/advanced_rag/agentic_rag.py index 8914d5f68f..4d8d854b78 100644 --- a/rag/advanced_rag/agentic_rag.py +++ b/rag/advanced_rag/agentic_rag.py @@ -33,9 +33,11 @@ the fast non-tool-calling path. import logging import re from collections.abc import Callable -from typing import Any, List +from typing import Any import json_repair + +from api.db.db_models import Document, Knowledgebase from api.db.services.doc_metadata_service import DocMetadataService from api.db.services.document_service import DocumentService from api.db.services.knowledgebase_service import KnowledgebaseService @@ -44,6 +46,7 @@ from common import settings from common.misc_utils import thread_pool_exec from common.token_utils import num_tokens_from_string from rag.advanced_rag.agentic_rag_graph import _split_think_stream +from rag.advanced_rag.harness.stats import CountingChatModel, LLMUsageStats, in_phase, using_stats from rag.app.tag import label_question from rag.llm.tool_decorator import tool from rag.prompts.generator import ( @@ -55,10 +58,8 @@ from rag.prompts.generator import ( multi_queries_gen, sufficiency_select, ) -from api.db.db_models import Document, Knowledgebase from rag.utils.web_search_conn import WebSearchProvider - # Tokens held back from the model's context when fitting retrieved evidence # into the sufficiency / follow-up prompts. The evidence sits in the MIDDLE of # those templates (question first, JSON output rules last), so if the combined @@ -87,9 +88,65 @@ _RAG_CACHE_MIN_SHARED = 2 # Lightweight stopwords for the cross-`rag`-call dedup only. Never reused for # retrieval/answer quality. _RAG_CACHE_STOPWORDS = frozenset( - "the a an is was were what which when where who how of in to for and or but on at by be as it that this" - " about with their its have has had been being from over under do does did not no yes can could should would" - " also only very much more most some any".split() + [ + "the", + "a", + "an", + "is", + "was", + "were", + "what", + "which", + "when", + "where", + "who", + "how", + "of", + "in", + "to", + "for", + "and", + "or", + "but", + "on", + "at", + "by", + "be", + "as", + "it", + "that", + "this", + "about", + "with", + "their", + "its", + "have", + "has", + "had", + "been", + "being", + "from", + "over", + "under", + "do", + "does", + "did", + "not", + "no", + "yes", + "can", + "could", + "should", + "would", + "also", + "only", + "very", + "much", + "more", + "most", + "some", + "any", + ] ) @@ -141,11 +198,11 @@ class RAGTools: tenant_ids: list[str], chat_mdl: LLMBundle, embed_mdl: LLMBundle | None = None, - kb_ids: List[str] | None = None, + kb_ids: list[str] | None = None, kbs: list[Knowledgebase] | None = None, web_search: WebSearchProvider | None = None, meta_data_filter: dict | None = None, - doc_scope: List[str] | None = None, + doc_scope: list[str] | None = None, user_defined_prompts: dict | None = None, empty_response: str = "", do_refer: bool | None = True, @@ -153,7 +210,11 @@ class RAGTools: text_attachments_content: str = "", ): self.tenant_ids = tenant_ids - self.chat_mdl = chat_mdl.clone() + # P0 instrumentation: count LLM calls / token usage per harness phase. + # The wrapper proxies every ``async_chat*`` entry point (and ``clone``) + # of the bundle, keeping the rest of the harness untouched. + self.llm_stats = LLMUsageStats() + self.chat_mdl = CountingChatModel(chat_mdl.clone(), self.llm_stats) self.embed_mdl = embed_mdl self.thinking_mode = thinking_mode self.field_map = {} @@ -228,7 +289,7 @@ class RAGTools: def has_llm(self) -> bool: return self.chat_mdl is not None - def scoped_doc_ids(self, doc_scope: List[str] | None = None) -> List[str] | None: + def scoped_doc_ids(self, doc_scope: list[str] | None = None) -> list[str] | None: if self.doc_scope is None: return doc_scope if not doc_scope: @@ -274,7 +335,8 @@ class RAGTools: # ------------------------------------------------------------------ # # Graph node helpers (plain async methods — never exposed as tools) # ------------------------------------------------------------------ # - async def formalize(self, messages: List[Any]) -> tuple[str, str]: + @in_phase("formalize") + async def formalize(self, messages: list[Any]) -> tuple[str, str]: """Rewrite the latest user message into a standalone question AND derive its search keywords (each with close synonyms), in one LLM call. @@ -347,7 +409,7 @@ class RAGTools: keywords = str(keywords).strip() return question, keywords - async def pick_documents(self, question: str) -> List[str] | None: + async def pick_documents(self, question: str) -> list[str] | None: """Narrow the search to a document subset for ``question``. Uses document metadata when the bound KBs carry any (mirrors the old @@ -367,7 +429,7 @@ class RAGTools: ids = await self._select_by_titles(question) return ids or None - async def _filter_by_metadata(self, question: str, metas: dict) -> List[str]: + async def _filter_by_metadata(self, question: str, metas: dict) -> list[str]: filters = await gen_meta_filter(self.chat_mdl, metas, question) logging.debug(f"Metadata filter(auto) generated: {filters}") conditions = filters.get("conditions") or [] @@ -386,7 +448,7 @@ class RAGTools: return [] return doc_ids or [] - async def _select_by_titles(self, question: str, max_docs: int = 512) -> List[str]: + async def _select_by_titles(self, question: str, max_docs: int = 512) -> list[str]: docs = await thread_pool_exec(self._collect_doc_titles, max_docs) if not docs: return [] @@ -448,7 +510,7 @@ class RAGTools: self, question: str, keywords: str | list = "", - doc_scope: List[str] | None = None, + doc_scope: list[str] | None = None, top_n: int = 6, similarity_threshold: float = 0.2, using_embedding: bool = False, @@ -565,6 +627,7 @@ class RAGTools: _, msg = message_fit_in(form_message(question, evidence_md), budget) return msg[-1]["content"] + @in_phase("sufficiency") async def judge_sufficiency(self, question: str, evidence_md: str) -> dict: """Judge whether ``evidence_md`` answers ``question`` and pick useful chunks. @@ -580,7 +643,8 @@ class RAGTools: logging.exception("judge_sufficiency failed") return {} - async def gen_followups(self, question: str, query: str, missing: List[str], evidence_md: str) -> List[dict]: + @in_phase("sufficiency") + async def gen_followups(self, question: str, query: str, missing: list[str], evidence_md: str) -> list[dict]: """Generate complementary follow-up (question, query) pairs for gaps.""" evidence_md = self._fit_evidence(question, evidence_md) try: @@ -653,36 +717,40 @@ class RAGTools: if self.tool_started_sink is not None: self.tool_started_sink() - # P0: reuse a near-identical question's cached answer instead of re-running - # the whole agentic graph. Significant-keyword overlap (>= min_overlap AND - # >=2 shared words, and matching numbers) collapses the re-ask pattern - # while leaving genuinely different questions untouched. Attachments bypass - # the cache (their content is appended to the question message below). - if question and not self.text_attachments_content: - qk = _question_keywords(question) - if self._rag_cache: - for cached_q, (cached_answer, cached_gram) in list(self._rag_cache.items()): - if cached_gram and _cache_similar(qk, cached_gram): - shared = len(qk[0] & cached_gram[0]) - _LOG.info("[rag] Reusing cached answer for near-identical question %r (%d shared words); skipping re-research.", question, shared) - return cached_answer + # Per-call instrumentation: each `rag` invocation gets its own stats + # object (bound through the per-task ContextVar), so parallel calls + # report independent usage instead of one shared cumulative sink. + with using_stats(LLMUsageStats()) as call_stats: + # P0: reuse a near-identical question's cached answer instead of re-running + # the whole agentic graph. Significant-keyword overlap (>= min_overlap AND + # >=2 shared words, and matching numbers) collapses the re-ask pattern + # while leaving genuinely different questions untouched. Attachments bypass + # the cache (their content is appended to the question message below). + if question and not self.text_attachments_content: + qk = _question_keywords(question) + if self._rag_cache: + for cached_q, (cached_answer, cached_gram) in list(self._rag_cache.items()): + if cached_gram and _cache_similar(qk, cached_gram): + shared = len(qk[0] & cached_gram[0]) + _LOG.info("[Agentic RAG] Cache hit — reused prior answer for near-identical question %r (%d shared words); skipped research.", question, shared) + return cached_answer - messages = [{"role": "user", "content": question}] if question else [] - if self.text_attachments_content and messages: - messages[-1]["content"] += self.text_attachments_content - final = "" - async for kind, delta in _split_think_stream(run_agentic_rag(self, messages)): - if kind == "answer": - final += delta - if self.answer_sink is not None: - self.answer_sink(delta, kind == "think") - for p, r in [(r"\(\**(ID:\d)\**\)", "[\1]")]: - final = re.sub(p, r, final) + messages = [{"role": "user", "content": question}] if question else [] + if self.text_attachments_content and messages: + messages[-1]["content"] += self.text_attachments_content + final = "" + async for kind, delta in _split_think_stream(run_agentic_rag(self, messages)): + if kind == "answer": + final += delta + if self.answer_sink is not None: + self.answer_sink(delta, kind == "think") + final = re.sub(r"\(\**(ID:\d+)\**\)", r"[\1]", final) - # Cache the freshly produced answer for later near-identical questions. - if question and final and not self.text_attachments_content: - self._rag_cache[question] = (final, _question_keywords(question)) - return final + # Cache the freshly produced answer for later near-identical questions. + if question and final and not self.text_attachments_content: + self._rag_cache[question] = (final, _question_keywords(question)) + call_stats.log() + return final @tool async def summarize_document(self, doc_id: str) -> list[str]: diff --git a/rag/advanced_rag/agentic_rag_graph.py b/rag/advanced_rag/agentic_rag_graph.py index d9ac823445..743f486112 100644 --- a/rag/advanced_rag/agentic_rag_graph.py +++ b/rag/advanced_rag/agentic_rag_graph.py @@ -33,12 +33,13 @@ from __future__ import annotations import asyncio import json -import re import logging +import re from typing import Any, TypedDict from langgraph.graph import END, START, StateGraph +from rag.advanced_rag.harness.stats import in_phase from rag.prompts.generator import form_message, kb_prompt, message_fit_in _LOG = logging.getLogger(__name__) @@ -277,6 +278,7 @@ def build_agentic_graph(tools, token_queue: asyncio.Queue, gen_conf: dict | None answer_conf = dict(gen_conf) if gen_conf else {"temperature": 0.3} # ── Node: formalize_question ── + @in_phase("formalize") async def formalize_question(state: AgenticState) -> dict: msgs = state.get("messages") or [] _LOG.info("[Formalizing the question] Reading the conversation (%d message(s)) to work out the standalone question...", len(msgs)) @@ -294,6 +296,7 @@ def build_agentic_graph(tools, token_queue: asyncio.Queue, gen_conf: dict | None } # ── Node: route ── + @in_phase("route") async def route(state: AgenticState) -> dict: from rag.advanced_rag.harness.route import route_node @@ -330,18 +333,21 @@ def build_agentic_graph(tools, token_queue: asyncio.Queue, gen_conf: dict | None return {"seed_chunks": chunks} # ── Node: planner ── + @in_phase("planner") async def planner(state: AgenticState) -> dict: from rag.advanced_rag.harness.planner import planner_node return await planner_node(state, tools) # ── Node: orchestrator_loop ── + @in_phase("orchestrator") async def orchestrator_loop(state: AgenticState) -> dict: from rag.advanced_rag.harness.orchestrator import orchestrator_loop as _run return await _run(state, tools) # ── Node: formalize_answer ── + @in_phase("finalize") async def formalize_answer(state: AgenticState) -> dict: kbinfos = state.get("kbinfos") or {"chunks": [], "doc_aggs": []} question = state.get("question") or "" diff --git a/rag/advanced_rag/harness/agent.py b/rag/advanced_rag/harness/agent.py index 9aa43b8215..eaa59263a1 100644 --- a/rag/advanced_rag/harness/agent.py +++ b/rag/advanced_rag/harness/agent.py @@ -27,6 +27,7 @@ from rag.advanced_rag.harness.prompts.research_agent_prompt import ( RESEARCH_AGENT_PROMPT, RESEARCH_AGENT_TEXT_PROMPT, ) +from rag.advanced_rag.harness.stats import in_phase _LOG = logging.getLogger(__name__) @@ -138,6 +139,7 @@ def _build_tool_schemas(gated_defs: list[dict]) -> list[dict]: return schemas +@in_phase("claim_research") async def research_agent_loop( claim: ClaimTarget, tools, diff --git a/rag/advanced_rag/harness/orchestrator/__init__.py b/rag/advanced_rag/harness/orchestrator/__init__.py index f2cd291b4d..d4efbbadd6 100644 --- a/rag/advanced_rag/harness/orchestrator/__init__.py +++ b/rag/advanced_rag/harness/orchestrator/__init__.py @@ -3,13 +3,15 @@ import logging from rag.advanced_rag.harness.config import get_mode -from rag.advanced_rag.harness.orchestrator.direct import direct_search -from rag.advanced_rag.harness.orchestrator.decompose import decompose_and_search from rag.advanced_rag.harness.orchestrator.agentic import agentic_research +from rag.advanced_rag.harness.orchestrator.decompose import decompose_and_search +from rag.advanced_rag.harness.orchestrator.direct import direct_search +from rag.advanced_rag.harness.stats import in_phase _LOG = logging.getLogger(__name__) +@in_phase("orchestrator") async def orchestrator_loop(state: dict, tools) -> dict: """Main orchestrator — dispatch to strategy based on thinking mode.""" route = state.get("route") @@ -20,7 +22,7 @@ async def orchestrator_loop(state: dict, tools) -> dict: mode_label = route.thinking_mode if isinstance(route, dict) else route.thinking_mode mode = get_mode(mode_label) - _LOG.info("[Orchestrator] Researching with the \"%s\" approach (%s thinking).", mode.execution_strategy, mode_label) + _LOG.info('[Orchestrator] Researching with the "%s" approach (%s thinking).', mode.execution_strategy, mode_label) if mode.execution_strategy == "direct_search": return await direct_search(state, tools) diff --git a/rag/advanced_rag/harness/orchestrator/agentic.py b/rag/advanced_rag/harness/orchestrator/agentic.py index b29dacf35a..4eac560b84 100644 --- a/rag/advanced_rag/harness/orchestrator/agentic.py +++ b/rag/advanced_rag/harness/orchestrator/agentic.py @@ -3,20 +3,21 @@ import asyncio import logging -from rag.advanced_rag.harness.types import ( - ClaimTarget, - AgentResult, - OrchestratorContext, -) -from rag.advanced_rag.harness.config import get_mode -from rag.advanced_rag.harness.pipeline import Pipeline from rag.advanced_rag.harness.agent import research_agent_loop +from rag.advanced_rag.harness.config import get_mode +from rag.advanced_rag.harness.orchestrator.sufficiency_llm import llm_sufficiency_boost +from rag.advanced_rag.harness.pipeline import Pipeline +from rag.advanced_rag.harness.stats import in_phase, record_round, record_round_claims from rag.advanced_rag.harness.sufficiency import ( - cross_check_claim, compute_fusion_score, + cross_check_claim, route_sufficiency_verdict, ) -from rag.advanced_rag.harness.orchestrator.sufficiency_llm import llm_sufficiency_boost +from rag.advanced_rag.harness.types import ( + AgentResult, + ClaimTarget, + OrchestratorContext, +) _LOG = logging.getLogger(__name__) CLAIM_RESEARCH_TIMEOUT_SECONDS = 180 @@ -51,6 +52,7 @@ def _discovered_entity(tools) -> str | None: return None +@in_phase("orchestrator") async def agentic_research(state: dict, tools) -> dict: """Two-level loop for high/ultra modes.""" question = state.get("question", "") @@ -76,12 +78,16 @@ async def agentic_research(state: dict, tools) -> dict: _STAGNATION_CYCLES = 2 # rounds with no meaningful gain before giving up _STAGNATION_GAIN = 0.05 # minimum fusion-score improvement to count + rounds_run = 0 for cycle in range(mode.max_orchestrator_cycles): + rounds_run = cycle + 1 ctx.iteration = cycle + record_round("orchestrator") _LOG.info("[Agentic research] Research round %d of %d — %d step(s) still unanswered.", cycle + 1, mode.max_orchestrator_cycles, sum(1 for c in ctx.claims if not c.is_verified)) # ── Step A: Research unverified claims (parallel if mode allows) ── unverified = [c for c in ctx.claims if not c.is_verified] + record_round_claims("claim_research", len(unverified)) if unverified: # Consume Phase-2 follow-up queries (missing-pieces feedback) ONCE for @@ -262,12 +268,12 @@ async def agentic_research(state: dict, tools) -> dict: _LOG.info("[Agentic research] Round %d: evidence looks %s (confidence %.0f%%) — next: %s", cycle + 1, verdict.status, verdict.score * 100, action) if action == "ANSWER": - return _finalize(ctx, tools, partial=False) + return _finalize(ctx, tools, partial=False, loop=rounds_run) if action == "ANSWER_PARTIAL": - return _finalize(ctx, tools, partial=True) + return _finalize(ctx, tools, partial=True, loop=rounds_run) if action == "ABSTAIN": tools.kbinfos["chunks"] = [] - return {"verdict": verdict.__dict__, "abstain": True} + return {"verdict": verdict.__dict__, "abstain": True, "loop": rounds_run} if action == "REPLAN": # Ultra: re-plan on low score. Ground the new plan on the evidence # gathered so far, and carry still-valid verified claims over so a @@ -289,10 +295,10 @@ async def agentic_research(state: dict, tools) -> dict: seen = {c.description for c in verified} ctx.claims = verified + [c for c in new_by_desc.values() if c.description not in seen] if action == "FALLBACK_LLM": - return _finalize(ctx, tools, partial=True, fallback=True) + return _finalize(ctx, tools, partial=True, fallback=True, loop=rounds_run) # Max cycles reached - return _finalize(ctx, tools, partial=True) + return _finalize(ctx, tools, partial=True, loop=rounds_run) async def _run_claim_research( @@ -312,7 +318,7 @@ async def _run_claim_research( ) except asyncio.CancelledError: raise - except asyncio.TimeoutError: + except TimeoutError: _LOG.warning( '[Agentic research] Gave up on "%s" — it took longer than %ss.', _snip(claim.description), @@ -348,12 +354,13 @@ async def _run_claim_research( return result -def _finalize(ctx: OrchestratorContext, tools, partial: bool = False, fallback: bool = False) -> dict: +def _finalize(ctx: OrchestratorContext, tools, partial: bool = False, fallback: bool = False, loop: int = 0) -> dict: """Merge agent results into kbinfos and return.""" _merge_agent_results(ctx, tools) return { "verdict": ctx.verdict.__dict__ if ctx.verdict else None, "partial_answer": partial or fallback, + "loop": loop, "kbinfos": tools.kbinfos, } @@ -452,8 +459,8 @@ async def _add_template_group_compilations(comps: set[str], parser_config: dict, if not tenant_id: return try: - from common.misc_utils import thread_pool_exec from api.db.services.compilation_template_group_service import CompilationTemplateGroupService + from common.misc_utils import thread_pool_exec from rag.svr.task_executor_refactor.chunk_post_processor import ( _parser_config_compilation_template_group_ids, ) diff --git a/rag/advanced_rag/harness/orchestrator/decompose.py b/rag/advanced_rag/harness/orchestrator/decompose.py index 33af96f393..4e28cee33a 100644 --- a/rag/advanced_rag/harness/orchestrator/decompose.py +++ b/rag/advanced_rag/harness/orchestrator/decompose.py @@ -13,6 +13,7 @@ from rag.advanced_rag.harness.sufficiency import ( route_sufficiency_verdict, ) from rag.advanced_rag.harness.orchestrator.sufficiency_llm import llm_sufficiency_boost +from rag.advanced_rag.harness.stats import in_phase from rag.advanced_rag.harness.tools.search import hybrid_search _LOG = logging.getLogger(__name__) @@ -56,6 +57,7 @@ Return JSON: Only list in grounded the facts you actually SAW in the evidence; prior-knowledge guesses go in gaps. If the claim is numerical or multi-hop and the evidence has multiple close-but-different figures, disclose all of them in numbers rather than silently picking one.""" +@in_phase("decompose") async def decompose_and_search(state: dict, tools) -> dict: """Decompose, retrieve, analyze evidence, then iterate with next-hop queries.""" question = state.get("question", "") diff --git a/rag/advanced_rag/harness/orchestrator/direct.py b/rag/advanced_rag/harness/orchestrator/direct.py index d7c1fc43b3..ed2b3126c5 100644 --- a/rag/advanced_rag/harness/orchestrator/direct.py +++ b/rag/advanced_rag/harness/orchestrator/direct.py @@ -2,11 +2,13 @@ import logging +from rag.advanced_rag.harness.stats import in_phase from rag.advanced_rag.harness.tools.search import hybrid_search _LOG = logging.getLogger(__name__) +@in_phase("direct") async def direct_search(state: dict, tools) -> dict: """Single hybrid search → merge into kbinfos.""" question = state.get("question", "") diff --git a/rag/advanced_rag/harness/orchestrator/grounded_llm.py b/rag/advanced_rag/harness/orchestrator/grounded_llm.py index fb53b19c88..6cd64e48f1 100644 --- a/rag/advanced_rag/harness/orchestrator/grounded_llm.py +++ b/rag/advanced_rag/harness/orchestrator/grounded_llm.py @@ -20,6 +20,7 @@ from __future__ import annotations import logging +from rag.advanced_rag.harness.stats import in_phase from rag.prompts.generator import PROMPT_JINJA_ENV, gen_json from rag.prompts.template import load_prompt @@ -43,6 +44,7 @@ def _render_reports(reports: list[tuple[str, str]]) -> str: return "\n".join(f"Claim {cid}: {rpt}" for cid, rpt in reports if rpt) +@in_phase("grounded") async def llm_grounded_verify( tools, question: str, diff --git a/rag/advanced_rag/harness/orchestrator/sufficiency_llm.py b/rag/advanced_rag/harness/orchestrator/sufficiency_llm.py index 70bdb74fe8..50756eb74c 100644 --- a/rag/advanced_rag/harness/orchestrator/sufficiency_llm.py +++ b/rag/advanced_rag/harness/orchestrator/sufficiency_llm.py @@ -21,6 +21,7 @@ from __future__ import annotations import logging import re +from rag.advanced_rag.harness.stats import in_phase from rag.advanced_rag.harness.types import SufficiencyVerdict _LOG = logging.getLogger(__name__) @@ -188,6 +189,7 @@ def _evidence_md(tools, evidence_ids=None, keywords: str | None = None) -> str: return "\n\n".join(blocks) +@in_phase("sufficiency") async def llm_sufficiency_boost( tools, question: str, diff --git a/rag/advanced_rag/harness/planner.py b/rag/advanced_rag/harness/planner.py index 51c40b4a98..126675aef7 100644 --- a/rag/advanced_rag/harness/planner.py +++ b/rag/advanced_rag/harness/planner.py @@ -4,15 +4,17 @@ import json import logging import re +from common.token_utils import num_tokens_from_string from rag.advanced_rag.agentic_rag_graph import _snip -from rag.advanced_rag.harness.types import ClaimTarget, WorkflowPlan, RouteDecision from rag.advanced_rag.harness.config import get_mode from rag.advanced_rag.harness.prompts.decompose_prompts import ( - DECOMPOSE_FACTUAL, DECOMPOSE_COMPARATIVE, - DECOMPOSE_PROCEDURAL, DECOMPOSE_EXPLORATORY, + DECOMPOSE_FACTUAL, + DECOMPOSE_PROCEDURAL, ) +from rag.advanced_rag.harness.stats import in_phase +from rag.advanced_rag.harness.types import ClaimTarget, RouteDecision, WorkflowPlan _LOG = logging.getLogger(__name__) @@ -32,6 +34,7 @@ def _extract_json(text: str) -> dict: return {} +@in_phase("planner") async def planner_node(state: dict, tools) -> dict: """Planner node — decompose question into claims based on question type.""" route: RouteDecision = state.get("route") @@ -123,13 +126,21 @@ async def planner_node(state: dict, tools) -> dict: def _format_seed_chunks(seed_chunks, tools) -> str: """Render preliminary-search chunks as grounding context for the planner.""" if not seed_chunks: + _LOG.info("[Planner] No preliminary passages — grounding the plan without seed chunks.") return "(no preliminary results)" try: from rag.prompts.generator import kb_prompt blocks = kb_prompt({"chunks": seed_chunks, "doc_aggs": []}, tools.chat_mdl.max_length) text = "\n".join(blocks).strip() - return text or "(no preliminary results)" + if not text: + return "(no preliminary results)" + _LOG.info( + "[Planner] Grounding the plan with %d preliminary passage(s) (~%d tokens of grounding context).", + len(seed_chunks), + num_tokens_from_string(text), + ) + return text except Exception: _LOG.exception("planner: failed to format seed chunks") return "(no preliminary results)" diff --git a/rag/advanced_rag/harness/route.py b/rag/advanced_rag/harness/route.py index 641c83b5ab..fa5193fe05 100644 --- a/rag/advanced_rag/harness/route.py +++ b/rag/advanced_rag/harness/route.py @@ -7,6 +7,7 @@ import re from rag.advanced_rag.harness.types import RouteDecision from rag.advanced_rag.harness.config import get_mode from rag.advanced_rag.harness.prompts.route_prompt import ROUTE_PROMPT +from rag.advanced_rag.harness.stats import in_phase _LOG = logging.getLogger(__name__) @@ -39,6 +40,7 @@ def _extract_json(text: str) -> dict: return parsed +@in_phase("route") async def route_node(state: dict, tools) -> dict: """Route node — analyze the question, produce RouteDecision.""" question = state.get("question", "") diff --git a/rag/advanced_rag/harness/stats.py b/rag/advanced_rag/harness/stats.py new file mode 100644 index 0000000000..12a5780b83 --- /dev/null +++ b/rag/advanced_rag/harness/stats.py @@ -0,0 +1,424 @@ +"""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))