From 0e34849b86998c73b8b39afac0aea414682d3cd4 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 29 Jul 2026 11:15:37 -0700 Subject: [PATCH 01/15] Recognize Gemini 3.x (and future families) as tool-calling - The agentic chat engine was silently disabled for Gemini 3.x models because the capability check only knew Gemini 1.5/2.x. Any Gemini is now treated as tool-calling except the legacy 1.0-era models, so future families work without a code change. Refs: GML-2171 --- common/llm_services/capabilities.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/common/llm_services/capabilities.py b/common/llm_services/capabilities.py index 653c3fd..4e34e74 100644 --- a/common/llm_services/capabilities.py +++ b/common/llm_services/capabilities.py @@ -109,7 +109,17 @@ def openai_rejects_temperature(model: str) -> bool: def _gemini_tool_calling(model: str) -> bool: # Gemini 1.5+ and 2.x support function calling. - return "gemini-1.5" in model or "gemini-2" in model or "gemini-exp" in model + # Every Gemini family from 1.5 onward supports function/tool calling, and + # future families (4.x, 5.x, ...) will too. Use a denylist instead of an + # allowlist so new models work without a code change: any Gemini is capable + # except the legacy 1.0-era models that predate function calling. + if "gemini" not in model: + return False + if "gemini-1.0" in model or "gemini-pro-vision" in model: + return False + if model.strip() == "gemini-pro": # bare 1.0 alias (versioned ids are fine) + return False + return True def _gemini_thinking(model: str) -> bool: From c7f184629429038b70f3d83345e10e25b54b1312 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 29 Jul 2026 11:15:37 -0700 Subject: [PATCH 02/15] Fix chat Stop button empty-send and light-theme icon size - Ignore empty submits so clicking Stop just as an answer finishes no longer sends an empty message. - Size the Stop icon in both light and dark themes; it was styled only for dark, so it rendered oversized and broke the layout in light mode. Refs: GML-2172 --- graphrag-ui/src/actions/MessageParser.tsx | 4 ++ graphrag-ui/src/index.css | 52 +++++++++++------------ 2 files changed, 30 insertions(+), 26 deletions(-) diff --git a/graphrag-ui/src/actions/MessageParser.tsx b/graphrag-ui/src/actions/MessageParser.tsx index 9645b11..a09f212 100644 --- a/graphrag-ui/src/actions/MessageParser.tsx +++ b/graphrag-ui/src/actions/MessageParser.tsx @@ -7,6 +7,10 @@ interface MessageParserProps { const MessageParser: React.FC = ({ children, actions }) => { const parse = (message: string) => { + // Ignore empty / whitespace-only submits. Clicking the Stop button just as + // an answer finishes can leak a send after the input was cleared, which + // otherwise fires an empty query at the backend. + if (!message || !message.trim()) return; actions.queryGraphragWs(message); }; diff --git a/graphrag-ui/src/index.css b/graphrag-ui/src/index.css index 22dc36e..560ce60 100755 --- a/graphrag-ui/src/index.css +++ b/graphrag-ui/src/index.css @@ -78,6 +78,32 @@ } } +/* Streaming Stop control — must apply in BOTH light and dark themes. While a + response streams, ActionProvider toggles `chat-streaming` on document.body; + Bot.tsx hides the native paper-plane and portals a red Stop icon into the + send button. These previously lived inside `.dark` only, which left the + stop icon unsized in light mode (it fell back to the SVG's default size and + broke the layout). The `.react-chatbot-kit-chat-btn-send` prefix keeps the + fill from being overridden by the base `... svg { fill: #999 }` rule. */ +body.chat-streaming .react-chatbot-kit-chat-input { + pointer-events: none; + opacity: 0.5; +} +body.chat-streaming .react-chatbot-kit-chat-btn-send-icon { + display: none; +} +.react-chatbot-kit-chat-btn-send .graphrag-stop-icon { + display: block; + width: 15px; + height: 15px; + margin: 0 auto; + fill: #dc2626; + cursor: pointer; +} +.react-chatbot-kit-chat-btn-send:hover .graphrag-stop-icon { + fill: #b91c1c; +} + .dark { a { @apply text-white; @@ -88,32 +114,6 @@ .react-chatbot-kit-chat-input-container { @apply !bg-background !border-[#3D3D3D]; } - /* While a response streams, lock the text input. The Send button keeps its - rounded cap (background) but its paper-plane icon is hidden and its click - disabled; a red Stop icon is overlaid in the exact same spot (Bot.tsx - portals it into the input-container). ActionProvider toggles - ``chat-streaming`` on ``document.body`` at stream start / end. */ - body.chat-streaming .react-chatbot-kit-chat-input { - pointer-events: none; - opacity: 0.5; - } - /* Replace the send icon with a red Stop icon IN PLACE — same button, same - position. Bot.tsx hides the native paper-plane and portals the stop icon - into the send button; the button's click is intercepted to stop. */ - body.chat-streaming .react-chatbot-kit-chat-btn-send-icon { - display: none; - } - .graphrag-stop-icon { - display: block; - width: 15px; - height: 15px; - margin: 0 auto; - fill: #dc2626; - cursor: pointer; - } - .react-chatbot-kit-chat-btn-send:hover .graphrag-stop-icon { - fill: #b91c1c; - } .open-dg { @apply bg-background; } From bd3aefbba31b070fed3125e82b8b6d16b66b61fe Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 29 Jul 2026 11:15:37 -0700 Subject: [PATCH 03/15] Give the answer-generation fallback a helpful next step - When an answer can't be generated, suggest retrying or rephrasing and point to the administrator, instead of a dead-end message. Refs: GML-2170 --- graphrag/app/agent/agent_generation.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py index 68074f3..d6dea51 100644 --- a/graphrag/app/agent/agent_generation.py +++ b/graphrag/app/agent/agent_generation.py @@ -83,7 +83,12 @@ def generate_answer(self, question: str, context: str | dict, query: str = "") - except Exception: logger.warning("generate_answer: generation failed") generation = GraphRAGAnswerOutput( - generated_answer="I wasn't able to generate an answer for this question.", + generated_answer=( + "I wasn't able to generate an answer for this question. " + "Try asking again, or rephrase it to be more specific or " + "focused on a single topic. If the problem continues, " + "contact your administrator for more details." + ), citation=[], ) From ec8e9ba1915b4a6f3796c1db72ebb919c4df9869 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 29 Jul 2026 11:41:33 -0700 Subject: [PATCH 04/15] Fix rebuild crash on document names with parentheses - Normalize vertex ids without truncating at "(", so names containing parentheses keep their chunk suffix and no longer corrupt chunk ids - Derive the chunk index from the chunk's position instead of parsing it back out of the id Refs: GML-2173 --- CHANGELOG.md | 5 +++++ ecc/app/graphrag/util.py | 9 ++++++--- ecc/app/graphrag/workers.py | 7 +++---- ecc/app/supportai/util.py | 8 ++++---- 4 files changed, 18 insertions(+), 11 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9a20cf2..a011f2e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,10 @@ # Changelog +## [2.0.2] + +### Fixed +- **Knowledge-graph rebuild no longer fails on document names containing parentheses.** Document names with `(` or `)` are normalized without dropping the rest of the name, so rebuilds complete instead of erroring while creating chunks. + ## [2.0.1] ### Changed diff --git a/ecc/app/graphrag/util.py b/ecc/app/graphrag/util.py index 60129e8..842d0a3 100644 --- a/ecc/app/graphrag/util.py +++ b/ecc/app/graphrag/util.py @@ -220,9 +220,12 @@ def map_attrs(attributes: dict): def process_id(v_id: str): - has_func = re.compile(r"(.*)\(").findall(v_id) - if len(has_func) > 0: - v_id = has_func[0] + # Strip parentheses in place — do NOT truncate at "(". The old `(.*)\(` + # truncation dropped everything from the first "(" onward, which mangled + # ids for document names containing parentheses (e.g. "Q3 (2026) report"): + # the "_chunk_{i}" suffix was lost, so chunk ids no longer ended in an + # integer and the rebuild crashed on int(chunk_id.split("_")[-1]). The + # replace() below removes the parens without dropping the rest. (GML-2173) v_id = v_id.replace(" ", "_").lower().replace("/", "_").replace("(", "").replace(")", "") if v_id == "''" or v_id == '""': return "" diff --git a/ecc/app/graphrag/workers.py b/ecc/app/graphrag/workers.py index 1dc7d1b..518a2d8 100644 --- a/ecc/app/graphrag/workers.py +++ b/ecc/app/graphrag/workers.py @@ -173,7 +173,7 @@ async def chunk_doc( # send chunks to be upserted (func, args) logger.debug("chunk writes to upsert_chan") - await upsert_chan.put((upsert_chunk, (conn, v_id, chunk_id, chunk))) + await upsert_chan.put((upsert_chunk, (conn, v_id, chunk_id, chunk, i))) # send chunks to have entities extracted logger.debug("chunk writes to extract_chan") @@ -208,7 +208,7 @@ async def upsert_doc(conn: AsyncTigerGraphConnection, doc_id, ctype, content_tex conn, "Document", doc_id, "HAS_CONTENT", "Content", doc_id ) -async def upsert_chunk(conn: AsyncTigerGraphConnection, doc_id, chunk_id, chunk): +async def upsert_chunk(conn: AsyncTigerGraphConnection, doc_id, chunk_id, chunk, idx): logger.debug(f"Upserting chunk {chunk_id}") date_added = int(time.time()) # Build the chunk's full vertex + edge bundle and enqueue atomically. @@ -220,7 +220,7 @@ async def upsert_chunk(conn: AsyncTigerGraphConnection, doc_id, chunk_id, chunk) ("DocumentChunk", chunk_id, { "epoch_added": date_added, "epoch_processed": date_added, - "idx": int(chunk_id.split("_")[-1]), + "idx": idx, }), ("Content", chunk_id, {"text": chunk, "epoch_added": date_added}), ] @@ -228,7 +228,6 @@ async def upsert_chunk(conn: AsyncTigerGraphConnection, doc_id, chunk_id, chunk) ("DocumentChunk", chunk_id, "HAS_CONTENT", "Content", chunk_id, None), ("Document", doc_id, "HAS_CHILD", "DocumentChunk", chunk_id, None), ] - idx = int(chunk_id.split("_")[-1]) if idx > 0: edges.append(( "DocumentChunk", chunk_id, "IS_AFTER", diff --git a/ecc/app/supportai/util.py b/ecc/app/supportai/util.py index 0d62c66..630bab9 100644 --- a/ecc/app/supportai/util.py +++ b/ecc/app/supportai/util.py @@ -2,7 +2,6 @@ import base64 import json import logging -import re import traceback from glob import glob from typing import Callable @@ -182,9 +181,10 @@ def map_attrs(attributes: dict): def process_id(v_id: str): - has_func = re.compile(r"(.*)\(").findall(v_id) - if len(has_func) > 0: - v_id = has_func[0] + # Strip parentheses in place — do NOT truncate at "(". Truncating dropped + # the "_chunk_{i}" suffix for ids containing parentheses, corrupting chunk + # ids. The replace() below removes the parens without dropping the rest. + # (GML-2173) v_id = v_id.replace(" ", "_").lower().replace("/", "_").replace("(", "").replace(")", "") if v_id == "''" or v_id == '""': return "" From c911cc15dd8d9309ad0af385b01f32cd102aab51 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 29 Jul 2026 11:46:46 -0700 Subject: [PATCH 05/15] Bump version to 2.0.2 - Set release version to 2.0.2 - Complete the 2.0.2 changelog with the answer-fallback message improvement, Gemini 3.x agentic support, and chat Stop button fixes --- CHANGELOG.md | 5 +++++ VERSION | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a011f2e..2c03ce1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,8 +2,13 @@ ## [2.0.2] +### Changed +- **Clearer message when an answer can't be generated.** When the assistant is unable to produce an answer, it now suggests rephrasing the question to be more specific and contacting an administrator if the problem continues, instead of a bare failure notice. + ### Fixed - **Knowledge-graph rebuild no longer fails on document names containing parentheses.** Document names with `(` or `)` are normalized without dropping the rest of the name, so rebuilds complete instead of erroring while creating chunks. +- **Newer Gemini models are available in agentic chat.** Gemini 3.x and later families are recognized as tool-calling models, so agentic mode works with them instead of silently falling back to classic chat. +- **Chat Stop button behaves correctly.** Stopping a response just as it finishes no longer sends a stray empty message, and the Stop button is the right size in the light theme. ## [2.0.1] diff --git a/VERSION b/VERSION index 10bf840..e9307ca 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -2.0.1 \ No newline at end of file +2.0.2 From 4d1d752acf67ba802d1cf6b75cd0ac920adedbc0 Mon Sep 17 00:00:00 2001 From: prinskumar-tigergraph Date: Fri, 7 Aug 2026 02:10:36 +0530 Subject: [PATCH 06/15] Improve PDF extract cleanup and wire hybrid expand for agent search (#55) * Improve PDF extract cleanup and wire hybrid expand for agent search. Normalize picture-text and numeric cleanup, recover mojibake pages, keep chart labels in chunks, and enable keyword+vector hybrid expand for agent search without eval-shaped age-band post-processing. * Restore _CJK_CHAR_CLASS name; keep fullwidth-digit exclusion. Avoids an unnecessary rename that noisied the PR diff. * Document hybrid expand knobs in docs; drop Toppan server_config from PR. Restore the original OCR figure comment in the chunker to avoid noisy comment-only diff. * Remove unused hybrid_expand/hybrid_method wiring from agent tools and docs. No measured accuracy gain from these settings in Toppan runs; keep extract/chunk product fixes only. --------- Co-authored-by: Prins Kumar --- common/chunkers/structured.py | 39 +++- .../embeddings/tigergraph_embedding_store.py | 3 +- common/utils/image_data_extractor.py | 9 +- common/utils/text_extractors.py | 209 +++++++++++++++++- 4 files changed, 240 insertions(+), 20 deletions(-) diff --git a/common/chunkers/structured.py b/common/chunkers/structured.py index aa3ec26..0980140 100644 --- a/common/chunkers/structured.py +++ b/common/chunkers/structured.py @@ -140,15 +140,28 @@ class Element: # pymupdf4llm artifacts: # • "==> picture [WxH] intentionally omitted <==" — image dropped (skip line) # • "----- Start of picture text -----" / "----- End of picture text -----" -# bracket OCR'd content inside an image; we fold the body into the figure -# so chart-internal labels stay with the image chunk. +# (markdown) or ```` (HTML-comment form) +# bracket pymupdf4llm picture-text inside a figure; we fold the body into +# the figure so chart-internal labels stay with the image chunk. _MD_PICTURE_OMITTED = re.compile(r"^\s*\*+\s*==>\s*picture\b.*intentionally omitted\s*<==\s*\*+.*$", re.IGNORECASE) -_MD_PICTURE_TEXT_START = re.compile(r"^\s*\*+\s*-+\s*Start of picture text\s*-+\s*\*+\s*()?\s*$", re.IGNORECASE) -_MD_PICTURE_TEXT_END = re.compile(r"^\s*\*+\s*-+\s*End of picture text\s*-+\s*\*+\s*()?\s*$", re.IGNORECASE) +_MD_PICTURE_TEXT_START = re.compile( + r"^\s*(?:\*+\s*-+\s*Start of picture text\s*-+\s*\*+|" + r")\s*(?:)?\s*$", + re.IGNORECASE, +) +_MD_PICTURE_TEXT_END = re.compile( + r"^\s*(?:\*+\s*-+\s*End of picture text\s*-+\s*\*+|" + r")\s*(?:)?\s*$", + re.IGNORECASE, +) # Inline variant of the End marker: the picture-text body can arrive as a # single
-joined line with the marker on its tail, so it is not always # line-anchored. Searched anywhere in a line to terminate the block. -_MD_PICTURE_TEXT_END_INLINE = re.compile(r"\*+\s*-+\s*End of picture text\s*-+\s*\*+\s*(?:)?", re.IGNORECASE) +_MD_PICTURE_TEXT_END_INLINE = re.compile( + r"(?:\*+\s*-+\s*End of picture text\s*-+\s*\*+|" + r")\s*(?:)?", + re.IGNORECASE, +) def _flush_prose(buf: List[str], heading: Optional[str], page: Optional[int], out: List[Element]) -> None: @@ -312,18 +325,14 @@ def markdown_to_elements(md: str, page: Optional[int] = None) -> List[Element]: i += 1 continue - # 5b. Other HTML comments (chunk markers etc.) — skip. - if _MD_HTML_COMMENT.match(line): - i += 1 - continue - # 5c. pymupdf4llm "==> picture ... intentionally omitted <==" — drop. if _MD_PICTURE_OMITTED.match(line): i += 1 continue - # 5d. pymupdf4llm picture-text block: ----- Start ... End of picture - # text ----- wraps OCR'd content (chart axis labels, legends). + # 5d. pymupdf4llm picture-text block (markdown dash markers OR HTML + # comments). Must run before the generic HTML-comment skip so + # ```` is not dropped. # Fold the body into the immediately preceding figure when # present so chart-internal text travels with the image. if _MD_PICTURE_TEXT_START.match(line): @@ -361,6 +370,12 @@ def markdown_to_elements(md: str, page: Optional[int] = None) -> List[Element]: out.append(Element(kind="figure", text=body, heading=heading, page=page)) continue + # 5b. Other HTML comments (chunk markers etc.) — skip. + # Picture-text comments are handled above. + if _MD_HTML_COMMENT.match(line): + i += 1 + continue + # 6. Blank line — flush current prose paragraph. if not stripped: _flush_prose(prose_buf, heading, page, out) diff --git a/common/embeddings/tigergraph_embedding_store.py b/common/embeddings/tigergraph_embedding_store.py index bfd1978..47dbed8 100644 --- a/common/embeddings/tigergraph_embedding_store.py +++ b/common/embeddings/tigergraph_embedding_store.py @@ -566,13 +566,14 @@ def retrieve_similar_with_score(self, query_embedding, top_k=10, similarity_thre logger.info(f"Fetch {top_k} similar entries from {vertex_types} with filter {filter_expr}") start_time = time() + # GSQL declares STRING expr=""; pyTigerGraph rejects null. verts = self.conn.runInstalledQuery( "get_topk_similar", params={ "vertex_types": vertex_types, "query_vector": query_embedding, "top_k": top_k*2, - "expr": filter_expr, + "expr": filter_expr or "", } ) end_time = time() diff --git a/common/utils/image_data_extractor.py b/common/utils/image_data_extractor.py index 6be929c..51721f7 100644 --- a/common/utils/image_data_extractor.py +++ b/common/utils/image_data_extractor.py @@ -119,7 +119,14 @@ def describe_image_with_llm(file_path): "stacked bar with a time-period axis), TRANSCRIBE every " "(period, value) pair you can read in the format " "`period: value; period: value; …` — do not summarize " - "the trend in place of the values; (3) the entities, " + "the trend in place of the values. For categorical " + "bar/rank charts (categories such as regions, " + "companies, or age groups on one axis and a numeric " + "unit on the other), TRANSCRIBE every visible " + "category with its value as " + "`category: value; category: value; …` including the " + "unit — never list categories without their numbers; " + "(3) the entities, " "relationships, or process steps in any diagram or " "flowchart; (4) any logo or branding mark, identified by " "name. Do NOT describe layout, background color, " diff --git a/common/utils/text_extractors.py b/common/utils/text_extractors.py index 09399a9..41a441e 100644 --- a/common/utils/text_extractors.py +++ b/common/utils/text_extractors.py @@ -41,8 +41,11 @@ # which bloat tokens 3-5x and confuse retrieval embeddings. The CJK # Unicode ranges below cover CJK Unified Ideographs (U+4E00-U+9FFF), # Hiragana / Katakana / CJK Symbols (U+3000-U+30FF), and full-width -# / half-width forms (U+FF00-U+FFEF). -_CJK_CHAR_CLASS = r"[　-鿿＀-￯]" +# / half-width forms (U+FF00-U+FFEF) excluding fullwidth digits +# (U+FF10-U+FF19). Collapsing digit runs would glue distinct chart +# values such as 767 and 808 into ``767808``. +# One CJK char excluding fullwidth digits (U+FF10-U+FF19). +_CJK_CHAR_CLASS = r"(?:[　-鿿]|[＀-／]|[：-￯])" _VERTICAL_BOLD_CJK = re.compile( rf"(?:\*\*{_CJK_CHAR_CLASS}\*\*(?:)){{2,}}\*\*{_CJK_CHAR_CLASS}\*\*" ) @@ -60,6 +63,20 @@ _TABLE_LINE_RE = re.compile(r"^\s*\|") _BR_TAG_RE = re.compile(r"", re.IGNORECASE) +# pymupdf4llm picture-text blocks (HTML-comment or markdown-dash markers). +# These are figure-associated text from pymupdf4llm, not a separate OCR engine. +_PICTURE_TEXT_BLOCK_RE = re.compile( + r"(?:|\*{0,3}\s*-+\s*Start of picture text\s*-+\s*\*{0,3})" + r"(.*?)" + r"(?:|\*{0,3}\s*-+\s*End of picture text\s*-+\s*\*{0,3})", + re.IGNORECASE | re.DOTALL, +) + +# Adjacent comma-grouped numbers glued with no separator, e.g. ``1,5461,518``. +_GLUED_COMMA_NUMBERS_RE = re.compile( + r"(? list[dict]: def _strip_br_in_table_rows(text: str) -> str: - """Remove ``
`` tags inside markdown table rows. + """Replace ``
`` tags inside markdown table rows with spaces. - Rationale documented at _TABLE_LINE_RE. + Using a space (not empty string) keeps stacked chart values distinct + — ``|767
808|`` becomes ``|767 808|``, never ``|767808|``. """ out: list[str] = [] for line in text.split("\n"): if _TABLE_LINE_RE.match(line): line = _BR_TAG_RE.sub(" ", line) + # Collapse runs of whitespace left by consecutive
tags. + line = re.sub(r"[ \t]{2,}", " ", line) out.append(line) return "\n".join(out) +def _split_glued_comma_numbers(text: str) -> str: + """Insert a space between adjacent comma-grouped numbers. + + pymupdf4llm / chart extraction sometimes emits ``1,5461,518`` instead of + ``1,546 1,518``. Repeat until stable for longer glued runs. + """ + prev = None + while prev != text: + prev = text + text = _GLUED_COMMA_NUMBERS_RE.sub(r"\1 \2", text) + return text + + +def _normalize_picture_text_blocks(text: str) -> str: + """Normalize pymupdf4llm picture-text blocks for chunking + retrieval. + + - Rewrite HTML-comment markers to the markdown form StructuredChunker + already recognizes. + - Turn in-block ``
`` into newlines (not empty joins) so values like + ``767`` and ``808`` stay separable. + - Split glued comma-numbers inside the block. + """ + + def _rewrite(match: re.Match) -> str: + body = match.group(1) or "" + body = _BR_TAG_RE.sub("\n", body) + body = _split_glued_comma_numbers(body) + # Trim excess blank lines inside the block. + body = re.sub(r"\n{3,}", "\n\n", body).strip("\n") + return ( + "***----- Start of picture text -----***\n" + f"{body}\n" + "***----- End of picture text -----***" + ) + + return _PICTURE_TEXT_BLOCK_RE.sub(_rewrite, text) + + +_PAGE_MARKER_RE = re.compile(r"") + + +def _recover_mojibake_pages( + file_path, + markdown: str, + graphname=None, + max_pages: int = 3, +) -> str: + """Recover table/chart text from PDF pages with broken ToUnicode CMaps. + + When glyph mapping fails, embedded text often keeps numbers but corrupts + labels. Drawn (non-embedded) figures also skip the normal image-describe + pass. For page sections that look both corrupted and table-like, render + the page and multimodal-transcribe it, then append the result next to the + original page body. Capped by ``max_pages`` to bound cost. + """ + if not markdown or max_pages <= 0: + return markdown + + # Collect page numbers whose section text looks corrupted. + parts = _PAGE_MARKER_RE.split(markdown) + # parts: [pre, pageNo, body, pageNo, body, ...] + bad_pages: list[int] = [] + if len(parts) >= 3: + for i in range(1, len(parts), 2): + try: + page_no = int(parts[i]) + except (TypeError, ValueError): + continue + body = parts[i + 1] if i + 1 < len(parts) else "" + findings = _detect_mojibake(body, source_hint=f"{file_path}:p{page_no}") + # Prefer pages that look like broken tables (pipe rows + mojibake). + pipe_rows = sum(1 for ln in body.splitlines() if ln.strip().startswith("|")) + if len(findings) >= 3 and pipe_rows >= 3: + bad_pages.append(page_no) + if not bad_pages: + return markdown + + try: + import pymupdf + from common.utils.image_data_extractor import ( + describe_image_with_llm, + should_extract_images, + ) + except Exception as e: # noqa: BLE001 + logger.warning("mojibake page recovery unavailable: %s", e) + return markdown + + if not should_extract_images(graphname): + return markdown + + recovered: dict[int, str] = {} + try: + doc = pymupdf.open(str(file_path)) + except Exception as e: # noqa: BLE001 + logger.warning("mojibake recovery: cannot open %s: %s", file_path, e) + return markdown + + try: + for page_no in bad_pages[:max_pages]: + idx = page_no - 1 + if idx < 0 or idx >= doc.page_count: + continue + try: + page = doc[idx] + # ~150 dpi — enough to read table cells without huge payloads. + pix = page.get_pixmap(matrix=pymupdf.Matrix(2.0, 2.0), alpha=False) + tmp = Path(tempfile.mkdtemp(prefix="mojibake_page_")) / f"p{page_no}.png" + pix.save(str(tmp)) + desc = describe_image_with_llm(str(tmp)) + try: + shutil.rmtree(tmp.parent, ignore_errors=True) + except Exception: + pass + if not desc or "decorative image" in desc.lower(): + continue + # Skip if the transcription itself looks glyph-broken. + if len(_detect_mojibake(desc)) >= 3: + continue + recovered[page_no] = desc.strip() + logger.info( + "mojibake page recovery: %s page %s recovered %s chars", + file_path, + page_no, + len(desc), + ) + except Exception as e: # noqa: BLE001 + logger.warning( + "mojibake page recovery failed for %s p%s: %s", + file_path, + page_no, + e, + ) + finally: + doc.close() + + if not recovered: + return markdown + + # Append recovered transcription under each page marker body. + out_parts: list[str] = [parts[0]] + for i in range(1, len(parts), 2): + page_no_s = parts[i] + body = parts[i + 1] if i + 1 < len(parts) else "" + out_parts.append(f"") + out_parts.append(body) + try: + page_no = int(page_no_s) + except (TypeError, ValueError): + continue + if page_no in recovered: + out_parts.append( + "\n\n\n" + + recovered[page_no] + + "\n" + ) + return "\n".join(out_parts) + + def _collapse_vertical_cjk(text: str) -> str: """Collapse pymupdf4llm's per-character vertical-CJK runs back into a single token. Bold runs ``**X**
**Y**
**Z**`` become ``**XYZ**``; @@ -158,11 +336,21 @@ def _clean_pdf_markdown(markdown: str, source_hint: str = "") -> str: with ``
`` separators and per-character bold markers. The run is collapsed back into a single token so embedding and retrieval see the intended word (e.g. ``**個別信用購入あっせん**``) rather than ten - fragments. + fragments. Fullwidth digits are excluded so chart values are not glued. + + 4. **Picture-text blocks** — figure text wrapped in + ```` (or the dash-marker form) by + pymupdf4llm is rewritten so StructuredChunker keeps the block atomic, + ``
`` becomes newlines, and glued comma-numbers like ``1,5461,518`` + are split. """ # --- Pass 1: remove ColN placeholders --- markdown = _coln_pattern.sub('', markdown) + # --- Pass 1b: normalize pymupdf4llm picture-text blocks before CJK/table + # passes so chart
stacks become newlines rather than empty joins. + markdown = _normalize_picture_text_blocks(markdown) + # --- Pass 2: collapse vertical-CJK runs (do this BEFORE row dedup so # rows that differ only by the collapsed form aren't treated as # distinct rows). @@ -171,7 +359,10 @@ def _clean_pdf_markdown(markdown: str, source_hint: str = "") -> str: # --- Pass 2b: strip
inside markdown table rows --- markdown = _strip_br_in_table_rows(markdown) - # --- Pass 2c: log lines that look like mojibake (failed glyph decode). + # --- Pass 2c: split glued comma-grouped numbers globally --- + markdown = _split_glued_comma_numbers(markdown) + + # --- Pass 2d: log lines that look like mojibake (failed glyph decode). # We don't repair these — the underlying glyphs aren't recoverable # from the markdown — but logging gives operators a grep target. findings = _detect_mojibake(markdown, source_hint) @@ -697,6 +888,12 @@ def _to_markdown_paged(strategy: str | None = None): # Clean up artefacts common in form PDFs (duplicate rows, ColN headers) markdown_content = _clean_pdf_markdown(markdown_content, source_hint=str(file_path)) + # Pages with broken CMaps (mojibake row labels, intact numbers) need a + # page-screenshot multimodal pass — embedded images are often absent. + markdown_content = _recover_mojibake_pages( + file_path, markdown_content, graphname=graphname + ) + # Rename image files that contain spaces to avoid path-parsing issues markdown_content = _sanitize_image_filenames(image_output_folder, markdown_content) From caf83af8fed3a4c26beda523f15419bd87731756 Mon Sep 17 00:00:00 2001 From: prinskumar-tigergraph Date: Fri, 7 Aug 2026 02:11:15 +0530 Subject: [PATCH 07/15] Fix Est. Cost showing $0 for models missing from LangChain pricing. (#57) * Fix Est. Cost showing $0 for models missing from LangChain pricing. Fall back to LiteLLM catalog rates when LangChain total_cost is 0, for the configured model only. * Use configured llm_service for LiteLLM pricing lookup. Avoid scanning all provider prefixes when the server config already names the provider. * Replace LiteLLM cost fallback with user-configured rates from LLM Config. When input/output USD-per-1M rates are set, always use them for Est. Cost; otherwise keep LangChain pricing. --------- Co-authored-by: Prins Kumar --- common/llm_services/base_llm.py | 61 +++++- graphrag-ui/src/pages/TraceLogs.tsx | 6 +- graphrag-ui/src/pages/setup/LLMConfig.tsx | 239 ++++++++++++++++++++-- graphrag/tests/test_token_cost_config.py | 74 +++++++ 4 files changed, 348 insertions(+), 32 deletions(-) create mode 100644 graphrag/tests/test_token_cost_config.py diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index e2f2ad3..fe77ac1 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -15,6 +15,7 @@ import os import re import logging +from typing import Optional from langchain_core.output_parsers import BaseOutputParser, PydanticOutputParser from langchain_core.exceptions import OutputParserException from langchain_core.prompts import BasePromptTemplate @@ -126,7 +127,53 @@ def reset_usage_collection(): _usage_collector.set(None) -def _record_usage(caller_name: str, usage_data: dict): +def _parse_cost_rate(value) -> Optional[float]: + """Return a non-negative float rate, or None if unset/invalid.""" + if value is None or value == "": + return None + try: + rate = float(value) + except (TypeError, ValueError): + return None + if rate < 0: + return None + return rate + + +def estimate_cost_from_config( + config: Optional[dict], + input_tokens: int, + output_tokens: int, +) -> Optional[float]: + """USD cost from user-configured per-1M rates, or None if not configured. + + Both ``input_cost_per_1m`` and ``output_cost_per_1m`` must be set on + ``config`` (USD per 1M tokens). When present they always override + LangChain's built-in cost. + """ + if not config: + return None + inp_rate = _parse_cost_rate(config.get("input_cost_per_1m")) + out_rate = _parse_cost_rate(config.get("output_cost_per_1m")) + if inp_rate is None or out_rate is None: + return None + return ( + max(0, int(input_tokens or 0)) * inp_rate / 1_000_000.0 + + max(0, int(output_tokens or 0)) * out_rate / 1_000_000.0 + ) + + +def _record_usage(caller_name: str, usage_data: dict, config: Optional[dict] = None): + # User-configured rates replace LangChain cost entirely; otherwise keep + # whatever LangChain reported (may be 0 for unknown models). + user_cost = estimate_cost_from_config( + config, + usage_data.get("input_tokens", 0), + usage_data.get("output_tokens", 0), + ) + if user_cost is not None: + usage_data["cost"] = user_cost + logger.info(f"{caller_name} usage: {usage_data}") bucket = _usage_collector.get() if bucket is not None: bucket.append({"caller_name": caller_name, **usage_data}) @@ -500,8 +547,7 @@ def invoke_with_parser( usage_data["output_tokens"] = cb.completion_tokens usage_data["total_tokens"] = cb.total_tokens usage_data["cost"] = cb.total_cost - logger.info(f"{caller_name} usage: {usage_data}") - _record_usage(caller_name, usage_data) + _record_usage(caller_name, usage_data, self.config) raw_text = self._message_text(raw_output) @@ -545,8 +591,7 @@ def invoke_with_tools( usage_data["output_tokens"] = cb.completion_tokens usage_data["total_tokens"] = cb.total_tokens usage_data["cost"] = cb.total_cost - logger.info(f"{caller_name} usage: {usage_data}") - _record_usage(caller_name, usage_data) + _record_usage(caller_name, usage_data, self.config) return resp def invoke_structured( @@ -579,8 +624,7 @@ def invoke_structured( usage_data["output_tokens"] = cb.completion_tokens usage_data["total_tokens"] = cb.total_tokens usage_data["cost"] = cb.total_cost - logger.info(f"{caller_name} usage: {usage_data}") - _record_usage(caller_name, usage_data) + _record_usage(caller_name, usage_data, self.config) return result async def ainvoke_with_parser( @@ -608,8 +652,7 @@ async def ainvoke_with_parser( usage_data["output_tokens"] = cb.completion_tokens usage_data["total_tokens"] = cb.total_tokens usage_data["cost"] = cb.total_cost - logger.info(f"{caller_name} usage: {usage_data}") - _record_usage(caller_name, usage_data) + _record_usage(caller_name, usage_data, self.config) raw_text = self._message_text(raw_output) diff --git a/graphrag-ui/src/pages/TraceLogs.tsx b/graphrag-ui/src/pages/TraceLogs.tsx index a038655..185b86f 100644 --- a/graphrag-ui/src/pages/TraceLogs.tsx +++ b/graphrag-ui/src/pages/TraceLogs.tsx @@ -669,7 +669,7 @@ const TokenOverviewPanel: FC<{ trace: TraceData }> = ({ trace }) => { - Cost is estimated based on the model's published per-token pricing. Actual billing may differ. + Cost uses rates from LLM Config when set; otherwise LangChain pricing. Actual billing may differ. @@ -705,7 +705,7 @@ const TokenOverviewPanel: FC<{ trace: TraceData }> = ({ trace }) => { - Cost is estimated based on the model's published per-token pricing. Actual billing may differ. + Cost uses rates from LLM Config when set; otherwise LangChain pricing. Actual billing may differ. @@ -750,7 +750,7 @@ const TokenOverviewPanel: FC<{ trace: TraceData }> = ({ trace }) => { {formatNumber(usage.total_tokens)} - + {formatCost(usage.cost)} diff --git a/graphrag-ui/src/pages/setup/LLMConfig.tsx b/graphrag-ui/src/pages/setup/LLMConfig.tsx index 382bf06..7737623 100644 --- a/graphrag-ui/src/pages/setup/LLMConfig.tsx +++ b/graphrag-ui/src/pages/setup/LLMConfig.tsx @@ -140,6 +140,9 @@ const LLMConfig = () => { const [completionProvider, setCompletionProvider] = useState("openai"); const [completionConfig, setCompletionConfig] = useState>({}); const [completionDefaultModel, setCompletionDefaultModel] = useState(""); + // Optional USD per 1M tokens — when both set, override LangChain cost + const [completionInputCostPer1m, setCompletionInputCostPer1m] = useState(""); + const [completionOutputCostPer1m, setCompletionOutputCostPer1m] = useState(""); const [embeddingProvider, setEmbeddingProvider] = useState("openai"); const [embeddingConfig, setEmbeddingConfig] = useState>({}); @@ -148,6 +151,8 @@ const LLMConfig = () => { const [multimodalProvider, setMultimodalProvider] = useState("openai"); const [multimodalConfig, setMultimodalConfig] = useState>({}); const [multimodalModelName, setMultimodalModelName] = useState(""); + const [multimodalInputCostPer1m, setMultimodalInputCostPer1m] = useState(""); + const [multimodalOutputCostPer1m, setMultimodalOutputCostPer1m] = useState(""); const isChatbotOnlyMode = llmConfigAccess === "chatbot_only"; // Per-graph chatbot config state (chatbot_only mode) @@ -156,6 +161,8 @@ const LLMConfig = () => { const [chatbotProviderConfig, setChatbotProviderConfig] = useState>({}); const [chatbotModelName, setChatbotModelName] = useState(""); const [chatbotTemperature, setChatbotTemperature] = useState("0"); + const [chatbotInputCostPer1m, setChatbotInputCostPer1m] = useState(""); + const [chatbotOutputCostPer1m, setChatbotOutputCostPer1m] = useState(""); const [globalChatInfo, setGlobalChatInfo] = useState({ llm_service: "", llm_model: "" }); // Superadmin scope: "global" edits global config, "graph" edits per-graph overrides @@ -274,6 +281,16 @@ const LLMConfig = () => { setChatbotProvider(data.chatbot_config.llm_service?.toLowerCase() || defaultProv); setChatbotModelName(data.chatbot_config.llm_model || ""); setChatbotTemperature(String(data.chatbot_config.model_kwargs?.temperature ?? "0")); + setChatbotInputCostPer1m( + data.chatbot_config.input_cost_per_1m != null + ? String(data.chatbot_config.input_cost_per_1m) + : "" + ); + setChatbotOutputCostPer1m( + data.chatbot_config.output_cost_per_1m != null + ? String(data.chatbot_config.output_cost_per_1m) + : "" + ); // Resolve chatbot config: base config + chatbot overrides setChatbotProviderConfig(loadServiceConfigResolved(data.chatbot_config)); } else { @@ -297,18 +314,40 @@ const LLMConfig = () => { setChatbotProvider(chatProv || defaultProv); setChatbotModelName(llmConfig.chat_service.llm_model || ""); setChatbotTemperature(String(llmConfig.chat_service.model_kwargs?.temperature ?? "0")); + setChatbotInputCostPer1m( + llmConfig.chat_service.input_cost_per_1m != null + ? String(llmConfig.chat_service.input_cost_per_1m) + : "" + ); + setChatbotOutputCostPer1m( + llmConfig.chat_service.output_cost_per_1m != null + ? String(llmConfig.chat_service.output_cost_per_1m) + : "" + ); setChatbotProviderConfig(loadServiceConfigResolved(llmConfig.chat_service)); } else { setUseCustomChatbot(false); setChatbotProvider(defaultProv); setChatbotModelName(""); setChatbotTemperature("0"); + setChatbotInputCostPer1m(""); + setChatbotOutputCostPer1m(""); setChatbotProviderConfig({ ...baseConfig }); } // Canonical per-service state — both single and multi-provider UIs read these setCompletionProvider(completionProv || "openai"); setCompletionDefaultModel(llmConfig.completion_service?.llm_model || ""); + setCompletionInputCostPer1m( + llmConfig.completion_service?.input_cost_per_1m != null + ? String(llmConfig.completion_service.input_cost_per_1m) + : "" + ); + setCompletionOutputCostPer1m( + llmConfig.completion_service?.output_cost_per_1m != null + ? String(llmConfig.completion_service.output_cost_per_1m) + : "" + ); setCompletionConfig(loadServiceConfigResolved(llmConfig.completion_service)); setEmbeddingProvider(embeddingProv || completionProv || "openai"); @@ -318,6 +357,16 @@ const LLMConfig = () => { setMultimodalProvider(multimodalProv || completionProv || "openai"); const mmModel = llmConfig.multimodal_service?.llm_model || ""; setMultimodalModelName(mmModel); + setMultimodalInputCostPer1m( + llmConfig.multimodal_service?.input_cost_per_1m != null + ? String(llmConfig.multimodal_service.input_cost_per_1m) + : "" + ); + setMultimodalOutputCostPer1m( + llmConfig.multimodal_service?.output_cost_per_1m != null + ? String(llmConfig.multimodal_service.output_cost_per_1m) + : "" + ); setMultimodalConfig(loadServiceConfigResolved(llmConfig.multimodal_service)); setUseCustomMultimodal(!!mmModel || !!multimodalProv); } catch (error: any) { @@ -393,6 +442,70 @@ const LLMConfig = () => { return serviceConfig; }; + /** Attach optional USD-per-1M rates when both fields are filled. */ + const attachTokenCosts = (svc: Record, input: string, output: string) => { + const inp = input.trim(); + const out = output.trim(); + if (inp === "" || out === "") return; + const inpN = Number(inp); + const outN = Number(out); + if (Number.isNaN(inpN) || Number.isNaN(outN) || inpN < 0 || outN < 0) return; + svc.input_cost_per_1m = inpN; + svc.output_cost_per_1m = outN; + }; + + const renderTokenCostFields = ( + inputValue: string, + outputValue: string, + setInput: (v: string) => void, + setOutput: (v: string) => void, + ) => ( +
+ +
+
+ + { + setInput(e.target.value); + clearTestResults(); + }} + /> +
+
+ + { + setOutput(e.target.value); + clearTestResults(); + }} + /> +
+
+

+ When both are set, Est. Cost always uses these rates (overrides LangChain). Leave empty to use LangChain pricing. +

+
+ ); + /** * Build the candidate LLM config payload. * Used by both test-connection and save — same structure, single source of truth. @@ -410,6 +523,11 @@ const LLMConfig = () => { prompt_path: `./common/prompts/${getPromptPath(completionProvider)}/`, ...buildServiceConfig(completionProvider, completionConfig) }; + attachTokenCosts( + completionServiceConfig, + completionInputCostPer1m, + completionOutputCostPer1m, + ); llmConfigData = { graphname: selectedGraph || undefined, @@ -430,6 +548,11 @@ const LLMConfig = () => { model_kwargs: { temperature: 0 }, ...buildServiceConfig(multimodalProvider, multimodalConfig) }; + attachTokenCosts( + llmConfigData.multimodal_service, + multimodalInputCostPer1m, + multimodalOutputCostPer1m, + ); } else { llmConfigData.multimodal_service = null; } @@ -442,6 +565,11 @@ const LLMConfig = () => { model_kwargs: { temperature: parseFloat(chatbotTemperature) || 0 }, ...buildServiceConfig(chatbotProvider, chatbotProviderConfig), }; + attachTokenCosts( + llmConfigData.chat_service, + chatbotInputCostPer1m, + chatbotOutputCostPer1m, + ); } else { llmConfigData.chat_service = null; } @@ -453,6 +581,11 @@ const LLMConfig = () => { prompt_path: `./common/prompts/${getPromptPath(completionProvider)}/`, ...buildServiceConfig(completionProvider, completionConfig) }; + attachTokenCosts( + completionServiceConfig, + completionInputCostPer1m, + completionOutputCostPer1m, + ); llmConfigData = { graphname: selectedGraph || undefined, @@ -468,6 +601,11 @@ const LLMConfig = () => { llmConfigData.multimodal_service = { llm_model: multimodalModelName, }; + attachTokenCosts( + llmConfigData.multimodal_service, + multimodalInputCostPer1m, + multimodalOutputCostPer1m, + ); } else { llmConfigData.multimodal_service = null; } @@ -478,6 +616,11 @@ const LLMConfig = () => { ...(chatbotModelName.trim() ? { llm_model: chatbotModelName } : {}), model_kwargs: { temperature: chatTemp }, }; + attachTokenCosts( + llmConfigData.chat_service, + chatbotInputCostPer1m, + chatbotOutputCostPer1m, + ); } else { llmConfigData.chat_service = null; } @@ -509,6 +652,7 @@ const LLMConfig = () => { model_kwargs: { temperature: parseFloat(chatbotTemperature) || 0 }, ...buildServiceConfig(chatbotProvider, chatbotProviderConfig), }; + attachTokenCosts(chatService, chatbotInputCostPer1m, chatbotOutputCostPer1m); llmConfigData = { graphname: selectedGraph || undefined, chat_service: chatService }; } else { // Revert to inherit: send null chat_service @@ -916,6 +1060,13 @@ const LLMConfig = () => { /> + {renderTokenCostFields( + chatbotInputCostPer1m, + chatbotOutputCostPer1m, + setChatbotInputCostPer1m, + setChatbotOutputCostPer1m, + )} +
+ {renderTokenCostFields( + completionInputCostPer1m, + completionOutputCostPer1m, + setCompletionInputCostPer1m, + setCompletionOutputCostPer1m, + )} +
@@ -1171,16 +1329,26 @@ const LLMConfig = () => {
{useCustomChatbot && ( - { - setChatbotModelName(e.target.value); - clearTestResults(); - }} - /> + <> + { + setChatbotModelName(e.target.value); + clearTestResults(); + }} + /> +
+ {renderTokenCostFields( + chatbotInputCostPer1m, + chatbotOutputCostPer1m, + setChatbotInputCostPer1m, + setChatbotOutputCostPer1m, + )} +
+ )}

Used by the chatbot for answering user questions @@ -1233,16 +1401,26 @@ const LLMConfig = () => {

)} {useCustomMultimodal && ( - { - setMultimodalModelName(e.target.value); - clearTestResults(); - }} - /> + <> + { + setMultimodalModelName(e.target.value); + clearTestResults(); + }} + /> +
+ {renderTokenCostFields( + multimodalInputCostPer1m, + multimodalOutputCostPer1m, + setMultimodalInputCostPer1m, + setMultimodalOutputCostPer1m, + )} +
+ )}

Used for processing images and multimodal content @@ -1324,6 +1502,13 @@ const LLMConfig = () => { Model for entity extraction and community summarization.

+ + {renderTokenCostFields( + completionInputCostPer1m, + completionOutputCostPer1m, + setCompletionInputCostPer1m, + setCompletionOutputCostPer1m, + )} @@ -1394,6 +1579,13 @@ const LLMConfig = () => { /> + {renderTokenCostFields( + chatbotInputCostPer1m, + chatbotOutputCostPer1m, + setChatbotInputCostPer1m, + setChatbotOutputCostPer1m, + )} +
+ + {renderTokenCostFields( + multimodalInputCostPer1m, + multimodalOutputCostPer1m, + setMultimodalInputCostPer1m, + setMultimodalOutputCostPer1m, + )} )} diff --git a/graphrag/tests/test_token_cost_config.py b/graphrag/tests/test_token_cost_config.py new file mode 100644 index 0000000..932cf31 --- /dev/null +++ b/graphrag/tests/test_token_cost_config.py @@ -0,0 +1,74 @@ +# Copyright (c) 2024-2026 TigerGraph, Inc. +# +# This program may be redistributed and/or modified under the terms of the GNU +# Affero General Public License as published by the Free Software Foundation, +# either version 3 of the License, or (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, but WITHOUT +# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS +# FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more +# details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +import unittest + +from common.llm_services.base_llm import estimate_cost_from_config, _record_usage + + +class TestTokenCostConfig(unittest.TestCase): + def test_none_when_rates_missing(self): + self.assertIsNone(estimate_cost_from_config({}, 1000, 100)) + self.assertIsNone( + estimate_cost_from_config({"input_cost_per_1m": 2.5}, 1000, 100) + ) + self.assertIsNone( + estimate_cost_from_config({"output_cost_per_1m": 15.0}, 1000, 100) + ) + + def test_computes_from_per_1m_rates(self): + # 1M in @ $2.50 + 1M out @ $15.00 = $17.50 + cost = estimate_cost_from_config( + {"input_cost_per_1m": 2.5, "output_cost_per_1m": 15.0}, + 1_000_000, + 1_000_000, + ) + self.assertAlmostEqual(cost, 17.5, places=9) + + def test_small_token_counts(self): + # 1000 in @ $2.50/1M + 100 out @ $15/1M = 0.0025 + 0.0015 = 0.004 + cost = estimate_cost_from_config( + {"input_cost_per_1m": "2.5", "output_cost_per_1m": "15"}, + 1000, + 100, + ) + self.assertAlmostEqual(cost, 0.004, places=9) + + def test_record_usage_overrides_langchain(self): + usage = { + "input_tokens": 1000, + "output_tokens": 100, + "total_tokens": 1100, + "cost": 0.123, # LangChain non-zero — still overridden + } + _record_usage( + "test", + usage, + {"input_cost_per_1m": 2.5, "output_cost_per_1m": 15.0}, + ) + self.assertAlmostEqual(usage["cost"], 0.004, places=9) + + def test_record_usage_keeps_langchain_when_unconfigured(self): + usage = { + "input_tokens": 1000, + "output_tokens": 100, + "total_tokens": 1100, + "cost": 0.123, + } + _record_usage("test", usage, {"llm_model": "gpt-5.4"}) + self.assertEqual(usage["cost"], 0.123) + + +if __name__ == "__main__": + unittest.main() From 07f63c7dff2d88825a838f17433687254ab96839 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 6 Aug 2026 13:44:42 -0700 Subject: [PATCH 08/15] Fix token-cost rate inheritance and document 2.0.2 changes - Completion input/output token-cost rates are no longer inherited by a same-provider embedding service, which is priced separately - Add changelog entries for configurable token cost rates and for the PDF chart/table extraction fidelity improvements --- CHANGELOG.md | 2 ++ common/config.py | 9 +++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c03ce1..ab3f2f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,9 +4,11 @@ ### Changed - **Clearer message when an answer can't be generated.** When the assistant is unable to produce an answer, it now suggests rephrasing the question to be more specific and contacting an administrator if the problem continues, instead of a bare failure notice. +- **Token cost rates can be set per service.** LLM Config accepts optional input/output USD-per-1M-token rates for the completion, chatbot, and multimodal services; when set, Est. Cost uses them so cost is accurate for models the pricing library doesn't recognize (otherwise it keeps the library's pricing). ### Fixed - **Knowledge-graph rebuild no longer fails on document names containing parentheses.** Document names with `(` or `)` are normalized without dropping the rest of the name, so rebuilds complete instead of erroring while creating chunks. +- **PDF chart labels and table numbers survive extraction.** Picture-text blocks are kept intact, stacked and comma-grouped chart values are no longer glued together, and pages with unreadable embedded fonts are recovered from a page image — so figures and tables reach the graph with their numbers intact. - **Newer Gemini models are available in agentic chat.** Gemini 3.x and later families are recognized as tool-calling models, so agentic mode works with them instead of silently falling back to classic chat. - **Chat Stop button behaves correctly.** Stopping a response just as it finishes no longer sends a stray empty message, and the Stop button is the right size in the light theme. diff --git a/common/config.py b/common/config.py index d229d69..d6144ca 100644 --- a/common/config.py +++ b/common/config.py @@ -187,8 +187,13 @@ def resolve_llm_services(llm_cfg: dict) -> dict: embedding_provider = embedding.get("embedding_model_service", "").lower() completion_provider = completion.get("llm_service", "").lower() if embedding_provider and embedding_provider == completion_provider: - # Identity/schema keys that belong to the embedding service itself - embedding_own_keys = {"embedding_model_service", "model_name", "authentication_configuration", "token_limit"} + # Identity/schema keys that belong to the embedding service itself. + # Token-cost rates stay per-service too — completion's per-1M rates + # must not be applied to embedding usage, which is priced separately. + embedding_own_keys = { + "embedding_model_service", "model_name", "authentication_configuration", + "token_limit", "input_cost_per_1m", "output_cost_per_1m", + } for k, v in completion.items(): if k not in embedding_own_keys and k not in embedding: embedding[k] = v From bf60c207ba369ad77fcee706248a183203fa6d49 Mon Sep 17 00:00:00 2001 From: prinskumar-tigergraph Date: Thu, 27 Aug 2026 01:52:18 +0530 Subject: [PATCH 09/15] Add Recall@5 regression metric with Multihop30 dataset (#59) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add Recall@5 regression metric with Multihop30 dataset - New recall_evaluator.py: computes Recall@K by querying GraphRAG and matching retrieved chunks against ground-truth chunks using embedding cosine similarity (threshold=0.70) or LLM judge - New run_recall.sh: Docker wrapper to run recall evaluation - New Multihop30 dataset: 30 multi-hop questions (10 each from HotpotQA, 2WikiMultiHopQA, MuSiQue) with raw text corpus, answers, and ground_truth_chunks.csv for retrieval scoring - Updated evaluator.py, run_eval.sh: unified --mode parameter replacing --agent + --search-type - Updated run_setup.sh, setup_graph.py: support for Multihop30 graph setup Co-authored-by: Cursor * Clean up Multihop30 README — remove stale build scripts and fix file references Co-authored-by: Cursor * Increase evaluator.py httpx timeout from 120s to 600s for complex agentic queries Co-authored-by: Cursor --------- Co-authored-by: Prins Kumar Co-authored-by: Cursor --- graphrag/tests/regression/evaluator.py | 2 +- graphrag/tests/regression/recall_evaluator.py | 936 ++++++++++++++++++ graphrag/tests/regression/run_eval.sh | 28 +- graphrag/tests/regression/run_recall.sh | 66 ++ graphrag/tests/regression/run_setup.sh | 2 +- graphrag/tests/regression/setup_graph.py | 30 +- .../tests/test_questions/Multihop30/README.md | 60 ++ .../test_questions/Multihop30/answers.csv | 31 + .../data/2wikimultihopqa_corpus.txt | 289 ++++++ .../Multihop30/data/hotpotqa_corpus.txt | 300 ++++++ .../Multihop30/data/musique_corpus.txt | 453 +++++++++ .../Multihop30/ground_truth_chunks.csv | 91 ++ .../test_questions/Multihop30/questions.csv | 31 + 13 files changed, 2297 insertions(+), 22 deletions(-) create mode 100644 graphrag/tests/regression/recall_evaluator.py create mode 100755 graphrag/tests/regression/run_recall.sh mode change 100644 => 100755 graphrag/tests/regression/run_setup.sh create mode 100644 graphrag/tests/test_questions/Multihop30/README.md create mode 100644 graphrag/tests/test_questions/Multihop30/answers.csv create mode 100644 graphrag/tests/test_questions/Multihop30/data/2wikimultihopqa_corpus.txt create mode 100644 graphrag/tests/test_questions/Multihop30/data/hotpotqa_corpus.txt create mode 100644 graphrag/tests/test_questions/Multihop30/data/musique_corpus.txt create mode 100644 graphrag/tests/test_questions/Multihop30/ground_truth_chunks.csv create mode 100644 graphrag/tests/test_questions/Multihop30/questions.csv diff --git a/graphrag/tests/regression/evaluator.py b/graphrag/tests/regression/evaluator.py index 3638fd4..528a40a 100644 --- a/graphrag/tests/regression/evaluator.py +++ b/graphrag/tests/regression/evaluator.py @@ -434,7 +434,7 @@ def _query_graphrag( "include_fields": "query_sources", }, auth=(username, password), - timeout=120.0, + timeout=600.0, # agentic planned mode can take >120s on complex multi-hop questions ) resp.raise_for_status() data = resp.json() diff --git a/graphrag/tests/regression/recall_evaluator.py b/graphrag/tests/regression/recall_evaluator.py new file mode 100644 index 0000000..45d30b8 --- /dev/null +++ b/graphrag/tests/regression/recall_evaluator.py @@ -0,0 +1,936 @@ +"""GraphRAG Regression — Recall@K Evaluator + +Measures retrieval recall against a labelled dataset using: + + Recall@K = (ground-truth chunks found in top-K retrieved chunks) + / (total ground-truth chunks for that question) + +Averaged over all questions to produce a final Avg Recall@K score. + +Dataset layout expected: + test_questions// + ├── data/_corpus.txt — one raw text file per source (GraphRAG chunks naturally) + ├── questions.csv — single column: question + └── ground_truth_chunks.csv — columns: question_index, chunk_index, context + Recall denominator = GT chunks per question (not total dataset chunks). + +Matching strategy — two options via --match: + + embedding (default): + Each GT chunk and each retrieved chunk are embedded with the model from + server_config.json (e.g. gemini-embedding-001). A GT chunk is "found" when + cosine similarity ≥ threshold (default 0.75). Handles paraphrasing and + contextual retrievers correctly. + + llm: + An LLM (same model as completions in server_config.json) judges whether the + retrieved chunk contains the same factual information as the GT chunk. + Most accurate but slowest and costliest. + +Run via: + ./graphrag/tests/regression/run_recall.sh --dataset Multihop30 --graphname +""" + +from __future__ import annotations + +import argparse +import csv +import json +import logging +import math +import os +import re +import sys +import threading +import time +import warnings +from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import asdict, dataclass, field +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import httpx + +if "/code" not in sys.path: + sys.path.insert(0, "/code") + +warnings.filterwarnings("ignore") +warnings.showwarning = lambda *a, **k: None +logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(name)s: %(message)s") + +# ── ANSI colours (mirrors evaluator.py) ────────────────────────────────────── +if os.environ.get("NO_COLOR"): + _USE_COLOR = False +elif os.environ.get("FORCE_COLOR"): + _USE_COLOR = True +else: + _USE_COLOR = sys.stdout.isatty() + +_G = "\033[32m" if _USE_COLOR else "" +_R = "\033[31m" if _USE_COLOR else "" +_Y = "\033[33m" if _USE_COLOR else "" +_C = "\033[36m" if _USE_COLOR else "" +_B = "\033[1m" if _USE_COLOR else "" +_X = "\033[0m" if _USE_COLOR else "" + +# ── Constants ───────────────────────────────────────────────────────────────── +_DEFAULT_K = 5 +_MAX_CHUNKS = 50 # max chunks extracted from GraphRAG response +_MAX_CHUNK_CH = 5000 # truncation per chunk before matching — keep higher than GT chunk size (4000) +_EMBED_THRESHOLD = 0.70 # cosine similarity threshold for embedding match + # Calibrated on Multihop30: genuine misses cluster at 0.54–0.64; + # true matches cluster at 0.70–0.85. Gap is clear and stable. +_EMBED_MAX_CHARS = 8000 # truncate text before embedding (API limit safety) +_EMBED_CONCURRENCY = 4 # max parallel embedding API calls +_LLM_CHUNK_CHARS = 4000 # truncate chunk text sent to LLM judge + +# Corpus passages are stored in the graph with a "===== Title =====" header prepended +# by the ingestion pipeline. GT chunks never have this header, so we strip it from +# retrieved text before embedding / LLM comparison to avoid a spurious similarity drop. +_TITLE_HEADER_RE = re.compile(r"^\s*={2,}[^=\n]+={2,}\s*\n?", re.MULTILINE) + +# ── Shared caches (embedding vectors, LLM verdicts) ─────────────────────────── +_embed_cache: Dict[str, List[float]] = {} +_llm_cache: Dict[str, bool] = {} +_embed_lock = threading.Lock() +_llm_lock = threading.Lock() +_embed_sem = threading.Semaphore(_EMBED_CONCURRENCY) + + +# ─── Data types ─────────────────────────────────────────────────────────────── + +@dataclass +class RecallQuestion: + index: int # 1-based, matches question_index in ground_truth_contexts.csv + question: str + gt_contexts: List[str] # GT chunk text (from ground_truth_chunks.csv) + gt_titles: List[Optional[str]] = field(default_factory=list) # Wikipedia title (may be None) + + +@dataclass +class RecallResult: + # ── written to CSV ──────────────────────────────────────────────────────── + question_index: int + question: str + gt_context_count: int = 0 + retrieved_chunk_count: int = 0 + matched_count: int = 0 + recall_at_k: Optional[float] = None + agent_mode: Optional[str] = None + search_type_used: Optional[str] = None + response_time_seconds: float = 0.0 + # ── CLI-only ────────────────────────────────────────────────────────────── + error: Optional[str] = field(default=None, repr=False) + answered_question: bool = field(default=False, repr=False) + matched_indices: List[int] = field(default_factory=list, repr=False) + unmatched_contexts: List[str] = field(default_factory=list, repr=False) + embed_scores: List[float] = field(default_factory=list, repr=False) # max cosine per GT chunk + + +# ─── Loaders ───────────────────────────────────────────────────────────────── + +def _detect_encoding(path: str) -> str: + with open(path, "rb") as f: + raw = f.read(4) + if raw[:3] == b"\xef\xbb\xbf": + return "utf-8-sig" + try: + import chardet + with open(path, "rb") as f: + detected = chardet.detect(f.read()) + enc = detected.get("encoding") or "utf-8" + return enc if (detected.get("confidence") or 0) >= 0.5 else "utf-8" + except ImportError: + return "utf-8" + + +def _read_single_column(path: str, column: str) -> List[str]: + encoding = _detect_encoding(path) + with open(path, newline="", encoding=encoding) as f: + reader = csv.reader(f) + rows = list(reader) + if not rows: + sys.exit(f"ERROR: {os.path.basename(path)} is empty.") + header = rows[0] + try: + col_idx = header.index(column) + except ValueError: + sys.exit( + f"ERROR: {os.path.basename(path)} must have a '{column}' column. " + f"Found: {header}" + ) + results = [] + for row in rows[1:]: + if not row: + continue + value = ",".join(row[col_idx:]).strip() + if value: + results.append(value) + return results + + +def load_recall_questions(dataset_dir: str) -> List[RecallQuestion]: + """Load questions.csv and ground_truth_chunks.csv. + + ground_truth_chunks.csv columns: + question_index — 1-based int + chunk_index — chunk number within the question (1-based) + context — chunk text (~400-500 chars) + + Denominator for Recall@K is the number of GT chunks for THAT question, + not the total chunks in the dataset. + Run build_gt_chunks.py once after setup to generate this file. + """ + q_path = os.path.join(dataset_dir, "questions.csv") + chunks_path = os.path.join(dataset_dir, "ground_truth_chunks.csv") + + if not os.path.exists(q_path): + sys.exit(f"ERROR: questions.csv not found at {q_path}") + if not os.path.exists(chunks_path): + sys.exit( + f"ERROR: ground_truth_chunks.csv not found in {dataset_dir}.\n" + f"Run build_gt_chunks.py first:\n" + f" python {dataset_dir}/build_gt_chunks.py" + ) + + questions_raw = _read_single_column(q_path, "question") + + encoding = _detect_encoding(chunks_path) + gt_map: Dict[int, List[str]] = {} + with open(chunks_path, newline="", encoding=encoding) as f: + reader = csv.DictReader(f) + for row in reader: + try: + idx = int(row["question_index"]) + except (KeyError, ValueError): + continue + ctx = row.get("context", "").strip() + if ctx: + gt_map.setdefault(idx, []).append(ctx) + + result: List[RecallQuestion] = [] + for i, question in enumerate(questions_raw, start=1): + chunks = gt_map.get(i, []) + if not chunks: + print(f" {_Y}WARNING: no GT chunks for Q{i} — " + f"recall will be 0 for this question{_X}", flush=True) + result.append(RecallQuestion( + index=i, question=question, + gt_contexts=chunks, gt_titles=[], + )) + + return result + + +# ─── Embedding cosine similarity matching ──────────────────────────────────── + +def _cosine(a: List[float], b: List[float]) -> float: + dot = sum(x * y for x, y in zip(a, b)) + norm_a = math.sqrt(sum(x * x for x in a)) + norm_b = math.sqrt(sum(x * x for x in b)) + return dot / (norm_a * norm_b) if norm_a and norm_b else 0.0 + + +def _get_embedding_service(): + """Return the GraphRAG embedding service (already configured in the container).""" + from common.config import get_embedding_service + return get_embedding_service() + + +def _embed(text: str, **_kwargs) -> Optional[List[float]]: + """Return embedding for text using GraphRAG's configured embedding service (cached). + Extra kwargs (api_key, model) are accepted but ignored — config comes from + the container's server_config.json via common.config.get_embedding_service(). + Returns None on any error. + """ + cache_key = text[:200] + with _embed_lock: + if cache_key in _embed_cache: + return _embed_cache[cache_key] + try: + with _embed_sem: + svc = _get_embedding_service() + vecs = svc.embed_documents([text[:_EMBED_MAX_CHARS]]) + vec = vecs[0] if vecs else None + if vec: + with _embed_lock: + _embed_cache[cache_key] = vec + return vec + except Exception as e: + logging.warning("Embedding error: %s", e) + return None + + +def _check_embedding_api(**_kwargs) -> Optional[str]: + """Return None if the embedding service works, else an error string.""" + vec = _embed("ping") + if vec is None: + return "embedding service unavailable — check container logs" + return None + + +def _chunk_found_embedding(gt_chunk: str, chunks: List[str], + threshold: float = _EMBED_THRESHOLD, + **_kwargs) -> Tuple[bool, float]: + """Check whether any retrieved chunk is similar enough to the GT chunk. + + Returns: + (found, max_cosine_similarity) + found — True if max_cosine_similarity >= threshold + max_cosine_similarity — highest cosine score observed (0.0 if no embeddings) + """ + gt_vec = _embed(gt_chunk) + if gt_vec is None: + return False, 0.0 + max_sim = 0.0 + for retrieved in chunks: + retrieved_vec = _embed(retrieved) + if retrieved_vec is not None: + sim = _cosine(gt_vec, retrieved_vec) + if sim > max_sim: + max_sim = sim + return max_sim >= threshold, max_sim + + +# ─── LLM judge matching ──────────────────────────────────────────────────────── + +_LLM_PROMPT = """\ +You are evaluating information retrieval quality for a RAG system. + +Ground-truth passage (the relevant source text): +{gt} + +Retrieved chunk (returned by the retriever): +{chunk} + +Task: Does the retrieved chunk contain information that is RELEVANT to the ground-truth passage? +The retrieved chunk does NOT need to be identical or cover everything — answer YES if it overlaps +on the same topic, entity, or key facts. Answer NO only if the retrieved chunk is clearly about +a completely different topic with no meaningful overlap. + +Answer with a single word: YES or NO.""" + + +def _llm_judge(gt_chunk: str, retrieved_chunk: str, **_kwargs) -> bool: + """Ask GraphRAG's configured LLM whether the retrieved chunk contains the GT chunk's information. + Uses the same get_llm_service / get_chat_config pattern as evaluator.py. + Returns True for YES, False for NO or on any error. + """ + cache_key = f"{gt_chunk[:100]}|||{retrieved_chunk[:100]}" + with _llm_lock: + if cache_key in _llm_cache: + return _llm_cache[cache_key] + try: + from common.config import get_chat_config, get_llm_service + from langchain_core.messages import HumanMessage + + cfg = get_chat_config() + llm = get_llm_service(cfg) + lc = getattr(llm, "llm", None) or llm + prompt = _LLM_PROMPT.format( + gt=gt_chunk[:_LLM_CHUNK_CHARS], + chunk=retrieved_chunk[:_LLM_CHUNK_CHARS], + ) + resp = lc.invoke([HumanMessage(content=prompt)]) + text = (resp.content if hasattr(resp, "content") else str(resp)).strip().upper() + verdict = text.startswith("YES") + with _llm_lock: + _llm_cache[cache_key] = verdict + return verdict + except Exception as e: + logging.warning("LLM judge error: %s", e) + return False + + +def _chunk_found_llm(gt_chunk: str, retrieved_chunks: List[str]) -> bool: + """True if the LLM judges any retrieved chunk as relevant to the GT chunk.""" + for retrieved in retrieved_chunks: + if _llm_judge(gt_chunk, retrieved): + return True + return False + + +# ─── Unified recall computation ─────────────────────────────────────────────── + +def _compute_recall( + gt_contexts: List[str], + top_k_chunks: List[str], + match_cfg: Optional[Dict[str, Any]] = None, + embed_threshold: float = _EMBED_THRESHOLD, + gt_titles: Optional[List[Optional[str]]] = None, # unused, kept for compat +) -> Tuple[float, int, List[int], List[str], List[float]]: + """Compute Recall@K for one question. + + match_cfg keys: + strategy — "embedding" (default) or "llm" + + Returns: + recall, matched_count, matched_indices, unmatched_chunks, embed_scores + embed_scores — max cosine similarity per GT chunk (empty list when strategy="llm") + """ + if not gt_contexts: + return 0.0, 0, [], [], [] + + strategy = (match_cfg or {}).get("strategy", "embedding") + + matched_idxs: List[int] = [] + unmatched: List[str] = [] + embed_scores: List[float] = [] + + for i, gt in enumerate(gt_contexts): + if strategy == "llm": + found = _chunk_found_llm(gt, top_k_chunks) + embed_scores.append(-1.0) # not applicable + else: + found, score = _chunk_found_embedding(gt, top_k_chunks, threshold=embed_threshold) + embed_scores.append(score) + + if found: + matched_idxs.append(i) + else: + unmatched.append(gt) + + recall = len(matched_idxs) / len(gt_contexts) + return recall, len(matched_idxs), matched_idxs, unmatched, embed_scores + + +# ─── GraphRAG query ─────────────────────────────────────────────────────────── + +_RETRIEVER_LABELS = { + "similaritysearch": "Similarity Search", + "contextualsearch": "Contextual Search", + "hybridsearch": "Hybrid Search", + "communitysearch": "Community Search", +} + + +def _query_graphrag( + url: str, + graphname: str, + username: str, + password: str, + question: str, + mode: str, + rag_pattern: str, +) -> dict: + resp = httpx.get( + f"{url}/ui/{graphname}/query", + params={ + "q": question, + "mode": mode, + "rag_pattern": rag_pattern, + "include_fields": "query_sources", + }, + auth=(username, password), + timeout=600.0, # agentic mode can take >300s on complex multi-hop questions + ) + resp.raise_for_status() + data = resp.json() + if isinstance(data, str): + data = json.loads(data) + return data + + +def _extract_chunk_texts(final_retrieval: dict) -> List[str]: + """Extract text strings from a final_retrieval dict. + + Each value is either: + - list[str] → one text per item (common case after ingestion) + - str → single text value + - dict → extract 'content' or 'text' sub-key + """ + texts: List[str] = [] + for v in final_retrieval.values(): + if isinstance(v, list): + for item in v: + if isinstance(item, str) and item.strip(): + texts.append(item) + elif isinstance(item, dict): + t = item.get("content") or item.get("text") or "" + if t: + texts.append(str(t)) + elif isinstance(v, str) and v.strip(): + texts.append(v) + elif isinstance(v, dict): + t = v.get("content") or v.get("text") or "" + if t: + texts.append(str(t)) + return texts + + +def _extract_top_k_chunks(query_sources: dict, k: int, mode: str = "classic") -> List[str]: + """Extract retrieved text chunks from the GraphRAG response. + + Classic mode (similarity/hybrid/contextual/community search): + Chunks are ordered by relevance score — most relevant first. + We return only the first k so Recall@K is honest and cannot + benefit from chunks beyond rank k. + + Agentic mode (planned/reactive/auto agent): + The agent calls multiple tools in arbitrary sequence — there is no + single relevance ranking across tool calls. Applying a k-cutoff + here would silently drop relevant chunks the agent DID retrieve but + that happen to appear after position k in dict-iteration order. + We therefore return ALL chunks the agent retrieved (up to _MAX_CHUNKS) + so the recall score reflects what the agent actually found, not an + arbitrary ordering artefact. + + Chunk extraction strategy: + 1. Top-level result.final_retrieval — populated in classic mode. + 2. Per-step result.unstructured[*].result.final_retrieval — populated + in agentic mode for each vector/hybrid search tool call. + 3. NO general fallback walker: agentic structural-only answers produce + no text chunks; falling back to walking the whole response would + extract Cypher query strings and reasoning text as fake chunks, leading + to artificially high LLM recall (entity name matching) and artificially + low embedding recall (near-zero cosine similarity with real passages). + + Title-header stripping: + The ingestion pipeline prepends "===== Title =====" to each stored chunk. + GT chunks contain only the passage body, so we strip headers before + comparison to avoid spurious similarity penalties. + """ + if not query_sources: + return [] + + result = query_sources.get("result") or {} + raw: List[str] = [] + + # ── Path 1: top-level final_retrieval (classic mode, and some agentic configs) ── + top_fr = result.get("final_retrieval") or {} + if isinstance(top_fr, dict) and top_fr: + raw.extend(_extract_chunk_texts(top_fr)) + + # ── Path 2: per-step unstructured results (agentic mode) ──────────────────────── + if not raw: + for u_item in result.get("unstructured") or []: + u_result = u_item.get("result") if isinstance(u_item, dict) else None + if not isinstance(u_result, dict): + continue + inner_fr = u_result.get("final_retrieval") or {} + if isinstance(inner_fr, dict) and inner_fr: + raw.extend(_extract_chunk_texts(inner_fr)) + + # ── Deduplicate, strip title headers, apply length guard, cap at _MAX_CHUNKS ──── + seen: set = set() + out: List[str] = [] + for text in raw: + # Strip "===== Title =====" header added by ingestion pipeline + text = _TITLE_HEADER_RE.sub("", text, count=1).strip() + if len(text) < 5 or text in seen: + continue + seen.add(text) + out.append(text[:_MAX_CHUNK_CH]) + if len(out) >= _MAX_CHUNKS: + break + + # For classic mode honour the K limit (ranked by relevance). + # For agentic mode return everything the agent retrieved — no meaningful rank exists. + if mode == "agentic": + return out + return out[:k] + + +# ─── Per-question evaluation ────────────────────────────────────────────────── + +def _eval_one_recall( + rq: RecallQuestion, + graphname: str, + url: str, + username: str, + password: str, + mode: str, + rag_pattern: str, + k: int, + match_cfg: Optional[Dict[str, Any]] = None, + embed_threshold: float = _EMBED_THRESHOLD, +) -> RecallResult: + r = RecallResult( + question_index=rq.index, + question=rq.question, + gt_context_count=len(rq.gt_contexts), + ) + t0 = time.monotonic() + + try: + data = _query_graphrag(url, graphname, username, password, rq.question, mode, rag_pattern) + r.answered_question = bool(data.get("answered_question", False)) + query_sources = data.get("query_sources") or {} + _chosen = (query_sources.get("chosen_retriever") or "").lower().replace(" ", "") + r.search_type_used = _RETRIEVER_LABELS.get(_chosen) or data.get("response_type") or None + r.agent_mode = f"{mode}/{rag_pattern}" + + top_k = _extract_top_k_chunks(query_sources, k, mode=mode) + r.retrieved_chunk_count = len(top_k) + except httpx.HTTPStatusError as e: + r.error = f"HTTP {e.response.status_code}: {e.response.text[:200]}" + r.response_time_seconds = time.monotonic() - t0 + return r + except Exception as e: + r.error = f"{type(e).__name__}: {e}" + r.response_time_seconds = time.monotonic() - t0 + return r + finally: + r.response_time_seconds = time.monotonic() - t0 + + if not rq.gt_contexts: + r.recall_at_k = 0.0 + return r + + if not top_k: + # No text chunks retrieved — likely a structural-only (graph traversal) answer. + # Retrieval recall cannot be measured without text chunks. + r.recall_at_k = 0.0 + r.matched_count = 0 + r.unmatched_contexts = list(rq.gt_contexts) + return r + + recall, matched, matched_idxs, unmatched, embed_scores = _compute_recall( + rq.gt_contexts, top_k, + match_cfg=match_cfg, + embed_threshold=embed_threshold, + ) + r.recall_at_k = round(recall, 4) + r.matched_count = matched + r.matched_indices = matched_idxs + r.unmatched_contexts = unmatched + r.embed_scores = embed_scores + + return r + + +# ─── Compact / detailed printer ────────────────────────────────────────────── + +def _recall_cell(r: RecallResult) -> str: + if r.recall_at_k is None: + return f"{_Y} n/a{_X}" + pct = r.recall_at_k * 100 + col = _G if pct >= 60 else _Y if pct >= 30 else _R + bar_filled = int(pct / 10) + bar = "█" * bar_filled + "░" * (10 - bar_filled) + return f"{col}{pct:5.1f}% {bar}{_X}" + + +def _print_compact(r: RecallResult, total: int, k: int) -> None: + n = r.question_index + if r.error: + print(f" Q{n:>3}/{total} [{r.response_time_seconds:5.1f}s] " + f"{_R}error: {r.error[:100]}{_X}", flush=True) + return + if r.gt_context_count == 0: + print(f" Q{n:>3}/{total} [{r.response_time_seconds:5.1f}s] " + f"{_Y}no GT contexts — skipped{_X}", flush=True) + return + if r.retrieved_chunk_count == 0: + print(f" Q{n:>3}/{total} [{r.response_time_seconds:5.1f}s] " + f"{_Y}no text chunks (structural-only answer) — recall=0{_X}", flush=True) + return + scope = f"top-{k}" if r.retrieved_chunk_count <= k else f"all-{r.retrieved_chunk_count}" + frac = f"({r.matched_count}/{r.gt_context_count} GT chunks in {scope})" + print( + f" Q{n:>3}/{total} [{r.response_time_seconds:5.1f}s] " + f"Recall@{k}={_recall_cell(r)} {frac}", + flush=True, + ) + + +def _print_detailed(r: RecallResult, total: int, k: int) -> None: + n = r.question_index + q_short = (r.question[:70] + "…") if len(r.question) > 71 else r.question + print(f" {'─' * 70}", flush=True) + print(f" {_B}Q{n}/{total}{_X} [{r.response_time_seconds:.1f}s] {q_short}", flush=True) + if r.error: + print(f" {_R}error: {r.error}{_X}", flush=True) + return + if r.gt_context_count == 0: + print(f" {_Y}no ground-truth contexts — skipped{_X}", flush=True) + return + print( + f" {_C}Recall@{k}{_X} = {_recall_cell(r)} " + f"matched {r.matched_count}/{r.gt_context_count} " + f"(retrieved {r.retrieved_chunk_count} chunks)", + flush=True, + ) + if r.unmatched_contexts: + strategy = "embedding" # default; score list is non-negative only for embedding + print(f" {_Y}Unmatched GT chunks:{_X}", flush=True) + # Build a map: gt_text → embed score (for embedding runs only) + score_map: Dict[str, float] = {} + if r.embed_scores and any(s >= 0 for s in r.embed_scores): + # embed_scores is aligned with gt_contexts; reconstruct unmatched scores + for idx, (score, _matched) in enumerate( + zip(r.embed_scores, [i in r.matched_indices for i in range(len(r.embed_scores))]) + ): + pass + # Simpler: zip embed_scores with all GT chunks, show score for unmatched + all_gt = r.unmatched_contexts # only unmatched are in this list + unmatched_scores = [s for i, s in enumerate(r.embed_scores) if i not in r.matched_indices] + for uc, sc in zip(all_gt, unmatched_scores if unmatched_scores else [None]*len(all_gt)): + score_str = f" max_cos={sc:.3f}" if sc is not None and sc >= 0 else "" + print(f" - {uc[:100]}{'…' if len(uc) > 100 else ''}{score_str}", flush=True) + else: + for uc in r.unmatched_contexts: + print(f" - {uc[:120]}{'…' if len(uc) > 120 else ''}", flush=True) + + +# ─── Main eval loop ─────────────────────────────────────────────────────────── + +def _parse_mode(mode_str: str) -> Tuple[str, str]: + """Parse --mode into (api_mode, rag_pattern). + + Agentic styles (mode=agentic, rag_pattern=