From ad5ad258bba903ef59cde6294b5f91ec7f4cd005 Mon Sep 17 00:00:00 2001 From: prinskumar-tigergraph Date: Fri, 25 Sep 2026 01:32:18 +0530 Subject: [PATCH 01/15] GML-2191 Planner and react: multi-clause routing, conversation context resolution, clarifying question fallback (#64) * Planner and react: multi-clause routing, conversation context resolution, clarifying question fallback - When a question refers to something from earlier in the conversation, resolve it from ## Conversation and use the full name in retrieval args. If history has multiple candidates and it is unclear which is meant, ask the user one short clarifying question instead of guessing. - For each unstructured retrieval step, set the question to that clause only as a standalone search query. Do not pass the full multi-part question, a clause already covered by another step, or an unresolved reference from conversation history. - When a question has multiple clauses, assign each clause to its own retrieval step; plan additional RAG steps for any clause not yet covered. - Synthesizer: answer each part of a multi-part question from its matching source; do not mix structured notes with document-passage answers. Co-authored-by: Cursor * Planner/react: resolve underspecified questions from conversation history before any retrieval * Synthesizer: generic multi-part answer rule, remove dataset-specific wording --------- Co-authored-by: Prins Kumar Co-authored-by: Cursor --- common/llm_services/base_llm.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index fe77ac1..3dfcb9d 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -1038,6 +1038,8 @@ def select_retriever_prompt(self): # Operator-customizable retrieval strategy for the react agent: the first # action, then each next action driven by what the previous result returned. _AGENTIC_AGENT_USER_DEFAULT = """\ +- Before calling any retrieval tool, check whether the question is self-contained: can it be fully understood without reading ## Conversation? If the subject, entity, or topic is not named explicitly in the question, look up the most recent relevant entity from ## Conversation and substitute its full name into every retrieval call. If ## Conversation contains multiple candidates and it is genuinely unclear which one the user means, ask one short clarifying question instead of guessing — do not call any retrieval tool until clarified. +- When calling an unstructured retriever (hybrid, contextual, similarity, community), pass only the sub-question for that specific part as a standalone search query in the user's language. Do not pass the full multi-part question, a part already covered by another step, or an unresolved reference from conversation history. - For most questions, make your FIRST action a vector search (graphrag__hybrid_search or graphrag__contextual_search) — it gives the broadest grounding. Skip it only when you are highly confident the question is a pure structured-data request (an exact count, an attribute/id lookup, a relationship traversal, or an aggregation over typed graph data) that a generated graph query fully answers on its own. - Let each observation drive the next action: if the passages you got back name specific entities or relationships you still need hard facts about, follow up with a structural query; if a result is thin, empty, or off-target, widen its parameters (top_k, num_hops) or switch method rather than repeating the same call. - Before answering, check that every part of the question is covered with the specific facts and figures it asks for; if a required value, table, or entity is still missing, retrieve again (widen top_k / num_hops or switch method) rather than answering vaguely or partially. @@ -1074,6 +1076,7 @@ def agentic_agent_prompt(self): Plan mechanics (fixed): - A later step may depend on an earlier one: set depends_on and use arg_bindings to pull a value from a prior step's result, e.g. {"question": "S1.context.result"}. - Retrieval params (top_k, num_hops, community_level) are optional; omit them to use defaults, or set higher values when you expect a broad answer. +- For each unstructured step, set args.question to that clause only as a standalone search query in the user's language. Do not pass the full multi-part question, a clause already assigned to another step, or an unresolved reference from conversation history. - The final step MUST have kind="answer" and tool="" (the orchestrator synthesizes the answer from gathered context); it should depend_on all retrieval steps. Decide which retrievals to include, how many, and in what order using the "Retrieval Strategy" below. Return ONLY the structured plan. @@ -1088,8 +1091,10 @@ def agentic_agent_prompt(self): # Strategy (operator-customizable) — moved out of the fixed rules so it can # be tuned without touching the role / act model / plan mechanics. _AGENTIC_PLANNER_USER_DEFAULT = """\ +- Before building the plan, check whether the question is self-contained: can it be fully understood without reading ## Conversation? If the subject, entity, or topic is not named explicitly in the question, find the most recent relevant entity from ## Conversation and substitute its full name in every step's args. If ## Conversation has multiple candidates and it is genuinely unclear which one the user means, plan only a final answer step (no retrieval) that asks the user one short clarifying question. +- When a question has multiple clauses, assign each clause to its own retrieval step. If another clause still needs passages or typed graph facts after one is covered, plan hybrid/community/structural for that clause too. - Prioritize including at least one vector search step (graphrag__hybrid_search or graphrag__contextual_search) unless you are highly confident the question is a pure structured-data request — an exact count, an attribute/id lookup, a relationship traversal, or an aggregation over typed graph data — that a generated graph query fully answers on its own. Whenever the answer could plausibly live in document text (what/why/how/describe/summarize, definitions, explanations, figures), include a vector search step. When unsure, include vector search. -- Use BOTH kinds when a question needs facts from the graph AND supporting text; you may run several of each, in any order. When you use STRUCTURAL, pair it with a vector search step unless the question is a pure structured-data request. +- Use BOTH structural and unstructured kinds when a question needs facts from the graph AND supporting text; you may run several of each, in any order. When you use STRUCTURAL, pair it with a vector search step unless the question is a pure structured-data request. - Prefer the smallest plan that will work. Trivial/greeting questions need only the final answer step. - Tabular / numeric questions (a specific value, a row, a column total, a ranking, or a year-over-year comparison from a table or chart): prefer graphrag__contextual_search or graphrag__hybrid_search with top_k>=10 (these return atomic table chunks that preserve full row/column structure); avoid graphrag__similarity_search alone; quote any specific table label, column header, year, or unit from the question (e.g. "ROE 2023"); for "compare X across years/regions/categories" set top_k>=15.""" @@ -1200,6 +1205,7 @@ def hyde_prompt(self): - **Quote exact values from the source.** Numbers, units, time periods, and named entities must appear verbatim — do not round, approximate, or translate units. Keep units in their original format, script, and language. For example, if the source says `1,234 km`, write `1,234 km`, not `767 miles` or `about 1,200 km`. - **For comparison or "which is the highest" questions, list each candidate's value before stating the conclusion.** Show the working — do not jump directly to a one-line answer. - **Score** each context for relevance and use only the high-scoring ones; do not invent additional logic. +- **Multi-part questions:** answer each part from its matching retrieval context. Do not mix structured-query results with document-passage results when answering different parts. If retrieved context is off-topic for a part, say that part is not covered by the retrieved information. - **Cover** the relevant information, especially image references that carry critical visual information. - **Format** the answer in Markdown — titles, paragraphs, bulleted / numbered lists, images, and tables. Place images and tables below the related text section. - **Tables**: every row, including the header, starts on a new line. From a5f33d7f9fcf38e3ee978734416ecb7be7751274 Mon Sep 17 00:00:00 2001 From: prinskumar-tigergraph Date: Fri, 25 Sep 2026 01:32:49 +0530 Subject: [PATCH 02/15] GML-2192 fix(eval): extend LLM judge to support gemini-3.5-flash output format (#66) Regression eval: fix LLM judge errors with gemini-3.5-flash - evaluator.py: pre-load common.config, hallucination_check, openai.types.chat at module level to prevent race conditions when ThreadPoolExecutor workers all try to import the same partially-initialized module simultaneously - hallucination_check.py: replace PydanticOutputParser format instructions (full JSON schema) with a simple inline example so models do not echo the schema definition back instead of filling in values - recall_evaluator.py: handle list content blocks in LLM judge response (gemini-3.5-flash returns resp.content as a list, not a string) Co-authored-by: Prins Kumar --- graphrag/tests/regression/evaluator.py | 38 ++++++++++++++++++- .../tests/regression/hallucination_check.py | 8 ++-- graphrag/tests/regression/recall_evaluator.py | 8 +++- 3 files changed, 46 insertions(+), 8 deletions(-) diff --git a/graphrag/tests/regression/evaluator.py b/graphrag/tests/regression/evaluator.py index 528a40a..e80e19e 100644 --- a/graphrag/tests/regression/evaluator.py +++ b/graphrag/tests/regression/evaluator.py @@ -50,6 +50,19 @@ if "/code" not in sys.path: sys.path.insert(0, "/code") +# Pre-load modules that are used inside per-question worker threads. +# Importing them here (single-threaded, at startup) prevents race conditions +# in Python's import lock when multiple ThreadPoolExecutor workers all try to +# initialize the same partially-loaded module simultaneously (manifests as +# ImportError: cannot import name '...' from partially initialized module). +try: + import common.config # noqa: F401 + from common.config import get_chat_config, get_llm_service # noqa: F401 + from hallucination_check import TigerGraphAgentHallucinationCheck # noqa: F401 + import openai.types.chat # noqa: F401 +except Exception: + pass # container may not have all deps at parse time; threads will surface real errors + logging.basicConfig(level=logging.WARNING, format="%(levelname)s %(name)s: %(message)s") # ANSI colours for clean CLI output. NO_COLOR disables; FORCE_COLOR forces on @@ -218,6 +231,27 @@ def _build_deepeval_llm(): model_name = cfg.get("llm_model", "unknown") lc_model = getattr(llm_provider, "llm", None) or llm_provider + # Gemini 3.x / thinking models return AIMessage.content as a list of + # typed blocks, not a string. DeepEval then does re.search() on it → + # AttributeError: 'list' object has no attribute 'find'. + _to_text = getattr(llm_provider, "_message_text", None) + + def _as_text(resp) -> str: + if callable(_to_text): + return _to_text(resp) or "" + content = resp.content if hasattr(resp, "content") else resp + if isinstance(content, str): + return content + if isinstance(content, list): + parts = [] + for b in content: + if isinstance(b, str): + parts.append(b) + elif isinstance(b, dict) and isinstance(b.get("text"), str): + parts.append(b["text"]) + return "".join(parts) + return str(content) + def _parse(text, schema): try: m = _re.search(r"\{.*\}", text, _re.DOTALL) @@ -236,12 +270,12 @@ def load_model(self): def generate(self, prompt: str, schema=None): resp = lc_model.invoke([HumanMessage(content=prompt)]) - text = resp.content if hasattr(resp, "content") else str(resp) + text = _as_text(resp) return _parse(text, schema) if schema else text async def a_generate(self, prompt: str, schema=None): resp = await lc_model.ainvoke([HumanMessage(content=prompt)]) - text = resp.content if hasattr(resp, "content") else str(resp) + text = _as_text(resp) return _parse(text, schema) if schema else text def get_model_name(self) -> str: diff --git a/graphrag/tests/regression/hallucination_check.py b/graphrag/tests/regression/hallucination_check.py index 8bb232b..3319969 100644 --- a/graphrag/tests/regression/hallucination_check.py +++ b/graphrag/tests/regression/hallucination_check.py @@ -111,12 +111,10 @@ def check_hallucination( to the wrong subject; (2) a claim directly contradicts the context; OR \ (3) the answer introduces topics entirely unrelated to the context. -Provide your verdict as JSON. -Format: {format_instructions}""", +Respond ONLY with a valid JSON object — no prose, no markdown fences: +{{"confidence": , "reason": ""}}""", input_variables=["generation", "context", "question_block"], - partial_variables={ - "format_instructions": hallucination_parser.get_format_instructions() - }, + partial_variables={}, ) prediction = self.llm.invoke_with_parser( diff --git a/graphrag/tests/regression/recall_evaluator.py b/graphrag/tests/regression/recall_evaluator.py index 45d30b8..be7c1f0 100644 --- a/graphrag/tests/regression/recall_evaluator.py +++ b/graphrag/tests/regression/recall_evaluator.py @@ -332,7 +332,13 @@ def _llm_judge(gt_chunk: str, retrieved_chunk: str, **_kwargs) -> bool: chunk=retrieved_chunk[:_LLM_CHUNK_CHARS], ) resp = lc.invoke([HumanMessage(content=prompt)]) - text = (resp.content if hasattr(resp, "content") else str(resp)).strip().upper() + content = resp.content if hasattr(resp, "content") else str(resp) + if isinstance(content, list): + content = "".join( + b["text"] if isinstance(b, dict) and "text" in b else str(b) + for b in content + ) + text = content.strip().upper() verdict = text.startswith("YES") with _llm_lock: _llm_cache[cache_key] = verdict From 738c0ba0cea425b8af68d9f3dff5409a8170135c Mon Sep 17 00:00:00 2001 From: Chengbiao Jin <45186596+chengbiao-jin@users.noreply.github.com> Date: Wed, 7 Oct 2026 12:08:49 -0700 Subject: [PATCH 03/15] GML-2182, GML-2194..2198, GML-2301..2303, GML-2311..2313: fixes for 2.0.3 (#67) * Take query-with-history conversation context from the request - Conversation history is no longer retained on the server between calls, so concurrent callers no longer receive each other's previous questions and answers as context - Callers supply prior turns on the request; the most recent are used - Requests that send no turns are answered without prior context - Update the demo notebook to carry its own conversation Refs: GML-2194 * Report unreadable files as failed instead of ingesting the error - A PDF or image that cannot be extracted no longer becomes a document whose text is the failure notice, so no placeholder content reaches the knowledge graph or turns up in answers - Such a file is reported as failed with the reason; other files in the same upload still process normally - Failure messages name the file and the remedy, keeping parser detail in the server log Refs: GML-2195 * Extract tables from Word documents instead of dropping them - Tables in a .docx are read and rendered in place, so the figures they hold reach the knowledge graph and can be cited; previously only the paragraphs around them were ingested and the table vanished silently - Tables nested inside a cell are included, and a pipe or line break in a cell no longer breaks the rendered row - Legacy .doc is reported as unreadable with a pointer to .docx, instead of being listed as supported and ingested as a note about itself - An empty spreadsheet or an unreadable file type is reported as a failure rather than ingested as placeholder text Refs: GML-2196 * Require valid credentials for document retrieval - The document-retrieval endpoint returned matching content without checking the caller; it now rejects anyone who cannot access the graph named in the request - Search runs against the graph named in the request instead of the configured default graph - A request with no credentials is rejected rather than failing as a server error Refs: GML-2197 * Stop swapping the two contextual-search options - Contextual (sibling-chunk) search received the hypothetical-answer and question-expansion options in each other's place on every call, so asking for one applied the other; they select different search strategies rather than variations of one - Requests that set neither option, which is all chat traffic, were unaffected Refs: GML-2198 * Fix knowledge-graph rebuild for token sign-ins - Rebuilds started from a token sign-in no longer fail as if the graph did not exist - An invalid token is rejected as an authentication failure instead of a server error Refs: GML-2182 * Save MCP server deletes and stop Cancel from deleting a server - Deleting an MCP server asks for confirmation and is saved immediately, instead of being lost on reload - Cancelling an edit restores the server's saved settings instead of removing it from the list Refs: GML-2301 * Request agentic plans in a form OpenAI accepts - Agentic planning no longer makes a request OpenAI rejects before every plan, removing the extra latency and token cost from each answer - Other model providers and other structured requests keep their current behavior Refs: GML-2302 * Delete chat conversations one at a time, in a group, or all at once - Each conversation in the chat sidebar has a delete button, with a confirmation - Select mode lets several conversations, or all shown ones, be deleted together - A Clear all action removes the whole history, shown whenever there is any; Clear older is unchanged - Deleting the open conversation starts a new chat - A delete that fails is reported instead of silently leaving the entry Refs: GML-2303 * Fit retrieved context to the chat model's input limit - Default to the chat model's own input limit when no token limit is set - Reserve room for the full answer prompt instead of a fixed amount - Decide on trimming by token count, so CJK context is trimmed correctly - Send each retrieved passage to the answer model once across a plan's steps - Log the cause when answer generation fails - Report and log how many chunks each search retrieved, and their size - Show oversized step output in the trace as shortened tool data with a note Refs: GML-2311 * Repair only graphs whose schema the current queries can run on - The Migration Assistant check reports whether a graph's schema matches the current base schema, as graphs created with 1.4 or later do - Repair is offered and accepted only for such graphs; older graphs are told to be recreated and their documents ingested again - Repair is withheld when the schema cannot be read Refs: GML-2312 * List each contextual-search passage once, at its correct position - Neighbor offsets are measured from the match each chunk belongs to - A chunk near several matches is listed once, under the nearest one - Every match appears under itself at offset zero Refs: GML-2313 --- CHANGELOG.md | 28 ++ README.md | 20 +- common/db/connections.py | 38 ++- common/db/migrate.py | 121 ++++++++ .../Chunk_Sibling_Vector_Search.gsql | 72 +++-- common/llm_services/azure_openai_service.py | 2 + common/llm_services/base_llm.py | 48 +++- common/llm_services/openai_service.py | 2 + common/py_schemas/schemas.py | 9 + common/utils/retrieval_stats.py | 87 ++++++ common/utils/text_extractors.py | 231 ++++++++++++--- common/utils/token_calculator.py | 49 ++++ docs/notebooks/TransactionFraud_demo.ipynb | 14 +- ecc/app/main.py | 13 +- graphrag-ui/src/components/SideMenu.tsx | 269 ++++++++++++++++-- graphrag-ui/src/pages/setup/IngestGraph.tsx | 33 ++- graphrag-ui/src/pages/setup/KGAdmin.tsx | 28 +- .../src/pages/setup/McpServersConfig.tsx | 87 +++++- graphrag/app/agent/agent_generation.py | 29 +- graphrag/app/agent/agentic_executor.py | 91 +++++- graphrag/app/agent/agentic_graph.py | 8 +- graphrag/app/agent/agentic_planner.py | 13 +- graphrag/app/agent/agentic_synthesizer.py | 51 +++- graphrag/app/routers/inquiryai.py | 69 +++-- graphrag/app/routers/supportai.py | 4 + graphrag/app/routers/ui.py | 79 ++++- .../app/supportai/retrievers/BaseRetriever.py | 33 ++- .../retrievers/CommunityRetriever.py | 1 + .../supportai/retrievers/HybridRetriever.py | 1 + .../supportai/retrievers/SiblingRetriever.py | 7 +- .../retrievers/SimilarityRetriever.py | 1 + graphrag/app/supportai/supportai.py | 29 +- graphrag/app/tools/graphrag_tools.py | 17 +- .../test_agentic_plan_structured_output.py | 156 ++++++++++ graphrag/tests/test_answer_context_dedup.py | 72 +++++ .../tests/test_async_id_token_connection.py | 231 +++++++++++++++ graphrag/tests/test_context_token_limit.py | 179 ++++++++++++ graphrag/tests/test_docx_table_extraction.py | 176 ++++++++++++ .../test_extraction_failure_not_a_document.py | 148 ++++++++++ .../test_extraction_failures_reported.py | 210 ++++++++++++++ .../tests/test_migration_schema_compat.py | 216 ++++++++++++++ .../test_query_with_history_isolation.py | 103 +++++++ graphrag/tests/test_retrieval_step_summary.py | 117 ++++++++ graphrag/tests/test_retrieve_docs_auth.py | 207 ++++++++++++++ graphrag/tests/test_retriever_flag_order.py | 134 +++++++++ graphrag/tests/test_trace_output_note.py | 100 +++++++ graphrag/tests/test_v1_4_1_uploads_schema.py | 60 +++- 47 files changed, 3502 insertions(+), 191 deletions(-) create mode 100644 common/utils/retrieval_stats.py create mode 100644 graphrag/tests/test_agentic_plan_structured_output.py create mode 100644 graphrag/tests/test_answer_context_dedup.py create mode 100644 graphrag/tests/test_async_id_token_connection.py create mode 100644 graphrag/tests/test_context_token_limit.py create mode 100644 graphrag/tests/test_docx_table_extraction.py create mode 100644 graphrag/tests/test_extraction_failure_not_a_document.py create mode 100644 graphrag/tests/test_extraction_failures_reported.py create mode 100644 graphrag/tests/test_migration_schema_compat.py create mode 100644 graphrag/tests/test_query_with_history_isolation.py create mode 100644 graphrag/tests/test_retrieval_step_summary.py create mode 100644 graphrag/tests/test_retrieve_docs_auth.py create mode 100644 graphrag/tests/test_retriever_flag_order.py create mode 100644 graphrag/tests/test_trace_output_note.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e323bec..4402981 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,33 @@ # Changelog +## [2.0.3] + +### Added +- **Chat conversations can be deleted one at a time, in a selected group, or all at once.** Each conversation in the chat sidebar has a delete button, *Select* lets several conversations be picked (or all shown ones at once) and deleted together, and a new *Clear all* action removes the whole history. The only option before, *Clear older*, always kept the ten most recent conversations and appeared only once there were more than ten. + +### Fixed +- **Contextual search returns each passage once, with its correct position.** When two matched chunks were close together in a document, the passages around them were returned under both matches with wrong, inconsistent positions (a match could appear several places away from itself), so the answer saw duplicated context. Each chunk is now listed once, under the nearest match, with its exact offset. Existing graphs pick this up by running Repair in the Migration Assistant. +- **Knowledge-graph rebuild works when signed in with a token.** A rebuild started from a token sign-in could fail at once, reported as the graph not existing or not being accessible. It now runs normally, and an invalid token is rejected as an authentication failure rather than a server error. +- **Contextual document search honours the search options it was given.** The option to use a hypothetical-answer embedding and the option to expand the question were swapped before reaching contextual (sibling-chunk) search, so asking for one applied the other — and each selects a different search strategy, not a variation of one. Requests that set neither option, including all chat traffic, were unaffected. +- **Document search now requires valid credentials.** The document-retrieval endpoint returned matching content without checking who was asking, and searched the configured default graph rather than the one named in the request. It now rejects callers who cannot access the requested graph, and searches that graph. +- **Tables in Word documents are no longer dropped on upload.** Only the paragraphs of a `.docx` were read, so every table vanished while the text around it was ingested normally — a document could state that limits are defined below and contain none of them. Tables are now extracted in place, so the figures they hold can be found and cited. +- **A document that fails to extract is no longer ingested as its own error message.** When a PDF or image could not be read, the extractor produced a document whose text was the failure notice — which was then chunked, embedded and returned by search like ordinary content. Nothing is ingested for such a file now, so the knowledge graph contains no placeholder text and answers can't be drawn from one. +- **Files that fail to extract are no longer counted as ingested.** An unreadable PDF or image was previously counted as a successfully ingested document, so the upload looked clean. It is now left out, the rest of the upload still processes, and the file and its reason are returned with the upload result. When none of the files can be read, the upload stops and gives the reasons. +- **Legacy `.doc` files are no longer accepted as if they could be read.** The format was listed as supported but could never be read, so such a file was ingested as a one-line note saying it was unsupported. The file picker now marks `.doc` as unsupported, and one sent through the API is reported as unreadable, with a pointer to `.docx`. +- **An empty spreadsheet is reported rather than ingested as a note about itself.** +- **Conversation context on the query-with-history endpoint no longer crosses callers.** `/{graphname}/query_with_history` kept a single conversation history for the whole service, so concurrent callers — on the same graph or different ones — had each other's previous questions and answers fed into their prompt. Context is now taken from the request instead of being retained on the server. +- **Deleting an MCP server takes effect, after a confirmation.** The delete button removed a server only from the page, without asking, and there was no way to save the change, so the server reappeared on reload. Deleting now asks first and is saved immediately. +- **Cancelling an MCP server edit no longer deletes the server.** Cancel removed an existing server from the list, so the next save of any server deleted it. Cancel now restores the saved settings. +- **Agentic answers no longer wait on a rejected planning request.** With an OpenAI chat model, each plan was first requested in a form OpenAI rejects and then retried, adding latency and token cost to every agentic answer. Plans are now requested in a form it accepts. +- **Answers no longer fail when retrieved context exceeds the model's input limit.** With no `token_limit` configured — the default in the sample configs — retrieved context was sent untrimmed, so a large retrieval could be rejected by the model and the user got a generic failure reply. When `token_limit` is unset, the chat model's own input limit now applies. Trimming also reserves room for the full prompt and counts tokens rather than characters, so Japanese and other CJK context is trimmed correctly. +- **Agentic answers no longer pay for the same passage twice.** When several search steps of a plan returned the same document passage, the answer model received every copy. Each passage is now included once, in the step that found it first. +- **Search steps in the trace report what they retrieved.** An agentic search step's summary showed a fixed "2 item(s)" regardless of results; it now gives the number of chunks retrieved (and any other entries), and a search that finds nothing is reported as empty. Each search also logs its chunk count and total size. A step output too large to keep in full in the trace is now shown as the tool's own data, shortened, with a note giving its full and shown size and each shortened text ending in how much of it was left out — previously it was replaced by a preview marker that read as if the tool's output itself had been cut. + +### Changed +- **The Migration Assistant repairs only graphs it can repair in place.** Its check now verifies that the graph has the schema GraphRAG's queries are written against, which every graph created with GraphRAG 1.4 or later has. A graph created with an earlier version is reported as needing to be recreated and its documents ingested again, and Repair is not offered — previously Repair would replace its working queries with ones its schema cannot run. If the schema cannot be read, Repair is withheld until the check succeeds. +- **`token_limit: 0` no longer means unlimited.** A `token_limit` of `0` or left unset now uses the chat model's own input limit (less room for the answer), so retrieved context is trimmed to what the model accepts. Models whose limit is not known keep the previous untrimmed behavior. Set a positive `token_limit` to choose a lower cap. +- **Callers of `/{graphname}/query_with_history` supply their own conversation context.** The request body accepts an optional `history` list of `{"query", "response"}` turns, oldest first, and the most recent turns are used. A request that omits it is answered without prior context; the endpoint no longer remembers anything between calls. Callers that relied on the server accumulating history must now send the turns they want considered. + ## [2.0.2] ### Added diff --git a/README.md b/README.md index fd83f9f..a54f0a6 100644 --- a/README.md +++ b/README.md @@ -22,6 +22,7 @@ - [Use TigerGraph GraphRAG](#use-tigergraph-graphrag) - [Run Demo with Preloaded GraphRAG](#run-demo-with-preloaded-graphrag) - [Manually Build GraphRAG From Scratch](#manually-build-graphrag-from-scratch) + - [Upgrade an Existing Graph](#upgrade-an-existing-graph) - [Chat Engines and Agents](#chat-engines-and-agents) - [Agentic](#agentic) - [Classic](#classic) @@ -364,6 +365,19 @@ The script will: 1. Init the GraphRAG based on the graph and install required queries 1. Ask a question via Python to get answer from GraphRAG +### Upgrade an Existing Graph + +Upgrading GraphRAG does not change graphs that already exist: their installed queries stay as they were until they are repaired. After upgrading, open **Knowledge Graph Admin → Migration Assistant** and check each graph. The check reports: + +- **Schema compatibility** — whether the graph's schema supports the current release. Graphs created with GraphRAG 1.4 or later are compatible. +- **Queries** — installed queries that are outdated or missing. **Repair** reinstalls them in place, without rebuilding the knowledge graph. +- **Prompt overrides** — saved custom prompts in a format the current release no longer uses. +- **Data health** — vertices missing embeddings and communities missing summaries, each with its own regenerate action. + +A graph created with a version earlier than 1.4 can't be repaired in place; the Migration Assistant reports this and does not offer **Repair**. Create a new graph and ingest its documents again. + +Repair can't run while the graph is being rebuilt or ingested. + [Go back to top](#top) --- @@ -666,7 +680,7 @@ In the `llm_config` section of `configs/server_config.json` file, copy JSON conf | Parameter | Type | Default | Description | | --- | --- | --- | --- | | `authentication_configuration` | object | — | Shared authentication credentials for all services. Service-level values take precedence. | -| `token_limit` | int | — | Hard cap on token count for retrieved context sent to the LLM. Context exceeding this limit is truncated. Inherited by all services if not set at service level. `0` or omitted means unlimited. | +| `token_limit` | int | — | Hard cap on token count for a prompt carrying retrieved context, prompt included. Context beyond it is truncated. Inherited by all services if not set at service level. `0` or omitted uses the chat model's own input limit when known, otherwise no cap. | **`completion_service` parameters:** @@ -678,7 +692,7 @@ In the `llm_config` section of `configs/server_config.json` file, copy JSON conf | `model_kwargs` | object | No | `{}` | Additional model parameters (e.g., `{"temperature": 0}`). | | `prompt_path` | string | No | `"./common/prompts/openai_gpt4/"` | Path to prompt template files. | | `base_url` | string | No | — | Custom API endpoint URL. | -| `token_limit` | int | No | inherited from top-level | Hard cap on token count for retrieved context sent to the LLM. Context exceeding this limit is truncated. `0` or omitted means unlimited. | +| `token_limit` | int | No | inherited from top-level | Hard cap on token count for a prompt carrying retrieved context, prompt included. Context beyond it is truncated. `0` or omitted uses the model's own input limit when known, otherwise no cap. | **`embedding_service` parameters:** @@ -701,7 +715,7 @@ Chatbot LLM override. If not configured, inherits from `completion_service`. Con | `model_kwargs` | object | No | inherited from completion | Additional model parameters (e.g., `{"temperature": 0}`). | | `prompt_path` | string | No | inherited from completion | Path to prompt template files. | | `base_url` | string | No | inherited from completion | Custom API endpoint URL. | -| `token_limit` | int | No | inherited from completion | Hard cap on token count for retrieved context sent to the chatbot LLM. Context exceeding this limit is truncated. `0` or omitted means unlimited. | +| `token_limit` | int | No | inherited from completion | Hard cap on token count for a prompt carrying retrieved context to the chatbot LLM, prompt included. Context beyond it is truncated. `0` or omitted uses the model's own input limit when known, otherwise no cap. | **`multimodal_service` parameters (optional):** diff --git a/common/db/connections.py b/common/db/connections.py index c9dc628..860dba1 100644 --- a/common/db/connections.py +++ b/common/db/connections.py @@ -16,6 +16,7 @@ import asyncio from typing import Annotated +import aiohttp from fastapi import Depends, HTTPException, status from fastapi.security import HTTPBasicCredentials, HTTPAuthorizationCredentials from pyTigerGraph import TigerGraphConnection, AsyncTigerGraphConnection @@ -33,6 +34,21 @@ consistency_checkers = {} +async def _verify_async_connection(conn, graphname: str) -> None: + """Validate the token, then release the HTTP session in the same loop. + + This runs under ``asyncio.run()``, which closes its loop on return. A session + left open stays bound to that closed loop, and with pyTigerGraph 2.0.4 the + next request on another loop fails with "Event loop is closed" (GML-2182). + Closing it here makes the connection open a fresh session on the caller's + loop at first use. + """ + try: + await conn.gsql("USE GRAPH " + graphname) + finally: + await conn.aclose() + + def get_db_connection_id_token( graphname: str, credentials: Annotated[HTTPAuthorizationCredentials, Depends(security)], @@ -65,9 +81,29 @@ def get_db_connection_id_token( try: if async_conn: - asyncio.run(conn.gsql("USE GRAPH " + graphname)) + asyncio.run(_verify_async_connection(conn, graphname)) else: conn.gsql("USE GRAPH " + graphname) + except aiohttp.ClientResponseError as e: + # The async client raises aiohttp's error type, not requests.HTTPError. + if e.status in (401, 403): + LogWriter.error("Failed to connect to TigerGraph. Incorrect ID Token.") + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect token", + headers={"WWW-Authenticate": "Bearer"}, + ) + LogWriter.error(f"Failed to verify ID token: {e}") + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Failed to get token - is the database running?" + ) + except (aiohttp.ClientConnectionError, asyncio.TimeoutError) as e: + LogWriter.error(f"Failed to reach TigerGraph to verify ID token: {e}") + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail="Failed to get token - is the database running?" + ) except HTTPError: LogWriter.error("Failed to connect to TigerGraph. Incorrect ID Token.") raise HTTPException( diff --git a/common/db/migrate.py b/common/db/migrate.py index 307ab51..9b23af9 100644 --- a/common/db/migrate.py +++ b/common/db/migrate.py @@ -13,6 +13,9 @@ against the body currently installed on TigerGraph and re-creates + re-installs only the ones whose body has actually drifted. +* ``check_schema_compatibility`` decides whether a graph can be repaired + in place at all: it must already have the shipped base schema. + * ``check_and_apply_schema`` (still a stub — see TODO below) is the schema counterpart: detect missing attributes on existing vertex / edge types and emit ``ALTER VERTEX ... ADD ATTRIBUTE …`` statements. @@ -28,6 +31,7 @@ import logging import os import re +from dataclasses import dataclass, field from typing import Iterable logger = logging.getLogger(__name__) @@ -230,6 +234,123 @@ async def filter_queries_needing_update_async( return out +# --------------------------------------------------------------------------- +# Schema compatibility +# --------------------------------------------------------------------------- +# +# Query repair reinstalls the shipped queries in place, which only works when +# the graph already has the base schema those queries are written against. +# That base schema has been unchanged since 1.4.0; graphs created earlier +# differ (1.3.x: IS_HEAD_OF / HAS_TAIL attach to Entity rather than +# EntityType; 1.1-1.2: different types altogether) and are not repairable in +# place. Rather than infer a version, compare the live schema with the shipped +# base schema file: every base vertex type, edge type, edge endpoint pair and +# attribute must exist. Extra types (domain types, images) are allowed. Vector +# attributes live in separate files and depend on the deployment's vector +# setup, so they are not part of the check. + +BASE_SCHEMA_PATH = "common/gsql/supportai/SupportAI_Schema.gsql" + +_ADD_VERTEX_RE = re.compile( + r"ADD\s+VERTEX\s+(\w+)\s*\((.*?)\)\s*WITH", re.IGNORECASE | re.DOTALL +) +_ADD_EDGE_RE = re.compile( + r"ADD\s+(?:UN)?DIRECTED\s+EDGE\s+(\w+)\s*\((.*?)\)\s*(?:WITH|;)", + re.IGNORECASE | re.DOTALL, +) +_EDGE_PAIR_RE = re.compile(r"FROM\s+(\w+)\s*,\s*TO\s+(\w+)", re.IGNORECASE) + + +@dataclass +class SchemaCompatibility: + compatible: bool + differences: list = field(default_factory=list) + + +def _attr_names(segment: str) -> set: + """Attribute names from a comma-separated ``name TYPE`` list.""" + names = set() + for part in segment.split(","): + words = part.split() + if words and words[0].upper() != "PRIMARY_ID": + names.add(words[0]) + return names + + +def parse_base_schema(text: str) -> tuple[dict, dict]: + """Expected schema from the shipped base schema GSQL. + + Returns ``(vertices, edges)``: ``{vertex: {attr, ...}}`` and + ``{edge: {"pairs": {(from, to), ...}, "attrs": {attr, ...}}}``. + """ + vertices = {m.group(1): _attr_names(m.group(2)) for m in _ADD_VERTEX_RE.finditer(text)} + edges = {} + for m in _ADD_EDGE_RE.finditer(text): + body = m.group(2) + pairs = {(f, t) for f, t in _EDGE_PAIR_RE.findall(body)} + edges[m.group(1)] = { + "pairs": pairs, + "attrs": _attr_names(_EDGE_PAIR_RE.sub("", body).replace("|", ",")), + } + return vertices, edges + + +def _live_attr_names(meta: dict) -> set: + return {a.get("AttributeName") for a in (meta or {}).get("Attributes", []) or []} + + +def _edge_connects(meta: dict, from_vt: str, to_vt: str) -> bool: + """Whether the live edge accepts ``from_vt -> to_vt``. + + Multi-pair edges list their pairs under ``EdgePairs``. An edge whose + endpoints are reported as ``*`` with no pair list accepts any type at + that end; TigerGraph reports ``IN_COMMUNITY`` this way once domain types + have been added to it. + """ + pairs = [(ep.get("From"), ep.get("To")) for ep in meta.get("EdgePairs", []) or []] + if pairs: + return (from_vt, to_vt) in pairs + f, t = meta.get("FromVertexTypeName"), meta.get("ToVertexTypeName") + return f in ("*", from_vt) and t in ("*", to_vt) + + +def check_schema_compatibility(conn, base_schema_path: str = BASE_SCHEMA_PATH) -> SchemaCompatibility: + """Whether the graph has every element of the shipped base schema. + + Raises when the live schema cannot be read, so callers can tell "not + compatible" from "could not check". + """ + with open(base_schema_path, encoding="utf-8") as f: + exp_vertices, exp_edges = parse_base_schema(f.read()) + + live_vertices = set(conn.getVertexTypes() or []) + live_edges = set(conn.getEdgeTypes() or []) + differences = [] + + for vt, attrs in exp_vertices.items(): + if vt not in live_vertices: + differences.append(f"missing vertex type {vt}") + continue + missing = attrs - _live_attr_names(conn.getVertexType(vt)) + if missing: + differences.append(f"vertex type {vt} is missing attribute(s) {', '.join(sorted(missing))}") + + for et, spec in exp_edges.items(): + if et not in live_edges: + differences.append(f"missing edge type {et}") + continue + meta = conn.getEdgeType(et) or {} + missing_pairs = {(f, t) for f, t in spec["pairs"] if not _edge_connects(meta, f, t)} + if missing_pairs: + pairs = ", ".join(f"{a}->{b}" for a, b in sorted(missing_pairs)) + differences.append(f"edge type {et} does not connect {pairs}") + missing = spec["attrs"] - _live_attr_names(meta) + if missing: + differences.append(f"edge type {et} is missing attribute(s) {', '.join(sorted(missing))}") + + return SchemaCompatibility(compatible=not differences, differences=differences) + + # --------------------------------------------------------------------------- # Schema migration — TODO # --------------------------------------------------------------------------- diff --git a/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql b/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql index 5487049..2a05aa9 100644 --- a/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql +++ b/common/gsql/supportai/retrievers/Chunk_Sibling_Vector_Search.gsql @@ -21,7 +21,11 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, TYPEDEF tuple Similarity_Results; HeapAccum(top_k, score DESC) @@topk_set; MapAccum> @@verbose_info; - MapAccum @distances; + // Offset of this chunk from each seed whose window reaches it: negative + // before the seed, positive after, 0 for the seed itself. + MapAccum> @offsets; + // Rank of the nearest seed (see below); the chunk is listed under it only. + MinAccum @best; SumAccum @content; MapAccum>> @@final_retrieval; OrAccum @touched; @@ -29,7 +33,7 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, SetAccum @@start_set_type; SetAccum @@sibling_set_type; SetAccum @@edges; - + all_chunks = {v_type}; result = SELECT v FROM all_chunks:v WHERE v.embedding.size() > 0 POST-ACCUM @@topk_set += Similarity_Results(v, 1 - gds.vector.distance(query_vector, v.embedding, "COSINE")); @@ -38,41 +42,52 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, END; start = {@@start_set}; - start = SELECT s FROM start:s ACCUM s.@distances += (s -> 0), s.@touched += TRUE, @@start_set_type += VertexTypes(s, s.type); + start = SELECT s FROM start:s ACCUM s.@offsets += (s -> 0), s.@touched += TRUE, @@start_set_type += VertexTypes(s, s.type); @@verbose_info += ("start_set" -> @@start_set_type); - FOREACH i IN RANGE[1, lookback] DO - start = SELECT tgt FROM start:s -(IS_AFTER)-> :tgt - ACCUM - tgt.@touched += TRUE, @@sibling_set_type += VertexTypes(tgt, tgt.type), - @@edges += EdgeTypes(s, tgt), - FOREACH (key, val) IN s.@distances DO - tgt.@distances += (key -> -1*i) - END; - END; - - start = {@@start_set}; - FOREACH i IN RANGE[1, lookahead] DO - start = SELECT tgt FROM start:s -(reverse_IS_AFTER)-> :tgt - ACCUM - tgt.@touched += TRUE, @@sibling_set_type += VertexTypes(tgt, tgt.type), - @@edges += EdgeTypes(s, tgt), - FOREACH (key, val) IN s.@distances DO - tgt.@distances += (key -> i) - END; + // Walk each seed's window on its own, so an offset is always measured from + // the seed it belongs to (IS_AFTER points to the previous chunk). + FOREACH seed IN @@start_set DO + frontier = {seed}; + FOREACH i IN RANGE[1, lookback] DO + frontier = SELECT tgt FROM frontier:s -(IS_AFTER)-> :tgt + ACCUM + tgt.@touched += TRUE, tgt.@offsets += (seed -> -1*i), + @@sibling_set_type += VertexTypes(tgt, tgt.type), + @@edges += EdgeTypes(s, tgt); + END; + frontier = {seed}; + FOREACH i IN RANGE[1, lookahead] DO + frontier = SELECT tgt FROM frontier:s -(reverse_IS_AFTER)-> :tgt + ACCUM + tgt.@touched += TRUE, tgt.@offsets += (seed -> i), + @@sibling_set_type += VertexTypes(tgt, tgt.type), + @@edges += EdgeTypes(s, tgt); + END; END; - + + // List each chunk once, under its nearest seed. Rank = 2*|offset|, plus 1 + // when the chunk is not after the seed, so a chunk exactly between two + // seeds goes to the earlier one and every seed ranks itself first. + res = SELECT s FROM all_chunks:s WHERE s.@touched == TRUE + POST-ACCUM + FOREACH (key, val) IN s.@offsets DO + s.@best += 2 * abs(val) + (CASE WHEN val > 0 THEN 0 ELSE 1 END) + END; + res = SELECT s FROM all_chunks:s -(HAS_CONTENT)- Content:t WHERE s.@touched == TRUE ACCUM s.@content += t.text; - + res = SELECT s FROM all_chunks:s WHERE s.@touched == TRUE POST-ACCUM - FOREACH (key, val) IN s.@distances DO - @@final_retrieval += (key -> (s -> ("distance" -> to_string(val)))), - @@final_retrieval += (key -> (s -> ("content" -> s.@content))) + FOREACH (key, val) IN s.@offsets DO + IF 2 * abs(val) + (CASE WHEN val > 0 THEN 0 ELSE 1 END) == s.@best THEN + @@final_retrieval += (key -> (s -> ("distance" -> to_string(val)))), + @@final_retrieval += (key -> (s -> ("content" -> s.@content))) + END END; - + @@verbose_info += ("selected_set" -> @@sibling_set_type); PRINT @@final_retrieval as final_retrieval, @@edges as edges; @@ -80,5 +95,4 @@ CREATE OR REPLACE DISTRIBUTED QUERY Chunk_Sibling_Vector_Search(STRING v_type, IF verbose THEN PRINT @@verbose_info as verbose; END; - } diff --git a/common/llm_services/azure_openai_service.py b/common/llm_services/azure_openai_service.py index 6baaa80..026a98c 100644 --- a/common/llm_services/azure_openai_service.py +++ b/common/llm_services/azure_openai_service.py @@ -9,6 +9,8 @@ class AzureOpenAI(LLM_Model): + strict_structured_output = True + def __init__(self, config): super().__init__(config) for auth_detail in config["authentication_configuration"].keys(): diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index 3dfcb9d..0626015 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -179,12 +179,25 @@ def _record_usage(caller_name: str, usage_data: dict, config: Optional[dict] = N bucket.append({"caller_name": caller_name, **usage_data}) +# Defaults for LLM_Model.context_token_limit when no token_limit is configured. +# Prompts are counted with tiktoken, which can undercount a provider's own +# tokenizer: Gemini counted ~23% more on a real 524K-token context. +_TOKENIZER_MARGIN = 0.6 +_DEFAULT_ANSWER_RESERVE = 4096 +_warned_no_token_limit: set = set() + + class LLM_Model: """Base LLM_Model Class Used to connect to external LLM API services, and retrieve customized prompts for the tools. """ + # Whether the provider's native structured output enforces strict JSON + # schema, which rejects free-form maps. True for the OpenAI family + # (GML-2302); callers with such maps then ask for tool calling instead. + strict_structured_output = False + def __init__(self, config): self.llm = None self.config = config @@ -192,6 +205,32 @@ def __init__(self, config): self._graphname = validate_graphname(config.get("graphname")) self.prompt_path = config.get("prompt_path", "") + def context_token_limit(self) -> int: + """Token budget for an answer prompt, retrieved context included. + + The configured ``token_limit`` when set. Otherwise the chat model's + input limit from its LangChain profile, less room for the answer and a + margin for tokenizer differences (counts use tiktoken, not the + provider's tokenizer). 0 means no limit: a model with no profile, such + as an Ollama or custom model name, keeps the old untrimmed behavior. + """ + configured = self.config.get("token_limit") + if configured: + return int(configured) + profile = getattr(self.llm, "profile", None) or {} + max_input = profile.get("max_input_tokens") + if not max_input: + model = self.config.get("llm_model") + if model not in _warned_no_token_limit: + _warned_no_token_limit.add(model) + logger.warning( + f"No token_limit configured and no known input limit for model " + f"{model!r}; retrieved context will not be trimmed" + ) + return 0 + max_output = profile.get("max_output_tokens") or _DEFAULT_ANSWER_RESERVE + return max(int(max_input * _TOKENIZER_MARGIN) - int(max_output), 0) + def _read_prompt_file(self, path): """Read a prompt file with per-graph override support. @@ -599,17 +638,24 @@ def invoke_structured( messages: list, schema, caller_name: str = "unknown", + method: str = None, ): """Invoke the chat model with native structured output. Returns an instance of ``schema`` (a pydantic class). Used by the planner to get a typed ``Plan`` back. Falls back to a JSON-extraction parse when the provider's structured-output path returns text. + + ``method`` is passed to ``with_structured_output`` when given; ``None`` + keeps the provider's default. """ usage_data = {} with get_openai_callback() as cb: try: - structured = self.llm.with_structured_output(schema) + if method is None: + structured = self.llm.with_structured_output(schema) + else: + structured = self.llm.with_structured_output(schema, method=method) result = structured.invoke(messages) except Exception as exc: logger.warning( diff --git a/common/llm_services/openai_service.py b/common/llm_services/openai_service.py index a81cf21..bd7f9ef 100644 --- a/common/llm_services/openai_service.py +++ b/common/llm_services/openai_service.py @@ -26,6 +26,8 @@ class OpenAI(LLM_Model): + strict_structured_output = True + def __init__(self, config): super().__init__(config) for auth_detail in config["authentication_configuration"].keys(): diff --git a/common/py_schemas/schemas.py b/common/py_schemas/schemas.py index a07b9da..6ee296f 100644 --- a/common/py_schemas/schemas.py +++ b/common/py_schemas/schemas.py @@ -18,6 +18,11 @@ from pydantic import BaseModel, Field +class ConversationTurn(BaseModel): + query: str + response: str + + class NaturalLanguageQuery(BaseModel): query: str # Engine: "agentic" | "classic" | None (defer to graph config). @@ -29,6 +34,10 @@ class NaturalLanguageQuery(BaseModel): # name fields (e.g. "query_sources") or "all" to include the supporting # sources / trace in the response. include_fields: Optional[List[str]] = Field(default=None) + # Prior turns of the caller's own conversation, oldest first. The server + # holds no conversation state, so a caller wanting follow-up context must + # send it; only the most recent turns are given to the agent. + history: Optional[List[ConversationTurn]] = Field(default=None) class SupportAIQuestion(BaseModel): diff --git a/common/utils/retrieval_stats.py b/common/utils/retrieval_stats.py new file mode 100644 index 0000000..e336035 --- /dev/null +++ b/common/utils/retrieval_stats.py @@ -0,0 +1,87 @@ +# 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 . + +"""Counts and sizes of what a retrieval query returned in ``final_retrieval``.""" + +import re +from dataclasses import dataclass + +# DocumentChunk ids are "_chunk_"; every other entry (entity +# text, community summaries, whole documents, the community +# ``Similarity_Context``) is counted separately. +_CHUNK_ID = re.compile(r"_chunk_\d+$") + + +@dataclass +class RetrievalStats: + chunks: int = 0 + chunk_chars: int = 0 + others: int = 0 + other_chars: int = 0 + + +def _chars(value) -> int: + if isinstance(value, str): + return len(value) + if isinstance(value, (list, tuple, set)): + return sum(_chars(v) for v in value) + return len(str(value)) + + +def is_grouped(value) -> bool: + """Whether a ``final_retrieval`` value is a contextual-search group, + ``{chunk: {"distance", "content"}}``, rather than an entry's own text.""" + return isinstance(value, dict) and bool(value) and all( + isinstance(v, dict) and "content" in v for v in value.values() + ) + + +def retrieved_entries(final_retrieval: dict) -> dict: + """``{id: text}`` for everything a retrieval returned, each id once. + + Most retrievals map an id to its text. Contextual (sibling) search nests + one level deeper, ``{seed: {chunk: {"distance", "content"}}}``, grouping + each match with the chunks around it; its entries are the inner chunks. + """ + entries = {} + for key, value in final_retrieval.items(): + if is_grouped(value): + for inner_key, inner in value.items(): + entries.setdefault(str(inner_key), inner.get("content", "")) + else: + entries.setdefault(str(key), value) + return entries + + +def retrieval_stats(final_retrieval: dict) -> RetrievalStats: + stats = RetrievalStats() + for key, value in retrieved_entries(final_retrieval).items(): + if _CHUNK_ID.search(key): + stats.chunks += 1 + stats.chunk_chars += _chars(value) + else: + stats.others += 1 + stats.other_chars += _chars(value) + return stats + + +def describe_retrieval(query_name: str, final_retrieval: dict) -> str: + """One line for the log: e.g. " retrieved 20 chunk(s), 16,102,331 + chars; 3 other entries, 12,345 chars".""" + s = retrieval_stats(final_retrieval) + line = f"{query_name} retrieved {s.chunks} chunk(s), {s.chunk_chars:,} chars" + if s.others: + line += (f"; {s.others} other entr{'y' if s.others == 1 else 'ies'}, " + f"{s.other_chars:,} chars") + return line diff --git a/common/utils/text_extractors.py b/common/utils/text_extractors.py index 41a441e..a5afb7e 100644 --- a/common/utils/text_extractors.py +++ b/common/utils/text_extractors.py @@ -17,6 +17,18 @@ logger = logging.getLogger(__name__) + +class ExtractionError(Exception): + """A file could not be turned into document content. + + Raised instead of returning a document whose text is the error message + (GML-2195): a placeholder like that reaches the graph, gets embedded and + becomes retrievable, and the agent cannot tell it from real content. The + message is surfaced to the caller, so it describes the file and what to do + about it — never exception internals or filesystem paths, which belong in + the log line at the raise site. + """ + # Global lock for pymupdf4llm calls (not thread-safe) _pymupdf4llm_lock = threading.Lock() @@ -494,6 +506,8 @@ def __init__(self): '.md': 'text/markdown', '.pdf': 'application/pdf', '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document', + # Listed so a .doc is attempted and reported as unreadable with a + # pointer to .docx; unlisted extensions are skipped silently. '.doc': 'application/msword', '.html': 'text/html', '.htm': 'text/html', @@ -552,12 +566,24 @@ async def _process_file_async(self, file_path, graphname, temp_folder): } except FileNotFoundError: - return {'success': False, 'file_path': str(file_path), 'error': 'File not found'} + return {'success': False, 'file_path': str(file_path), + 'error': f"{Path(file_path).name} could not be found on the server."} except PermissionError: - return {'success': False, 'file_path': str(file_path), 'error': 'Permission denied'} - except Exception as e: + return {'success': False, 'file_path': str(file_path), + 'error': f"{Path(file_path).name} could not be opened on the server."} + except ExtractionError as e: + # Written for the person who uploaded the file. logger.warning(f"Failed to process file {file_path}: {e}") return {'success': False, 'file_path': str(file_path), 'error': str(e)} + except Exception as e: + # Anything else can carry paths, URLs or service errors; those stay + # in the log, and the caller gets a message about the file. + logger.warning(f"Failed to process file {file_path}: {e}", exc_info=True) + return { + 'success': False, + 'file_path': str(file_path), + 'error': f"Could not process {Path(file_path).name}.", + } def _write_to_jsonl(self, jsonl_file, doc_entries): """ @@ -864,12 +890,14 @@ def _to_markdown_paged(strategy: str | None = None): # Cleanup folder if it was created if image_output_folder.exists(): shutil.rmtree(image_output_folder, ignore_errors=True) - return [{ - "doc_id": base_doc_id, - "doc_type": "markdown", - "content": f"[PDF extraction failed: {e}]", - "position": 0 - }] + # Surface the failure instead of ingesting the error text as + # the document's content (GML-2195). The raw exception stays + # in the log above; callers get a message safe to display. + raise ExtractionError( + f"Could not read text from {file_path.name}. " + "The file may be corrupt, password-protected or an " + "unsupported PDF variant." + ) from e if not markdown_content or not markdown_content.strip(): logger.warning( @@ -878,12 +906,10 @@ def _to_markdown_paged(strategy: str | None = None): ) if image_output_folder.exists(): shutil.rmtree(image_output_folder, ignore_errors=True) - return [{ - "doc_id": base_doc_id, - "doc_type": "markdown", - "content": f"[Scanned PDF — no text layer extracted: {file_path.name}]", - "position": 0 - }] + raise ExtractionError( + f"No text could be extracted from {file_path.name}. If it is a " + "scanned document, upload a version with selectable text." + ) # Clean up artefacts common in form PDFs (duplicate rows, ColN headers) markdown_content = _clean_pdf_markdown(markdown_content, source_hint=str(file_path)) @@ -1036,17 +1062,20 @@ def _describe_and_encode(img_ref: dict) -> dict: result.extend(image_entries) return result + except ExtractionError: + # Already reported at the raise site with a caller-safe message. + if image_output_folder.exists(): + shutil.rmtree(image_output_folder, ignore_errors=True) + raise except ImportError as import_err: logger.error(f"Required library missing: {import_err}") # Cleanup on import error if image_output_folder.exists(): shutil.rmtree(image_output_folder, ignore_errors=True) - return [{ - "doc_id": base_doc_id, - "doc_type": "markdown", - "content": "[PDF extraction requires pymupdf4llm and PyMuPDF]", - "position": 0 - }] + raise ExtractionError( + f"Could not read {file_path.name}: PDF support is unavailable on " + "this server." + ) from import_err except Exception as e: logger.error(f"Error extracting PDF: {e}") # Cleanup on any other error @@ -1112,14 +1141,104 @@ def _extract_standalone_image_as_doc(file_path, base_doc_id, graphname=None): } ] + except ExtractionError: + raise except Exception as e: logger.error(f"Error extracting image: {e}") - return [{ - "doc_id": base_doc_id, - "doc_type": "markdown", - "content": f"[Image extraction failed: {str(e)}]", - "position": 0 - }] + # This also covers the image-description service failing, so the + # message doesn't blame the file. + raise ExtractionError(f"Could not process image {Path(file_path).name}.") from e + + +def _docx_cell_text(cell): + """Text of one table cell, including any table nested inside it. + + ``_Cell.text`` joins the cell's paragraphs only, so a nested table would + be dropped the same way top-level tables were (GML-2196). + """ + return "\n".join(_docx_blocks_to_text(cell)).strip() + + +def _docx_render_table(table): + """Render a Word table as a markdown pipe table. + + Matches the shape the spreadsheet branch emits, so the chunker sees one + representation of tabular data. Pipes and newlines inside a cell are + neutralised; either would otherwise break the row. + """ + rows = [] + for row in table.rows: + cells = [ + _docx_cell_text(cell).replace("|", "\\|").replace("\n", " ").strip() + for cell in row.cells + ] + if any(cells): + rows.append(cells) + if not rows: + return "" + + width = max(len(r) for r in rows) + rows = [r + [""] * (width - len(r)) for r in rows] + header, body = rows[0], rows[1:] + lines = [ + "| " + " | ".join(header) + " |", + "| " + " | ".join("---" for _ in header) + " |", + ] + lines += ["| " + " | ".join(r) + " |" for r in body] + return "\n".join(lines) + + +def _docx_flatten_table(table): + """Render a table nested inside a cell as plain text. + + Markdown has no nested tables, and rendering one inside a cell would only + escape its pipes into an unreadable line, so the values are joined instead. + """ + rows = [] + for row in table.rows: + cells = [c for c in (_docx_cell_text(cell) for cell in row.cells) if c] + if cells: + rows.append(" — ".join(cells)) + return "; ".join(rows) + + +def _docx_blocks_to_text(parent): + """Yield the text of ``parent``'s paragraphs and tables, in document order. + + Word stores tables as siblings of paragraphs in the body. Reading + ``doc.paragraphs`` alone skips them entirely, so a table's contents never + reached the graph while the prose around it did (GML-2196). Walking the + XML children keeps each table anchored to the text that introduces it. + """ + from docx.document import Document as _DocxDocument + from docx.oxml.table import CT_Tbl + from docx.oxml.text.paragraph import CT_P + from docx.table import Table, _Cell + from docx.text.paragraph import Paragraph + + if isinstance(parent, _DocxDocument): + parent_elm = parent.element.body + render = _docx_render_table + elif isinstance(parent, _Cell): + parent_elm = parent._tc + render = _docx_flatten_table + else: + return + + for child in parent_elm.iterchildren(): + if isinstance(child, CT_P): + text = Paragraph(child, parent).text.strip() + if text: + yield text + elif isinstance(child, CT_Tbl): + rendered = render(Table(child, parent)) + if rendered: + yield rendered + + +def _docx_to_markdown(doc): + """Extract a .docx as text, with tables rendered in place.""" + return "\n\n".join(_docx_blocks_to_text(doc)).strip() def extract_text_from_file(file_path, graphname=None): @@ -1157,8 +1276,13 @@ def extract_text_from_file(file_path, graphname=None): return json.dumps(data, indent=2, ensure_ascii=False) elif extension == '.docx': import docx - doc = docx.Document(file_path) - return "\n".join(p.text for p in doc.paragraphs if p.text.strip()) + return _docx_to_markdown(docx.Document(file_path)) + elif extension == '.doc': + # Legacy OLE format; python-docx reads only the OOXML .docx form. + raise ExtractionError( + f"Could not read {file_path.name}: the legacy .doc format is " + "not supported. Save it as .docx and upload it again." + ) elif extension in ['.xlsx', '.xls']: import pandas as pd engine = 'openpyxl' if extension == '.xlsx' else 'xlrd' @@ -1189,7 +1313,12 @@ def extract_text_from_file(file_path, graphname=None): df.columns = [f"Column {i + 1}" for i in range(len(df.columns))] sheet_md = df.to_markdown(index=False) sheet_texts.append(f"## Sheet: {sheet_name}\n\n{sheet_md}") - return "\n\n".join(sheet_texts) if sheet_texts else "[Excel file is empty or contains no data]" + if not sheet_texts: + raise ExtractionError( + f"No data found in {file_path.name}. The spreadsheet " + "appears to be empty." + ) + return "\n\n".join(sheet_texts) elif extension == '.xml': import xml.etree.ElementTree as ET tree = ET.parse(file_path) @@ -1207,11 +1336,45 @@ def extract_text_from_element(element): import re return re.sub(r'\s+', ' ', content).strip() else: - return f"[Unsupported file type: {extension}]" + raise ExtractionError( + f"Could not read {file_path.name}: {extension} files are not " + "supported." + ) + except ExtractionError: + # Already carries a caller-safe message; don't re-wrap it with the + # raw exception text (GML-2196). + raise except Exception as e: logger.error(f"Error extracting text from {file_path}: {e}") - raise Exception(f"Text extraction failed: {e}") + raise ExtractionError( + f"Could not read {file_path.name}. The file may be corrupt or " + "not match its extension." + ) from e + + +def failed_extractions(result: dict) -> dict: + """Map each file that failed in a folder-processing *result* to its reason. + + ``_process_folder_async`` records a failed file as a ``status: failed`` + entry rather than raising, so the rest of the folder still converts. Its + ``error`` is the caller-safe message the extractor raised. + """ + failures = {} + for entry in (result or {}).get("files", []): + if entry.get("status") == "failed": + name = os.path.basename(entry.get("file_path") or "") + if name: + failures[name] = entry.get("error") or "The file could not be read." + return failures + + +def describe_failures(failures: dict) -> str: + """One line per failed file, for a message shown to the person uploading.""" + return " ".join( + reason if name in reason else f"{name}: {reason}" + for name, reason in sorted(failures.items()) + ) def get_doc_type_from_extension(extension): @@ -1229,7 +1392,9 @@ def get_doc_type_from_extension(extension): def get_supported_extensions(): """Get list of supported file extensions.""" - return {'.txt', '.md', '.html', '.htm', '.csv', '.json', '.pdf', '.docx', '.doc', '.xml', '.jpeg', '.jpg', '.png', '.gif', '.xlsx', '.xls', '.jsonl'} + # '.doc' (legacy OLE) is deliberately absent: python-docx reads only the + # OOXML '.docx' form, so it was never actually supported (GML-2196). + return {'.txt', '.md', '.html', '.htm', '.csv', '.json', '.pdf', '.docx', '.xml', '.jpeg', '.jpg', '.png', '.gif', '.xlsx', '.xls', '.jsonl'} def is_supported_file(file_path): """Check if a file is supported for text extraction.""" diff --git a/common/utils/token_calculator.py b/common/utils/token_calculator.py index bc58cee..e664a84 100644 --- a/common/utils/token_calculator.py +++ b/common/utils/token_calculator.py @@ -19,6 +19,10 @@ logger = logging.getLogger(__name__) +# Extra tokens kept free beyond the measured prompt, for chat-message framing +# and small differences between tiktoken and the provider's tokenizer. +PROMPT_RESERVE_MARGIN = 200 + # Cache for TokenCalculator instances to avoid re-initialization _token_calculator_cache: dict[tuple[str, int], 'TokenCalculator'] = {} @@ -113,6 +117,51 @@ def count_tokens(self, text: str | dict) -> int: # Fallback: rough estimation (1 token ≈ 4 characters for English text) return len(text) // 4 + def fit_context(self, context: str | dict, prompt_without_context: str) -> str | dict: + """Trim ``context`` so the whole prompt stays within the token limit. + + ``prompt_without_context`` is the prompt rendered with an empty context; + its size, plus a small margin, is reserved for the rest of the prompt. + Decides on token counts, never character counts, which undercount + tokens for CJK text. + """ + if self.is_unlimited_tokens(): + return context + budget = max( + self.max_context_tokens + - self.count_tokens(prompt_without_context) + - PROMPT_RESERVE_MARGIN, + 0, + ) + context_tokens = self.count_tokens(context) + if context_tokens <= budget: + return context + logger.info(f"Truncating context from {context_tokens} to {budget} tokens") + if isinstance(context, list): + return self._truncate_list(context, budget) + if budget == 0: + # The dict truncation treats a limit of 0 as "no limit". + return {} if isinstance(context, dict) else "" + return self.truncate_to_token_limit(context, budget) + + def _truncate_list(self, items: list, max_tokens: int) -> list: + """Keep items in order until ``max_tokens``; a string that only partly + fits is cut to the remaining budget.""" + kept, used = [], 0 + for item in items: + remaining = max_tokens - used + if remaining <= 0: + break + item_tokens = self.count_tokens(item) + if item_tokens <= remaining: + kept.append(item) + used += item_tokens + else: + if isinstance(item, str): + kept.append(self.truncate_text_to_token_limit(item, remaining)) + break + return kept + def truncate_dict_to_token_limit(self, sources_dict: dict, max_tokens: Optional[int] = None) -> dict: """ Truncate dictionary to fit within the token limit by keeping original values diff --git a/docs/notebooks/TransactionFraud_demo.ipynb b/docs/notebooks/TransactionFraud_demo.ipynb index c7ce334..79b8f8d 100644 --- a/docs/notebooks/TransactionFraud_demo.ipynb +++ b/docs/notebooks/TransactionFraud_demo.ipynb @@ -116,17 +116,27 @@ "\n", "for run_number in range(1, num_runs + 1):\n", " print(f\"Run {run_number}:\")\n", + "\n", + " # The endpoint keeps no conversation state, so the caller carries it.\n", + " # Each run starts a fresh conversation; later questions refer back to\n", + " # earlier ones (\"in this transaction\"), so the turns so far are sent\n", + " # along with every request.\n", + " history = []\n", + "\n", " for question in questions:\n", " start_time = time.time()\n", " payload = {\n", - " \"query\": question\n", + " \"query\": question,\n", + " \"history\": history\n", " }\n", "\n", " res = requests.post(\"http://GRAPHRAG_ADDRESS/Transaction_Fraud/query_with_history\", \n", " json=payload, \n", " auth=HTTPBasicAuth(conn.username, conn.password)\n", " )\n", - " print (json.loads(res.text).get(\"natural_language_response\", \"no answer\"))\n", + " answer = json.loads(res.text).get(\"natural_language_response\", \"no answer\")\n", + " print (answer)\n", + " history.append({\"query\": question, \"response\": answer})\n", " end_time = time.time()\n", "\n", " execution_time = end_time - start_time\n", diff --git a/ecc/app/main.py b/ecc/app/main.py index 9624cb3..3df56b0 100644 --- a/ecc/app/main.py +++ b/ecc/app/main.py @@ -287,8 +287,17 @@ async def run_with_tracking(task_key: str, run_func, graphname: str, conn): # Verify the graph still exists before doing any work try: await conn.getVertexTypes() - except Exception: - raise Exception(f"Graph '{graphname}' does not exist or is not accessible") + except Exception as e: + # Log the real cause (GML-2182). This message once hid an unrelated + # failure — "Event loop is closed" from a connection reused across + # event loops — and sent debugging the wrong way. The cause stays out + # of the message itself, which is shown to operators. + LogWriter.error( + f"Graph check failed for '{graphname}': {type(e).__name__}: {e}" + ) + raise Exception( + f"Graph '{graphname}' does not exist or is not accessible" + ) from e # Reload config at the start of each job to ensure latest settings are used LogWriter.info("Reloading configuration for new job...") diff --git a/graphrag-ui/src/components/SideMenu.tsx b/graphrag-ui/src/components/SideMenu.tsx index 02041cb..ba317ce 100644 --- a/graphrag-ui/src/components/SideMenu.tsx +++ b/graphrag-ui/src/components/SideMenu.tsx @@ -3,7 +3,7 @@ import { IoDocumentTextOutline } from "react-icons/io5"; import { FiTerminal } from "react-icons/fi"; import { FiLoader } from "react-icons/fi"; import { IoCartOutline } from "react-icons/io5"; -import { FiKey } from "react-icons/fi"; +import { FiKey, FiTrash2 } from "react-icons/fi"; import { IoIosHelpCircleOutline } from "react-icons/io"; import { HiOutlineChatBubbleOvalLeft } from "react-icons/hi2"; import { MdKeyboardArrowDown, MdKeyboardArrowUp } from "react-icons/md"; @@ -55,6 +55,7 @@ import { FaPaperclip } from "react-icons/fa6"; import { useCallback } from "react"; import { conversationManager } from "../actions/ActionProvider"; import { useConfirm } from "@/hooks/useConfirm"; +import { useAlert } from "@/hooks/useAlert"; import { useNavigate } from "react-router-dom"; // TODO make dynamic @@ -86,7 +87,13 @@ const SideMenu = ({ const [loadedCount, setLoadedCount] = useState(0); const [loadingMore, setLoadingMore] = useState(false); const [clearing, setClearing] = useState(false); + const [clearingAll, setClearingAll] = useState(false); + // Multi-select: pick several shown conversations and delete them together. + const [selecting, setSelecting] = useState(false); + const [selected, setSelected] = useState>(new Set()); + const [deletingSelected, setDeletingSelected] = useState(false); const [confirm, confirmDialog] = useConfirm(); + const [alert, alertDialog] = useAlert(); // Fade + disable the side menu (conversation list + New Chat) while // the chat is streaming an answer, so the user can't unmount Chat by // switching conversations mid-response. @@ -171,9 +178,9 @@ const SideMenu = ({ const fetchHistory2 = useCallback(async () => { try { const list = await fetchConvList(); - setConvList(list); const firstBatch = list.slice(0, PAGE_SIZE); const details = await loadDetails(firstBatch); + setConvList(list); setConversationId(details as any); setLoadedCount(firstBatch.length); } catch (error) { @@ -242,6 +249,137 @@ const SideMenu = ({ } }; + // Delete conversations by id, a few at a time so a long history can't flood + // the browser. Returns the ids the server confirmed deleted. + const deleteConversations = async (ids: string[]) => { + const creds = sessionStorage.getItem("auth"); + const settings = { + method: "DELETE", + headers: { Authorization: creds!, "Content-Type": "application/json" }, + }; + const deleted = new Set(); + const BATCH = 5; + for (let i = 0; i < ids.length; i += BATCH) { + const results = await Promise.all( + ids.slice(i, i + BATCH).map((id) => + fetch(`${WS_CONVO_URL}/${id}`, settings) + .then((r) => (r.ok ? id : null)) + .catch(() => null) + ) + ); + results.forEach((id) => { if (id) deleted.add(id); }); + } + return deleted; + }; + + // Drop deleted conversations from the sidebar. If the open conversation was + // one of them, start a new chat rather than keep showing it. + const forgetConversations = (deleted: Set) => { + if (deleted.size === 0) return; + const loadedGone = convList + .slice(0, loadedCount) + .filter((c: any) => deleted.has(c.conversation_id)).length; + setConvList((prev: any[]) => prev.filter((c: any) => !deleted.has(c.conversation_id))); + setConversationId((prev: any[]) => prev.filter((c: any) => !deleted.has(c.conversation_id))); + setLoadedCount((c) => c - loadedGone); + const openId = conversationManager.getCurrentConversationId(); + if ((openId && deleted.has(openId)) || (activeConversationId && deleted.has(activeConversationId))) { + handleNewChat(); + } + }; + + const deleteOne = async (id: string) => { + const ok = await confirm( + "Delete this conversation?\n\nThis permanently removes it and cannot be undone." + ); + if (!ok) return; + const deleted = await deleteConversations([id]); + forgetConversations(deleted); + if (deleted.size === 0) { + await alert("The conversation could not be deleted. Please try again."); + } + }; + + const clearAll = async () => { + const n = convList.length; + if (n === 0) return; + const ok = await confirm( + `Delete all ${n} conversation${n === 1 ? "" : "s"}?\n\n` + + `This permanently removes them and cannot be undone.` + ); + if (!ok) return; + setClearing(true); + setClearingAll(true); + let failed = 0; + try { + const deleted = await deleteConversations(convList.map((c: any) => c.conversation_id)); + forgetConversations(deleted); + failed = n - deleted.size; + } finally { + setClearing(false); + setClearingAll(false); + } + if (failed > 0) { + await alert( + `${failed} of ${n} conversation${n === 1 ? "" : "s"} could not be deleted. Please try again.` + ); + } + }; + + const visibleIds: string[] = newSet.map((c: any) => c.conversation_id); + // Only conversations still shown count as selected, so a refresh that drops + // one can't leave it selected out of sight. + const selectedIds = visibleIds.filter((id) => selected.has(id)); + const allSelected = visibleIds.length > 0 && selectedIds.length === visibleIds.length; + + const exitSelecting = () => { + setSelecting(false); + setSelected(new Set()); + }; + + const toggleSelected = (id: string) => { + setSelected((prev) => { + const next = new Set(prev); + if (next.has(id)) next.delete(id); + else next.add(id); + return next; + }); + }; + + const toggleSelectAll = () => { + setSelected(allSelected ? new Set() : new Set(visibleIds)); + }; + + const deleteSelected = async () => { + const ids = selectedIds; + const n = ids.length; + if (n === 0) return; + const ok = await confirm( + `Delete ${n} selected conversation${n === 1 ? "" : "s"}?\n\n` + + `This permanently removes them and cannot be undone.` + ); + if (!ok) return; + setClearing(true); + setDeletingSelected(true); + let failed = 0; + try { + const deleted = await deleteConversations(ids); + forgetConversations(deleted); + failed = n - deleted.size; + // Keep only the ones that failed selected, so they can be retried. + setSelected(new Set(ids.filter((id) => !deleted.has(id)))); + if (failed === 0) setSelecting(false); + } finally { + setClearing(false); + setDeletingSelected(false); + } + if (failed > 0) { + await alert( + `${failed} of ${n} conversation${n === 1 ? "" : "s"} could not be deleted. Please try again.` + ); + } + }; + // eslint-disable-next-line // @ts-ignore const resumeConvo = async (id):any => { @@ -303,7 +441,10 @@ const SideMenu = ({ } const renderConvoHistory = () => { - if (newSet.length === 0) { + // Only truly empty history gets the empty state; if every loaded + // conversation was deleted but older ones remain, keep the footer so + // "more…" and "Clear all" stay reachable. + if (newSet.length === 0 && convList.length === 0) { return (

@@ -332,6 +473,23 @@ const SideMenu = ({ return (

+ {selecting && ( + + )} {sortedDates.map(([date, conversations]) => { return (
@@ -353,12 +511,23 @@ const SideMenu = ({
{ e.preventDefault(); - resumeConvo(item.conversation_id); + if (selecting) toggleSelected(item.conversation_id); + else resumeConvo(item.conversation_id); }} > + {selecting && ( + + )}
{previewText}
{userMessages.length > 1 && ( @@ -377,6 +546,21 @@ const SideMenu = ({ )} )} + {!selecting && ( + + )} {isExpanded && userMessages.length > 1 && (
@@ -386,7 +570,8 @@ const SideMenu = ({ className="py-2 text-sm text-gray-600 dark:text-gray-400 truncate cursor-pointer hover:text-gray-900 dark:hover:text-gray-200" onClick={(e) => { e.preventDefault(); - resumeConvo(item.conversation_id); + if (selecting) toggleSelected(item.conversation_id); + else resumeConvo(item.conversation_id); }} > {msg.content} @@ -402,22 +587,63 @@ const SideMenu = ({
); })} - {loadedCount < convList.length && ( + {convList.length > 0 && (
- - + {loadedCount < convList.length ? ( + + ) : ( + + )} + {selecting ? ( +
+ + +
+ ) : ( +
+ + {loadedCount < convList.length && ( + + )} + +
+ )}
)}
@@ -675,6 +901,7 @@ const SideMenu = ({ {renderConvoHistory()} {confirmDialog} + {alertDialog} {/*
= ({ isModal = false }) => { ); const [selectedFiles, setSelectedFiles] = useState(null); const fileInputRef = useRef(null); + // Files the last ingest preparation could not read, per tab. They are + // skipped, so each success message names them rather than reading as clean. + type Unreadable = Array<{ file: string; error: string }>; + const unreadableRef = useRef>({ + uploaded: [], + downloaded: [], + }); + const unreadableNote = (source: "uploaded" | "downloaded") => { + const failed = unreadableRef.current[source]; + if (!failed.length) return ""; + return ` ⚠️ ${failed.length} file(s) could not be read and were skipped: ${failed.map((f) => f.error).join(" ")}`; + }; const [uploadedFiles, setUploadedFiles] = useState([]); const [isUploading, setIsUploading] = useState(false); const [uploadMessage, setUploadMessage] = useState(""); @@ -328,6 +340,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { } ); const data = await safeJson(response); + unreadableRef.current.uploaded = unreadableRef.current.uploaded.filter((f) => f.file !== filename); setUploadMessage(`✅ ${data.message}`); await fetchUploadedFiles(); } catch (error: any) { @@ -351,6 +364,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { headers: { Authorization: creds! }, }); const data = await safeJson(response); + unreadableRef.current.uploaded = []; setUploadMessage(`✅ ${data.message}`); await fetchUploadedFiles(); } catch (error: any) { @@ -488,6 +502,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { } ); const data = await safeJson(response); + unreadableRef.current.downloaded = unreadableRef.current.downloaded.filter((f) => f.file !== filename); setDownloadMessage(`✅ ${data.message}`); await fetchDownloadedFiles(); } catch (error: any) { @@ -511,6 +526,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { headers: { Authorization: creds! }, }); const data = await safeJson(response); + unreadableRef.current.downloaded = []; setDownloadMessage(`✅ ${data.message}`); await fetchDownloadedFiles(); } catch (error: any) { @@ -561,6 +577,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { ); } const createData = await safeJson(createResp); + unreadableRef.current[sourceType] = createData.failed_files || []; jobData = { load_job_id: createData.load_job_id, data_source_id: createData.data_source_id, @@ -591,7 +608,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { const ingestData = await safeJson(ingestResponse); console.log("Ingest response:", ingestData); - setIngestMessage(`✅ Ingestion completed successfully!`); + setIngestMessage(`✅ Ingestion completed successfully!${unreadableNote(sourceType)}`); setUploadMessage(""); } catch (error: any) { console.error("Error during ingestion:", error); @@ -642,6 +659,8 @@ const IngestGraph: React.FC = ({ isModal = false }) => { const createData = await safeJson(createResponse); console.log("Create ingest response:", createData); + unreadableRef.current[sourceType] = createData.failed_files || []; + const readyCount = createData.ready_files ?? fileCount; // Store ingest job data for later use setIngestJobData({ @@ -651,7 +670,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { }); if (!directIngestion) { - setIngestMessage(`✅ ${fileCount} file(s) ready for ingestion.`); + setIngestMessage(`✅ ${readyCount} file(s) ready for ingestion.${unreadableNote(sourceType)}`); setIsIngesting(false); } else { setIngestMessage("Step 2/2: Running document ingest..."); @@ -678,7 +697,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { const ingestData = await safeJson(ingestResponse); console.log("Ingest response:", ingestData); - setIngestMessage(`✅ Data ingested successfully! Processed documents from ${folderPath}/`); + setIngestMessage(`✅ Data ingested successfully! Processed documents from ${folderPath}/${unreadableNote(sourceType)}`); setIsIngesting(false); } } catch (error: any) { @@ -735,6 +754,8 @@ const IngestGraph: React.FC = ({ isModal = false }) => { const createData = await safeJson(createResponse); console.log("create_ingest response data:", createData); + unreadableRef.current[sourceType] = createData.failed_files || []; + const readyCount = createData.ready_files ?? fileCount; setIngestJobData({ load_job_id: createData.load_job_id, @@ -753,9 +774,9 @@ const IngestGraph: React.FC = ({ isModal = false }) => { await handleRunIngest(sourceType); } else { if (sourceType === "uploaded") { - setUploadMessage(`✅ ${fileCount} file(s) ready for ingestion.`); + setUploadMessage(`✅ ${readyCount} file(s) ready for ingestion.${unreadableNote(sourceType)}`); } else { - setDownloadMessage(`✅ ${fileCount} file(s) ready for ingestion.`); + setDownloadMessage(`✅ ${readyCount} file(s) ready for ingestion.${unreadableNote(sourceType)}`); } } } catch (error: any) { @@ -1188,7 +1209,7 @@ const IngestGraph: React.FC = ({ isModal = false }) => { Maximum upload per request: {MAX_UPLOAD_SIZE_MB} MB.

{selectedFiles && (() => { - const SUPPORTED_EXTENSIONS = new Set([".txt", ".md", ".pdf", ".docx", ".doc", ".html", ".htm", ".json", ".csv", ".xlsx", ".xls", ".xml", ".jpeg", ".jpg", ".png", ".gif", ".jsonl"]); + const SUPPORTED_EXTENSIONS = new Set([".txt", ".md", ".pdf", ".docx", ".html", ".htm", ".json", ".csv", ".xlsx", ".xls", ".xml", ".jpeg", ".jpg", ".png", ".gif", ".jsonl"]); const files = Array.from(selectedFiles); const unsupported = files.filter((f) => !SUPPORTED_EXTENSIONS.has(f.name.slice(f.name.lastIndexOf(".")).toLowerCase())); const hasCsvExcel = files.some((f) => [".csv", ".xlsx", ".xls"].includes(f.name.slice(f.name.lastIndexOf(".")).toLowerCase())); diff --git a/graphrag-ui/src/pages/setup/KGAdmin.tsx b/graphrag-ui/src/pages/setup/KGAdmin.tsx index 2cf23b5..69da191 100644 --- a/graphrag-ui/src/pages/setup/KGAdmin.tsx +++ b/graphrag-ui/src/pages/setup/KGAdmin.tsx @@ -79,6 +79,8 @@ const KGAdmin = () => { missing_files: string[]; }; needs_repair?: boolean; + repairable?: boolean; + schema?: { compatible: boolean; check_failed: boolean }; embeddings?: { by_type: Record; total_missing: number; @@ -154,7 +156,15 @@ const KGAdmin = () => { return; } setMigrationStatus(data); - if ( + if (data.schema?.check_failed) { + setMigrationMessage( + "⚠️ The graph's schema could not be checked, so repair is unavailable. Try again shortly." + ); + } else if (data.schema && !data.schema.compatible) { + setMigrationMessage( + "❌ This graph was created with an older version of GraphRAG and can't be repaired in place. Create a new graph and ingest its documents again." + ); + } else if ( !data.needs_repair && !data.embeddings_incomplete && !data.community_summaries_incomplete @@ -808,9 +818,15 @@ const KGAdmin = () => { // Capture the composite fingerprint (files + hint chips) so the // Extract button stays disabled until something changes. setExtractedFingerprint(sampleFingerprint); + // Samples that could not be read were left out of the draft; say which. + const unreadable: Array<{ file: string; error: string }> = convertData.failed_files || []; + const skippedNote = unreadable.length + ? ` ${unreadable.length} sample file(s) could not be read and were left out: ${unreadable.map((f) => f.error).join(" ")}` + : ""; setStatusMessage( `Draft schema ready (${data.summary?.vertex_count ?? "?"} vertex types, ` + - `${data.summary?.edge_count ?? "?"} edge types). Review/edit below, then click Initialize.` + `${data.summary?.edge_count ?? "?"} edge types). Review/edit below, then click Initialize.` + + skippedNote ); setStatusType("success"); } catch (error: any) { @@ -1460,7 +1476,7 @@ const KGAdmin = () => { Migration Assistant

- Check an existing graph against the current release — repair drifted queries and review prompt-override compatibility. + Check an existing graph against the current release — schema compatibility, outdated queries, prompt overrides and data health.

@@ -2857,8 +2873,8 @@ const KGAdmin = () => { Migration Assistant - Check an existing graph against the current release after upgrading graphrag: - repair drifted GSQL queries and review prompt-override compatibility — without recreating the graph. + Check an existing graph against the current release after upgrading GraphRAG: confirm its schema + is compatible, repair outdated queries and review prompt overrides — without recreating the graph. @@ -3073,7 +3089,7 @@ const KGAdmin = () => { > Close - {migrationStatus?.needs_repair && ( + {migrationStatus?.needs_repair && migrationStatus?.repairable && ( -
@@ -763,7 +836,7 @@ const McpServersConfig: React.FC = () => { server={s} onPatch={(patch) => patchRow(idx, patch)} onSave={async () => { const ok = await handleSave(); if (ok) setEditingIndex(null); }} - onCancel={() => removeRow(idx)} + onCancel={() => cancelEdit(idx)} isSaving={isSaving} testPassed={!!testResults[idx]?.ok} /> diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py index d6dea51..bf70978 100644 --- a/graphrag/app/agent/agent_generation.py +++ b/graphrag/app/agent/agent_generation.py @@ -29,7 +29,12 @@ class TigerGraphAgentGenerator: def __init__(self, llm_service): self.llm = llm_service svc_config = getattr(llm_service, "config", {}) - self.token_calculator = get_token_calculator(token_limit=svc_config.get("token_limit"), model_name=svc_config.get("llm_model")) + token_limit = ( + llm_service.context_token_limit() + if hasattr(llm_service, "context_token_limit") + else svc_config.get("token_limit") + ) + self.token_calculator = get_token_calculator(token_limit=token_limit, model_name=svc_config.get("llm_model")) def generate_answer(self, question: str, context: str | dict, query: str = "") -> dict: """Generate an answer based on the question and context. @@ -51,17 +56,6 @@ def generate_answer(self, question: str, context: str | dict, query: str = "") - if isinstance(context, dict): context = json.dumps(context, ensure_ascii=False) - # Truncate context to fit within token limit - if not self.token_calculator.is_unlimited_tokens(): - # Reserve tokens for question, query, and format instructions (approximately 1000 tokens) - max_context_tokens = self.token_calculator.get_max_context_tokens() - 1000 - - if len(context) > max_context_tokens: - context_tokens = self.token_calculator.count_tokens(context) - if context_tokens > max_context_tokens: - context = self.token_calculator.truncate_to_token_limit(context, max_context_tokens) - logger.info(f"Truncated context from {context_tokens} to {max_context_tokens} tokens") - answer_parser = PydanticOutputParser(pydantic_object=GraphRAGAnswerOutput) prompt = PromptTemplate( template=self.llm.chatbot_response_prompt, @@ -71,6 +65,12 @@ def generate_answer(self, question: str, context: str | dict, query: str = "") - } ) + # Trim the context so the full prompt fits the model's token limit. + if not self.token_calculator.is_unlimited_tokens(): + context = self.token_calculator.fit_context( + context, prompt.format(question=question, context="", query=query) + ) + try: generation = self.llm.invoke_with_parser( prompt, answer_parser, @@ -80,8 +80,9 @@ def generate_answer(self, question: str, context: str | dict, query: str = "") - # from the raw model output. on_parse_error=self.llm._salvage_answer_output, ) - except Exception: - logger.warning("generate_answer: generation failed") + except Exception as e: + logger.warning(f"generate_answer: generation failed: {type(e).__name__}: {e}") + logger.debug("generate_answer: generation failure traceback", exc_info=True) generation = GraphRAGAnswerOutput( generated_answer=( "I wasn't able to generate an answer for this question. " diff --git a/graphrag/app/agent/agentic_executor.py b/graphrag/app/agent/agentic_executor.py index f4f660d..9878003 100644 --- a/graphrag/app/agent/agentic_executor.py +++ b/graphrag/app/agent/agentic_executor.py @@ -28,6 +28,7 @@ from common.llm_services.base_llm import get_collected_usage from common.py_schemas import StepResult +from common.utils.retrieval_stats import retrieved_entries from tools import tool_registry as registry logger = logging.getLogger(__name__) @@ -37,16 +38,80 @@ _TRACE_FIELD_CAP = 12000 -def cap_for_trace(obj, limit: int = _TRACE_FIELD_CAP): - """Return ``obj`` unchanged if its JSON is under ``limit`` chars; else a - truncated-preview marker (kept JSON-valid for the Trace Logs UI).""" +def _json_size(obj) -> int: + # Count characters as written, not as \uXXXX escapes, so non-English + # text is neither overcounted nor squeezed out of the preview. + return len(json.dumps(obj, ensure_ascii=False, default=str)) + + +def _cap_strings(obj, max_len: int): + if isinstance(obj, str): + if len(obj) <= max_len: + return obj + return f"{obj[:max_len]}…(+{len(obj) - max_len:,} chars)" + if isinstance(obj, dict): + return {k: _cap_strings(v, max_len) for k, v in obj.items()} + if isinstance(obj, list): + return [_cap_strings(v, max_len) for v in obj] + return obj + + +def _cap_items(obj, max_items: int): + if isinstance(obj, dict): + return {k: _cap_items(v, max_items) for k, v in list(obj.items())[:max_items]} + if isinstance(obj, list): + return [_cap_items(v, max_items) for v in obj[:max_items]] + return obj + + +def fit_for_trace(obj, limit: int = _TRACE_FIELD_CAP): + """Return ``(value, shortened)`` for storing ``obj`` in the trace. + + ``value`` is the tool's own data with its shape intact — long strings cut + first, then long lists and maps — so the Trace Logs view still shows + real tool output, never a wrapper. ``shortened`` is ``None`` when nothing + was cut, else ``{"full_chars": N, "shown_chars": M}`` for the trace to + report alongside the value. + """ try: - s = json.dumps(obj, default=str) + plain = json.loads(json.dumps(obj, default=str)) except Exception: - s = str(obj) - if len(s) <= limit: - return obj - return {"_truncated": True, "chars": len(s), "preview": s[:limit]} + plain = str(obj) + full = _json_size(plain) + if full <= limit: + return obj, None + value = plain + for max_len in (4000, 2000, 1000, 500, 250, 120, 60): + value = _cap_strings(plain, max_len) + if _json_size(value) <= limit: + break + else: + for max_items in (50, 20, 10, 5, 2, 1): + value = _cap_items(_cap_strings(plain, 60), max_items) + if _json_size(value) <= limit: + break + return value, {"full_chars": full, "shown_chars": _json_size(value)} + + +def trace_note(cut: dict) -> str: + """The note recorded next to a value ``fit_for_trace`` shortened.""" + return ( + f"Truncated from {cut['full_chars']:,} to {cut['shown_chars']:,} " + "characters due to the trace log size limit." + ) + + +def with_trace_note(value, cut): + """``value`` with a ``note`` key when it was shortened and is a map.""" + if cut and isinstance(value, dict): + return {**value, "note": trace_note(cut)} + return value + + +def cap_for_trace(obj, limit: int = _TRACE_FIELD_CAP): + """The trace-sized value of ``obj``, for callers with nowhere to record + that it was shortened.""" + return fit_for_trace(obj, limit)[0] def retrieved_chunk_ids(context) -> list: @@ -64,7 +129,7 @@ def retrieved_chunk_ids(context) -> list: fr = inner.get("final_retrieval") if isinstance(inner, dict) else None if not isinstance(fr, dict): return [] - return [k for k in fr.keys() if k != "Similarity_Context"] + return [k for k in retrieved_entries(fr) if k != "Similarity_Context"] def _usage_since(start_idx: int) -> dict: @@ -134,15 +199,19 @@ def _run_step(step, args, ctx, results, traces): # Trace output carries the one-line summary AND the actual result, so # the Trace Logs detail view shows what each step returned (not just a # status line). Input is the resolved tool args. + # Input and output keep the tool's own data; a value shortened to fit the + # trace gets a ``note`` saying so next to it. trace_output = {"summary": out.get("summary", "")} if out.get("context") is not None: - trace_output["result"] = cap_for_trace(out.get("context")) + trace_output["result"], cut = fit_for_trace(out.get("context")) + trace_output = with_trace_note(trace_output, cut) + trace_input = with_trace_note(*fit_for_trace(args)) traces.append({ "node": f"{step.id}: {step.tool}", "kind": step.kind, "tool": step.tool, "duration_s": duration, - "input": cap_for_trace(args), + "input": trace_input, "output": trace_output, "rationale": step.rationale or "", "usage": _usage_since(usage_start), diff --git a/graphrag/app/agent/agentic_graph.py b/graphrag/app/agent/agentic_graph.py index 7a88c02..46aa0b5 100644 --- a/graphrag/app/agent/agentic_graph.py +++ b/graphrag/app/agent/agentic_graph.py @@ -24,7 +24,7 @@ import logging import time -from agent.agentic_executor import cap_for_trace, execute_plan, _run_step, _usage_since +from agent.agentic_executor import execute_plan, fit_for_trace, with_trace_note, _run_step, _usage_since from agent.agentic_planner import plan_question from agent.agentic_synthesizer import _gather, has_context, synthesize from common.llm_services.base_llm import get_collected_usage @@ -138,12 +138,12 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: _u0 = len(get_collected_usage() or []) _t0 = time.time() resp = synthesize(llm, question, results, plan=plan, conversation=conversation) + # Input is the combined context fed to the answer LLM (what actually + # grounded the answer); output is the answer + citations. agent_steps.append({ "node": "synthesize", "kind": "answer", "duration_s": round(time.time() - _t0, 3), - # Input is the combined context fed to the answer LLM (what actually - # grounded the answer); output is the answer + citations. - "input": cap_for_trace(_gather(results)), + "input": with_trace_note(*fit_for_trace(_gather(results))), "output": { "answer": resp.natural_language_response, "citations": (resp.query_sources or {}).get("citations", []), diff --git a/graphrag/app/agent/agentic_planner.py b/graphrag/app/agent/agentic_planner.py index 69d58d0..8d6c578 100644 --- a/graphrag/app/agent/agentic_planner.py +++ b/graphrag/app/agent/agentic_planner.py @@ -126,7 +126,18 @@ def plan_question(llm, question, conversation=None, schema_rep="", prior_results # editable "Additional Instructions" portion); the default lives in base_llm. messages = [("system", llm.agentic_planner_prompt), ("user", "\n\n".join(user_parts))] try: - plan = llm.invoke_structured(messages, Plan, caller_name="agentic_plan") + # Plan steps carry free-form ``args`` maps. Where a provider's default + # structured output is strict JSON schema (the OpenAI family), it + # rejects them, so every plan paid for a failed request before falling + # back to parsing (GML-2302). Tool calling accepts them there; other + # providers keep their default. + plan = llm.invoke_structured( + messages, Plan, caller_name="agentic_plan", + method=( + "function_calling" + if getattr(llm, "strict_structured_output", False) else None + ), + ) except Exception as exc: logger.warning(f"planner failed ({exc}); falling back to single hybrid step") from common.py_schemas import PlanStep diff --git a/graphrag/app/agent/agentic_synthesizer.py b/graphrag/app/agent/agentic_synthesizer.py index 7b8e241..af1d057 100644 --- a/graphrag/app/agent/agentic_synthesizer.py +++ b/graphrag/app/agent/agentic_synthesizer.py @@ -25,25 +25,68 @@ from agent.agent_generation import TigerGraphAgentGenerator from agent.agentic_executor import retrieved_chunk_ids from common.py_schemas import GraphRAGResponse +from common.utils.retrieval_stats import is_grouped, retrieved_entries logger = logging.getLogger(__name__) -def _gather(results: dict) -> dict: - """Collect non-empty step contexts into a combined context block.""" +def _without_seen(ctx, seen: set): + """``ctx`` minus the retrieved entries already in ``seen`` (which it then + extends), or ``None`` when nothing new is left. ``ctx`` is not modified.""" + result = ctx.get("result") if isinstance(ctx, dict) else None + fr = result.get("final_retrieval") if isinstance(result, dict) else None + if not isinstance(fr, dict): + return ctx + kept = {} + for key, value in fr.items(): + if is_grouped(value): + fresh = {k: v for k, v in value.items() if k not in seen} + seen.update(fresh) + if fresh: + kept[key] = fresh + elif key == "Similarity_Context" or key not in seen: + seen.add(key) + kept[key] = value + if not kept: + return None + return {**ctx, "result": {**result, "final_retrieval": kept}} + + +def _gather(results: dict, log: bool = False) -> dict: + """Collect non-empty step contexts into a combined context block. + + Document passages already returned by an earlier step are left out of + later ones, so the answer model reads (and is billed for) each passage + once; a step left with nothing new is dropped. + """ structural, unstructured = [], [] + seen: set = set() + skipped = 0 for sr in results.values(): if not sr.ok or sr.context is None: continue ctx = sr.context fc = ctx.get("function_call") if isinstance(ctx, dict) else None if fc and "Vector_Search" in str(fc): - unstructured.append(ctx) + before = len(retrieved_entries(_final_retrieval(ctx))) + ctx = _without_seen(ctx, seen) + after = len(retrieved_entries(_final_retrieval(ctx))) if ctx else 0 + skipped += before - after + if ctx is not None: + unstructured.append(ctx) else: structural.append(ctx) + if skipped and log: + logger.info(f"synthesize: left out {skipped} passage(s) already retrieved by an earlier step") return {"structural": structural, "unstructured": unstructured} +def _final_retrieval(ctx) -> dict: + result = ctx.get("result") if isinstance(ctx, dict) else None + fr = result.get("final_retrieval") if isinstance(result, dict) else None + return fr if isinstance(fr, dict) else {} + + def has_context(results: dict) -> bool: g = _gather(results) return bool(g["structural"] or g["unstructured"]) @@ -51,7 +94,7 @@ def has_context(results: dict) -> bool: def synthesize(llm, question, results: dict, plan=None, conversation=None) -> GraphRAGResponse: """Produce the final answer from gathered step contexts.""" - combined = _gather(results) + combined = _gather(results, log=True) generator = TigerGraphAgentGenerator(llm) answer = generator.generate_answer(question, combined) diff --git a/graphrag/app/routers/inquiryai.py b/graphrag/app/routers/inquiryai.py index 7d37211..c71bba2 100644 --- a/graphrag/app/routers/inquiryai.py +++ b/graphrag/app/routers/inquiryai.py @@ -5,6 +5,7 @@ from typing import Annotated, List, Union from agent.agent import make_agent +import requests from fastapi import (APIRouter, Depends, HTTPException, Request, status) from fastapi.security.http import HTTPBase @@ -42,6 +43,42 @@ def _caller_is_superadmin(credentials) -> bool: return False +def _require_graph_access(request: Request, graphname: str): + """Return the caller's connection, having proven it can reach *graphname*. + + ``auth_middleware`` builds the connection object lazily: on the Basic-auth + path nothing contacts the database, so bad credentials are not rejected + until the connection is first used. A route that reads only + service-credentialed resources would therefore never authenticate its + caller at all (GML-2197). + + Reading the graph's schema forces that check and is scoped to this graph — + unlike ``conn.echo()``, which pings ``/echo`` and succeeds for any valid + user regardless of which graph the request names. + """ + conn = getattr(request.state, "conn", None) + if conn is None: + # No Authorization header: the middleware never built a connection. + raise HTTPException(status_code=401, detail="Invalid credentials") + try: + conn.getVertexTypes() + except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e: + # The database is unreachable — not a verdict on the credentials. + raise HTTPException( + status_code=503, detail="The database is unavailable. Try again shortly." + ) from e + except requests.exceptions.HTTPError as e: + # A 5xx (nginx answers 502/503 while GSQL restarts) is an outage too. + if getattr(e.response, "status_code", 0) >= 500: + raise HTTPException( + status_code=503, detail="The database is unavailable. Try again shortly." + ) from e + raise HTTPException(status_code=401, detail="Invalid credentials") from e + except Exception as e: + raise HTTPException(status_code=401, detail="Invalid credentials") from e + return conn + + # Response fields returned only when the caller opts in via include_fields. # The answer envelope (natural_language_response, answered_question, # response_type) is always returned; query_sources carries the heavier @@ -132,7 +169,8 @@ def retrieve_answer( return _apply_field_selection(resp, query.include_fields) -conversation_history = [] +# Most recent caller-supplied turns passed to the agent (GML-2194). +_MAX_HISTORY_TURNS = 3 # TODO: This could be merged with /{graphname}/query endpoints, all agents can be refactored in seperated function or file @@ -143,8 +181,6 @@ def retrieve_answer_with_chathistory( conn: Request, credentials: Annotated[HTTPBase, Depends(security)], ) -> GraphRAGResponse: - global conversation_history - conn = conn.state.conn logger.debug_pii( f"/{graphname}/query_with_history request_id={req_id_cv.get()} question={query.query}" @@ -160,17 +196,17 @@ def retrieve_answer_with_chathistory( natural_language_response="", answered_question=False, response_type="inquiryai" ) try: - # Retrieve latest 3 Q&A pairs in full conversation history - latest_history = conversation_history[-3:] + # Conversation context is supplied by the caller: this endpoint keeps + # no server-side history, so each caller sees only its own turns. latest_history_query = [ - { - "query": interaction.get("query", ""), - "response": interaction.get("response", ""), - } - for interaction in latest_history + {"query": turn.query, "response": turn.response} + for turn in (query.history or [])[-_MAX_HISTORY_TURNS:] ] - logger.info(f"latest 3 pairs of queries: {latest_history_query}") + logger.debug_pii( + f"/{graphname}/query_with_history request_id={req_id_cv.get()} " + f"history_turns={len(latest_history_query)} {latest_history_query}" + ) resp = agent.question_for_agent(query.query, latest_history_query) @@ -185,10 +221,6 @@ def retrieve_answer_with_chathistory( pmetrics.llm_success_response_total.labels(get_embedding_service().model_name).inc() - conversation_history.append( - {"query": query.query, "response": resp.natural_language_response} - ) - except MapQuestionToSchemaException: resp.natural_language_response = ( "A schema mapping error occurred. Please try rephrasing your question." @@ -516,14 +548,19 @@ def delete_docs( def retrieve_docs( graphname, query: NaturalLanguageQuery, + request: Request, credentials: Annotated[HTTPBase, Depends(security)], top_k: int = 3, ): logger.debug_pii( f"/{graphname}/retrieve_docs request_id={req_id_cv.get()} top_k={top_k} question={query.query}" ) + _require_graph_access(request, graphname) check_embedding_store_status() - return get_embedding_store().retrieve_similar( + # Search the graph named in the path. Without an argument this returns the + # default store, so results came from the configured default graph no + # matter which graph the caller asked for (GML-2197). + return get_embedding_store(graphname).retrieve_similar( get_embedding_service().embed_query(query.query), top_k=top_k ) diff --git a/graphrag/app/routers/supportai.py b/graphrag/app/routers/supportai.py index 196fc2c..cf93bb9 100644 --- a/graphrag/app/routers/supportai.py +++ b/graphrag/app/routers/supportai.py @@ -44,6 +44,7 @@ SupportAIMethod, SupportAIQuestion, ) +from common.utils.text_extractors import ExtractionError logger = logging.getLogger(__name__) router = APIRouter(tags=["SupportAI"]) @@ -95,6 +96,9 @@ def create_ingest( return supportai.create_ingest(graphname, cfg, conn) except HTTPException: raise + except ExtractionError as e: + # None of the files could be read — a problem with the input. + raise HTTPException(status_code=400, detail=str(e)) except Exception as e: logger.error(f"create_ingest failed for graph '{graphname}': {e}", exc_info=True) raise HTTPException(status_code=500, detail=f"Ingest preparation failed: {str(e)}") diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py index a482e90..29717e0 100644 --- a/graphrag/app/routers/ui.py +++ b/graphrag/app/routers/ui.py @@ -56,7 +56,12 @@ from common.db.connections import get_db_connection_pwd_manual from common.db import schema_utils as schema_utils_mod from common.db import schema_extraction as schema_extraction_mod -from common.utils.text_extractors import TextExtractor +from common.utils.text_extractors import ( + ExtractionError, + TextExtractor, + describe_failures, + failed_extractions, +) from common.logs.log import req_id_cv from common.logs.logwriter import LogWriter from common.metrics.prometheus_metrics import metrics as pmetrics @@ -1284,8 +1289,10 @@ def migration_status( Read-only — does NOT modify the graph. Pair with ``POST /migration/apply`` to actually repair. - Schema-attribute drift is reported as ``{}`` for now; the detection - is stubbed in ``common.db.migrate.check_and_apply_schema``. + ``schema.compatible`` says whether the graph has the base schema the + shipped queries are written against; query repair is offered only when + it does (graphs created before 1.4 do not and need recreating). + ``schema.check_failed`` means the schema could not be read. """ from common.db.migrate import _gsql_hash, get_installed_query_names, get_installed_query_body @@ -1294,6 +1301,8 @@ def migration_status( cred_obj = creds[1] conn = get_db_connection_pwd_manual(graphname, cred_obj.username, cred_obj.password) + schema_compatible, schema_check_failed = _schema_compatibility(conn, graphname) + # Read live domain schema once for templated-retriever rendering. domain_vts, domain_edges, include_entity = _get_domain_schema_for_render(conn, graphname) @@ -1407,10 +1416,8 @@ def migration_status( "missing_files": missing_files, }, "schema": { - # Populated from common.db.migrate.check_and_apply_schema once - # attribute additions are tracked; empty until then. - "missing_attributes": {}, - "schema_change_required": False, + "compatible": schema_compatible, + "check_failed": schema_check_failed, }, "prompts": prompt_issues, "embeddings": { @@ -1423,9 +1430,29 @@ def migration_status( community_summaries.get("needs_resummarize", 0) ), "needs_repair": bool(outdated) or bool(not_installed) or bool(prompt_issues), + "repairable": schema_compatible, } +def _schema_compatibility(conn, graphname) -> tuple[bool, bool]: + """``(compatible, check_failed)`` for the graph's schema; differences are + logged, not returned, since they are only useful to an administrator + reading the server log.""" + from common.db.migrate import check_schema_compatibility + + try: + result = check_schema_compatibility(conn) + except Exception as e: + logger.warning(f"migration: schema check for {graphname} failed: {e}") + return False, True + if not result.compatible: + logger.warning( + f"migration: graph {graphname} does not match the current base schema " + f"and cannot be repaired in place: {'; '.join(result.differences)}" + ) + return result.compatible, False + + async def _proxy_regen(graphname, creds, action): """Forward a targeted regenerate action to ECC (synchronous) and return its counts. Refuses while a rebuild is in flight (both write embeddings).""" @@ -1485,6 +1512,10 @@ def migration_apply( per-category opt-in. When neither list is provided the endpoint detects the repair set itself. Only shipped query names are honored. + Refuses with 409 a graph whose schema lacks the shipped base schema + (created before 1.4), and with 503 when the schema cannot be read, so + queries are never replaced with ones the graph cannot run. + Acquires the per-graph lock for the duration of the repair so that a concurrent ingest / rebuild / schema-extraction on the same graph cannot race against the CREATE OR REPLACE + INSTALL QUERY @@ -1531,6 +1562,14 @@ def migration_apply( try: cred_obj = creds[1] conn = get_db_connection_pwd_manual(graphname, cred_obj.username, cred_obj.password) + compatible, check_failed = _schema_compatibility(conn, graphname) + if check_failed: + raise HTTPException( + status_code=503, + detail="The graph's schema could not be checked. Try again shortly.", + ) + if not compatible: + raise HTTPException(status_code=409, detail=_INCOMPATIBLE_SCHEMA_MESSAGE) return _migration_apply_inner( graphname, conn, queries_outdated=queries_outdated, @@ -1541,6 +1580,12 @@ def migration_apply( release_graph_lock(graphname, "migration") +_INCOMPATIBLE_SCHEMA_MESSAGE = ( + "This graph was created with an older version of GraphRAG and can't be " + "repaired in place. Create a new graph and ingest its documents again." +) + + def _migration_apply_inner( graphname: str, conn, @@ -2116,15 +2161,28 @@ async def convert_sample_files( detail=f"Text extraction failed: {exc}", ) + # A file that failed to convert has no JSONL, so naming it in the + # schema-extraction step would fail that whole step. Report it here and + # hand on only the files that converted. + failures = failed_extractions(result) + converted = [n for n in saved_basenames if n not in failures] + if failures and not converted: + raise HTTPException(status_code=400, detail=describe_failures(failures)) + LogWriter.info( f"Converted sample files for {graphname}: " - f"{len(accepted)} uploaded, {result.get('num_documents', 0)} docs in JSONL" + f"{len(accepted)} uploaded, {result.get('num_documents', 0)} docs in JSONL, " + f"{len(failures)} failed" ) return { "status": "success", "graphname": graphname, - "saved_files": list(saved_basenames), + "saved_files": converted, "skipped_files": sorted(skip_set), + "failed_files": [ + {"file": name, "error": reason} + for name, reason in sorted(failures.items()) + ], "num_documents": result.get("num_documents", 0), } finally: @@ -2493,6 +2551,9 @@ def create_ingest( except HTTPException: raise + except ExtractionError as e: + # None of the uploaded files could be read — a problem with the input. + raise HTTPException(status_code=400, detail=str(e)) except Exception as e: LogWriter.error(f"Error creating ingest configuration for graph {graphname}: {str(e)}") raise HTTPException( diff --git a/graphrag/app/supportai/retrievers/BaseRetriever.py b/graphrag/app/supportai/retrievers/BaseRetriever.py index e6e8697..e782af7 100644 --- a/graphrag/app/supportai/retrievers/BaseRetriever.py +++ b/graphrag/app/supportai/retrievers/BaseRetriever.py @@ -4,6 +4,7 @@ from common.llm_services.base_llm import LLM_Model from common.py_schemas import CandidateScore, CandidateGenerator, GraphRAGAnswerOutput, CommunityAnswer from common.utils.token_calculator import get_token_calculator +from common.utils.retrieval_stats import describe_retrieval from common.config import get_chat_config, get_embedding_store, get_graphrag_config from langchain_core.output_parsers import StrOutputParser, PydanticOutputParser @@ -36,7 +37,19 @@ def __init__( # Use llm_service's own config when available (chatbot path); # fall back to get_chat_config() (direct supportai API path). llm_cfg = getattr(llm_service, "config", None) or get_chat_config() - self.token_calculator = get_token_calculator(token_limit=llm_cfg.get("token_limit"), model_name=llm_cfg.get("llm_model")) + token_limit = ( + llm_service.context_token_limit() + if hasattr(llm_service, "context_token_limit") + else llm_cfg.get("token_limit") + ) + self.token_calculator = get_token_calculator(token_limit=token_limit, model_name=llm_cfg.get("llm_model")) + + def _log_retrieval(self, query_name, res): + """Log how many chunks a retrieval query returned and how big they are.""" + result = res[0] if isinstance(res, list) and res else res + retrieved = result.get("final_retrieval") if isinstance(result, dict) else None + if isinstance(retrieved, dict): + self.logger.info(describe_retrieval(query_name, retrieved)) def _install_query(self, query_name): self.logger.info(f"Installing query {query_name}") @@ -144,23 +157,19 @@ def _score_candidates(self, question, contexts, top_k=None): return [{"candidate_answer": x.answer, "score": x.quality_score} for x in results] def _generate_response(self, question, retrieved, query = "", verbose = False): - # Truncate retrieved sources to fit within token limit - if not self.token_calculator.is_unlimited_tokens(): - # Reserve tokens for question, query, and format instructions (approximately 1000 tokens) - max_context_tokens = self.token_calculator.get_max_context_tokens() - 1000 - - if len(retrieved) > max_context_tokens: - retrieved_tokens = self.token_calculator.count_tokens(retrieved) - if retrieved_tokens > max_context_tokens: - retrieved = self.token_calculator.truncate_to_token_limit(retrieved, max_context_tokens) - self.logger.info(f"Truncated retrieved text from {retrieved_tokens} to {max_context_tokens} tokens") - response_parser = PydanticOutputParser(pydantic_object=GraphRAGAnswerOutput) # {format_instructions} lives in the (hardcoded) system prompt; bind it # as a partial, consistent with the other prompts (see base_llm A1b). prompt = ChatPromptTemplate.from_template( self.llm_service.chatbot_response_prompt ).partial(format_instructions=response_parser.get_format_instructions()) + + # Trim retrieved sources so the full prompt fits the model's token limit. + if not self.token_calculator.is_unlimited_tokens(): + retrieved = self.token_calculator.fit_context( + retrieved, prompt.format(question=question, context="", query=query) + ) + input_vars = { "question": question, "context": retrieved, "query": query, } diff --git a/graphrag/app/supportai/retrievers/CommunityRetriever.py b/graphrag/app/supportai/retrievers/CommunityRetriever.py index 2582274..eb1a4f3 100644 --- a/graphrag/app/supportai/retrievers/CommunityRetriever.py +++ b/graphrag/app/supportai/retrievers/CommunityRetriever.py @@ -87,6 +87,7 @@ def search(self, question, community_level: int, top_k: int = 5, similarity_thre self.logger.info(f"Retrived GraphRAG query verbose info: {verbose_info}") if expand: res[1]["verbose"]["expanded_questions"] = questions + self._log_retrieval("GraphRAG_Community_Vector_Search", res) return res def retrieve_answer(self, diff --git a/graphrag/app/supportai/retrievers/HybridRetriever.py b/graphrag/app/supportai/retrievers/HybridRetriever.py index 140bf5f..462cbb8 100644 --- a/graphrag/app/supportai/retrievers/HybridRetriever.py +++ b/graphrag/app/supportai/retrievers/HybridRetriever.py @@ -97,6 +97,7 @@ def search(self, question, indices, top_k=1, similarity_threshold=0.90, num_hops self.logger.info(f"Retrived HybridSearch query verbose info: {verbose_info}") if expand: res[1]["verbose"]["expanded_questions"] = questions + self._log_retrieval("GraphRAG_Hybrid_Vector_Search", res) return res def retrieve_answer(self, question, index, top_k=1, similarity_threshold=0.90, num_hops=2, num_seen_min=1, expand: bool = False, method: str = "similarity", chunk_only: bool = False, doc_only: bool = False, combine: bool = False, verbose: bool = False, max_results: int = 0): diff --git a/graphrag/app/supportai/retrievers/SiblingRetriever.py b/graphrag/app/supportai/retrievers/SiblingRetriever.py index 2e66f01..8657db6 100644 --- a/graphrag/app/supportai/retrievers/SiblingRetriever.py +++ b/graphrag/app/supportai/retrievers/SiblingRetriever.py @@ -13,7 +13,11 @@ def __init__( ): super().__init__(embedding_service, embedding_store, llm_service, connection) - def search(self, question, index, top_k=1, lookback=3, lookahead=3, expand=False, withHyDE=False, verbose=False): + # Flag order matches SimilarityRetriever.search and retrieve_answer below. + # These two were declared reversed here while both positional callers passed + # SimilarityRetriever's order, so every call transposed them (GML-2198) — + # and they select different installed queries, not just a ranking tweak. + def search(self, question, index, top_k=1, lookback=3, lookahead=3, withHyDE=False, expand=False, verbose=False): if expand: questions = self._expand_question(question, top_k, verbose) verbose and self.logger.info(f"Expanded questions to use: {questions}") @@ -57,6 +61,7 @@ def search(self, question, index, top_k=1, lookback=3, lookahead=3, expand=False self.logger.info(f"Retrived SiblingSearch query verbose info: {verbose_info}") if expand: res[1]["verbose"]["expanded_questions"] = questions + self._log_retrieval("Chunk_Sibling_Vector_Search", res) return res def retrieve_answer( diff --git a/graphrag/app/supportai/retrievers/SimilarityRetriever.py b/graphrag/app/supportai/retrievers/SimilarityRetriever.py index e5f948c..3b0244b 100644 --- a/graphrag/app/supportai/retrievers/SimilarityRetriever.py +++ b/graphrag/app/supportai/retrievers/SimilarityRetriever.py @@ -53,6 +53,7 @@ def search(self, question, index, top_k=1, withHyDE=False, expand=False, verbose self.logger.info(f"Retrived SimilaritySearch query verbose info: {verbose_info}") if expand: res[1]["verbose"]["expanded_questions"] = questions + self._log_retrieval("Content_Similarity_Vector_Search", res) return res def retrieve_answer(self, question, index, top_k=1, withHyDE=False, expand=False, combine=False, verbose=False): diff --git a/graphrag/app/supportai/supportai.py b/graphrag/app/supportai/supportai.py index 89bb6c1..548f79a 100644 --- a/graphrag/app/supportai/supportai.py +++ b/graphrag/app/supportai/supportai.py @@ -17,7 +17,12 @@ # SupportAIMethod, # SupportAIQuestion, ) -from common.utils.text_extractors import TextExtractor +from common.utils.text_extractors import ( + ExtractionError, + TextExtractor, + describe_failures, + failed_extractions, +) logger = logging.getLogger(__name__) def init_supportai(conn: TigerGraphConnection, graphname: str) -> tuple[dict, dict]: @@ -510,6 +515,10 @@ def create_ingest( raise Exception(f"Server folder processing failed: {server_processing_result}") doc_count = server_processing_result.get("num_documents", 0) + failures = failed_extractions(server_processing_result) + if failures and not doc_count: + # Nothing converted, so ingest would find no documents; say why. + raise ExtractionError(describe_failures(failures)) logger.info(f"Server folder processing completed: {server_processing_result.get('message')}") res_ingest_config["data_path"] = temp_folder @@ -518,6 +527,21 @@ def create_ingest( # Use a placeholder path to indicate temp storage res["data_path"] = "in_temp_storage" res["data_source_id"] = res_ingest_config + res["failed_files"] = [ + {"file": name, "error": reason} + for name, reason in sorted(failures.items()) + ] + # Ingest loads one JSONL per converted file across the whole folder, + # so count distinct JSONLs rather than the latest upload or the source + # files (two sources sharing a stem produce a single JSONL). + res["ready_files"] = len({ + f.get("jsonl_file") + for f in server_processing_result.get("files", []) + if f.get("status") == "success" and f.get("jsonl_file") + and f.get("num_documents", 0) > 0 # e.g. a decorative image yields none + }) + except ExtractionError: + raise except Exception as e: logger.error(f"Server folder processing failed for graph '{graphname}', path '{data_path}': {e}", exc_info=True) raise Exception(f"Error during server folder processing: {e}") @@ -704,7 +728,8 @@ def ingest( # Get all JSONL files from temp folder jsonl_files = [f for f in os.listdir(data_path) if f.endswith('.jsonl')] if not jsonl_files: - raise Exception(f"No JSONL files found in: {data_path}") + logger.error(f"No JSONL files found in: {data_path}") + raise Exception("No readable documents were found to ingest.") logger.info(f"Found {len(jsonl_files)} JSONL files to ingest from: {data_path}") # Ensure loading job exists — recreate if missing (e.g. after schema drop) diff --git a/graphrag/app/tools/graphrag_tools.py b/graphrag/app/tools/graphrag_tools.py index ca4b9f9..cf60fd9 100644 --- a/graphrag/app/tools/graphrag_tools.py +++ b/graphrag/app/tools/graphrag_tools.py @@ -32,6 +32,7 @@ from dataclasses import dataclass, field from typing import Any, Callable, Dict, Optional +from common.utils.retrieval_stats import retrieval_stats from tools import tool_guards as guards from tools.validation_utils import MapQuestionToSchemaException @@ -294,11 +295,19 @@ def community_search( def _unstructured_result(query_name: str, step) -> dict: result = step[0] if isinstance(step, (list, tuple)) and step else step - if _result_is_empty(result): + retrieved = result.get("final_retrieval") if isinstance(result, dict) else None + if _result_is_empty(result) or retrieved == {}: return _empty(f"{query_name} returned no chunks") - n = len(result) if hasattr(result, "__len__") else "?" - return _ok(f"{query_name} returned {n} item(s)", - {"function_call": query_name, "result": result}) + if not isinstance(retrieved, dict): + return _ok(f"{query_name} returned results", + {"function_call": query_name, "result": result}) + # Count what the answer will see; non-chunk entries (entity text, + # community summaries, whole documents) are reported separately. + stats = retrieval_stats(retrieved) + summary = f"{query_name} returned {stats.chunks} chunk(s)" + if stats.others: + summary += f" and {stats.others} other entr{'y' if stats.others == 1 else 'ies'}" + return _ok(summary, {"function_call": query_name, "result": result}) # -------------------------------------------------------------------------- diff --git a/graphrag/tests/test_agentic_plan_structured_output.py b/graphrag/tests/test_agentic_plan_structured_output.py new file mode 100644 index 0000000..97c3c28 --- /dev/null +++ b/graphrag/tests/test_agentic_plan_structured_output.py @@ -0,0 +1,156 @@ +import ast +import os +import types +import unittest + +_ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "..")) +_PLANNER = os.path.join(_ROOT, "graphrag", "app", "agent", "agentic_planner.py") +_AGENT = os.path.join(_ROOT, "graphrag", "app", "agent", "agentic_agent.py") + + +def _structured_calls(path): + """(schema name, method keyword or None) for each invoke_structured call.""" + with open(path, encoding="utf-8") as handle: + tree = ast.parse(handle.read()) + calls = [] + for node in ast.walk(tree): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "invoke_structured" + ): + schema = ast.unparse(node.args[1]) if len(node.args) > 1 else None + method = next( + (ast.unparse(k.value) for k in node.keywords if k.arg == "method"), + None, + ) + calls.append((schema, method)) + return calls + + +class TestPlannerRequestsToolCalling(unittest.TestCase): + """``Plan`` steps carry free-form ``args`` maps, which OpenAI's default + strict JSON-schema output rejects — so every plan made a failed request + before falling back to parsing (GML-2302).""" + + def test_planner_asks_for_tool_calling_only_where_output_is_strict(self): + """Tool calling is verified for the OpenAI family only; other providers + keep their default rather than risk a new failure there.""" + methods = [m for schema, m in _structured_calls(_PLANNER) if schema == "Plan"] + self.assertEqual(len(methods), 1) + self.assertIn("function_calling", methods[0]) + self.assertIn("strict_structured_output", methods[0]) + + def test_triage_keeps_the_default(self): + """Its schema satisfies strict mode, which enforces it exactly.""" + calls = [c for c in _structured_calls(_AGENT) if c[0] == "_Triage"] + self.assertTrue(calls, "triage call not found") + self.assertTrue(all(method is None for _, method in calls)) + + +class TestStrictProviders(unittest.TestCase): + """Only providers whose structured output is strict JSON schema opt in.""" + + @classmethod + def setUpClass(cls): + try: + from common.llm_services.base_llm import LLM_Model + from common.llm_services.openai_service import OpenAI + from common.llm_services.azure_openai_service import AzureOpenAI + except ImportError as exc: + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.base, cls.openai, cls.azure = LLM_Model, OpenAI, AzureOpenAI + + def test_openai_family_is_strict(self): + self.assertTrue(self.openai.strict_structured_output) + self.assertTrue(self.azure.strict_structured_output) + + def test_default_is_not_strict(self): + self.assertFalse(self.base.strict_structured_output) + + def test_other_providers_keep_the_default(self): + import importlib + for module, cls in ( + ("common.llm_services.google_genai_service", "GoogleGenAI"), + ("common.llm_services.google_vertexai_service", "GoogleVertexAI"), + ("common.llm_services.aws_bedrock_service", "AWSBedrock"), + ("common.llm_services.ollama", "Ollama"), + ("common.llm_services.groq_llm_service", "Groq"), + ): + with self.subTest(provider=cls): + try: + provider = getattr(importlib.import_module(module), cls) + except ImportError: + continue # optional provider SDK not installed + self.assertFalse(provider.strict_structured_output) + + +class TestInvokeStructuredPassesMethod(unittest.TestCase): + """``invoke_structured`` forwards ``method`` only when one is given.""" + + @classmethod + def setUpClass(cls): + try: + from common.llm_services.base_llm import LLM_Model + from common.py_schemas.schemas import Plan + except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.invoke = staticmethod(LLM_Model.invoke_structured) + cls.Plan = Plan + + def _run(self, method): + seen = {} + plan = self.Plan() + + class FakeLLM: + def with_structured_output(self, schema, **kwargs): + seen.update(kwargs) + return types.SimpleNamespace(invoke=lambda messages: plan) + + host = types.SimpleNamespace(llm=FakeLLM(), config={}) + result = self.invoke(host, [("user", "q")], self.Plan, caller_name="t", method=method) + self.assertIs(result, plan) + return seen + + def test_no_method_keeps_the_provider_default(self): + self.assertEqual(self._run(None), {}) + + def test_method_is_forwarded(self): + self.assertEqual(self._run("function_calling"), {"method": "function_calling"}) + + +class TestOpenAIRequestShape(unittest.TestCase): + """Built offline with a dummy key: the request that would go to OpenAI.""" + + @classmethod + def setUpClass(cls): + try: + from langchain_openai import ChatOpenAI + from common.py_schemas.schemas import Plan + except ImportError as exc: + raise unittest.SkipTest(f"langchain-openai unavailable: {exc}") + cls.llm = ChatOpenAI(model="gpt-4o-mini", api_key="sk-test") + cls.Plan = Plan + + @staticmethod + def _bound_kwargs(runnable): + first = getattr(runnable, "first", runnable) + return getattr(first, "kwargs", {}) + + def test_function_calling_sends_no_strict_schema(self): + kwargs = self._bound_kwargs( + self.llm.with_structured_output(self.Plan, method="function_calling") + ) + tools = kwargs.get("tools") or [] + self.assertTrue(tools, f"expected a bound tool, got {sorted(kwargs)}") + self.assertNotEqual(tools[0].get("function", {}).get("strict"), True) + self.assertNotIn("response_format", kwargs) + + def test_default_takes_the_strict_response_format_path(self): + """The path that rejects free-form maps — what the planner now avoids.""" + kwargs = self._bound_kwargs(self.llm.with_structured_output(self.Plan)) + self.assertIn("response_format", kwargs) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_answer_context_dedup.py b/graphrag/tests/test_answer_context_dedup.py new file mode 100644 index 0000000..3ad7e28 --- /dev/null +++ b/graphrag/tests/test_answer_context_dedup.py @@ -0,0 +1,72 @@ +import copy +import unittest + +try: + from agent.agentic_synthesizer import _gather + from common.py_schemas import StepResult +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +def _step(sid, query, final_retrieval): + return StepResult(step_id=sid, ok=True, summary="", citations=[], context={ + "function_call": query, "result": {"final_retrieval": final_retrieval, "edges": []}, + }) + + +def _fr(ctx): + return ctx["result"]["final_retrieval"] + + +class TestAnswerContextDedup(unittest.TestCase): + """Each passage reaches the answer model once, whichever step found it.""" + + def test_overlapping_searches_keep_the_first_copy(self): + results = { + "S1": _step("S1", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_1": ["one"], "a_chunk_2": ["two"]}), + "S2": _step("S2", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_2": ["two"], "a_chunk_3": ["three"]}), + } + unstructured = _gather(results)["unstructured"] + self.assertEqual(_fr(unstructured[0]), {"a_chunk_1": ["one"], "a_chunk_2": ["two"]}) + self.assertEqual(_fr(unstructured[1]), {"a_chunk_3": ["three"]}) + + def test_step_with_nothing_new_is_dropped(self): + results = { + "S1": _step("S1", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_1": ["one"]}), + "S2": _step("S2", "Content_Similarity_Vector_Search", {"a_chunk_1": ["one"]}), + } + self.assertEqual(len(_gather(results)["unstructured"]), 1) + + def test_grouped_results_lose_only_repeated_chunks(self): + results = { + "S1": _step("S1", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_14": ["fourteen"]}), + "S2": _step("S2", "Chunk_Sibling_Vector_Search", { + "a_chunk_14": { + "a_chunk_13": {"distance": "-1", "content": "thirteen"}, + "a_chunk_14": {"distance": "0", "content": "fourteen"}, + }, + "a_chunk_30": {"a_chunk_14": {"distance": "-16", "content": "fourteen"}}, + }), + } + grouped = _fr(_gather(results)["unstructured"][1]) + self.assertEqual(grouped, {"a_chunk_14": {"a_chunk_13": {"distance": "-1", "content": "thirteen"}}}) + + def test_step_results_are_not_modified(self): + results = { + "S1": _step("S1", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_1": ["one"]}), + "S2": _step("S2", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_1": ["one"], "a_chunk_2": ["two"]}), + } + before = copy.deepcopy(results["S2"].context) + _gather(results) + self.assertEqual(results["S2"].context, before) + + def test_structural_results_are_untouched(self): + structural = StepResult(step_id="S1", ok=True, summary="", citations=[], + context={"function_call": "tg__run_query", "result": [{"n": 1}]}) + results = {"S1": structural, "S2": copy.deepcopy(structural)} + results["S2"].step_id = "S2" + self.assertEqual(len(_gather(results)["structural"]), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_async_id_token_connection.py b/graphrag/tests/test_async_id_token_connection.py new file mode 100644 index 0000000..6591a9c --- /dev/null +++ b/graphrag/tests/test_async_id_token_connection.py @@ -0,0 +1,231 @@ +import ast +import asyncio +import json +import os +import threading +import unittest + +_ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "..")) +_CONNECTIONS = os.path.join(_ROOT, "common", "db", "connections.py") +_ECC_MAIN = os.path.join(_ROOT, "ecc", "app", "main.py") + + +def _tree(path): + with open(path, encoding="utf-8") as handle: + return ast.parse(handle.read()) + + +def _function(tree, name): + return next( + n + for n in ast.walk(tree) + if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef)) and n.name == name + ) + + +def _load_verify_helper(): + """Compile ``_verify_async_connection`` straight from source. + + Importing ``common.db.connections`` pulls in the app config and FastAPI; the + helper needs neither, so it is exercised on its own. + """ + fn = _function(_tree(_CONNECTIONS), "_verify_async_connection") + module = ast.Module(body=[fn], type_ignores=[]) + namespace = {} + exec(compile(module, _CONNECTIONS, "exec"), namespace) + return namespace["_verify_async_connection"] + + +class _FakeConnection: + def __init__(self, fail=False): + self.fail = fail + self.statements = [] + self.closed = False + + async def gsql(self, statement): + self.statements.append(statement) + if self.fail: + raise PermissionError("rejected") + + async def aclose(self): + self.closed = True + + +class TestVerifyReleasesItsSession(unittest.TestCase): + """The async token check must not leave a session bound to a dead loop. + + It runs under ``asyncio.run()``, which closes its loop on return. With + pyTigerGraph 2.0.4 a session opened there stays cached, and the connection's + next request on another loop — ECC's rebuild — failed with "Event loop is + closed", reported as "graph does not exist" (GML-2182). + """ + + def setUp(self): + self.verify = _load_verify_helper() + + def test_checks_the_graph(self): + conn = _FakeConnection() + asyncio.run(self.verify(conn, "MyGraph")) + self.assertEqual(conn.statements, ["USE GRAPH MyGraph"]) + + def test_releases_the_session_after_success(self): + conn = _FakeConnection() + asyncio.run(self.verify(conn, "MyGraph")) + self.assertTrue(conn.closed) + + def test_releases_the_session_when_the_check_fails(self): + """A rejected token must still release the session, and still raise.""" + conn = _FakeConnection(fail=True) + with self.assertRaises(PermissionError): + asyncio.run(self.verify(conn, "MyGraph")) + self.assertTrue(conn.closed) + + +class TestConnectionSurvivesTheCheck(unittest.TestCase): + """End to end with the real client against a keep-alive HTTP server.""" + + @classmethod + def setUpClass(cls): + try: + from aiohttp import web + from pyTigerGraph import AsyncTigerGraphConnection # noqa: F401 + except ImportError as exc: + raise unittest.SkipTest(f"needs aiohttp and pyTigerGraph: {exc}") + + async def ok(_request): + body = json.dumps({"error": False, "message": "ok", "results": []}) + return web.Response(text=body, content_type="application/json") + + ready = threading.Event() + cls._port = None + + def serve(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + app = web.Application() + app.router.add_route("*", "/{tail:.*}", ok) + runner = web.AppRunner(app) + loop.run_until_complete(runner.setup()) + site = web.TCPSite(runner, "127.0.0.1", 0) + loop.run_until_complete(site.start()) + cls._port = site._server.sockets[0].getsockname()[1] + cls._loop = loop + ready.set() + loop.run_forever() + + threading.Thread(target=serve, daemon=True).start() + ready.wait(10) + + @classmethod + def tearDownClass(cls): + loop = getattr(cls, "_loop", None) + if loop is not None: + loop.call_soon_threadsafe(loop.stop) + + def test_connection_is_usable_on_a_new_loop_after_the_check(self): + from pyTigerGraph import AsyncTigerGraphConnection + + port = str(self._port) + conn = AsyncTigerGraphConnection( + host="http://127.0.0.1", graphname="g", apiToken="t", + gsPort=port, restppPort=port, + ) + # The check runs on one throwaway loop; the connection is then reused + # on another, which is how ECC drives it. + asyncio.run(_load_verify_helper()(conn, "g")) + self.assertIsNone(conn._async_client) + asyncio.run(conn.gsql("ls")) # raised "Event loop is closed" before + + +class TestIdTokenPathKeepsValidation(unittest.TestCase): + """Skipping the check would also skip token validation on this path.""" + + def setUp(self): + self.fn = _function(_tree(_CONNECTIONS), "get_db_connection_id_token") + + def test_async_branch_runs_the_releasing_check(self): + calls = [ + ast.unparse(n) + for n in ast.walk(self.fn) + if isinstance(n, ast.Call) and ast.unparse(n.func) == "asyncio.run" + ] + self.assertTrue( + any("_verify_async_connection" in c for c in calls), + f"async branch must validate via the releasing helper: {calls}", + ) + self.assertFalse( + any("conn.gsql" in c for c in calls), + "a bare conn.gsql under asyncio.run leaves the session on a dead loop", + ) + + def test_async_auth_failure_maps_to_401(self): + """The async client raises aiohttp's error, not requests.HTTPError, so the + sync-only handler never caught it and a bad token surfaced as a 500.""" + handled = [ + ast.unparse(h.type) + for n in ast.walk(self.fn) + if isinstance(n, ast.Try) + for h in n.handlers + if h.type is not None + ] + self.assertIn("aiohttp.ClientResponseError", handled) + + +class TestIdTokenOutcomes(unittest.TestCase): + """Run the real function against a port nothing listens on.""" + + @classmethod + def setUpClass(cls): + try: + from fastapi import HTTPException + from common.db import connections + except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.HTTPException = HTTPException + cls.connections = connections + + def test_unreachable_database_is_503(self): + """A down database raises aiohttp's connection error, which was not + handled, so the rebuild request failed with a 500.""" + from unittest.mock import patch + + with patch.dict(self.connections.db_config, {"hostname": "http://127.0.0.1"}): + with self.assertRaises(self.HTTPException) as caught: + self.connections.get_db_connection_id_token("g", "token", async_conn=True) + self.assertEqual(caught.exception.status_code, 503) + + +class TestEccGraphCheckMessage(unittest.TestCase): + """The failure message is shown to operators; the cause belongs in the log.""" + + def setUp(self): + self.fn = _function(_tree(_ECC_MAIN), "run_with_tracking") + + def _graph_check_handler(self): + for node in ast.walk(self.fn): + if isinstance(node, ast.Try) and "getVertexTypes" in ast.unparse(node.body[0]): + return node.handlers[0] + self.fail("graph check not found in run_with_tracking") + + def test_message_does_not_carry_the_raw_exception(self): + handler = self._graph_check_handler() + raised = [n for n in ast.walk(handler) if isinstance(n, ast.Raise)] + self.assertTrue(raised) + for node in raised: + names = { + n.id for n in ast.walk(node.exc) if isinstance(n, ast.Name) + } + self.assertNotIn(handler.name, names, ast.unparse(node.exc)) + + def test_cause_is_logged(self): + handler = self._graph_check_handler() + logged = [ + ast.unparse(n) + for n in ast.walk(handler) + if isinstance(n, ast.Call) and ast.unparse(n.func).startswith("LogWriter.") + ] + self.assertTrue(logged, "the underlying error must be logged") + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_context_token_limit.py b/graphrag/tests/test_context_token_limit.py new file mode 100644 index 0000000..933ccad --- /dev/null +++ b/graphrag/tests/test_context_token_limit.py @@ -0,0 +1,179 @@ +import types +import unittest + +try: + from common.utils.token_calculator import PROMPT_RESERVE_MARGIN, TokenCalculator +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +def _calc(limit): + return TokenCalculator(token_limit=limit, model_name="gpt-4o-mini") + + +class TestFitContext(unittest.TestCase): + """The context is trimmed so the whole prompt fits, by token count (GML-2311).""" + + PROMPT = "Answer the question from the context below.\nQuestion: q\nContext:\n" + + def _budget(self, calc): + return calc.max_context_tokens - calc.count_tokens(self.PROMPT) - PROMPT_RESERVE_MARGIN + + def test_no_limit_leaves_context_untouched(self): + text = "word " * 5000 + self.assertEqual(_calc(0).fit_context(text, self.PROMPT), text) + + def test_context_within_budget_is_unchanged(self): + self.assertEqual(_calc(2000).fit_context("short context", self.PROMPT), "short context") + + def test_reserve_is_the_measured_prompt_not_a_flat_amount(self): + calc = _calc(2000) + long_prompt = self.PROMPT + ("rule. " * 600) + trimmed = calc.fit_context("word " * 5000, long_prompt) + total = calc.count_tokens(trimmed) + calc.count_tokens(long_prompt) + self.assertLessEqual(total, 2000 - PROMPT_RESERVE_MARGIN) + self.assertGreater(calc.count_tokens(trimmed), 0) + + def test_cjk_context_is_trimmed_by_tokens_not_characters(self): + """Fewer characters than the limit, but more tokens than the budget. + Non-OpenAI models count with cl100k_base, where Japanese runs over one + token per character.""" + calc = TokenCalculator(token_limit=1500, model_name="gemini-3.5-flash") + text = "東京都の書類管理における情報資産" * 80 + self.assertLess(len(text), calc.max_context_tokens) + self.assertGreater(calc.count_tokens(text), self._budget(calc)) + trimmed = calc.fit_context(text, self.PROMPT) + self.assertLessEqual(calc.count_tokens(trimmed), self._budget(calc)) + + def test_list_of_passages_is_trimmed_in_order(self): + calc = _calc(1500) + passages = [f"passage {i}: " + "detail " * 200 for i in range(10)] + trimmed = calc.fit_context(passages, self.PROMPT) + self.assertIsInstance(trimmed, list) + self.assertLess(len(trimmed), len(passages)) + self.assertEqual(trimmed[0], passages[0]) + self.assertLessEqual(sum(calc.count_tokens(p) for p in trimmed), self._budget(calc)) + + def test_prompt_larger_than_limit_leaves_no_context(self): + calc = _calc(300) + self.assertEqual(calc.fit_context("word " * 500, "rule " * 400), "") + self.assertEqual(calc.fit_context(["word " * 500], "rule " * 400), []) + self.assertEqual(calc.fit_context({"a": "word " * 500}, "rule " * 400), {}) + + +class TestContextTokenLimit(unittest.TestCase): + """Without a configured token_limit, the model's own input limit applies.""" + + @classmethod + def setUpClass(cls): + try: + from common.llm_services import base_llm + except ImportError as exc: + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.base_llm = base_llm + + def _service(self, config, profile=None): + svc = object.__new__(self.base_llm.LLM_Model) + svc.config = config + svc.llm = types.SimpleNamespace(profile=profile) + return svc + + def test_configured_limit_wins(self): + svc = self._service({"token_limit": 50000}, {"max_input_tokens": 1_000_000}) + self.assertEqual(svc.context_token_limit(), 50000) + + def test_profile_limit_less_answer_room_and_margin(self): + svc = self._service( + {"llm_model": "m"}, {"max_input_tokens": 128000, "max_output_tokens": 16384} + ) + self.assertEqual( + svc.context_token_limit(), + int(128000 * self.base_llm._TOKENIZER_MARGIN) - 16384, + ) + + def test_unknown_model_keeps_no_limit_and_warns_once(self): + self.base_llm._warned_no_token_limit.discard("custom-model") + svc = self._service({"llm_model": "custom-model", "token_limit": 0}, None) + with self.assertLogs(self.base_llm.logger, level="WARNING") as logs: + self.assertEqual(svc.context_token_limit(), 0) + self.assertEqual(svc.context_token_limit(), 0) + self.base_llm.logger.warning("sentinel") + self.assertEqual(sum("custom-model" in line for line in logs.output), 1) + + def test_real_openai_profile_is_read_offline(self): + try: + from langchain_openai import ChatOpenAI + except ImportError as exc: + self.skipTest(f"langchain-openai unavailable: {exc}") + svc = self._service({"llm_model": "gpt-4o-mini"}) + svc.llm = ChatOpenAI(model="gpt-4o-mini", api_key="sk-test") + self.assertGreater(svc.context_token_limit(), 50000) + + +class TestAnswerGenerationTrimsContext(unittest.TestCase): + """The answer prompt sent to the model fits the limit, prompt included.""" + + @classmethod + def setUpClass(cls): + try: + from agent.agent_generation import TigerGraphAgentGenerator + except ImportError as exc: + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.Generator = TigerGraphAgentGenerator + + def test_sent_context_fits_with_the_prompt(self): + seen = {} + + class FakeService: + config = {"llm_model": "gpt-4o-mini"} + chatbot_response_prompt = ( + "Rules " * 300 + + "\n{format_instructions}\nQuestion: {question}\nQuery: {query}\nContext: {context}" + ) + + def context_token_limit(self): + return 3000 + + @staticmethod + def _salvage_answer_output(raw): + return raw + + def invoke_with_parser(self, prompt, parser, inputs, **kwargs): + seen["prompt"] = prompt.format(**inputs) + seen["context"] = inputs["context"] + return "ok" + + gen = self.Generator(FakeService()) + gen.generate_answer("What is X?", "fact " * 10000) + calc = gen.token_calculator + self.assertLess(calc.count_tokens(seen["context"]), calc.count_tokens("fact " * 10000)) + self.assertLessEqual(calc.count_tokens(seen["prompt"]), 3000) + + def test_failure_log_names_the_provider_error(self): + class FailingService: + config = {"llm_model": "gpt-4o-mini"} + chatbot_response_prompt = "{format_instructions}\n{question}\n{query}\n{context}" + + def context_token_limit(self): + return 0 + + @staticmethod + def _salvage_answer_output(raw): + return raw + + def invoke_with_parser(self, *args, **kwargs): + raise ValueError("400 INVALID_ARGUMENT: request contains an invalid argument") + + import logging + gen = self.Generator(FailingService()) + with self.assertLogs("agent.agent_generation", level=logging.WARNING) as logs: + answer = gen.generate_answer("q", "ctx") + self.assertIn("I wasn't able to generate an answer", answer.generated_answer) + self.assertTrue(any( + "ValueError: 400 INVALID_ARGUMENT: request contains an invalid argument" in line + for line in logs.output + )) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_docx_table_extraction.py b/graphrag/tests/test_docx_table_extraction.py new file mode 100644 index 0000000..1c504ce --- /dev/null +++ b/graphrag/tests/test_docx_table_extraction.py @@ -0,0 +1,176 @@ +import os +import tempfile +import unittest + +import docx + +from common.utils.text_extractors import ( + ExtractionError, + _docx_to_markdown, + extract_text_from_file, + get_supported_extensions, +) + + +def _saved(build): + """Build a .docx with ``build(document)`` and return its path.""" + handle = tempfile.NamedTemporaryFile(suffix=".docx", delete=False) + handle.close() + document = docx.Document() + build(document) + document.save(handle.name) + return handle.name + + +class TestDocxTablesSurvive(unittest.TestCase): + """Tables in a .docx must reach the extracted text. + + Extraction previously joined ``doc.paragraphs`` only (GML-2196). That + yields just the direct children of the document body, so every table was + dropped — silently, because the prose around it still came through. + """ + + def setUp(self): + self.paths = [] + + def tearDown(self): + for path in self.paths: + try: + os.unlink(path) + except OSError: + pass + + def _build(self, build): + path = _saved(build) + self.paths.append(path) + return path + + @staticmethod + def _thresholds(document): + document.add_paragraph("Approval thresholds are defined below.") + table = document.add_table(rows=3, cols=2) + rows = [("Tier", "Limit"), ("Manager", "$10,000"), ("VP", "$250,000")] + for index, (left, right) in enumerate(rows): + table.rows[index].cells[0].text = left + table.rows[index].cells[1].text = right + document.add_paragraph("Requests above the VP limit escalate to the CFO.") + + def test_cell_values_are_extracted(self): + """The figures a reader would ask about are present.""" + text = extract_text_from_file(self._build(self._thresholds)) + for value in ("Tier", "Limit", "Manager", "$10,000", "VP", "$250,000"): + self.assertIn(value, text) + + def test_table_stays_between_its_paragraphs(self): + """A table keeps its position, so the prose introducing it still refers + to something.""" + text = extract_text_from_file(self._build(self._thresholds)) + intro = text.index("Approval thresholds") + table = text.index("$250,000") + outro = text.index("escalate to the CFO") + self.assertLess(intro, table) + self.assertLess(table, outro) + + def test_rendered_as_a_markdown_table(self): + """Tabular data is rendered the way the spreadsheet branch renders it.""" + text = extract_text_from_file(self._build(self._thresholds)) + self.assertIn("| Tier | Limit |", text) + self.assertIn("| VP | $250,000 |", text) + + def test_pipe_and_newline_in_a_cell_do_not_break_the_row(self): + """A cell containing a pipe or a line break must not split the row.""" + + def build(document): + table = document.add_table(rows=2, cols=2) + table.rows[0].cells[0].text = "Region" + table.rows[0].cells[1].text = "Rule" + table.rows[1].cells[0].text = "EU|UK" + cell = table.rows[1].cells[1] + cell.text = "line one" + cell.add_paragraph("line two") + + text = extract_text_from_file(self._build(build)) + rows = [line for line in text.splitlines() if line.startswith("|")] + self.assertEqual(len(rows), 3, f"expected header, rule and one row: {rows}") + self.assertIn(r"EU\|UK", text) + self.assertIn("line one line two", text) + + def test_nested_table_contents_are_reached(self): + """A table inside a cell is flattened rather than dropped.""" + + def build(document): + outer = document.add_table(rows=1, cols=1) + inner = outer.rows[0].cells[0].add_table(rows=1, cols=2) + inner.rows[0].cells[0].text = "NESTED_KEY" + inner.rows[0].cells[1].text = "NESTED_VAL" + + text = extract_text_from_file(self._build(build)) + self.assertIn("NESTED_KEY", text) + self.assertIn("NESTED_VAL", text) + + def test_merged_cells_do_not_raise(self): + """Merged cells repeat rather than producing a ragged row.""" + + def build(document): + table = document.add_table(rows=2, cols=3) + table.rows[0].cells[0].merge(table.rows[0].cells[1]) + table.rows[0].cells[0].text = "Merged" + table.rows[0].cells[2].text = "C" + for index, value in enumerate(["x", "y", "z"]): + table.rows[1].cells[index].text = value + + text = extract_text_from_file(self._build(build)) + self.assertIn("Merged", text) + self.assertIn("| x | y | z |", text) + + def test_document_without_tables_is_unchanged(self): + """Prose-only documents keep their previous shape.""" + + def build(document): + document.add_paragraph("First.") + document.add_paragraph("Second.") + + self.assertEqual( + _docx_to_markdown(docx.Document(self._build(build))), + "First.\n\nSecond.", + ) + + +class TestUnreadableFilesRaise(unittest.TestCase): + """Formats that cannot be read must raise, not become a document.""" + + def _write(self, suffix, payload): + handle = tempfile.NamedTemporaryFile(suffix=suffix, delete=False) + handle.write(payload) + handle.close() + self.addCleanup(lambda: os.path.exists(handle.name) and os.unlink(handle.name)) + return handle.name + + def test_legacy_doc_raises(self): + """.doc is the OLE format python-docx cannot read.""" + path = self._write(".doc", b"\xd0\xcf\x11\xe0 legacy") + with self.assertRaises(ExtractionError) as caught: + extract_text_from_file(path) + self.assertIn(".docx", str(caught.exception)) + + def test_legacy_doc_is_not_advertised_as_supported(self): + self.assertNotIn(".doc", get_supported_extensions()) + self.assertIn(".docx", get_supported_extensions()) + + def test_unsupported_extension_raises(self): + path = self._write(".rtf", b"{}") + with self.assertRaises(ExtractionError): + extract_text_from_file(path) + + def test_corrupt_docx_raises_without_leaking_parser_detail(self): + """The caller-facing message must not carry the parser's wording.""" + path = self._write(".docx", b"not a zip at all") + with self.assertRaises(ExtractionError) as caught: + extract_text_from_file(path) + message = str(caught.exception) + self.assertNotIn("zip", message.lower()) + self.assertIn("corrupt", message.lower()) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_extraction_failure_not_a_document.py b/graphrag/tests/test_extraction_failure_not_a_document.py new file mode 100644 index 0000000..e61b0c8 --- /dev/null +++ b/graphrag/tests/test_extraction_failure_not_a_document.py @@ -0,0 +1,148 @@ +import ast +import os +import re +import unittest + +_EXTRACTORS = os.path.normpath( + os.path.join( + os.path.dirname(__file__), "..", "..", "common", "utils", "text_extractors.py" + ) +) + + +def _source(): + with open(_EXTRACTORS, encoding="utf-8") as handle: + return handle.read() + + +class TestNoPlaceholderDocuments(unittest.TestCase): + """A file that cannot be extracted must not become a document. + + Extraction failures used to return a document whose text was the error + message (GML-2195). That placeholder reached the graph, was chunked and + embedded, and became retrievable — indistinguishable to the agent from + real document content — while the ingest was reported as successful. + """ + + def test_no_bracketed_placeholder_is_returned_as_content(self): + """No extractor returns an error string as a document's content.""" + offenders = [ + line.strip() + for line in _source().splitlines() + if re.search(r'"content":\s*f?"\[', line) + ] + self.assertEqual( + offenders, + [], + "extraction failures must raise, not return placeholder content", + ) + + def test_no_bare_placeholder_string_is_returned(self): + """No function returns a bracketed notice as its extracted text. + + GML-2195 caught dict literals carrying a `content` key. Bare returns + of the same shape survived it — `[Unsupported file type: .doc]` and + `[Excel file is empty or contains no data]` were still ingested as a + document's entire text (GML-2196). + """ + def leading_literals(value): + """Leading string literal of every branch this expression can yield.""" + if isinstance(value, ast.Constant) and isinstance(value.value, str): + yield value.value + elif isinstance(value, ast.JoinedStr) and value.values: + first = value.values[0] + if isinstance(first, ast.Constant) and isinstance(first.value, str): + yield first.value + elif isinstance(value, ast.IfExp): + # `return x if cond else "[...]"` hides a placeholder in a branch + yield from leading_literals(value.body) + yield from leading_literals(value.orelse) + + offenders = [] + for node in ast.walk(ast.parse(_source())): + if not isinstance(node, ast.Return) or node.value is None: + continue + for literal in leading_literals(node.value): + if literal.startswith("["): + offenders.append(f"line {node.lineno}: {literal[:50]}") + self.assertEqual( + offenders, + [], + "placeholder notices must be raised, not returned as content", + ) + + def test_failure_paths_raise_extraction_error(self): + """Each failure branch in the PDF/image extractors raises.""" + tree = ast.parse(_source()) + for name in ( + "_extract_pdf_with_images_as_docs", + "_extract_standalone_image_as_doc", + ): + fn = next( + ( + n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == name + ), + None, + ) + self.assertIsNotNone(fn, f"{name} not found") + raises = [ + n + for n in ast.walk(fn) + if isinstance(n, ast.Raise) + and isinstance(n.exc, ast.Call) + and getattr(n.exc.func, "id", "") == "ExtractionError" + ] + self.assertTrue(raises, f"{name} never raises ExtractionError") + + def test_extraction_error_is_defined(self): + """The error type exists and derives from Exception.""" + tree = ast.parse(_source()) + cls = next( + ( + n + for n in ast.walk(tree) + if isinstance(n, ast.ClassDef) and n.name == "ExtractionError" + ), + None, + ) + self.assertIsNotNone(cls, "ExtractionError not defined") + self.assertIn("Exception", [ast.unparse(b) for b in cls.bases]) + + +class TestMessagesCarryNoInternals(unittest.TestCase): + """Messages reach the caller in the per-file `error` field, so they must + describe the file and the remedy — not exception text or paths.""" + + def _messages(self): + tree = ast.parse(_source()) + for node in ast.walk(tree): + if ( + isinstance(node, ast.Raise) + and isinstance(node.exc, ast.Call) + and getattr(node.exc.func, "id", "") == "ExtractionError" + and node.exc.args + ): + yield node.lineno, ast.unparse(node.exc.args[0]) + + def test_no_raw_exception_interpolated(self): + """`{e}` / `{import_err}` belong in the log line, not the message.""" + for lineno, msg in self._messages(): + for leak in ("{e}", "{str(e)}", "{import_err}", "{exc}"): + self.assertNotIn( + leak, msg, f"exception text in message at line {lineno}: {msg}" + ) + + def test_no_full_path_interpolated(self): + """Only the bare filename may appear, never the full path.""" + for lineno, msg in self._messages(): + self.assertNotIn( + "{file_path}", + msg, + f"filesystem path in message at line {lineno}: {msg}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_extraction_failures_reported.py b/graphrag/tests/test_extraction_failures_reported.py new file mode 100644 index 0000000..1a1c6b6 --- /dev/null +++ b/graphrag/tests/test_extraction_failures_reported.py @@ -0,0 +1,210 @@ +import asyncio +import os +import tempfile +import unittest +from unittest.mock import MagicMock, patch + +from common.utils.text_extractors import ( + ExtractionError, + TextExtractor, + describe_failures, + failed_extractions, +) + + +class TestFailedExtractions(unittest.TestCase): + """Folder processing records a failed file instead of raising, so callers + read the failures back out of its result.""" + + def test_maps_each_failed_file_to_its_reason(self): + result = {"files": [ + {"file_path": "/up/g/good.pdf", "status": "success"}, + {"file_path": "/up/g/bad.pdf", "status": "failed", "error": "unreadable"}, + ]} + self.assertEqual(failed_extractions(result), {"bad.pdf": "unreadable"}) + + def test_tolerates_a_result_without_files(self): + self.assertEqual(failed_extractions({"num_documents": 1}), {}) + self.assertEqual(failed_extractions(None), {}) + + def test_description_names_each_file_once(self): + text = describe_failures({ + "a.pdf": "Could not read text from a.pdf.", + "b.doc": "legacy format", + }) + self.assertEqual(text.count("a.pdf"), 1) + self.assertIn("b.doc: legacy format", text) + + +class TestRealFolderProcessing(unittest.TestCase): + """End to end through the real extractor: failures come back with reasons + fit to show the person who uploaded the files.""" + + def test_unreadable_files_are_reported_and_readable_ones_convert(self): + with tempfile.TemporaryDirectory() as src, tempfile.TemporaryDirectory() as out: + with open(os.path.join(src, "good.txt"), "w") as handle: + handle.write("Approval thresholds are defined below.") + with open(os.path.join(src, "corrupt.docx"), "wb") as handle: + handle.write(b"not a zip archive") + with open(os.path.join(src, "legacy.doc"), "wb") as handle: + handle.write(b"\xd0\xcf\x11\xe0 legacy") + + result = asyncio.run( + TextExtractor()._process_folder_async(src, "g", out) + ) + failures = failed_extractions(result) + + self.assertEqual(sorted(failures), ["corrupt.docx", "legacy.doc"]) + self.assertTrue(os.path.exists(os.path.join(out, "good.jsonl"))) + self.assertIn(".docx", failures["legacy.doc"]) + for reason in failures.values(): + self.assertNotIn(src, reason, "no filesystem path in the reason") + self.assertNotIn("zip", reason.lower(), "no parser wording") + + +class TestOnlyCallerSafeReasons(unittest.TestCase): + """Reasons now reach the person uploading, so unexpected errors — which can + carry paths, URLs or service messages — must not pass through verbatim.""" + + def test_unexpected_error_is_replaced_by_a_message_about_the_file(self): + leak = RuntimeError("POST http://llm.internal/v1 failed; see /code/uploads/g/bad.pdf") + with tempfile.TemporaryDirectory() as src, tempfile.TemporaryDirectory() as out: + with open(os.path.join(src, "bad.pdf"), "wb") as handle: + handle.write(b"%PDF-1.4") + with patch( + "common.utils.text_extractors.extract_text_from_file_with_images_as_docs", + side_effect=leak, + ): + result = asyncio.run(TextExtractor()._process_folder_async(src, "g", out)) + reason = failed_extractions(result)["bad.pdf"] + self.assertEqual(reason, "Could not process bad.pdf.") + + def test_server_side_reasons_name_the_file(self): + """The UI shows only the reason, so it must say which file, and that + the cause is on the server rather than in the upload.""" + for error in (FileNotFoundError("gone"), PermissionError("denied")): + with self.subTest(error=type(error).__name__): + with tempfile.TemporaryDirectory() as src, tempfile.TemporaryDirectory() as out: + with open(os.path.join(src, "report.pdf"), "wb") as handle: + handle.write(b"%PDF-1.4") + with patch( + "common.utils.text_extractors.extract_text_from_file_with_images_as_docs", + side_effect=error, + ): + result = asyncio.run(TextExtractor()._process_folder_async(src, "g", out)) + reason = failed_extractions(result)["report.pdf"] + self.assertIn("report.pdf", reason) + self.assertIn("server", reason) + + def test_image_failure_does_not_blame_the_file(self): + """The image path also fails when the description service does, so the + message must not tell the user their image is corrupt.""" + with tempfile.TemporaryDirectory() as src, tempfile.TemporaryDirectory() as out: + with open(os.path.join(src, "photo.png"), "wb") as handle: + handle.write(b"not really a png") + result = asyncio.run(TextExtractor()._process_folder_async(src, "g", out)) + reason = failed_extractions(result)["photo.png"] + self.assertIn("photo.png", reason) + self.assertNotIn("corrupt", reason.lower()) + + +class TestCreateIngestReportsFailures(unittest.TestCase): + """``create_ingest`` converts the uploaded folder for server ingest.""" + + @classmethod + def setUpClass(cls): + try: + from common.py_schemas.schemas import CreateIngestConfig + from supportai import supportai + except ImportError as exc: # needs the full app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.config_type = CreateIngestConfig + cls.supportai = supportai + + def _config(self, data_path): + return self.config_type( + data_source="server", + file_format="multi", + data_source_config={"data_path": data_path}, + ) + + def _run(self, result): + with tempfile.TemporaryDirectory() as tmp, patch.object( + TextExtractor, "process_folder", return_value=result + ): + return self.supportai.create_ingest( + "g", self._config(os.path.join(tmp, "uploads")), MagicMock() + ) + + def test_partial_failure_returns_the_failed_files(self): + res = self._run({"statusCode": 200, "num_documents": 1, "files": [ + {"file_path": "/u/good.pdf", "status": "success", "jsonl_file": "good.jsonl", "num_documents": 1}, + {"file_path": "/u/bad.pdf", "status": "failed", "error": "unreadable"}, + ]}) + self.assertEqual(res["failed_files"], [{"file": "bad.pdf", "error": "unreadable"}]) + # Folder-wide, matching what ingest loads — not the latest upload's size. + self.assertEqual(res["ready_files"], 1) + + def test_ready_count_matches_the_jsonl_files_ingest_loads(self): + """Two sources sharing a stem convert to one JSONL, so ingest loads one.""" + res = self._run({"statusCode": 200, "num_documents": 2, "files": [ + {"file_path": "/u/report.pdf", "status": "success", "jsonl_file": "report.jsonl", "num_documents": 1}, + {"file_path": "/u/report.docx", "status": "success", "jsonl_file": "report.jsonl", "num_documents": 1}, + {"file_path": "/u/notes.txt", "status": "success", "jsonl_file": "notes.jsonl", "num_documents": 1}, + ]}) + self.assertEqual(res["ready_files"], 2) + + def test_file_with_no_documents_is_not_counted_ready(self): + """A skipped decorative image converts to nothing, so ingest loads nothing + for it — counting it would announce files ingest then can't find.""" + res = self._run({"statusCode": 200, "num_documents": 1, "files": [ + {"file_path": "/u/notes.txt", "status": "success", + "jsonl_file": "notes.jsonl", "num_documents": 1}, + {"file_path": "/u/logo.png", "status": "success", + "jsonl_file": "logo.jsonl", "num_documents": 0}, + ]}) + self.assertEqual(res["ready_files"], 1) + + def test_nothing_readable_raises_the_reasons(self): + """Otherwise ingest later fails with an error naming an internal path.""" + with self.assertRaises(ExtractionError) as caught: + self._run({"statusCode": 200, "num_documents": 0, "files": [ + {"file_path": "/u/bad.pdf", "status": "failed", "error": "unreadable"}, + ]}) + self.assertIn("bad.pdf", str(caught.exception)) + self.assertNotIn("/u/", str(caught.exception)) + + +class TestPublicCreateIngestRoute(unittest.TestCase): + """The API route pyTigerGraph uses must answer unreadable input with 400, + matching the UI route, rather than a 500.""" + + @classmethod + def setUpClass(cls): + try: + from fastapi.testclient import TestClient + from app.main import app + except ImportError as exc: + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.client = TestClient(app) + + def test_nothing_readable_is_400_with_the_reasons(self): + failed = {"statusCode": 200, "num_documents": 0, "files": [ + {"file_path": "/u/bad.pdf", "status": "failed", + "error": "Could not read text from bad.pdf."}, + ]} + with tempfile.TemporaryDirectory() as tmp, patch.object( + TextExtractor, "process_folder", return_value=failed + ): + resp = self.client.post( + "/g/supportai/create_ingest", + json={"data_source": "server", "file_format": "multi", + "data_source_config": {"data_path": os.path.join(tmp, "uploads")}}, + auth=("user", "pass"), + ) + self.assertEqual(resp.status_code, 400, resp.text) + self.assertIn("bad.pdf", resp.json()["detail"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_migration_schema_compat.py b/graphrag/tests/test_migration_schema_compat.py new file mode 100644 index 0000000..a61364b --- /dev/null +++ b/graphrag/tests/test_migration_schema_compat.py @@ -0,0 +1,216 @@ +import copy +import os +import unittest + +from common.db.migrate import BASE_SCHEMA_PATH, check_schema_compatibility, parse_base_schema + +_ROOT = os.path.normpath(os.path.join(os.path.dirname(__file__), "..", "..")) +_BASE = os.path.join(_ROOT, BASE_SCHEMA_PATH) + + +class FakeConn: + """Schema metadata in the shapes pyTigerGraph returns.""" + + def __init__(self, vertices, edges): + self.vertices = vertices # {name: {attr, ...}} + self.edges = edges # {name: {"pairs": [(f, t)], "attrs": {...}}} + + def getVertexTypes(self): + return list(self.vertices) + + def getEdgeTypes(self): + return list(self.edges) + + def getVertexType(self, vt): + return {"Name": vt, "Attributes": [{"AttributeName": a} for a in self.vertices[vt]]} + + def getEdgeType(self, et): + spec = self.edges[et] + if "meta" in spec: # verbatim metadata, as TigerGraph reported it + return dict(spec["meta"], Attributes=[{"AttributeName": a} for a in spec["attrs"]]) + pairs = sorted(spec["pairs"]) + meta = {"Name": et, "Attributes": [{"AttributeName": a} for a in spec["attrs"]]} + if len(pairs) == 1: + meta["FromVertexTypeName"], meta["ToVertexTypeName"] = pairs[0] + else: # multi-pair edges report "*" plus an EdgePairs list + meta["FromVertexTypeName"] = meta["ToVertexTypeName"] = "*" + meta["EdgePairs"] = [{"From": f, "To": t} for f, t in pairs] + return meta + + +def _current_graph(): + """A graph created with the current base schema (1.4 or later).""" + with open(_BASE, encoding="utf-8") as f: + vertices, edges = parse_base_schema(f.read()) + return copy.deepcopy(vertices), { + k: {"pairs": set(v["pairs"]), "attrs": set(v["attrs"])} for k, v in edges.items() + } + + +def _check(vertices, edges): + return check_schema_compatibility(FakeConn(vertices, edges), base_schema_path=_BASE) + + +class TestParseBaseSchema(unittest.TestCase): + def test_reads_the_shipped_base_schema(self): + vertices, edges = _current_graph() + self.assertEqual( + set(vertices), + {"DocumentChunk", "Document", "Entity", "RelationshipType", "Content", + "EntityType", "Community"}, + ) + self.assertEqual(edges["IS_HEAD_OF"]["pairs"], {("EntityType", "RelationshipType")}) + self.assertEqual( + edges["HAS_CONTENT"]["pairs"], + {("Document", "Content"), ("DocumentChunk", "Content")}, + ) + self.assertEqual(edges["LINKS_TO"]["attrs"], {"weight"}) # lowercase from/to + self.assertIn("description", vertices["Entity"]) + self.assertNotIn("id", vertices["Entity"]) + + +class TestSchemaCompatibility(unittest.TestCase): + def test_current_schema_is_compatible(self): + result = _check(*_current_graph()) + self.assertTrue(result.compatible, result.differences) + + def test_extra_domain_and_image_types_are_allowed(self): + vertices, edges = _current_graph() + vertices["Company"] = {"name", "industry"} + vertices["Image"] = {"image_data", "image_format"} + edges["WORKS_FOR"] = {"pairs": {("Company", "Company")}, "attrs": set()} + edges["CONTAINS_ENTITY"]["pairs"].add(("DocumentChunk", "Company")) + self.assertTrue(_check(vertices, edges).compatible) + + def test_wildcard_source_accepts_every_type(self): + """IN_COMMUNITY on a graph with domain types, as TigerGraph reports it.""" + vertices, edges = _current_graph() + edges["IN_COMMUNITY"]["meta"] = { + "Name": "IN_COMMUNITY", "FromVertexTypeName": "*", + "ToVertexTypeName": "Community", "EdgePairs": [], "IsDirected": True, + } + self.assertTrue(_check(vertices, edges).compatible) + + def test_wildcard_with_the_wrong_target_is_incompatible(self): + vertices, edges = _current_graph() + edges["IN_COMMUNITY"]["meta"] = { + "Name": "IN_COMMUNITY", "FromVertexTypeName": "*", + "ToVertexTypeName": "EntityType", "EdgePairs": [], + } + self.assertFalse(_check(vertices, edges).compatible) + + def test_13_style_meta_edges_are_incompatible(self): + vertices, edges = _current_graph() + edges["IS_HEAD_OF"]["pairs"] = {("Entity", "RelationshipType")} + edges["HAS_TAIL"]["pairs"] = {("RelationshipType", "Entity")} + edges["RELATIONSHIP_TYPE"] = {"pairs": {("EntityType", "EntityType")}, "attrs": {"frequency"}} + result = _check(vertices, edges) + self.assertFalse(result.compatible) + self.assertTrue(any("IS_HEAD_OF" in d and "EntityType->RelationshipType" in d + for d in result.differences), result.differences) + self.assertTrue(any("HAS_TAIL" in d for d in result.differences)) + + def test_missing_vertex_type_is_incompatible(self): + vertices, edges = _current_graph() + del vertices["EntityType"] + result = _check(vertices, edges) + self.assertFalse(result.compatible) + self.assertIn("missing vertex type EntityType", result.differences) + + def test_missing_attribute_is_incompatible(self): + vertices, edges = _current_graph() + vertices["Entity"].discard("entity_type") + edges["RELATIONSHIP"]["attrs"].clear() + result = _check(vertices, edges) + self.assertFalse(result.compatible) + self.assertTrue(any("Entity" in d and "entity_type" in d for d in result.differences)) + self.assertTrue(any("RELATIONSHIP" in d and "relation_type" in d for d in result.differences)) + + def test_missing_edge_pair_is_incompatible(self): + vertices, edges = _current_graph() + edges["CONTAINS_ENTITY"]["pairs"].discard(("Document", "Entity")) + result = _check(vertices, edges) + self.assertFalse(result.compatible) + self.assertTrue(any("CONTAINS_ENTITY" in d and "Document->Entity" in d + for d in result.differences)) + + def test_unreadable_schema_raises(self): + class Broken(FakeConn): + def getVertexTypes(self): + raise ConnectionError("TigerGraph unreachable") + + with self.assertRaises(ConnectionError): + check_schema_compatibility(Broken({}, {}), base_schema_path=_BASE) + + +class TestRepairGate(unittest.TestCase): + """The repair endpoint refuses graphs it cannot repair in place.""" + + @classmethod + def setUpClass(cls): + try: + from routers import ui + except Exception as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.ui = ui + + def setUp(self): + from unittest import mock + ui = self.ui + self.inner = mock.Mock(return_value={"success": True}) + self.conn = None + patches = [ + mock.patch.object(ui, "get_rebuilding_graph", return_value=None), + mock.patch.object(ui, "acquire_graph_lock", return_value=True), + mock.patch.object(ui, "release_graph_lock"), + mock.patch.object(ui, "get_db_connection_pwd_manual", side_effect=lambda *a, **k: self.conn), + mock.patch.object(ui, "_migration_apply_inner", self.inner), + ] + for p in patches: + p.start() + self.addCleanup(p.stop) + # The shipped schema path is relative to the app directory; point the + # check at the repo copy, since tests run from elsewhere. + import common.db.migrate as migrate + real = migrate.check_schema_compatibility + p = mock.patch.object(migrate, "check_schema_compatibility", + side_effect=lambda conn: real(conn, base_schema_path=_BASE)) + p.start() + self.addCleanup(p.stop) + + def _apply(self): + creds = (["superuser"], type("C", (), {"username": "u", "password": "p"})()) + return self.ui.migration_apply("g", creds, {"outdated": ["Q"]}) + + def test_compatible_graph_is_repaired(self): + self.conn = FakeConn(*_current_graph()) + self.assertEqual(self._apply(), {"success": True}) + self.inner.assert_called_once() + + def test_incompatible_graph_is_refused(self): + from fastapi import HTTPException + vertices, edges = _current_graph() + edges["IS_HEAD_OF"]["pairs"] = {("Entity", "RelationshipType")} + self.conn = FakeConn(vertices, edges) + with self.assertRaises(HTTPException) as ctx: + self._apply() + self.assertEqual(ctx.exception.status_code, 409) + self.assertIn("older version of GraphRAG", ctx.exception.detail) + self.inner.assert_not_called() + + def test_unreadable_schema_is_refused(self): + from fastapi import HTTPException + + class Broken(FakeConn): + def getVertexTypes(self): + raise ConnectionError("TigerGraph unreachable") + + self.conn = Broken({}, {}) + with self.assertRaises(HTTPException) as ctx: + self._apply() + self.assertEqual(ctx.exception.status_code, 503) + self.inner.assert_not_called() + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_query_with_history_isolation.py b/graphrag/tests/test_query_with_history_isolation.py new file mode 100644 index 0000000..8d41a78 --- /dev/null +++ b/graphrag/tests/test_query_with_history_isolation.py @@ -0,0 +1,103 @@ +import ast +import os +import unittest + +from pydantic import ValidationError + +from common.py_schemas.schemas import NaturalLanguageQuery + +_INQUIRYAI = os.path.join( + os.path.dirname(__file__), "..", "app", "routers", "inquiryai.py" +) + + +class TestCallerSuppliedHistory(unittest.TestCase): + """History on /{graphname}/query_with_history comes from the caller. + + The endpoint previously kept a process-global list shared by every + caller and every graph, so one caller's questions and answers were + injected into another's prompt (GML-2194). Context is now supplied per + request, which makes that class of bug unrepresentable. + """ + + def test_history_defaults_to_none(self): + """A request that omits history is valid and carries none.""" + self.assertIsNone(NaturalLanguageQuery(query="hi").history) + + def test_omitted_history_yields_no_turns(self): + """The handler's `(history or [])` guard turns None into no turns.""" + query = NaturalLanguageQuery(query="hi") + self.assertEqual(list(query.history or []), []) + + def test_history_round_trips(self): + """Supplied turns are parsed in order.""" + query = NaturalLanguageQuery( + query="and his email?", + history=[ + {"query": "who is Ada?", "response": "an engineer"}, + {"query": "her ID?", "response": "42"}, + ], + ) + self.assertEqual( + [(t.query, t.response) for t in query.history], + [("who is Ada?", "an engineer"), ("her ID?", "42")], + ) + + def test_only_the_most_recent_turns_are_used(self): + """Older turns beyond the cap are dropped, newest kept.""" + turns = [{"query": f"q{i}", "response": f"r{i}"} for i in range(6)] + query = NaturalLanguageQuery(query="hi", history=turns) + self.assertEqual( + [t.query for t in (query.history or [])[-3:]], ["q3", "q4", "q5"] + ) + + def test_malformed_turn_is_rejected(self): + """A turn missing `response` fails validation rather than passing through.""" + with self.assertRaises(ValidationError): + NaturalLanguageQuery(query="hi", history=[{"query": "a"}]) + + +class TestNoProcessGlobalHistory(unittest.TestCase): + """Regression guard for GML-2194.""" + + def _module(self): + with open(os.path.normpath(_INQUIRYAI), encoding="utf-8") as handle: + return ast.parse(handle.read()) + + def test_handler_declares_no_global(self): + """The handler must not reach for process-wide mutable state.""" + for node in ast.walk(self._module()): + if ( + isinstance(node, ast.FunctionDef) + and node.name == "retrieve_answer_with_chathistory" + ): + globals_declared = [ + name + for child in ast.walk(node) + if isinstance(child, ast.Global) + for name in child.names + ] + self.assertEqual( + globals_declared, + [], + "conversation state must be per-request, not process-global", + ) + return + self.fail("retrieve_answer_with_chathistory not found") + + def test_no_module_level_conversation_store(self): + """No module-level list/dict is left to accumulate caller history.""" + stores = [ + target.id + for node in self._module().body + if isinstance(node, ast.Assign) + and isinstance(node.value, (ast.List, ast.Dict)) + for target in node.targets + if isinstance(target, ast.Name) + and "history" in target.id.lower() + ] + self.assertEqual(stores, [], f"module-level history store(s): {stores}") + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_retrieval_step_summary.py b/graphrag/tests/test_retrieval_step_summary.py new file mode 100644 index 0000000..9b061e1 --- /dev/null +++ b/graphrag/tests/test_retrieval_step_summary.py @@ -0,0 +1,117 @@ +import unittest + +try: + from tools.graphrag_tools import _unstructured_result +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +Q = "GraphRAG_Hybrid_Vector_Search" + + +def _step(final_retrieval, edges=()): + return [{"final_retrieval": final_retrieval, "edges": list(edges)}] + + +class TestRetrievalStepSummary(unittest.TestCase): + """A search step's summary counts what it retrieved, not the result's keys.""" + + def test_counts_chunks(self): + fr = {f"guide.pdf_chunk_{i}": [f"text {i}"] for i in range(20)} + out = _unstructured_result(Q, _step(fr, edges=[{"s": "a", "t": "b"}] * 7)) + self.assertTrue(out["ok"]) + self.assertEqual(out["summary"], f"{Q} returned 20 chunk(s)") + self.assertEqual(out["context"]["result"]["final_retrieval"], fr) + + def test_reports_non_chunk_entries_separately(self): + fr = {"ticket_42_chunk_0": ["t"], "ticket_42_chunk_1": ["t"], + "TigerGraph": ["Entity: TigerGraph ..."], "GPE": ["Entity: GPE ..."], + "Similarity_Context": ["..."]} + out = _unstructured_result(Q, _step(fr)) + self.assertEqual(out["summary"], f"{Q} returned 2 chunk(s) and 3 other entries") + + def test_one_other_entry_is_singular(self): + out = _unstructured_result(Q, _step({"a_chunk_0": ["t"], "community_1": ["s"]})) + self.assertEqual(out["summary"], f"{Q} returned 1 chunk(s) and 1 other entry") + + def test_name_merely_containing_chunk_is_not_a_chunk(self): + out = _unstructured_result(Q, _step({"chunk_size_setting": ["t"]})) + self.assertEqual(out["summary"], f"{Q} returned 0 chunk(s) and 1 other entry") + + def test_empty_retrieval_is_reported_as_empty(self): + out = _unstructured_result(Q, _step({})) + self.assertFalse(out["ok"]) + self.assertEqual(out["summary"], f"{Q} returned no chunks") + self.assertIsNone(out["context"]) + + def test_missing_result_is_reported_as_empty(self): + self.assertFalse(_unstructured_result(Q, [])["ok"]) + self.assertFalse(_unstructured_result(Q, None)["ok"]) + + +class TestRetrievalLog(unittest.TestCase): + """Each retrieval logs its chunk count and size right after the query.""" + + def test_line_counts_chunks_and_their_size(self): + from common.utils.retrieval_stats import describe_retrieval + fr = {"a_chunk_0": ["x" * 1000], "a_chunk_1": ["y" * 2500, "z" * 500]} + self.assertEqual(describe_retrieval(Q, fr), f"{Q} retrieved 2 chunk(s), 4,000 chars") + + def test_line_reports_other_entries_and_nested_text(self): + from common.utils.retrieval_stats import describe_retrieval + fr = {"a_chunk_0": ["x" * 10], "TigerGraph": ["e" * 300], + "Similarity_Context": [["s" * 20], ["t" * 30]]} + self.assertEqual( + describe_retrieval(Q, fr), + f"{Q} retrieved 1 chunk(s), 10 chars; 2 other entries, 350 chars", + ) + + def test_retriever_logs_after_the_query(self): + try: + from supportai.retrievers.BaseRetriever import BaseRetriever + except ImportError as exc: + self.skipTest(f"app environment unavailable: {exc}") + import logging + retriever = object.__new__(BaseRetriever) + retriever.logger = logging.getLogger("test.retrieval") + with self.assertLogs("test.retrieval", level="INFO") as logs: + retriever._log_retrieval(Q, _step({"a_chunk_0": ["x" * 42]})) + self.assertEqual(logs.output, [f"INFO:test.retrieval:{Q} retrieved 1 chunk(s), 42 chars"]) + + +class TestNestedSiblingResults(unittest.TestCase): + """Contextual search groups chunks under the match they surround.""" + + FR = { + "doc_chunk_14": { + "doc_chunk_13": {"distance": "-1", "content": "a" * 10}, + "doc_chunk_14": {"distance": "0", "content": "b" * 20}, + "doc_chunk_15": {"distance": "1", "content": "c" * 30}, + }, + "doc_chunk_21": { + "doc_chunk_21": {"distance": "0", "content": "d" * 40}, + "doc_chunk_15": {"distance": "-6", "content": "c" * 30}, # listed twice + }, + } + + def test_summary_counts_the_chunks_returned_not_the_matches(self): + out = _unstructured_result("Chunk_Sibling_Vector_Search", _step(self.FR)) + self.assertEqual(out["summary"], "Chunk_Sibling_Vector_Search returned 4 chunk(s)") + + def test_log_sizes_by_chunk_text(self): + from common.utils.retrieval_stats import describe_retrieval + self.assertEqual( + describe_retrieval("Q", self.FR), "Q retrieved 4 chunk(s), 100 chars" + ) + + def test_trace_records_every_fetched_chunk(self): + try: + from agent.agentic_executor import retrieved_chunk_ids + except ImportError as exc: + self.skipTest(f"app environment unavailable: {exc}") + ids = retrieved_chunk_ids({"result": {"final_retrieval": self.FR}}) + self.assertEqual(ids, ["doc_chunk_13", "doc_chunk_14", "doc_chunk_15", "doc_chunk_21"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_retrieve_docs_auth.py b/graphrag/tests/test_retrieve_docs_auth.py new file mode 100644 index 0000000..951e81f --- /dev/null +++ b/graphrag/tests/test_retrieve_docs_auth.py @@ -0,0 +1,207 @@ +import ast +import os +import unittest + +_INQUIRYAI = os.path.normpath( + os.path.join( + os.path.dirname(__file__), "..", "app", "routers", "inquiryai.py" + ) +) + + +def _tree(): + with open(_INQUIRYAI, encoding="utf-8") as handle: + return ast.parse(handle.read()) + + +def _route(tree, name): + return next( + ( + n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) and n.name == name + ), + None, + ) + + +class TestRetrieveDocsAuthenticatesTheCaller(unittest.TestCase): + """``/{graphname}/retrieve_docs`` must authenticate before returning content. + + The route previously took no ``Request``, so it never touched the + connection ``auth_middleware`` had built. On the Basic-auth path that + connection is constructed lazily — nothing contacts the database — so a + request with absent or invalid credentials was never rejected, and the + documents came back from a service-credentialed store (GML-2197). + """ + + def setUp(self): + self.tree = _tree() + self.fn = _route(self.tree, "retrieve_docs") + self.assertIsNotNone(self.fn, "retrieve_docs not found") + + def test_route_receives_the_request(self): + """Without a Request the route cannot reach the caller's connection.""" + annotations = [ + ast.unparse(a.annotation) for a in self.fn.args.args if a.annotation + ] + self.assertIn("Request", annotations) + + def test_route_checks_graph_access(self): + """The caller must be proven against this graph before anything is read.""" + calls = [ + n.func.id + for n in ast.walk(self.fn) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) + ] + self.assertIn("_require_graph_access", calls) + + def test_access_check_precedes_the_document_read(self): + """Ordering is the guarantee: authorize, then retrieve.""" + check_line = retrieve_line = None + for node in ast.walk(self.fn): + if isinstance(node, ast.Call): + if isinstance(node.func, ast.Name) and node.func.id == "_require_graph_access": + check_line = node.lineno + if isinstance(node.func, ast.Attribute) and node.func.attr == "retrieve_similar": + retrieve_line = node.lineno + self.assertIsNotNone(check_line, "no access check") + self.assertIsNotNone(retrieve_line, "no retrieval call") + self.assertLess(check_line, retrieve_line) + + def test_searches_the_graph_named_in_the_path(self): + """Authorizing graph A then reading graph B would make the check moot.""" + for node in ast.walk(self.fn): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "get_embedding_store" + ): + self.assertTrue( + node.args or node.keywords, + "get_embedding_store() with no graph returns the default " + "store, ignoring the graph in the path", + ) + return + self.fail("get_embedding_store not called") + + +class TestGraphAccessHelper(unittest.TestCase): + """The helper must be graph-scoped and must reject a missing connection.""" + + def setUp(self): + self.tree = _tree() + self.fn = _route(self.tree, "_require_graph_access") + self.assertIsNotNone(self.fn, "_require_graph_access not found") + + def test_rejects_when_no_connection_was_built(self): + """No Authorization header means the middleware built nothing — 401, + not an AttributeError surfacing as a 500.""" + source = ast.unparse(self.fn) + self.assertIn("getattr(request.state, 'conn', None)", source) + self.assertIn("401", source) + + def test_uses_a_graph_scoped_probe_not_echo(self): + """``echo()`` pings /echo and succeeds for any valid user on any graph, + so it cannot establish access to *this* graph.""" + attrs = [ + n.func.attr + for n in ast.walk(self.fn) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) + ] + self.assertIn("getVertexTypes", attrs) + self.assertNotIn("echo", attrs) + + def test_raises_http_401_rather_than_propagating(self): + raises = [ + ast.unparse(n.exc) + for n in ast.walk(self.fn) + if isinstance(n, ast.Raise) and n.exc is not None + ] + self.assertTrue(raises, "helper never raises") + self.assertTrue( + all("HTTPException" in r for r in raises), + f"must raise HTTPException, got {raises}", + ) + + +class TestGraphAccessOutcomes(unittest.TestCase): + """What the helper actually returns for each failure, run from source.""" + + @classmethod + def setUpClass(cls): + try: + import requests + except ImportError as exc: + raise unittest.SkipTest(f"needs requests: {exc}") + + class HTTPException(Exception): + def __init__(self, status_code, detail=None): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + fn = _route(_tree(), "_require_graph_access") + # ``Request`` only appears in the signature's annotation. + namespace = {"requests": requests, "HTTPException": HTTPException, "Request": object} + exec(compile(ast.Module(body=[fn], type_ignores=[]), _INQUIRYAI, "exec"), namespace) + cls.check = staticmethod(namespace["_require_graph_access"]) + cls.HTTPException = HTTPException + cls.requests = requests + + def _request(self, conn): + request = type("Request", (), {})() + request.state = type("State", (), {})() + if conn is not None: + request.state.conn = conn + return request + + def _conn(self, raises=None): + conn = type("Conn", (), {})() + + def get_vertex_types(): + if raises: + raise raises + return ["Document"] + + conn.getVertexTypes = get_vertex_types + return conn + + def _status(self, conn): + with self.assertRaises(self.HTTPException) as caught: + self.check(self._request(conn), "g") + return caught.exception.status_code + + def test_reachable_and_authorized_returns_the_connection(self): + conn = self._conn() + self.assertIs(self.check(self._request(conn), "g"), conn) + + def test_missing_connection_is_401(self): + self.assertEqual(self._status(None), 401) + + def test_rejected_credentials_are_401(self): + self.assertEqual(self._status(self._conn(Exception("User authentication failed"))), 401) + + def _http_error(self, status): + response = self.requests.Response() + response.status_code = status + return self.requests.exceptions.HTTPError(f"{status}", response=response) + + def test_database_5xx_is_503(self): + """nginx answers 502/503 while GSQL restarts; valid callers must not be + told their credentials are wrong.""" + self.assertEqual(self._status(self._conn(self._http_error(502))), 503) + + def test_database_4xx_is_401(self): + self.assertEqual(self._status(self._conn(self._http_error(401))), 401) + + def test_unreachable_database_is_503_not_401(self): + """An outage must not read as bad credentials.""" + down = self.requests.exceptions.ConnectionError("connection refused") + self.assertEqual(self._status(self._conn(down)), 503) + slow = self.requests.exceptions.Timeout("timed out") + self.assertEqual(self._status(self._conn(slow)), 503) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_retriever_flag_order.py b/graphrag/tests/test_retriever_flag_order.py new file mode 100644 index 0000000..03130fe --- /dev/null +++ b/graphrag/tests/test_retriever_flag_order.py @@ -0,0 +1,134 @@ +import ast +import glob +import os +import unittest + +_RETRIEVERS = os.path.normpath( + os.path.join(os.path.dirname(__file__), "..", "app", "supportai", "retrievers") +) +_ROUTER = os.path.normpath( + os.path.join(os.path.dirname(__file__), "..", "app", "routers", "supportai.py") +) +# Flags that select behaviour rather than tune it; transposing two of these +# swaps the installed query being run, silently. +_FLAGS = ("withHyDE", "expand", "combine", "chunk_only", "doc_only") + + +def _parse(path): + with open(path, encoding="utf-8") as handle: + return ast.parse(handle.read()) + + +def _params(tree, name): + fn = next( + (n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == name), + None, + ) + return [a.arg for a in fn.args.args if a.arg != "self"] if fn else None + + +def _flag_order(params): + return [p for p in params or [] if p in _FLAGS and p != "combine"] + + +class TestRetrieverFlagOrderIsConsistent(unittest.TestCase): + """A retriever's own two entry points must order their flags alike. + + ``SiblingRetriever.search`` declared ``expand`` before ``withHyDE`` while + ``retrieve_answer`` — and both positional callers — passed the reverse, so + every call transposed them (GML-2198). The flags pick different installed + GSQL queries, so the swap changed which query ran, with no error. + """ + + def test_search_and_retrieve_answer_agree_within_each_retriever(self): + for path in sorted(glob.glob(os.path.join(_RETRIEVERS, "*.py"))): + name = os.path.basename(path) + if name in ("__init__.py", "BaseRetriever.py"): + continue + tree = _parse(path) + search = _flag_order(_params(tree, "search")) + answer = _flag_order(_params(tree, "retrieve_answer")) + with self.subTest(retriever=name): + self.assertEqual( + search, + answer, + f"{name}: search{search} vs retrieve_answer{answer}", + ) + + def test_sibling_and_similarity_share_a_flag_order(self): + """These two take the same pair; disagreeing is what caused the swap.""" + sibling = _flag_order(_params(_parse(os.path.join(_RETRIEVERS, "SiblingRetriever.py")), "search")) + similarity = _flag_order(_params(_parse(os.path.join(_RETRIEVERS, "SimilarityRetriever.py")), "search")) + self.assertEqual(sibling, similarity) + self.assertEqual(sibling, ["withHyDE", "expand"]) + + +class TestPositionalCallsMatchSignatures(unittest.TestCase): + """Positional calls are what make a transposition possible at all.""" + + def _self_call(self, tree, caller, callee): + fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == caller) + for node in ast.walk(fn): + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == callee + ): + return [ast.unparse(a) for a in node.args] + return None + + def test_each_retriever_calls_its_own_search_in_order(self): + for path in sorted(glob.glob(os.path.join(_RETRIEVERS, "*.py"))): + name = os.path.basename(path) + if name in ("__init__.py", "BaseRetriever.py"): + continue + tree = _parse(path) + passed = self._self_call(tree, "retrieve_answer", "search") + signature = _params(tree, "search") + if not passed or not signature: + continue + # Compare only the flag positions. Parameter names may legitimately + # differ either side of the call (HybridRetriever passes `index` + # into `indices`); a flag landing in another flag's slot may not. + for position, param in enumerate(signature[: len(passed)]): + if param not in _FLAGS: + continue + with self.subTest(retriever=name, param=param): + self.assertEqual( + passed[position], + param, + f"{name}: {passed[position]!r} passed as {param!r}", + ) + + def test_router_passes_sibling_flags_in_signature_order(self): + """``supportai.py`` calls ``search`` positionally with method_params.""" + signature = _params( + _parse(os.path.join(_RETRIEVERS, "SiblingRetriever.py")), "search" + ) + flags = {"withHyDE", "expand"} + for node in ast.walk(_parse(_ROUTER)): + if not ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "search" + and len(node.args) == len(signature) + ): + continue + passed = [ast.unparse(a) for a in node.args] + named = [p for p in passed if any(f"'{f}'" in p for f in flags)] + if len(named) != 2: + continue + for arg, param in zip(passed, signature): + if param in flags: + with self.subTest(line=node.lineno, param=param): + self.assertIn( + f"'{param}'", + arg, + f"line {node.lineno}: {arg} passed as {param}", + ) + return + self.skipTest("no positional sibling search call found in the router") + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_trace_output_note.py b/graphrag/tests/test_trace_output_note.py new file mode 100644 index 0000000..4193fe8 --- /dev/null +++ b/graphrag/tests/test_trace_output_note.py @@ -0,0 +1,100 @@ +import json +import unittest +from unittest import mock + +try: + from agent import agentic_executor as ex +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +def _size(obj): + return len(json.dumps(obj, ensure_ascii=False)) + + +def _retrieval(n_chunks, chunk_chars, text="x"): + return { + "function_call": "GraphRAG_Hybrid_Vector_Search", + "result": { + "final_retrieval": { + f"doc_{i}_chunk_{i}": [text * chunk_chars] for i in range(n_chunks) + }, + "edges": [], + }, + } + + +class TestFitForTrace(unittest.TestCase): + """The trace keeps the tool's own data, shortened to fit when needed.""" + + def test_small_value_is_unchanged(self): + value, cut = ex.fit_for_trace({"question": "q"}) + self.assertEqual(value, {"question": "q"}) + self.assertIsNone(cut) + + def test_large_retrieval_keeps_its_shape(self): + ctx = _retrieval(20, 3400) + value, cut = ex.fit_for_trace(ctx) + self.assertEqual(cut["full_chars"], _size(ctx)) + self.assertLessEqual(cut["shown_chars"], ex._TRACE_FIELD_CAP) + fr = value["result"]["final_retrieval"] + self.assertEqual(set(fr), set(ctx["result"]["final_retrieval"])) # every chunk id kept + for cid, texts in fr.items(): + original = ctx["result"]["final_retrieval"][cid][0] + shown, _, marker = texts[0].partition("…(+") + self.assertTrue(original.startswith(shown)) + self.assertEqual(marker, f"{len(original) - len(shown):,} chars)") + self.assertNotIn("_truncated", json.dumps(value)) + + def test_many_entries_are_cut_to_fit(self): + value, cut = ex.fit_for_trace(_retrieval(2000, 200)) + self.assertLessEqual(_size(value), ex._TRACE_FIELD_CAP) + self.assertIn("final_retrieval", value["result"]) + + def test_non_english_text_is_counted_as_written(self): + ctx = _retrieval(20, 1000, text="日") + _, cut = ex.fit_for_trace(ctx) + self.assertEqual(cut["full_chars"], _size(ctx)) + self.assertLess(cut["full_chars"], len(json.dumps(ctx))) # escaped form is larger + + def test_short_texts_get_no_marker(self): + ctx = _retrieval(20, 3400) + ctx["result"]["final_retrieval"]["short_chunk_0"] = ["brief"] + value, _ = ex.fit_for_trace(ctx) + self.assertEqual(value["result"]["final_retrieval"]["short_chunk_0"], ["brief"]) + + def test_note_wording(self): + self.assertEqual( + ex.trace_note({"full_chars": 68331, "shown_chars": 11958}), + "Truncated from 68,331 to 11,958 characters due to the trace log size limit.", + ) + + +class TestStepTrace(unittest.TestCase): + """A step's trace output is the tool's output, plus a note when shortened.""" + + def _run(self, context): + out = {"ok": True, "summary": "returned 20 chunk(s)", "context": context, "citations": []} + step = mock.Mock(id="S1", tool="graphrag__hybrid_search", kind="unstructured", rationale="") + results, traces = {}, [] + with mock.patch.object(ex.registry, "run", return_value=out): + ex._run_step(step, {"question": "q"}, mock.Mock(), results, traces) + self.assertEqual(results["S1"].context, context) # the answer still gets everything + return traces[0] + + def test_large_output_gets_a_note(self): + ctx = _retrieval(20, 3400) + entry = self._run(ctx) + self.assertEqual(set(entry["output"]), {"summary", "result", "note"}) + self.assertEqual(entry["output"]["summary"], "returned 20 chunk(s)") + self.assertIn(f"Truncated from {_size(ctx):,} to", entry["output"]["note"]) + self.assertIn("final_retrieval", entry["output"]["result"]["result"]) + self.assertEqual(entry["input"], {"question": "q"}) # small input, no note + + def test_small_output_has_no_note(self): + entry = self._run(_retrieval(2, 50)) + self.assertEqual(set(entry["output"]), {"summary", "result"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/graphrag/tests/test_v1_4_1_uploads_schema.py b/graphrag/tests/test_v1_4_1_uploads_schema.py index c4927ca..7ace401 100644 --- a/graphrag/tests/test_v1_4_1_uploads_schema.py +++ b/graphrag/tests/test_v1_4_1_uploads_schema.py @@ -261,7 +261,7 @@ def setUp(self): @patch("app.routers.ui.auth", return_value=_ok_auth()) def test_writes_to_flat_layout_and_returns_saved_files(self, _mock_auth): - async def _fake_process(self, folder, graphname, temp): + async def _fake_process(self, folder, graphname, temp, filenames=None): stem = "report" Path(temp).mkdir(parents=True, exist_ok=True) Path(temp, f"{stem}.jsonl").write_text('{"text":"hello"}\n') @@ -286,9 +286,65 @@ async def _fake_process(self, folder, graphname, temp): # No request-id field surfaces to the caller in the new contract self.assertNotIn("request_id", body) + @patch("app.routers.ui.auth", return_value=_ok_auth()) + def test_unreadable_sample_is_reported_and_not_passed_on(self, _mock_auth): + """A sample that fails to convert has no JSONL, so naming it in the + schema step would fail the whole step. It is reported instead.""" + async def _fake_process(self, folder, graphname, temp, filenames=None): + Path(temp).mkdir(parents=True, exist_ok=True) + Path(temp, "good.jsonl").write_text('{"text":"hello"}\n') + return {"num_documents": 1, "files": [ + {"file_path": str(Path(folder, "good.pdf")), "status": "success", + "num_documents": 1}, + {"file_path": str(Path(folder, "bad.pdf")), "status": "failed", + "error": "Could not read text from bad.pdf. The file may be corrupt."}, + ]} + + with _ChdirTempDir(), patch( + "app.routers.ui.TextExtractor._process_folder_async", new=_fake_process, + ): + resp = self.client.post( + f"/ui/{GRAPH}/convert_sample_files", + files=[ + ("files", ("good.pdf", io.BytesIO(b"x"), "application/pdf")), + ("files", ("bad.pdf", io.BytesIO(b"x"), "application/pdf")), + ], + auth=("testuser", "testpass"), + ) + self.assertEqual(resp.status_code, 200, resp.text) + body = resp.json() + self.assertEqual(body["saved_files"], ["good.pdf"]) + self.assertEqual( + body["failed_files"], + [{"file": "bad.pdf", + "error": "Could not read text from bad.pdf. The file may be corrupt."}], + ) + + @patch("app.routers.ui.auth", return_value=_ok_auth()) + def test_all_samples_unreadable_returns_the_reasons(self, _mock_auth): + async def _fake_process(self, folder, graphname, temp, filenames=None): + Path(temp).mkdir(parents=True, exist_ok=True) + return {"num_documents": 0, "files": [ + {"file_path": str(Path(folder, "bad.pdf")), "status": "failed", + "error": "Could not read text from bad.pdf. The file may be corrupt."}, + ]} + + with _ChdirTempDir(), patch( + "app.routers.ui.TextExtractor._process_folder_async", new=_fake_process, + ): + resp = self.client.post( + f"/ui/{GRAPH}/convert_sample_files", + files=[("files", ("bad.pdf", io.BytesIO(b"x"), "application/pdf"))], + auth=("testuser", "testpass"), + ) + self.assertEqual(resp.status_code, 400, resp.text) + detail = resp.json()["detail"] + self.assertIn("bad.pdf", detail) + self.assertNotIn("/", detail, "no filesystem path in the message") + @patch("app.routers.ui.auth", return_value=_ok_auth()) def test_sweeps_legacy_schema_subdirs_on_entry(self, _mock_auth): - async def _fake_process(self, folder, graphname, temp): + async def _fake_process(self, folder, graphname, temp, filenames=None): Path(temp).mkdir(parents=True, exist_ok=True) return {"num_documents": 0} From 001456654d567acccc7694e147d7d4a6babf9b83 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 7 Oct 2026 12:10:30 -0700 Subject: [PATCH 04/15] Prepare 2.0.3 release - Bump the version to 2.0.3 - Add the 2.0.3 entry to the README release list - Record the multi-part and follow-up question improvements in the changelog --- CHANGELOG.md | 1 + README.md | 1 + VERSION | 2 +- graphrag-ui/package.json | 2 +- 4 files changed, 4 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4402981..2119ab2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,7 @@ - **Search steps in the trace report what they retrieved.** An agentic search step's summary showed a fixed "2 item(s)" regardless of results; it now gives the number of chunks retrieved (and any other entries), and a search that finds nothing is reported as empty. Each search also logs its chunk count and total size. A step output too large to keep in full in the trace is now shown as the tool's own data, shortened, with a note giving its full and shown size and each shortened text ending in how much of it was left out — previously it was replaced by a preview marker that read as if the tool's output itself had been cut. ### Changed +- **Agentic answers handle multi-part and follow-up questions better.** Each part of a multi-part question is searched on its own and answered from its matching source, references to earlier turns are resolved from the conversation, and when a reference is ambiguous the assistant asks one short clarifying question instead of guessing. - **The Migration Assistant repairs only graphs it can repair in place.** Its check now verifies that the graph has the schema GraphRAG's queries are written against, which every graph created with GraphRAG 1.4 or later has. A graph created with an earlier version is reported as needing to be recreated and its documents ingested again, and Repair is not offered — previously Repair would replace its working queries with ones its schema cannot run. If the schema cannot be read, Repair is withheld until the check succeeds. - **`token_limit: 0` no longer means unlimited.** A `token_limit` of `0` or left unset now uses the chat model's own input limit (less room for the answer), so retrieved context is trimmed to what the model accepts. Models whose limit is not known keep the previous untrimmed behavior. Set a positive `token_limit` to choose a lower cap. - **Callers of `/{graphname}/query_with_history` supply their own conversation context.** The request body accepts an optional `history` list of `{"query", "response"}` turns, oldest first, and the most recent turns are used. A request that omits it is answered without prior context; the endpoint no longer remembers anything between calls. Callers that relied on the server accumulating history must now send the turns they want considered. diff --git a/README.md b/README.md index a54f0a6..6fd74fc 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ --- ## Releases +* **10/7/2026**: GraphRAG v2.0.3 released. Improved agentic answers to multi-part and follow-up questions, kept answers within the chat model's input limit, added deleting chat conversations individually, in groups, or all at once, made the Migration Assistant check a graph's schema before repairing it, and fixed document retrieval, ingestion, and contextual search issues. See [Release Notes](https://github.com/tigergraph/graphrag/releases/tag/v2.0.3) for details. * **8/28/2026**: GraphRAG v2.0.2 released. Added Migration Assistant data-integrity checks with targeted re-embedding of missing embeddings and re-summarization of incomplete community summaries, made agentic chat work with current and future tool-calling models automatically, and fixed rebuild, PDF extraction, and chat issues. See [Release Notes](https://github.com/tigergraph/graphrag/releases/tag/v2.0.2) for details. * **7/22/2026**: GraphRAG v2.0.1 released. Fixed a chat error with models that return structured content (such as Gemini 3), and hardened knowledge-graph rebuild so it no longer hangs when the language model is unreachable. Relicensed from Apache 2.0 to AGPL-3.0. See [Release Notes](https://github.com/tigergraph/graphrag/releases/tag/v2.0.1) for details. * **7/1/2026**: GraphRAG v2.0.0 released. Added an agentic chat engine that plans and runs its own retrieval (Planner and Reactive styles), external MCP tools, and structure-aware document chunking, along with additive prompt customization and many other improvements and bug fixes. See [Release Notes](https://github.com/tigergraph/graphrag/releases/tag/v2.0.0) for details. diff --git a/VERSION b/VERSION index e9307ca..50ffc5a 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -2.0.2 +2.0.3 diff --git a/graphrag-ui/package.json b/graphrag-ui/package.json index 936dba7..3fecc8a 100755 --- a/graphrag-ui/package.json +++ b/graphrag-ui/package.json @@ -1,7 +1,7 @@ { "name": "tg-cbot-v5", "private": true, - "version": "2.0.2", + "version": "2.0.3", "type": "module", "packageManager": "pnpm@9.15.0", "scripts": { From aff364e6723b5fc3f7007b715a9058293178dfbc Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 7 Oct 2026 13:59:21 -0700 Subject: [PATCH 05/15] Keep more context for reasoning models and fall back safely - Reserve at most 8K tokens for the answer when deriving the default limit, so models with large output budgets keep most of their context - Drop only text already sent when de-duplicating passages across steps, so a later step's different or extra text for the same entry is kept - Return the usual fallback answer when the answer prompt cannot be built Refs: GML-2311 --- common/llm_services/base_llm.py | 15 ++++--- graphrag/app/agent/agent_generation.py | 13 +++--- graphrag/app/agent/agentic_synthesizer.py | 44 +++++++++++++++++---- graphrag/tests/test_answer_context_dedup.py | 25 ++++++++++++ graphrag/tests/test_context_token_limit.py | 38 +++++++++++++++++- 5 files changed, 115 insertions(+), 20 deletions(-) diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index 0626015..19c1c22 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -183,7 +183,9 @@ def _record_usage(caller_name: str, usage_data: dict, config: Optional[dict] = N # Prompts are counted with tiktoken, which can undercount a provider's own # tokenizer: Gemini counted ~23% more on a real 524K-token context. _TOKENIZER_MARGIN = 0.6 -_DEFAULT_ANSWER_RESERVE = 4096 +# Room kept for the answer. Capped: reasoning models declare output budgets of +# 100K+ tokens, and reserving all of it left them a few thousand for context. +_ANSWER_RESERVE_CAP = 8192 _warned_no_token_limit: set = set() @@ -209,8 +211,8 @@ def context_token_limit(self) -> int: """Token budget for an answer prompt, retrieved context included. The configured ``token_limit`` when set. Otherwise the chat model's - input limit from its LangChain profile, less room for the answer and a - margin for tokenizer differences (counts use tiktoken, not the + input limit from its LangChain profile, less room for the answer (at + most ``_ANSWER_RESERVE_CAP``) and a margin for tokenizer differences (counts use tiktoken, not the provider's tokenizer). 0 means no limit: a model with no profile, such as an Ollama or custom model name, keeps the old untrimmed behavior. """ @@ -228,8 +230,11 @@ def context_token_limit(self) -> int: f"{model!r}; retrieved context will not be trimmed" ) return 0 - max_output = profile.get("max_output_tokens") or _DEFAULT_ANSWER_RESERVE - return max(int(max_input * _TOKENIZER_MARGIN) - int(max_output), 0) + budget = int(max_input * _TOKENIZER_MARGIN) + # Never more than half the budget, so small-window models keep context. + answer_room = min(int(profile.get("max_output_tokens") or _ANSWER_RESERVE_CAP), + _ANSWER_RESERVE_CAP, budget // 2) + return budget - answer_room def _read_prompt_file(self, path): """Read a prompt file with per-graph override support. diff --git a/graphrag/app/agent/agent_generation.py b/graphrag/app/agent/agent_generation.py index bf70978..4b34402 100644 --- a/graphrag/app/agent/agent_generation.py +++ b/graphrag/app/agent/agent_generation.py @@ -65,13 +65,14 @@ def generate_answer(self, question: str, context: str | dict, query: str = "") - } ) - # Trim the context so the full prompt fits the model's token limit. - if not self.token_calculator.is_unlimited_tokens(): - context = self.token_calculator.fit_context( - context, prompt.format(question=question, context="", query=query) - ) - try: + # Trim the context so the full prompt fits the model's token limit. + # Inside the try: rendering the prompt can fail (e.g. a custom + # prompt with stray braces), which must still yield the fallback. + if not self.token_calculator.is_unlimited_tokens(): + context = self.token_calculator.fit_context( + context, prompt.format(question=question, context="", query=query) + ) generation = self.llm.invoke_with_parser( prompt, answer_parser, {"question": question, "context": context, "query": query}, diff --git a/graphrag/app/agent/agentic_synthesizer.py b/graphrag/app/agent/agentic_synthesizer.py index af1d057..1ef9255 100644 --- a/graphrag/app/agent/agentic_synthesizer.py +++ b/graphrag/app/agent/agentic_synthesizer.py @@ -20,6 +20,7 @@ the out-of-corpus honesty match classic mode. """ +import json import logging from agent.agent_generation import TigerGraphAgentGenerator @@ -30,9 +31,23 @@ logger = logging.getLogger(__name__) -def _without_seen(ctx, seen: set): - """``ctx`` minus the retrieved entries already in ``seen`` (which it then - extends), or ``None`` when nothing new is left. ``ctx`` is not modified.""" +def _pieces(value) -> list: + """The separate texts in a ``final_retrieval`` value, as comparable strings.""" + if isinstance(value, str): + return [value] + if isinstance(value, dict) and "content" in value: + return [str(value["content"])] + if isinstance(value, list): + return [v if isinstance(v, str) else json.dumps(v, sort_keys=True, default=str) + for v in value] + return [json.dumps(value, sort_keys=True, default=str)] + + +def _without_seen(ctx, seen: dict): + """``ctx`` minus the texts already sent for the same id, or ``None`` when + nothing new is left. ``seen`` maps id -> texts sent so far and is extended; + ``ctx`` is not modified. Comparing text, not just ids, keeps a later step's + different or extra text for an id an earlier step also returned.""" result = ctx.get("result") if isinstance(ctx, dict) else None fr = result.get("final_retrieval") if isinstance(result, dict) else None if not isinstance(fr, dict): @@ -40,12 +55,25 @@ def _without_seen(ctx, seen: set): kept = {} for key, value in fr.items(): if is_grouped(value): - fresh = {k: v for k, v in value.items() if k not in seen} - seen.update(fresh) + fresh = {} + for inner_key, inner in value.items(): + text = str(inner.get("content", "")) + sent = seen.setdefault(inner_key, set()) + if text not in sent: + sent.add(text) + fresh[inner_key] = inner if fresh: kept[key] = fresh - elif key == "Similarity_Context" or key not in seen: - seen.add(key) + continue + sent = seen.setdefault(key, set()) + pieces = _pieces(value) + new = [p for p in pieces if p not in sent] + if not new: + continue + sent.update(new) + if isinstance(value, list) and len(new) < len(pieces): + kept[key] = [v for v, p in zip(value, pieces) if p in new] + else: kept[key] = value if not kept: return None @@ -60,7 +88,7 @@ def _gather(results: dict, log: bool = False) -> dict: once; a step left with nothing new is dropped. """ structural, unstructured = [], [] - seen: set = set() + seen: dict = {} skipped = 0 for sr in results.values(): if not sr.ok or sr.context is None: diff --git a/graphrag/tests/test_answer_context_dedup.py b/graphrag/tests/test_answer_context_dedup.py index 3ad7e28..06d4563 100644 --- a/graphrag/tests/test_answer_context_dedup.py +++ b/graphrag/tests/test_answer_context_dedup.py @@ -51,6 +51,31 @@ def test_grouped_results_lose_only_repeated_chunks(self): grouped = _fr(_gather(results)["unstructured"][1]) self.assertEqual(grouped, {"a_chunk_14": {"a_chunk_13": {"distance": "-1", "content": "thirteen"}}}) + def test_same_id_with_different_text_is_kept(self): + """Hybrid seeds and communities can return the same id with other text.""" + results = { + "S1": _step("S1", "GraphRAG_Community_Vector_Search", {"c_1": "summary v1"}), + "S2": _step("S2", "GraphRAG_Community_Vector_Search", {"c_1": "summary v2"}), + } + unstructured = _gather(results)["unstructured"] + self.assertEqual([_fr(c) for c in unstructured], [{"c_1": "summary v1"}, {"c_1": "summary v2"}]) + + def test_only_new_pieces_of_a_list_are_kept(self): + results = { + "S1": _step("S1", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_1": ["text", "Entity: A"]}), + "S2": _step("S2", "GraphRAG_Hybrid_Vector_Search", + {"a_chunk_1": ["text", "Entity: A", "Entity: B"]}), + } + self.assertEqual(_fr(_gather(results)["unstructured"][1]), {"a_chunk_1": ["Entity: B"]}) + + def test_flat_and_grouped_copies_of_one_chunk_dedupe(self): + results = { + "S1": _step("S1", "Chunk_Sibling_Vector_Search", + {"a_chunk_2": {"a_chunk_2": {"distance": "0", "content": "two"}}}), + "S2": _step("S2", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_2": ["two"], "a_chunk_3": ["three"]}), + } + self.assertEqual(_fr(_gather(results)["unstructured"][1]), {"a_chunk_3": ["three"]}) + def test_step_results_are_not_modified(self): results = { "S1": _step("S1", "GraphRAG_Hybrid_Vector_Search", {"a_chunk_1": ["one"]}), diff --git a/graphrag/tests/test_context_token_limit.py b/graphrag/tests/test_context_token_limit.py index 933ccad..95d46ef 100644 --- a/graphrag/tests/test_context_token_limit.py +++ b/graphrag/tests/test_context_token_limit.py @@ -88,9 +88,26 @@ def test_profile_limit_less_answer_room_and_margin(self): ) self.assertEqual( svc.context_token_limit(), - int(128000 * self.base_llm._TOKENIZER_MARGIN) - 16384, + int(128000 * self.base_llm._TOKENIZER_MARGIN) - self.base_llm._ANSWER_RESERVE_CAP, ) + def test_large_output_budgets_reserve_only_a_capped_answer(self): + """Reasoning models declare 100K+ output budgets; reserving all of it + left gpt-5 about 35K and o3 about 20K tokens for context.""" + for max_in, max_out in ((272000, 128000), (200000, 100000), (272000, 272000)): + svc = self._service({"llm_model": "m"}, + {"max_input_tokens": max_in, "max_output_tokens": max_out}) + self.assertEqual( + svc.context_token_limit(), + int(max_in * self.base_llm._TOKENIZER_MARGIN) - self.base_llm._ANSWER_RESERVE_CAP, + ) + + def test_small_window_keeps_at_least_half_for_context(self): + svc = self._service({"llm_model": "m"}, + {"max_input_tokens": 8192, "max_output_tokens": 8192}) + budget = int(8192 * self.base_llm._TOKENIZER_MARGIN) + self.assertEqual(svc.context_token_limit(), budget - budget // 2) + def test_unknown_model_keeps_no_limit_and_warns_once(self): self.base_llm._warned_no_token_limit.discard("custom-model") svc = self._service({"llm_model": "custom-model", "token_limit": 0}, None) @@ -149,6 +166,25 @@ def invoke_with_parser(self, prompt, parser, inputs, **kwargs): self.assertLess(calc.count_tokens(seen["context"]), calc.count_tokens("fact " * 10000)) self.assertLessEqual(calc.count_tokens(seen["prompt"]), 3000) + def test_prompt_that_cannot_render_gets_the_fallback_answer(self): + class BracesService: + config = {"llm_model": "gpt-4o-mini"} + # A custom prompt with stray braces cannot be rendered. + chatbot_response_prompt = "{format_instructions}\n{question}\n{query}\n{context}\nUse {oops}" + + def context_token_limit(self): + return 3000 + + @staticmethod + def _salvage_answer_output(raw): + return raw + + def invoke_with_parser(self, *args, **kwargs): + raise AssertionError("not reached") + + answer = self.Generator(BracesService()).generate_answer("q", "ctx") + self.assertIn("I wasn't able to generate an answer", answer.generated_answer) + def test_failure_log_names_the_provider_error(self): class FailingService: config = {"llm_model": "gpt-4o-mini"} From 006007098b12c3fd4e749cd8942a2fdaf83a81cb Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 7 Oct 2026 14:00:55 -0700 Subject: [PATCH 06/15] Plan safely with models that answer without calling tools - Fall back to parsing the planner's text when tool-calling structured output comes back empty instead of failing the request - Count a model as tool-calling only when it actually calls the tool, asking a second time with the call required before deciding Refs: GML-2302 --- CHANGELOG.md | 1 + common/llm_services/base_llm.py | 5 ++ common/llm_services/capabilities.py | 31 +++++++---- .../test_agentic_plan_structured_output.py | 18 +++++++ graphrag/tests/test_tool_calling_probe.py | 54 +++++++++++++++++++ 5 files changed, 99 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2119ab2..74d36b6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - **Chat conversations can be deleted one at a time, in a selected group, or all at once.** Each conversation in the chat sidebar has a delete button, *Select* lets several conversations be picked (or all shown ones at once) and deleted together, and a new *Clear all* action removes the whole history. The only option before, *Clear older*, always kept the ten most recent conversations and appeared only once there were more than ten. ### Fixed +- **Agentic mode is used only with models that actually call tools.** The capability check counted any accepted request as support, so an OpenAI-compatible server that ignores tools ran in Agentic mode and then failed to plan. The check now requires the model to call the tool, asking a second time with the call required, and such models use Classic mode. - **Contextual search returns each passage once, with its correct position.** When two matched chunks were close together in a document, the passages around them were returned under both matches with wrong, inconsistent positions (a match could appear several places away from itself), so the answer saw duplicated context. Each chunk is now listed once, under the nearest match, with its exact offset. Existing graphs pick this up by running Repair in the Migration Assistant. - **Knowledge-graph rebuild works when signed in with a token.** A rebuild started from a token sign-in could fail at once, reported as the graph not existing or not being accessible. It now runs normally, and an invalid token is rejected as an authentication failure rather than a server error. - **Contextual document search honours the search options it was given.** The option to use a hypothetical-answer embedding and the option to expand the question were swapped before reaching contextual (sibling-chunk) search, so asking for one applied the other — and each selects a different search strategy, not a variation of one. Requests that set neither option, including all chat traffic, were unaffected. diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index 19c1c22..b092b64 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -662,6 +662,11 @@ def invoke_structured( else: structured = self.llm.with_structured_output(schema, method=method) result = structured.invoke(messages) + if result is None: + # Tool-calling structured output returns None, not an + # error, when the model answers without calling the tool + # (common on OpenAI-compatible servers); parse instead. + raise ValueError("model returned no structured result") except Exception as exc: logger.warning( f"{caller_name}: structured output failed ({exc}); " diff --git a/common/llm_services/capabilities.py b/common/llm_services/capabilities.py index 0d5f326..36140c1 100644 --- a/common/llm_services/capabilities.py +++ b/common/llm_services/capabilities.py @@ -241,27 +241,35 @@ def _looks_like_no_tool_support(exc) -> bool: return any(m in msg for m in _NO_TOOL_SUPPORT_MARKERS) -def _invoke_probe(llm): +def _invoke_probe(llm, force: bool = False) -> bool: + """Ask the model to call a trivial tool; True when it actually did. + + ``force`` requires the call (``tool_choice``). A reply without a tool call + is not an error, so success means a tool call in the response, not just a + request the endpoint accepted. + """ from pydantic import BaseModel, Field class _ProbePing(BaseModel): """Acknowledge readiness by calling this tool.""" ok: bool = Field(default=True, description="always true") - bound = llm.bind_tools([_ProbePing]) - bound.invoke([ + bound = (llm.bind_tools([_ProbePing], tool_choice="any") if force + else llm.bind_tools([_ProbePing])) + resp = bound.invoke([ ("system", "You can call tools."), ("user", "Call the ProbePing tool with ok set to true."), ]) + return bool(getattr(resp, "tool_calls", None)) def _run_tool_calling_probe(llm_provider): - """Bind a trivial tool and make a minimal, time-bounded call. + """Check that the model calls a tool when asked, within a time bound. - Returns ``True`` (confirmed tool-calling), ``False`` (the model explicitly - rejects tools — a *confident* no-support signal), or ``None`` (unknown: - transient, ambiguous, timeout, or no usable provider). ``None`` must not - disable Agentic mode. + Returns ``True`` (it called the tool), ``False`` (confident no-support: the + model rejects tools, or answered without the tool call even when the call + was required), or ``None`` (unknown: transient, ambiguous, timeout, or no + usable provider). ``None`` must not disable Agentic mode. """ import concurrent.futures @@ -270,8 +278,11 @@ def _run_tool_calling_probe(llm_provider): return None try: with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex: - ex.submit(_invoke_probe, llm).result(timeout=_PROBE_TIMEOUT_S) - return True + for force in (False, True): + if ex.submit(_invoke_probe, llm, force).result(timeout=_PROBE_TIMEOUT_S): + return True + logger.info("tool-calling probe: model answered without calling the tool, even when required") + return False except concurrent.futures.TimeoutError: logger.warning("tool-calling probe timed out; leaving Agentic enabled") return None diff --git a/graphrag/tests/test_agentic_plan_structured_output.py b/graphrag/tests/test_agentic_plan_structured_output.py index 97c3c28..37ec47f 100644 --- a/graphrag/tests/test_agentic_plan_structured_output.py +++ b/graphrag/tests/test_agentic_plan_structured_output.py @@ -112,6 +112,24 @@ def with_structured_output(self, schema, **kwargs): self.assertIs(result, plan) return seen + def test_empty_structured_result_falls_back_to_parsing(self): + """A model that answers without the tool call yields None, not an error.""" + from common.llm_services.base_llm import LLM_Model + plan_json = self.Plan().model_dump_json() + + class NoToolCallLLM: + def with_structured_output(self, schema, **kwargs): + return types.SimpleNamespace(invoke=lambda messages: None) + + def invoke(self, messages): + return types.SimpleNamespace(content=plan_json) + + svc = object.__new__(LLM_Model) + svc.llm, svc.config = NoToolCallLLM(), {} + result = svc.invoke_structured([("user", "q")], self.Plan, caller_name="t", + method="function_calling") + self.assertIsInstance(result, self.Plan) + def test_no_method_keeps_the_provider_default(self): self.assertEqual(self._run(None), {}) diff --git a/graphrag/tests/test_tool_calling_probe.py b/graphrag/tests/test_tool_calling_probe.py index 67cf97c..4d27c66 100644 --- a/graphrag/tests/test_tool_calling_probe.py +++ b/graphrag/tests/test_tool_calling_probe.py @@ -77,3 +77,57 @@ def test_no_tool_support_error_is_confident(self): if __name__ == "__main__": unittest.main() + + +class _FakeChat: + """A chat model whose tool calls are scripted per attempt.""" + + def __init__(self, calls_tool_when, raise_on_force=None): + self.calls_tool_when = calls_tool_when # {"free": bool, "forced": bool} + self.raise_on_force = raise_on_force + self.attempts = [] + + def bind_tools(self, tools, tool_choice=None): + forced = tool_choice is not None + self.attempts.append("forced" if forced else "free") + if forced and self.raise_on_force: + raise self.raise_on_force + calls = self.calls_tool_when["forced" if forced else "free"] + + class _Bound: + def invoke(_self, messages): + tool_calls = [{"name": "_ProbePing", "args": {"ok": True}, "id": "1"}] if calls else [] + return mock.Mock(tool_calls=tool_calls) + + return _Bound() + + +class TestProbeRequiresAToolCall(unittest.TestCase): + """Accepting the request is not enough; the model must call the tool.""" + + def _probe(self, chat): + return cap._run_tool_calling_probe(mock.Mock(llm=chat)) + + def test_tool_call_on_first_attempt(self): + chat = _FakeChat({"free": True, "forced": True}) + self.assertIs(self._probe(chat), True) + self.assertEqual(chat.attempts, ["free"]) + + def test_tool_call_only_when_forced_still_counts(self): + chat = _FakeChat({"free": False, "forced": True}) + self.assertIs(self._probe(chat), True) + self.assertEqual(chat.attempts, ["free", "forced"]) + + def test_no_tool_call_even_when_forced_disables(self): + """An endpoint that accepts tools but ignores them.""" + chat = _FakeChat({"free": False, "forced": False}) + self.assertIs(self._probe(chat), False) + + def test_forcing_rejected_as_unsupported_disables(self): + chat = _FakeChat({"free": False, "forced": True}, + raise_on_force=ValueError("tool_choice is not supported")) + self.assertIs(self._probe(chat), False) + + def test_ambiguous_error_stays_unknown(self): + chat = _FakeChat({"free": False, "forced": True}, raise_on_force=RuntimeError("502 Bad Gateway")) + self.assertIsNone(self._probe(chat)) From e17612f4549adeec38dc741c1137b033ac07cce1 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 7 Oct 2026 14:02:45 -0700 Subject: [PATCH 07/15] Ask clarifying questions before planning - Unclear follow-ups get one short clarifying question from triage, which can reply to the user, instead of from the planner, which cannot - Triage reads the most recent turns of the conversation, up to a larger budget, so it sees the subjects a follow-up refers to - The planner still resolves a follow-up's subject from the conversation Refs: GML-2191 --- common/llm_services/base_llm.py | 6 +- graphrag/app/agent/agentic_agent.py | 41 ++++++++++-- graphrag/tests/test_clarify_at_triage.py | 81 ++++++++++++++++++++++++ 3 files changed, 122 insertions(+), 6 deletions(-) create mode 100644 graphrag/tests/test_clarify_at_triage.py diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index b092b64..1b9a8de 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -1147,7 +1147,7 @@ def agentic_agent_prompt(self): # Strategy (operator-customizable) — moved out of the fixed rules so it can # be tuned without touching the role / act model / plan mechanics. _AGENTIC_PLANNER_USER_DEFAULT = """\ -- Before building the plan, check whether the question is self-contained: can it be fully understood without reading ## Conversation? If the subject, entity, or topic is not named explicitly in the question, find the most recent relevant entity from ## Conversation and substitute its full name in every step's args. If ## Conversation has multiple candidates and it is genuinely unclear which one the user means, plan only a final answer step (no retrieval) that asks the user one short clarifying question. +- Before building the plan, check whether the question is self-contained: can it be fully understood without reading ## Conversation? If the subject, entity, or topic is not named explicitly in the question, find the most recent relevant entity from ## Conversation and substitute its full name in every step's args. If ## Conversation has multiple candidates, use the most recent one (a genuinely ambiguous reference is clarified with the user before planning). - When a question has multiple clauses, assign each clause to its own retrieval step. If another clause still needs passages or typed graph facts after one is covered, plan hybrid/community/structural for that clause too. - Prioritize including at least one vector search step (graphrag__hybrid_search or graphrag__contextual_search) unless you are highly confident the question is a pure structured-data request — an exact count, an attribute/id lookup, a relationship traversal, or an aggregation over typed graph data — that a generated graph query fully answers on its own. Whenever the answer could plausibly live in document text (what/why/how/describe/summarize, definitions, explanations, figures), include a vector search step. When unsure, include vector search. - Use BOTH structural and unstructured kinds when a question needs facts from the graph AND supporting text; you may run several of each, in any order. When you use STRUCTURAL, pair it with a vector search step unless the question is a pure structured-data request. @@ -1168,7 +1168,7 @@ def agentic_planner_prompt(self): You are the front desk for an agentic assistant. The agent behind you has tools: it retrieves from a TigerGraph knowledge base and may also have external tools attached (e.g. weather, web, or other data sources). Decide whether the user's latest message can be answered directly without any lookup, or needs the agent to retrieve or call a tool: -- needs_retrieval=false WITH a brief, friendly direct answer when the message is purely conversational per the routing policy below; +- needs_retrieval=false WITH a brief, friendly direct answer when the message is purely conversational per the routing policy below, or WITH one short clarifying question when the routing policy says the message is too unclear to look up; - needs_retrieval=true WITH an empty answer otherwise — the agent will then pick the right tool, or honestly report it cannot answer. When unsure, choose needs_retrieval=true. Match the user's language. @@ -1188,6 +1188,8 @@ def agentic_planner_prompt(self): - broad questions about what the data CONTAINS or is ABOUT — e.g. "what is this graph about?", "what data is in the graph?", "what topics are covered?", "summarize the documents"; - anything else a tool might fetch (weather, current events, a calculation, etc.). +- UNCLEAR FOLLOW-UP — the message refers back to something in ## Conversation ("it", "that one", "the company", "what about the other?") and ## Conversation offers more than one candidate, so it is genuinely unclear which one the user means. Answer with one short clarifying question that names the candidates. If ## Conversation makes the reference clear, the message is INFORMATIONAL, not this bucket. + Key distinction: a question about the ASSISTANT's capabilities is CONVERSATIONAL; a question about the DATA's contents (what is in the graph, or what it is about) is INFORMATIONAL — never deflect those. Do not deflect an informational question just because it looks outside the knowledge base — the agent may have a tool that answers it.""" @property diff --git a/graphrag/app/agent/agentic_agent.py b/graphrag/app/agent/agentic_agent.py index b585c6e..a319c39 100644 --- a/graphrag/app/agent/agentic_agent.py +++ b/graphrag/app/agent/agentic_agent.py @@ -51,23 +51,56 @@ class _Triage(BaseModel): """Front-desk classification of a user message before any DB/MCP work.""" needs_retrieval: bool = Field( description="True if answering requires looking up the user's data in the " - "knowledge graph. False for greetings, small talk, thanks/goodbye, or " - "questions about the assistant itself (who/what are you, what can you do)." + "knowledge graph. False for greetings, small talk, thanks/goodbye, " + "questions about the assistant itself (who/what are you, what can you do), " + "or a follow-up that is too unclear to look up without asking which " + "subject the user means." ) answer: str = Field( default="", description="When needs_retrieval is False, the complete direct answer to " - "give the user. Empty when needs_retrieval is True.", + "give the user, or one short clarifying question. Empty when " + "needs_retrieval is True.", ) +# Conversation budget for triage, in characters of the JSON sent. +_TRIAGE_CONVO_CHARS = 4000 + + +def _recent_conversation(convo, limit: int = _TRIAGE_CONVO_CHARS) -> str: + """The most recent turns that fit ``limit``, as JSON, oldest first. + + Follow-ups refer to the latest turns, so the budget is spent from the + end. If even the newest turn is too long, its response is shortened and + its question kept. + """ + def dumps(turns): + return json.dumps(turns, ensure_ascii=False) + + kept = [] + for turn in reversed(convo or []): + if len(dumps([turn] + kept)) > limit: + break + kept.insert(0, turn) + if not kept and convo: + newest = dict(convo[-1]) + room = limit - len(dumps([{**newest, "response": ""}])) + newest["response"] = str(newest.get("response", ""))[:max(room, 0)] + # JSON escaping can still push it over; trim until it fits. + while newest["response"] and len(dumps([newest])) > limit: + newest["response"] = newest["response"][: int(len(newest["response"]) * 0.9)] + kept = [newest] + return dumps(kept) + + def _triage_question(llm, question, convo): """One cheap classify-and-answer call. Returns a ``_Triage`` or ``None`` if triage itself fails (caller then proceeds with normal retrieval). Uses only the question + conversation — no schema, no MCP, no DB.""" try: user = ( - f"## Conversation\n{json.dumps(convo or [])[:2000]}\n\n" + f"## Conversation\n{_recent_conversation(convo)}\n\n" f"## Message\n{question}" ) # Customizable routing prompt (fixed contract + operator-editable diff --git a/graphrag/tests/test_clarify_at_triage.py b/graphrag/tests/test_clarify_at_triage.py new file mode 100644 index 0000000..f49fa27 --- /dev/null +++ b/graphrag/tests/test_clarify_at_triage.py @@ -0,0 +1,81 @@ +import json +import unittest + + +class TestClarifyBeforePlanning(unittest.TestCase): + """An unclear follow-up is clarified at triage, which can reply to the + user; the planner's output is a plan, so it cannot ask.""" + + @classmethod + def setUpClass(cls): + try: + from common.llm_services.base_llm import LLM_Model + from agent.agentic_agent import _Triage + except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.M, cls.Triage = LLM_Model, _Triage + + def test_triage_policy_asks_on_an_unclear_follow_up(self): + policy = self.M._AGENTIC_TRIAGE_USER_DEFAULT + self.assertIn("UNCLEAR FOLLOW-UP", policy) + self.assertIn("clarifying question", policy) + self.assertIn("clarifying question", self.M._AGENTIC_TRIAGE_SYSTEM) + + def test_triage_output_allows_a_question_instead_of_retrieval(self): + fields = self.Triage.model_fields + self.assertIn("clarifying question", fields["answer"].description) + self.assertIn("too unclear", fields["needs_retrieval"].description) + + def test_planner_no_longer_plans_a_question(self): + planner = self.M._AGENTIC_PLANNER_USER_DEFAULT + self.assertNotIn("plan only a final answer step (no retrieval)", planner) + self.assertIn("use the most recent one", planner) + + + +class TestTriageSeesTheLatestTurns(unittest.TestCase): + """Follow-ups refer to the latest turns, so those are what triage gets.""" + + @classmethod + def setUpClass(cls): + try: + from agent.agentic_agent import _recent_conversation, _TRIAGE_CONVO_CHARS + except ImportError as exc: + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.recent, cls.limit = staticmethod(_recent_conversation), _TRIAGE_CONVO_CHARS + + @staticmethod + def _turn(i, size=1500): + return {"query": f"question {i}", "response": f"answer {i} " + "x" * size} + + def test_short_conversation_is_sent_whole(self): + convo = [self._turn(1, 100), self._turn(2, 100)] + self.assertEqual(json.loads(self.recent(convo)), convo) + + def test_long_conversation_keeps_the_newest_turns_in_order(self): + convo = [self._turn(i) for i in range(1, 8)] + sent = self.recent(convo) + self.assertLessEqual(len(sent), self.limit) + kept = json.loads(sent) + self.assertEqual(kept, convo[-len(kept):]) + self.assertEqual(kept[-1]["query"], "question 7") + self.assertGreaterEqual(len(kept), 2) + + def test_oversized_newest_turn_keeps_its_question(self): + convo = [self._turn(1), {"query": "the latest question", "response": 'say "hi"\n' * 2000}] + sent = self.recent(convo) + self.assertLessEqual(len(sent), self.limit) + kept = json.loads(sent) + self.assertEqual([t["query"] for t in kept], ["the latest question"]) + + def test_non_english_text_is_counted_as_written(self): + convo = [{"query": "質問", "response": "日" * 1500}] + self.assertIn("日" * 1500, self.recent(convo)) + + def test_empty_conversation(self): + self.assertEqual(self.recent([]), "[]") + self.assertEqual(self.recent(None), "[]") + + +if __name__ == "__main__": + unittest.main() From d5111e4b5d3ea15977c2c9f24f021690f9513c4d Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 7 Oct 2026 14:03:24 -0700 Subject: [PATCH 08/15] Keep the conversation list consistent when it refreshes during a delete - Count deleted conversations against the list as it is when the delete lands, so a background refresh can't leave a duplicate or missing row - Disable the other delete actions while a single delete is in progress Refs: GML-2303 --- graphrag-ui/src/components/SideMenu.tsx | 41 +++++++++++++++++-------- 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/graphrag-ui/src/components/SideMenu.tsx b/graphrag-ui/src/components/SideMenu.tsx index ba317ce..e79f890 100644 --- a/graphrag-ui/src/components/SideMenu.tsx +++ b/graphrag-ui/src/components/SideMenu.tsx @@ -11,7 +11,7 @@ import { IoIosArrowForward } from "react-icons/io"; import { useTheme } from "@/components/ThemeProvider"; import { safeJson } from "@/utils/safeJson"; import { GoGear } from "react-icons/go"; -import { useState, useEffect } from "react"; +import { useState, useEffect, useRef } from "react"; import { Popover, PopoverContent, @@ -92,6 +92,13 @@ const SideMenu = ({ const [selecting, setSelecting] = useState(false); const [selected, setSelected] = useState>(new Set()); const [deletingSelected, setDeletingSelected] = useState(false); + const [deletingOne, setDeletingOne] = useState(false); + // Any delete in flight: keeps the other delete actions disabled meanwhile. + const busy = clearing || deletingOne; + // The list as last rendered, read when a delete lands: a background refresh + // may have replaced it since the delete was started. + const convListRef = useRef([]); + convListRef.current = convList; const [confirm, confirmDialog] = useConfirm(); const [alert, alertDialog] = useAlert(); // Fade + disable the side menu (conversation list + New Chat) while @@ -276,12 +283,14 @@ const SideMenu = ({ // one of them, start a new chat rather than keep showing it. const forgetConversations = (deleted: Set) => { if (deleted.size === 0) return; - const loadedGone = convList - .slice(0, loadedCount) - .filter((c: any) => deleted.has(c.conversation_id)).length; setConvList((prev: any[]) => prev.filter((c: any) => !deleted.has(c.conversation_id))); setConversationId((prev: any[]) => prev.filter((c: any) => !deleted.has(c.conversation_id))); - setLoadedCount((c) => c - loadedGone); + setLoadedCount((c) => { + const gone = convListRef.current + .slice(0, c) + .filter((x: any) => deleted.has(x.conversation_id)).length; + return Math.max(0, c - gone); + }); const openId = conversationManager.getCurrentConversationId(); if ((openId && deleted.has(openId)) || (activeConversationId && deleted.has(activeConversationId))) { handleNewChat(); @@ -293,8 +302,14 @@ const SideMenu = ({ "Delete this conversation?\n\nThis permanently removes it and cannot be undone." ); if (!ok) return; - const deleted = await deleteConversations([id]); - forgetConversations(deleted); + setDeletingOne(true); + let deleted: Set; + try { + deleted = await deleteConversations([id]); + forgetConversations(deleted); + } finally { + setDeletingOne(false); + } if (deleted.size === 0) { await alert("The conversation could not be deleted. Please try again."); } @@ -553,7 +568,7 @@ const SideMenu = ({ e.stopPropagation(); deleteOne(item.conversation_id); }} - disabled={clearing || loadingMore} + disabled={busy || loadingMore} aria-label="Delete conversation" title="Delete conversation" className="ml-1 flex-shrink-0 p-1 rounded text-gray-400 opacity-40 group-hover:opacity-100 focus:opacity-100 hover:text-red-600 hover:bg-gray-200 dark:hover:bg-gray-700 dark:hover:text-red-400 disabled:opacity-50" @@ -592,7 +607,7 @@ const SideMenu = ({ {loadedCount < convList.length ? (