From 797ff619c557a14df273662af91d76828734f581 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 14 Oct 2026 12:00:42 -0700 Subject: [PATCH 01/32] Keep the agentic planner from looping on its plan summary - Ask for one short sentence in the plan summary and step reasons, and shorten an overlong value instead of failing the plan - Cap the planner's output tokens, with a new planner_max_output_tokens setting and a higher floor for reasoning models Refs: GML-2324 --- CHANGELOG.md | 8 + README.md | 1 + common/llm_services/base_llm.py | 37 ++++- common/py_schemas/schemas.py | 57 ++++++- graphrag/app/agent/agentic_planner.py | 20 +++ .../tests/test_planner_repetition_loop.py | 147 ++++++++++++++++++ 6 files changed, 263 insertions(+), 7 deletions(-) create mode 100644 graphrag/tests/test_planner_repetition_loop.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 6bd9daf..955c865 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,13 @@ # Changelog +## [2.0.4] + +### Fixed +- **Agentic planning no longer gets stuck writing its plan summary.** The planner's summary and step reasons had no stated purpose or length, and the planner's output had no limit, so a model could repeat a phrase until it reached its own output maximum — costing thousands of tokens and tens of seconds, and sometimes the plan itself. These fields now ask for one short sentence, an overlong value is shortened instead of failing the plan, and the planner's output is capped. + +### Added +- **`planner_max_output_tokens` setting.** The most output tokens the planner may use for one plan (default `4000`; reasoning models use at least `16000`). + ## [2.0.3] ### Added diff --git a/README.md b/README.md index 1e70505..ab64e5a 100644 --- a/README.md +++ b/README.md @@ -546,6 +546,7 @@ Copy the below code into `configs/server_config.json`. You shouldn’t need to c | `agent_style` | string | `"planned"` | Default agentic engine style: `"planned"` (plan the whole retrieval up front) or `"reactive"` (decide each step from the last result). The chat menu can override per request. See [Chat Engines and Agents](#chat-engines-and-agents). | | `agent_max_iterations` | int | `30` | Reactive agent only: maximum reason-act-observe steps before it must answer. | | `agent_max_replans` | int | `3` | Planned agent only: how many times the planner may extend its plan when the gathered context is insufficient. | +| `planner_max_output_tokens` | int | `4000` | Planned agent only: the most output tokens the planner may use for one plan. Reasoning models use at least `16000`. | | `agent_max_total_steps` | int | `20` | Planned agent only: hard cap on executed retrieval steps across all replans. | | `chunk_only` | bool | `true` | If true, hybrid search only retrieves document chunks, excluding entity data. | | `doc_only` | bool | `false` | If true, hybrid search retrieves whole documents instead of chunks. Significantly increases context size. | diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index 1b9a8de..5c2e2e5 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -638,12 +638,38 @@ def invoke_with_tools( _record_usage(caller_name, usage_data, self.config) return resp + def with_output_limit(self, max_output_tokens: int): + """The chat model with its output capped at ``max_output_tokens``. + + A per-call copy; the service's own model is not changed. The limit + field differs by provider (``max_tokens``, ``max_output_tokens``, + ``num_predict``). An existing lower limit is kept. Unknown models are + returned unchanged. + """ + llm = self.llm + if not max_output_tokens or llm is None or not hasattr(llm, "model_copy"): + return llm + model_kwargs = getattr(llm, "model_kwargs", None) + if isinstance(model_kwargs, dict) and "max_tokens" in model_kwargs: + # Bedrock: the service resolves its limit into model_kwargs. + current = model_kwargs["max_tokens"] or max_output_tokens + return llm.model_copy(update={"model_kwargs": { + **model_kwargs, "max_tokens": min(current, max_output_tokens)}}) + fields = getattr(type(llm), "model_fields", {}) + for name in ("max_tokens", "max_output_tokens", "num_predict"): + if name in fields: + current = getattr(llm, name, None) + limit = min(current, max_output_tokens) if current else max_output_tokens + return llm.model_copy(update={name: limit}) + return llm + def invoke_structured( self, messages: list, schema, caller_name: str = "unknown", method: str = None, + max_output_tokens: int = None, ): """Invoke the chat model with native structured output. @@ -652,15 +678,17 @@ def invoke_structured( 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. + keeps the provider's default. ``max_output_tokens`` caps the output of + this call (see ``with_output_limit``). """ + llm = self.with_output_limit(max_output_tokens) if max_output_tokens else self.llm usage_data = {} with get_openai_callback() as cb: try: if method is None: - structured = self.llm.with_structured_output(schema) + structured = llm.with_structured_output(schema) else: - structured = self.llm.with_structured_output(schema, method=method) + structured = llm.with_structured_output(schema, method=method) result = structured.invoke(messages) if result is None: # Tool-calling structured output returns None, not an @@ -673,7 +701,7 @@ def invoke_structured( "falling back to parser" ) parser = PydanticOutputParser(pydantic_object=schema) - raw = self.llm.invoke(messages) + raw = llm.invoke(messages) raw_text = self._message_text(raw) result = self._parse_or_repair(parser, raw_text, caller_name) usage_data["input_tokens"] = cb.prompt_tokens @@ -1134,6 +1162,7 @@ def agentic_agent_prompt(self): - 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. +- Write every step first. Then set "strategy" to one short sentence (under 25 words). Give each step a one-sentence "rationale". Decide which retrievals to include, how many, and in what order using the "Retrieval Strategy" below. Return ONLY the structured plan. diff --git a/common/py_schemas/schemas.py b/common/py_schemas/schemas.py index 6ee296f..fde19f8 100644 --- a/common/py_schemas/schemas.py +++ b/common/py_schemas/schemas.py @@ -13,9 +13,10 @@ # along with this program. If not, see . import enum +import logging from typing import Dict, List, Optional, Union -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator class ConversationTurn(BaseModel): @@ -72,6 +73,37 @@ class GraphRAGResponse(BaseModel): # --- Agentic engine (v2.0 deep-thinking mode) ------------------------------ +logger = logging.getLogger(__name__) + +# Length bound for the planner's free-text fields. Without a purpose and a +# bound, a model can write these as an endless repetition loop (GML-2324). +_PLAN_TEXT_MAX = 200 + + +def _collapse_repeats(text: str, min_words: int = 4, max_words: int = 60) -> str: + """Cut ``text`` at the first phrase of ``min_words``–``max_words`` words + that repeats three times in a row, keeping one copy of it.""" + words = text.split() + for i in range(len(words)): + for n in range(min_words, max_words + 1): + if i + 3 * n > len(words): + break + if words[i:i + n] == words[i + n:i + 2 * n] == words[i + 2 * n:i + 3 * n]: + return " ".join(words[:i + n]) + return text + + +def _short_plan_text(value, field: str): + """Shorten an overlong plan text instead of rejecting the whole plan.""" + if not isinstance(value, str) or len(value) <= _PLAN_TEXT_MAX: + return value + text = _collapse_repeats(value) + if len(text) > _PLAN_TEXT_MAX: + text = text[:_PLAN_TEXT_MAX - 1].rstrip() + "…" + logger.warning(f"agentic plan: shortened {field} from {len(value):,} to {len(text):,} chars") + return text + + class PlanStep(BaseModel): """One step in an agentic plan DAG. @@ -86,12 +118,31 @@ class PlanStep(BaseModel): args: Dict = {} arg_bindings: Dict[str, str] = {} depends_on: List[str] = [] - rationale: str = "" + rationale: str = Field( + "", + description="One short sentence on why this step is needed.", + json_schema_extra={"maxLength": _PLAN_TEXT_MAX}, + ) + + @field_validator("rationale", mode="before") + @classmethod + def _short_rationale(cls, v): + return _short_plan_text(v, "rationale") class Plan(BaseModel): steps: List[PlanStep] = [] - strategy: str = "" # one-line, user-facing summary + strategy: str = Field( + "", + description="One short sentence (under 25 words) on the approach. " + "Write it after the steps.", + json_schema_extra={"maxLength": _PLAN_TEXT_MAX}, + ) + + @field_validator("strategy", mode="before") + @classmethod + def _short_strategy(cls, v): + return _short_plan_text(v, "strategy") class StepResult(BaseModel): diff --git a/graphrag/app/agent/agentic_planner.py b/graphrag/app/agent/agentic_planner.py index 8d6c578..5594429 100644 --- a/graphrag/app/agent/agentic_planner.py +++ b/graphrag/app/agent/agentic_planner.py @@ -25,6 +25,7 @@ import json import logging +from common.llm_services.capabilities import model_capabilities from common.py_schemas import Plan from tools import tool_registry as registry @@ -91,6 +92,24 @@ def _sanitize(plan: Plan, ctx=None) -> Plan: return plan +# A normal plan is a few hundred tokens; the cap only ends a runaway +# generation early (GML-2324). Reasoning models spend output tokens on +# thinking, so they get a higher floor. +_PLANNER_MAX_OUTPUT_TOKENS = 4000 +_REASONING_PLANNER_MIN_OUTPUT_TOKENS = 16000 + + +def _planner_output_limit(llm, ctx) -> int: + cfg = getattr(ctx, "graphrag_cfg", None) or {} + try: + limit = int(cfg.get("planner_max_output_tokens") or _PLANNER_MAX_OUTPUT_TOKENS) + except (TypeError, ValueError): + limit = _PLANNER_MAX_OUTPUT_TOKENS + if model_capabilities(getattr(llm, "config", None)).get("supports_thinking"): + limit = max(limit, _REASONING_PLANNER_MIN_OUTPUT_TOKENS) + return limit + + def plan_question(llm, question, conversation=None, schema_rep="", prior_results=None, ctx=None) -> Plan: """Draft (or extend) a plan for ``question``. @@ -137,6 +156,7 @@ def plan_question(llm, question, conversation=None, schema_rep="", prior_results "function_calling" if getattr(llm, "strict_structured_output", False) else None ), + max_output_tokens=_planner_output_limit(llm, ctx), ) except Exception as exc: logger.warning(f"planner failed ({exc}); falling back to single hybrid step") diff --git a/graphrag/tests/test_planner_repetition_loop.py b/graphrag/tests/test_planner_repetition_loop.py new file mode 100644 index 0000000..5e5f668 --- /dev/null +++ b/graphrag/tests/test_planner_repetition_loop.py @@ -0,0 +1,147 @@ +import types +import unittest + +try: + from common.py_schemas.schemas import Plan, PlanStep, _PLAN_TEXT_MAX, _collapse_repeats + from common.llm_services.base_llm import LLM_Model +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + +# The cycle seen in a real trace (GML-2324), repeated as the model wrote it. +_CYCLE = ("here cleanly built step by step DAG design plan configuration object " + "instance declaration format correctly structured now ready for " + "evaluation execution lifecycle flow initialization now complete ") +_LOOP = ("The user is asking for an exact count of Jira tickets. This is a pure " + "structured-data count query suitable for graphrag__structural_retrieve " + + _CYCLE * 60) + + +class TestPlanTextFields(unittest.TestCase): + """strategy and rationale carry a purpose and a bound.""" + + def test_schema_tells_the_model_what_to_write(self): + props = Plan.model_json_schema()["properties"] + self.assertIn("One short sentence", props["strategy"]["description"]) + self.assertEqual(props["strategy"]["maxLength"], _PLAN_TEXT_MAX) + step = Plan.model_json_schema()["$defs"]["PlanStep"]["properties"] + self.assertIn("One short sentence", step["rationale"]["description"]) + + def test_runaway_strategy_is_shortened_not_rejected(self): + plan = Plan(strategy=_LOOP, steps=[PlanStep(id="A", kind="answer", tool="")]) + self.assertLessEqual(len(plan.strategy), _PLAN_TEXT_MAX) + self.assertTrue(plan.strategy.startswith("The user is asking")) + self.assertEqual(len(plan.steps), 1) + + def test_runaway_rationale_is_shortened(self): + step = PlanStep(id="S1", tool="graphrag__hybrid_search", rationale=_CYCLE * 20) + self.assertLessEqual(len(step.rationale), _PLAN_TEXT_MAX) + + def test_short_text_is_unchanged(self): + self.assertEqual(Plan(strategy="Count tickets, then answer.").strategy, + "Count tickets, then answer.") + + def test_collapse_keeps_one_copy_of_the_cycle(self): + out = _collapse_repeats(_LOOP) + self.assertEqual(out.count("lifecycle flow initialization"), 1) + + def test_collapse_leaves_normal_text_alone(self): + text = ("Search the documents for the policy, then search the documents " + "for the exceptions, then answer from the documents.") + self.assertEqual(_collapse_repeats(text), text) + + +class TestPlannerPrompt(unittest.TestCase): + def test_fixed_rule_asks_for_one_short_sentence(self): + self.assertIn('Then set "strategy" to one short sentence', + LLM_Model._AGENTIC_PLANNER_SYSTEM) + + +class TestOutputLimit(unittest.TestCase): + """The planner call caps output on every provider's model class.""" + + def _svc(self, llm): + svc = object.__new__(LLM_Model) + svc.llm, svc.config = llm, {} + return svc + + def test_openai_uses_max_tokens(self): + from langchain_openai import ChatOpenAI + capped = self._svc(ChatOpenAI(model="gpt-4o-mini", api_key="x")).with_output_limit(4000) + self.assertEqual(capped.max_tokens, 4000) + + def test_existing_lower_limit_is_kept(self): + from langchain_openai import ChatOpenAI + llm = ChatOpenAI(model="gpt-4o-mini", api_key="x", max_tokens=1000) + self.assertEqual(self._svc(llm).with_output_limit(4000).max_tokens, 1000) + + def test_gemini_uses_max_output_tokens(self): + from langchain_google_genai import ChatGoogleGenerativeAI + llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", google_api_key="x") + self.assertEqual(self._svc(llm).with_output_limit(4000).max_output_tokens, 4000) + + def test_ollama_uses_num_predict(self): + from langchain_community.llms import Ollama + self.assertEqual(self._svc(Ollama(model="llama3.1")).with_output_limit(4000).num_predict, 4000) + + def test_limit_in_model_kwargs_is_capped_there(self): + """Bedrock: the service resolves its limit into model_kwargs.""" + from langchain_openai import ChatOpenAI + llm = ChatOpenAI(model="gpt-4o-mini", api_key="x") + llm = llm.model_copy(update={"model_kwargs": {"max_tokens": 8192}}) + capped = self._svc(llm).with_output_limit(4000) + self.assertEqual(capped.model_kwargs["max_tokens"], 4000) + + def test_service_model_is_not_changed(self): + from langchain_openai import ChatOpenAI + svc = self._svc(ChatOpenAI(model="gpt-4o-mini", api_key="x")) + svc.with_output_limit(4000) + self.assertIsNone(svc.llm.max_tokens) + + def test_unknown_model_is_returned_unchanged(self): + sentinel = types.SimpleNamespace() + self.assertIs(self._svc(sentinel).with_output_limit(4000), sentinel) + + def test_structured_call_uses_the_capped_model(self): + seen = {} + + class FakeLLM: + max_tokens = None + model_fields = {"max_tokens": None} + + def model_copy(self, update): + seen.update(update) + return self + + def with_structured_output(self, schema, **kw): + return types.SimpleNamespace(invoke=lambda m: Plan()) + + svc = self._svc(FakeLLM()) + svc.invoke_structured([("user", "q")], Plan, caller_name="t", max_output_tokens=4000) + self.assertEqual(seen, {"max_tokens": 4000}) + + +class TestPlannerLimit(unittest.TestCase): + @classmethod + def setUpClass(cls): + try: + from agent.agentic_planner import _planner_output_limit + except ImportError as exc: + raise unittest.SkipTest(f"app environment unavailable: {exc}") + cls.limit = staticmethod(_planner_output_limit) + + def _llm(self, service, model): + return types.SimpleNamespace(config={"llm_service": service, "llm_model": model}) + + def test_default(self): + self.assertEqual(self.limit(self._llm("openai", "gpt-4o-mini"), types.SimpleNamespace(graphrag_cfg={})), 4000) + + def test_setting_overrides(self): + ctx = types.SimpleNamespace(graphrag_cfg={"planner_max_output_tokens": 6000}) + self.assertEqual(self.limit(self._llm("openai", "gpt-4o-mini"), ctx), 6000) + + def test_reasoning_model_gets_room_to_think(self): + self.assertEqual(self.limit(self._llm("openai", "o3"), types.SimpleNamespace(graphrag_cfg={})), 16000) + + +if __name__ == "__main__": + unittest.main() From 2047c1d17362af5ba22746d257d5b3006e4e608c Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 14 Oct 2026 12:11:04 -0700 Subject: [PATCH 02/32] Let agentic answers reuse what the conversation already holds - Planner and reactive agent decide first whether the conversation answers the question fully, partly, or not at all, and retrieve only what is missing - Questions about the current or latest state still retrieve - Answer step reads the recent conversation alongside retrieved context - Planner and reactive agent see the most recent turns, not the oldest - Trace Logs mark a plan that answers without retrieval Refs: GML-2325 --- CHANGELOG.md | 4 + README.md | 2 +- common/llm_services/base_llm.py | 13 ++ graphrag/app/agent/agentic_agent.py | 31 +--- graphrag/app/agent/agentic_graph.py | 27 +++- graphrag/app/agent/agentic_planner.py | 4 +- graphrag/app/agent/agentic_react.py | 3 +- graphrag/app/agent/agentic_synthesizer.py | 20 ++- graphrag/app/agent/conversation_window.py | 55 +++++++ .../test_agentic_conversation_routing.py | 143 ++++++++++++++++++ 10 files changed, 260 insertions(+), 42 deletions(-) create mode 100644 graphrag/app/agent/conversation_window.py create mode 100644 graphrag/tests/test_agentic_conversation_routing.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 955c865..f0b5057 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,7 +2,11 @@ ## [2.0.4] +### Changed +- **Agentic answers reuse what the conversation already holds.** Before, the agentic engines retrieved again for a question the conversation had already answered, and the planned engine's answer step did not see the conversation at all. The planner and the reactive agent now first decide how much of the answer the conversation holds: all of it is answered without retrieval, part of it gets retrieval only for the missing parts, and questions about the current or latest state still retrieve. The answer step reads the recent conversation together with the retrieved context, and the retrieved context wins where they disagree. Trace Logs mark a plan that answers without retrieval. + ### Fixed +- **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. - **Agentic planning no longer gets stuck writing its plan summary.** The planner's summary and step reasons had no stated purpose or length, and the planner's output had no limit, so a model could repeat a phrase until it reached its own output maximum — costing thousands of tokens and tens of seconds, and sometimes the plan itself. These fields now ask for one short sentence, an overlong value is shortened instead of failing the plan, and the planner's output is capped. ### Added diff --git a/README.md b/README.md index ab64e5a..8629b9d 100644 --- a/README.md +++ b/README.md @@ -389,7 +389,7 @@ GraphRAG offers two chat engines, chosen from the chat menu (or set as a per-gra ### Agentic -The agent decides its **own** retrieval instead of following a fixed pipeline: it picks which methods to use (structural graph queries, vector search, community search), can call external [MCP tools](#mcp-servers-agentic-tools), answers greetings and questions about itself directly, and cites the chunks and queries it used. It comes in two styles: +The agent decides its **own** retrieval instead of following a fixed pipeline: it picks which methods to use (structural graph queries, vector search, community search), can call external [MCP tools](#mcp-servers-agentic-tools), answers greetings and questions about itself directly, reuses what the conversation already answers and retrieves only what is missing, and cites the chunks and queries it used. It comes in two styles: - **Planned** *(default)* — analyzes the question up front and lays out the whole retrieval plan (which methods, how many, in what order) as a small DAG, executes it, then synthesizes one grounded answer. Predictable and efficient; a strong fit for most questions, including multi-part ones. - **Reactive** — reasons and acts one step at a time, deciding each next retrieval from what the previous step returned, and keeps going until it can answer completely and accurately. More adaptive to what it finds, at the cost of more steps (and tokens) on complex questions. diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index 5c2e2e5..d167aa0 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -1106,6 +1106,12 @@ def select_retriever_prompt(self): - If not — a part is still unanswered, a needed value or table is missing, or the results were thin — take another action to close the gap (follow a lead, widen top_k / num_hops, or switch method). Do not settle for a partial or vague answer when more retrieval could complete it. Do not commit to a full multi-step plan up front; let each next step be driven by what is still missing for a complete answer. +Before your first tool call, decide how much of the answer ## Conversation already holds: +- All of it: give the final answer without calling a tool. +- Part of it: call tools only for the missing parts. Do not retrieve again what ## Conversation already states. Combine both in the final answer. +- None of it: retrieve as described below. +Do not use an earlier answer when the question asks for the current or latest state (for example "now", "current", "latest", "again"), or when the user says that an earlier answer was wrong or incomplete. + The graph schema is required for the structural and unstructured query tools: before your first structural query or vector/unstructured retrieval, call graphrag__get_schema once to load the graph's vertex and edge types. Questions answered without graph data (e.g. by an external tool) do not need the schema. Run independent tool calls in parallel within one response; chain dependent calls across iterations. Cite specific findings from tool results in your final answer. @@ -1153,6 +1159,12 @@ def agentic_agent_prompt(self): The graph schema is NOT provided here — the structural and unstructured query tools load it themselves at run time, so plan retrieval steps directly. A question that needs no graph data should not include any graph-retrieval step (plan only the final answer step, or the relevant non-graph tool). +Before you plan retrieval, decide how much of the answer ## Conversation already holds: +- All of it: plan only the final answer step. +- Part of it: plan retrieval steps only for the missing parts. Do not retrieve again what ## Conversation already states. The answer step combines both. +- None of it: plan the retrieval steps as described below. +Do not use an earlier answer when the question asks for the current or latest state (for example "now", "current", "latest", "again"), or when the user says that an earlier answer was wrong or incomplete. + You have two kinds of retrieval: - STRUCTURAL (graphrag__structural_retrieve): generates and runs a graph query. Best for counts, lookups by attribute/id, relationships, and aggregations over typed data. It depends on the LLM generating a correct query against the live schema — it can return nothing or the wrong rows when the question doesn't map cleanly to typed graph data, so it is NOT a safe sole source of context. - UNSTRUCTURED (graphrag__hybrid_search / similarity_search / contextual_search / community_search): vector search over document text. Best for "what/why/how/describe/summarize" questions answered from passages. community_search suits broad/overall questions. @@ -1264,6 +1276,7 @@ def hyde_prompt(self): ## Rules - The contexts arrive as JSON key-context pairs. **Combine and rephrase** them to answer the question. - **Preserve** image links exactly as `![description](url)` in the final answer when used. Do NOT modify or omit them. +- The context "earlier_conversation", when present, holds earlier questions and answers from this chat. Use it to resolve references, and to answer when it holds the answer. When it conflicts with another context, the other context is correct. ## Inputs - **Question**: {question} diff --git a/graphrag/app/agent/agentic_agent.py b/graphrag/app/agent/agentic_agent.py index a319c39..4dd4b9c 100644 --- a/graphrag/app/agent/agentic_agent.py +++ b/graphrag/app/agent/agentic_agent.py @@ -23,7 +23,6 @@ """ import logging -import json import time from typing import Dict, List @@ -64,34 +63,8 @@ class _Triage(BaseModel): ) -# 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) +from agent.conversation_window import CONVERSATION_CHARS as _TRIAGE_CONVO_CHARS +from agent.conversation_window import recent_conversation as _recent_conversation def _triage_question(llm, question, convo): diff --git a/graphrag/app/agent/agentic_graph.py b/graphrag/app/agent/agentic_graph.py index 46aa0b5..743f660 100644 --- a/graphrag/app/agent/agentic_graph.py +++ b/graphrag/app/agent/agentic_graph.py @@ -26,7 +26,7 @@ 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 agent.agentic_synthesizer import _gather, has_context, synthesize, with_conversation from common.llm_services.base_llm import get_collected_usage from common.py_schemas import GraphRAGResponse, PlanStep @@ -55,6 +55,21 @@ def _hybrid_fallback_step() -> PlanStep: ) +def _plan_output(plan) -> dict: + """A plan as shown in the trace.""" + out = {"strategy": plan.strategy, "steps": [s.model_dump() for s in plan.steps]} + if not _retrieves(plan): + out["routing"] = "answered without retrieval" + return out + + +def _retrieves(plan) -> bool: + """Whether the plan has a step that runs a tool. A plan with only the + answer step is the planner's decision to answer from the conversation + (or without data), not a retrieval that found nothing.""" + return any(s.kind != "answer" and s.tool for s in plan.steps) + + def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: """Run the agentic workflow for one question and return a response. @@ -83,15 +98,14 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: "node": "plan", "kind": "plan", "duration_s": round(time.time() - _t0, 3), "input": {"question": question, "conversation": conversation or []}, - "output": {"strategy": plan.strategy, - "steps": [s.model_dump() for s in plan.steps]}, + "output": _plan_output(plan), "usage": _usage_since(_u0), }) if plan.strategy: emit(plan.strategy) replans = 0 - while True: + while _retrieves(plan): new_results, step_traces = execute_plan(plan, ctx) results.update(new_results) agent_steps.extend(step_traces) @@ -129,8 +143,7 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: "node": f"replan {replans}", "kind": "plan", "duration_s": round(time.time() - _t0, 3), "input": {"results_so_far": prior}, - "output": {"strategy": plan.strategy, - "steps": [s.model_dump() for s in plan.steps]}, + "output": _plan_output(plan), "usage": _usage_since(_u0), }) @@ -143,7 +156,7 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: agent_steps.append({ "node": "synthesize", "kind": "answer", "duration_s": round(time.time() - _t0, 3), - "input": with_trace_note(*fit_for_trace(_gather(results))), + "input": with_trace_note(*fit_for_trace(with_conversation(_gather(results), conversation))), "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 5594429..501d926 100644 --- a/graphrag/app/agent/agentic_planner.py +++ b/graphrag/app/agent/agentic_planner.py @@ -22,11 +22,11 @@ orchestrator). """ -import json import logging from common.llm_services.capabilities import model_capabilities from common.py_schemas import Plan +from agent.conversation_window import recent_conversation from tools import tool_registry as registry logger = logging.getLogger(__name__) @@ -123,7 +123,7 @@ def plan_question(llm, question, conversation=None, schema_rep="", prior_results catalog = _catalog_text(ctx) user_parts = [ f"## Question\n{question}", - f"## Conversation\n{json.dumps(conversation or [])[:2000]}", + f"## Conversation\n{recent_conversation(conversation)}", ] # Schema is normally not pre-loaded (the query tools load it themselves); # include it only if a caller explicitly supplied one. diff --git a/graphrag/app/agent/agentic_react.py b/graphrag/app/agent/agentic_react.py index f838720..365232a 100644 --- a/graphrag/app/agent/agentic_react.py +++ b/graphrag/app/agent/agentic_react.py @@ -33,6 +33,7 @@ from langchain_core.output_parsers import PydanticOutputParser from agent.agentic_executor import cap_for_trace, retrieved_chunk_ids, _usage_since +from agent.conversation_window import recent_conversation from common.llm_services.base_llm import get_collected_usage from common.py_schemas import GraphRAGAnswerOutput, GraphRAGResponse from tools import tool_registry as registry @@ -70,7 +71,7 @@ def run_react(ctx, llm, question, conversation=None) -> GraphRAGResponse: # other tools (e.g. external MCP tools) skip the schema read entirely. user = ( f"## Question\n{question}\n\n" - f"## Conversation\n{json.dumps(conversation or [])[:2000]}" + f"## Conversation\n{recent_conversation(conversation)}" ) # Customizable system prompt (fixed rules + user "Additional Instructions"); # the default lives in base_llm. diff --git a/graphrag/app/agent/agentic_synthesizer.py b/graphrag/app/agent/agentic_synthesizer.py index 1ef9255..9df9077 100644 --- a/graphrag/app/agent/agentic_synthesizer.py +++ b/graphrag/app/agent/agentic_synthesizer.py @@ -24,6 +24,7 @@ import logging from agent.agent_generation import TigerGraphAgentGenerator +from agent.conversation_window import recent_turns from agent.agentic_executor import retrieved_chunk_ids from common.py_schemas import GraphRAGResponse from common.utils.retrieval_stats import is_grouped, retrieved_entries @@ -115,6 +116,16 @@ def _final_retrieval(ctx) -> dict: return fr if isinstance(fr, dict) else {} +def with_conversation(combined: dict, conversation) -> dict: + """The answer model's context: ``combined`` plus the recent turns. + + The turns go last: when the context is too large, trimming starts from + the end, so they are cut before any retrieved context. + """ + earlier = recent_turns(conversation) + return {**combined, "earlier_conversation": earlier} if earlier else combined + + def has_context(results: dict) -> bool: g = _gather(results) return bool(g["structural"] or g["unstructured"]) @@ -123,12 +134,17 @@ 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, log=True) + answer_ctx = with_conversation(combined, conversation) generator = TigerGraphAgentGenerator(llm) - answer = generator.generate_answer(question, combined) + answer = generator.generate_answer(question, answer_ctx) nl = getattr(answer, "generated_answer", None) or str(answer) citations = getattr(answer, "citation", []) or [] - answered = bool(combined["structural"] or combined["unstructured"]) + retrieved = bool(combined["structural"] or combined["unstructured"]) + # A plan with no retrieval steps answers from the conversation. + from_conversation = (plan is not None and "earlier_conversation" in answer_ctx + and not any(s.kind != "answer" and s.tool for s in plan.steps)) + answered = retrieved or from_conversation # Chunk ids the plan FETCHED (across all unstructured steps), de-duped in # order — the agent's record of what was retrieved, distinct from the diff --git a/graphrag/app/agent/conversation_window.py b/graphrag/app/agent/conversation_window.py new file mode 100644 index 0000000..493af48 --- /dev/null +++ b/graphrag/app/agent/conversation_window.py @@ -0,0 +1,55 @@ +# 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 . + +"""The part of the conversation the agentic engine shows a model. + +Triage, the planner, the reactive agent and the answer step all read the +conversation through this window, so they see the same, most recent turns. +""" + +import json + +# Conversation budget, in characters of the JSON sent. +CONVERSATION_CHARS = 4000 + + +def recent_turns(convo, limit: int = CONVERSATION_CHARS) -> list: + """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 kept + + +def recent_conversation(convo, limit: int = CONVERSATION_CHARS) -> str: + """``recent_turns`` as the JSON text placed in a prompt.""" + return json.dumps(recent_turns(convo, limit), ensure_ascii=False) diff --git a/graphrag/tests/test_agentic_conversation_routing.py b/graphrag/tests/test_agentic_conversation_routing.py new file mode 100644 index 0000000..d125305 --- /dev/null +++ b/graphrag/tests/test_agentic_conversation_routing.py @@ -0,0 +1,143 @@ +"""The agentic engines decide from the conversation before they retrieve +(GML-2325).""" +import json +import types +import unittest +from unittest import mock + +try: + import agent.agentic_graph as g + import agent.agentic_synthesizer as syn + from agent.agentic_planner import plan_question + from agent.conversation_window import CONVERSATION_CHARS, recent_conversation, recent_turns + from common.llm_services.base_llm import LLM_Model + from common.py_schemas import GraphRAGResponse, Plan, PlanStep, StepResult +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +def _turns(n, size=300): + return [{"query": f"question {i}", "response": f"answer {i} " + "x" * size} for i in range(n)] + + +class TestConversationWindow(unittest.TestCase): + def test_newest_turns_are_kept(self): + convo = _turns(40) + kept = recent_turns(convo) + self.assertEqual(kept[-1]["query"], "question 39") + self.assertLessEqual(len(json.dumps(kept, ensure_ascii=False)), CONVERSATION_CHARS) + self.assertEqual(kept, convo[-len(kept):]) + + def test_planner_sees_the_newest_turns(self): + seen = {} + + def invoke(messages, schema, **kw): + seen["user"] = messages[-1][1] + return Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) + + llm = types.SimpleNamespace(agentic_planner_prompt="sys", config={}, + invoke_structured=invoke) + plan_question(llm, "q", _turns(40)) + self.assertIn("question 39", seen["user"]) + self.assertNotIn('"question 0"', seen["user"]) + + def test_triage_name_still_works(self): + from agent.agentic_agent import _recent_conversation + self.assertEqual(_recent_conversation(_turns(3)), recent_conversation(_turns(3))) + + +class _Ctx: + def __init__(self): + self.graphrag_cfg = {} + self.emit = lambda *a, **k: None + + +def _synth(*a, **k): + return GraphRAGResponse(natural_language_response="ok", + answered_question=True, response_type="agentic") + + +class TestAnswerOnlyPlan(unittest.TestCase): + def _run(self, plans): + with mock.patch.object(g, "plan_question", side_effect=plans) as planner, \ + mock.patch.object(g, "execute_plan", return_value=({}, [])) as execute, \ + mock.patch.object(g, "synthesize", side_effect=_synth), \ + mock.patch.object(g, "_run_step") as run_step: + resp = g.run_agentic(_Ctx(), llm=object(), question="q", conversation=_turns(1)) + return resp, planner, execute, run_step + + def test_answer_only_plan_skips_retrieval_and_replanning(self): + answer_only = Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) + resp, planner, execute, run_step = self._run([answer_only]) + self.assertEqual(planner.call_count, 1) + execute.assert_not_called() + run_step.assert_not_called() + steps = resp.query_sources["agent_steps"] + self.assertEqual(steps[0]["output"]["routing"], "answered without retrieval") + self.assertIn("earlier_conversation", steps[-1]["input"]) + + def test_empty_retrieval_still_replans(self): + search = Plan(steps=[PlanStep(id="S1", kind="unstructured", tool=g._HYBRID_TOOL), + PlanStep(id="A", kind="answer", tool="", depends_on=["S1"])]) + _, planner, execute, _ = self._run([search] * 10) + self.assertGreater(planner.call_count, 1) + self.assertGreater(execute.call_count, 1) + + def test_trace_marks_the_decision(self): + answer_only = Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) + self.assertEqual(g._plan_output(answer_only)["routing"], "answered without retrieval") + search = Plan(steps=[PlanStep(id="S1", kind="unstructured", tool=g._HYBRID_TOOL)]) + self.assertNotIn("routing", g._plan_output(search)) + + +class TestAnswerStepSeesTheConversation(unittest.TestCase): + def _synthesize(self, results, plan, conversation): + seen = {} + + class Gen: + def __init__(self, llm): + pass + + def generate_answer(self, question, context): + seen["context"] = context + return types.SimpleNamespace(generated_answer="a", citation=[]) + + with mock.patch.object(syn, "TigerGraphAgentGenerator", Gen): + resp = syn.synthesize(object(), "q", results, plan=plan, conversation=conversation) + return resp, seen["context"] + + def test_conversation_goes_last(self): + results = {"S1": StepResult(step_id="S1", ok=True, summary="rows", + context={"result": [{"row": 1}]})} + plan = Plan(steps=[PlanStep(id="S1", kind="structural", tool=g._STRUCTURAL_TOOL)]) + resp, ctx = self._synthesize(results, plan, _turns(2)) + self.assertEqual(list(ctx)[-1], "earlier_conversation") + self.assertEqual(ctx["earlier_conversation"][-1]["query"], "question 1") + # The recorded result is the retrieved context only. + self.assertNotIn("earlier_conversation", resp.query_sources["result"]) + + def test_answer_only_plan_counts_as_answered(self): + plan = Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) + resp, _ = self._synthesize({}, plan, _turns(1)) + self.assertTrue(resp.answered_question) + + def test_no_conversation_no_key(self): + plan = Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) + resp, ctx = self._synthesize({}, plan, []) + self.assertNotIn("earlier_conversation", ctx) + self.assertFalse(resp.answered_question) + + +class TestPrompts(unittest.TestCase): + def test_planner_and_react_check_the_conversation_first(self): + for prompt in (LLM_Model._AGENTIC_PLANNER_SYSTEM, LLM_Model._AGENTIC_AGENT_SYSTEM): + self.assertIn("decide how much of the answer ## Conversation already holds", prompt) + self.assertIn("Part of it", prompt) + self.assertIn('"now"', prompt) + + def test_answer_prompt_explains_the_conversation_context(self): + self.assertIn('"earlier_conversation"', LLM_Model._CHATBOT_RESPONSE_SYSTEM) + + +if __name__ == "__main__": + unittest.main() From 6fddd8d81ab74ad38f6603b26a385104d4b6ae60 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 14 Oct 2026 16:45:53 -0700 Subject: [PATCH 03/32] Restore external MCP servers in agentic chat - Keep the MCP SDK below version 2, whose client API the add-on support does not use yet Refs: GML-2326 --- CHANGELOG.md | 1 + common/requirements.txt | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f0b5057..71286d8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ - **Agentic answers reuse what the conversation already holds.** Before, the agentic engines retrieved again for a question the conversation had already answered, and the planned engine's answer step did not see the conversation at all. The planner and the reactive agent now first decide how much of the answer the conversation holds: all of it is answered without retrieval, part of it gets retrieval only for the missing parts, and questions about the current or latest state still retrieve. The answer step reads the recent conversation together with the retrieved context, and the retrieved context wins where they disagree. Trace Logs mark a plan that answers without retrieval. ### Fixed +- **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. - **Agentic planning no longer gets stuck writing its plan summary.** The planner's summary and step reasons had no stated purpose or length, and the planner's output had no limit, so a model could repeat a phrase until it reached its own output maximum — costing thousands of tokens and tens of seconds, and sometimes the plan itself. These fields now ask for one short sentence, an overlong value is shortened instead of failing the plan, and the planner's output is capped. diff --git a/common/requirements.txt b/common/requirements.txt index 17892bf..0c7f03c 100644 --- a/common/requirements.txt +++ b/common/requirements.txt @@ -186,8 +186,8 @@ zipp>=3.23.0 # Agentic engine (v2.0) — MCP + tigergraph-mcp for in-process, per-user # tool execution. Requires the fastapi/starlette bump above (mcp pulls -# starlette>=0.49). -mcp>=1.27.1 +# starlette>=0.49). mcp 2.x changes the client API that mcp_addons uses. +mcp>=1.27.1,<2 tigergraph-mcp>=1.0.4 sse-starlette>=3.4.4 httpx-sse>=0.4.3 From ca3cf5f0b54bb13cb5f6bdfb1a8b53d7c2dd9d94 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 14 Oct 2026 17:15:08 -0700 Subject: [PATCH 04/32] Check plans that skip retrieval before answering from the conversation - Planner and reactive agent compare each fact the question asks for with the conversation - In planned mode, a plan without retrieval gets a second check - A question the conversation does not fully answer is looked up instead Refs: GML-2325 --- CHANGELOG.md | 2 +- common/llm_services/base_llm.py | 35 +++++-- common/py_schemas/schemas.py | 14 +++ graphrag/app/agent/agentic_graph.py | 53 ++++++++++- graphrag/app/agent/agentic_planner.py | 31 ++++++- .../test_agentic_conversation_routing.py | 92 ++++++++++++++++++- 6 files changed, 209 insertions(+), 18 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 71286d8..ad18e6a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,7 @@ ## [2.0.4] ### Changed -- **Agentic answers reuse what the conversation already holds.** Before, the agentic engines retrieved again for a question the conversation had already answered, and the planned engine's answer step did not see the conversation at all. The planner and the reactive agent now first decide how much of the answer the conversation holds: all of it is answered without retrieval, part of it gets retrieval only for the missing parts, and questions about the current or latest state still retrieve. The answer step reads the recent conversation together with the retrieved context, and the retrieved context wins where they disagree. Trace Logs mark a plan that answers without retrieval. +- **Agentic answers reuse what the conversation already holds.** Before, the agentic engines retrieved again for a question the conversation had already answered, and the planned engine's answer step did not see the conversation at all. The planner and the reactive agent now first decide how much of the answer the conversation holds: all of it is answered without retrieval, part of it gets retrieval only for the missing parts, and questions about the current or latest state still retrieve. The answer step reads the recent conversation together with the retrieved context, and the retrieved context wins where they disagree. In planned mode, a plan that skips retrieval gets one more check, so a question the conversation does not fully answer is still looked up. Trace Logs show the check and mark a plan that answers without retrieval. ### Fixed - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index d167aa0..ef9afbc 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -1106,10 +1106,11 @@ def select_retriever_prompt(self): - If not — a part is still unanswered, a needed value or table is missing, or the results were thin — take another action to close the gap (follow a lead, widen top_k / num_hops, or switch method). Do not settle for a partial or vague answer when more retrieval could complete it. Do not commit to a full multi-step plan up front; let each next step be driven by what is still missing for a complete answer. -Before your first tool call, decide how much of the answer ## Conversation already holds: -- All of it: give the final answer without calling a tool. -- Part of it: call tools only for the missing parts. Do not retrieve again what ## Conversation already states. Combine both in the final answer. -- None of it: retrieve as described below. +Before your first tool call, list the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. Then: +- Every fact is stated (for example, the same question was already asked and answered): give the final answer without calling a tool. +- Some facts are stated: call tools only for the missing facts. Do not retrieve again what ## Conversation already states. Combine both in the final answer. +- No fact is stated: retrieve as described below. +This decision comes before the "Retrieval Strategy" below. Do not use an earlier answer when the question asks for the current or latest state (for example "now", "current", "latest", "again"), or when the user says that an earlier answer was wrong or incomplete. The graph schema is required for the structural and unstructured query tools: before your first structural query or vector/unstructured retrieval, call graphrag__get_schema once to load the graph's vertex and edge types. Questions answered without graph data (e.g. by an external tool) do not need the schema. @@ -1159,10 +1160,11 @@ def agentic_agent_prompt(self): The graph schema is NOT provided here — the structural and unstructured query tools load it themselves at run time, so plan retrieval steps directly. A question that needs no graph data should not include any graph-retrieval step (plan only the final answer step, or the relevant non-graph tool). -Before you plan retrieval, decide how much of the answer ## Conversation already holds: -- All of it: plan only the final answer step. -- Part of it: plan retrieval steps only for the missing parts. Do not retrieve again what ## Conversation already states. The answer step combines both. -- None of it: plan the retrieval steps as described below. +Before you plan retrieval, list the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. Then: +- Every fact is stated (for example, the same question was already asked and answered): plan only the final answer step. Do not retrieve those facts again. +- Some facts are stated: plan retrieval steps only for the missing facts. Do not retrieve again what ## Conversation already states. The answer step combines both. +- No fact is stated: plan the retrieval steps as described below. +This decision comes before the "Retrieval Strategy" below. Do not use an earlier answer when the question asks for the current or latest state (for example "now", "current", "latest", "again"), or when the user says that an earlier answer was wrong or incomplete. You have two kinds of retrieval: @@ -1183,6 +1185,23 @@ def agentic_agent_prompt(self): ## Retrieval Strategy {user_prompt} +""" + + # Second opinion on a plan that skips retrieval (GML-2325). Fixed, not + # customizable: it only guards the planner's own conversation rule. + _AGENTIC_CONVERSATION_CHECK_SYSTEM = """\ +You check one decision of a question-answering agent. The agent plans to answer the question from ## Conversation only, without new retrieval. + +List the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. + +Set "needs_retrieval" to true when one of these is correct: +- ## Conversation does not state a fact that the question asks for. +- The question asks for the current or latest state (for example "now", "current", "latest", "again"). +- The user says that an earlier answer was wrong or incomplete. + +Set "needs_retrieval" to false when ## Conversation states every fact, or when the question needs no data (for example a greeting or thanks). + +In "missing", name the facts that ## Conversation does not state. Return ONLY the structured result. """ # Strategy (operator-customizable) — moved out of the fixed rules so it can diff --git a/common/py_schemas/schemas.py b/common/py_schemas/schemas.py index fde19f8..5c7bb65 100644 --- a/common/py_schemas/schemas.py +++ b/common/py_schemas/schemas.py @@ -145,6 +145,20 @@ def _short_strategy(cls, v): return _short_plan_text(v, "strategy") +class ConversationCheck(BaseModel): + """Whether a plan that skips retrieval can answer from the conversation.""" + missing: str = Field( + "", + description="The facts the question asks for that the conversation does " + "not state. Empty when it states all of them.", + ) + needs_retrieval: bool = Field( + False, + description="True when the answer needs a fact the conversation does not " + "state, or the question asks for the current or latest state.", + ) + + class StepResult(BaseModel): step_id: str ok: bool diff --git a/graphrag/app/agent/agentic_graph.py b/graphrag/app/agent/agentic_graph.py index 743f660..f66e0c9 100644 --- a/graphrag/app/agent/agentic_graph.py +++ b/graphrag/app/agent/agentic_graph.py @@ -25,10 +25,10 @@ import time 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_planner import check_conversation, plan_question from agent.agentic_synthesizer import _gather, has_context, synthesize, with_conversation from common.llm_services.base_llm import get_collected_usage -from common.py_schemas import GraphRAGResponse, PlanStep +from common.py_schemas import GraphRAGResponse, Plan, PlanStep logger = logging.getLogger(__name__) @@ -55,6 +55,20 @@ def _hybrid_fallback_step() -> PlanStep: ) +def _search_plan(question: str) -> Plan: + """One hybrid search for the question, then the answer.""" + return Plan( + strategy="Search the documents for the question, then answer.", + steps=[ + PlanStep(id="S1", kind="unstructured", tool=_HYBRID_TOOL, + args={"question": question}, + rationale="The conversation does not hold the answer."), + PlanStep(id="A", kind="answer", tool="", depends_on=["S1"], + rationale="Answer from the retrieved context."), + ], + ) + + def _plan_output(plan) -> dict: """A plan as shown in the trace.""" out = {"strategy": plan.strategy, "steps": [s.model_dump() for s in plan.steps]} @@ -101,6 +115,41 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: "output": _plan_output(plan), "usage": _usage_since(_u0), }) + # A plan that skips retrieval relies on the conversation alone. A second, + # narrow check catches the cases where the planner judged that wrongly + # (a fact the conversation does not state, or a request for fresh data). + # It also runs without a conversation: a data question planned without + # retrieval would otherwise be answered with nothing. + if not _retrieves(plan): + _u0 = len(get_collected_usage() or []) + _t0 = time.time() + check = check_conversation(llm, question, conversation, ctx=ctx) + agent_steps.append({ + "node": "conversation check", "kind": "plan", + "duration_s": round(time.time() - _t0, 3), + "input": {"question": question}, + "output": check.model_dump() if check else {"error": "check failed; plan kept"}, + "usage": _usage_since(_u0), + }) + if check is not None and check.needs_retrieval: + emit("Refining the plan") + _u0 = len(get_collected_usage() or []) + _t0 = time.time() + plan = plan_question(llm, question, conversation, ctx=ctx, + missing_from_conversation=check.missing or "some facts the question asks for") + agent_steps.append({ + "node": "replan (conversation check)", "kind": "plan", + "duration_s": round(time.time() - _t0, 3), + "input": {"missing_from_conversation": check.missing}, + "output": _plan_output(plan), + "usage": _usage_since(_u0), + }) + if not _retrieves(plan): + # The planner skipped retrieval twice for a question the check + # says needs data: search instead of answering with nothing. + logger.info("agentic: replan still has no retrieval; using a hybrid search") + plan = _search_plan(question) + if plan.strategy: emit(plan.strategy) diff --git a/graphrag/app/agent/agentic_planner.py b/graphrag/app/agent/agentic_planner.py index 501d926..a51d65e 100644 --- a/graphrag/app/agent/agentic_planner.py +++ b/graphrag/app/agent/agentic_planner.py @@ -25,7 +25,7 @@ import logging from common.llm_services.capabilities import model_capabilities -from common.py_schemas import Plan +from common.py_schemas import ConversationCheck, Plan from agent.conversation_window import recent_conversation from tools import tool_registry as registry @@ -110,7 +110,28 @@ def _planner_output_limit(llm, ctx) -> int: return limit -def plan_question(llm, question, conversation=None, schema_rep="", prior_results=None, ctx=None) -> Plan: +def check_conversation(llm, question, conversation, ctx=None): + """Ask whether a plan that skips retrieval can answer from the conversation. + + Returns a ``ConversationCheck``, or ``None`` when the check fails (the + caller then keeps the plan). + """ + try: + messages = [ + ("system", llm._AGENTIC_CONVERSATION_CHECK_SYSTEM), + ("user", f"## Question\n{question}\n\n## Conversation\n{recent_conversation(conversation)}"), + ] + return llm.invoke_structured( + messages, ConversationCheck, caller_name="agentic_conversation_check", + max_output_tokens=_planner_output_limit(llm, ctx), + ) + except Exception as exc: + logger.warning(f"conversation check failed ({exc}); keeping the plan") + return None + + +def plan_question(llm, question, conversation=None, schema_rep="", prior_results=None, ctx=None, + missing_from_conversation: str = "") -> Plan: """Draft (or extend) a plan for ``question``. ``prior_results`` (a list of ``StepResult``-like dicts) is supplied on @@ -130,6 +151,12 @@ def plan_question(llm, question, conversation=None, schema_rep="", prior_results if schema_rep: user_parts.append(f"## Graph schema\n{schema_rep[:6000]}") user_parts.append(f"## Tools\n{catalog}") + if missing_from_conversation: + user_parts.append( + "## Missing from the conversation\n" + f"## Conversation does not state: {missing_from_conversation}\n\n" + "Plan retrieval steps for these facts, then the final answer step." + ) if prior_results: summary = "\n".join( f"- {r.get('step_id')}: ok={r.get('ok')} — {r.get('summary')}" diff --git a/graphrag/tests/test_agentic_conversation_routing.py b/graphrag/tests/test_agentic_conversation_routing.py index d125305..c45e270 100644 --- a/graphrag/tests/test_agentic_conversation_routing.py +++ b/graphrag/tests/test_agentic_conversation_routing.py @@ -11,7 +11,7 @@ from agent.agentic_planner import plan_question from agent.conversation_window import CONVERSATION_CHARS, recent_conversation, recent_turns from common.llm_services.base_llm import LLM_Model - from common.py_schemas import GraphRAGResponse, Plan, PlanStep, StepResult + from common.py_schemas import ConversationCheck, GraphRAGResponse, Plan, PlanStep, StepResult except ImportError as exc: # needs the app environment raise unittest.SkipTest(f"app environment unavailable: {exc}") @@ -58,8 +58,9 @@ def _synth(*a, **k): class TestAnswerOnlyPlan(unittest.TestCase): - def _run(self, plans): - with mock.patch.object(g, "plan_question", side_effect=plans) as planner, \ + def _run(self, plans, check=ConversationCheck(needs_retrieval=False)): + with mock.patch.object(g, "check_conversation", return_value=check), \ + mock.patch.object(g, "plan_question", side_effect=plans) as planner, \ mock.patch.object(g, "execute_plan", return_value=({}, [])) as execute, \ mock.patch.object(g, "synthesize", side_effect=_synth), \ mock.patch.object(g, "_run_step") as run_step: @@ -90,6 +91,86 @@ def test_trace_marks_the_decision(self): self.assertNotIn("routing", g._plan_output(search)) +_ANSWER_ONLY = Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) +_IPAD_SEARCH = Plan(steps=[PlanStep(id="S1", kind="unstructured", tool="graphrag__hybrid_search", + args={"question": "iPad net sales"}), + PlanStep(id="A", kind="answer", tool="", depends_on=["S1"])]) +_FOUND = {"S1": StepResult(step_id="S1", ok=True, summary="rows", context={"result": [{"row": 1}]})} + + +class TestConversationCheck(unittest.TestCase): + """A plan that skips retrieval gets a second check.""" + + def _run(self, plans, check=None, conversation=None, check_side_effect=None): + with mock.patch.object(g, "check_conversation", return_value=check, + side_effect=check_side_effect) as checker, \ + mock.patch.object(g, "plan_question", side_effect=plans) as planner, \ + mock.patch.object(g, "execute_plan", return_value=(_FOUND, [])) as execute, \ + mock.patch.object(g, "synthesize", side_effect=_synth): + resp = g.run_agentic(_Ctx(), llm=object(), question="Add iPad.", + conversation=_turns(1) if conversation is None else conversation) + return resp, checker, planner, execute + + def test_missing_fact_replans_with_the_gap(self): + check = ConversationCheck(missing="iPad net sales", needs_retrieval=True) + resp, checker, planner, execute = self._run([_ANSWER_ONLY, _IPAD_SEARCH], check) + self.assertEqual(planner.call_count, 2) + self.assertEqual(planner.call_args.kwargs["missing_from_conversation"], "iPad net sales") + execute.assert_called_once() + nodes = [s["node"] for s in resp.query_sources["agent_steps"]] + self.assertEqual(nodes[:3], ["plan", "conversation check", "replan (conversation check)"]) + + def test_second_skip_searches_anyway(self): + check = ConversationCheck(missing="iPhone net sales", needs_retrieval=True) + _, _, planner, execute = self._run([_ANSWER_ONLY, _ANSWER_ONLY], check) + self.assertEqual(planner.call_count, 2) + execute.assert_called_once() + plan = execute.call_args[0][0] + self.assertEqual([s.tool for s in plan.steps], [g._HYBRID_TOOL, ""]) + self.assertEqual(plan.steps[0].args, {"question": "Add iPad."}) + + def test_confirmed_plan_answers_without_retrieval(self): + _, _, planner, execute = self._run([_ANSWER_ONLY], ConversationCheck(needs_retrieval=False)) + self.assertEqual(planner.call_count, 1) + execute.assert_not_called() + + def test_failed_check_keeps_the_plan(self): + _, _, planner, execute = self._run([_ANSWER_ONLY], None) + self.assertEqual(planner.call_count, 1) + execute.assert_not_called() + + def test_check_runs_without_conversation(self): + """A data question planned without retrieval still gets retrieval.""" + check = ConversationCheck(missing="iPhone net sales", needs_retrieval=True) + _, checker, planner, execute = self._run([_ANSWER_ONLY, _IPAD_SEARCH], check, conversation=[]) + checker.assert_called_once() + self.assertEqual(planner.call_count, 2) + execute.assert_called_once() + + def test_no_check_for_a_retrieval_plan(self): + _, checker, _, execute = self._run([_IPAD_SEARCH]) + checker.assert_not_called() + execute.assert_called_once() + + def test_check_errors_keep_the_plan(self): + from agent.agentic_planner import check_conversation + self.assertIsNone(check_conversation(object(), "q", _turns(1))) + + def test_planner_is_told_what_is_missing(self): + seen = {} + + def invoke(messages, schema, **kw): + seen["user"] = messages[-1][1] + return _IPAD_SEARCH + + llm = types.SimpleNamespace(agentic_planner_prompt="sys", config={}, invoke_structured=invoke) + plan_question(llm, "Add iPad.", _turns(1), missing_from_conversation="iPad net sales") + self.assertIn("## Conversation does not state: iPad net sales", seen["user"]) + + def test_check_prompt_covers_fresh_data_requests(self): + self.assertIn('"now"', LLM_Model._AGENTIC_CONVERSATION_CHECK_SYSTEM) + + class TestAnswerStepSeesTheConversation(unittest.TestCase): def _synthesize(self, results, plan, conversation): seen = {} @@ -131,8 +212,9 @@ def test_no_conversation_no_key(self): class TestPrompts(unittest.TestCase): def test_planner_and_react_check_the_conversation_first(self): for prompt in (LLM_Model._AGENTIC_PLANNER_SYSTEM, LLM_Model._AGENTIC_AGENT_SYSTEM): - self.assertIn("decide how much of the answer ## Conversation already holds", prompt) - self.assertIn("Part of it", prompt) + self.assertIn("list the facts the question asks for", prompt) + self.assertIn("Some facts are stated", prompt) + self.assertIn("A related topic or a similar question is not enough", prompt) self.assertIn('"now"', prompt) def test_answer_prompt_explains_the_conversation_context(self): From daa070c88235f7297501cc9b0ea2154df90b5ad3 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 14 Oct 2026 17:47:09 -0700 Subject: [PATCH 05/32] Show the structured query in Trace Logs and treat empty results as no rows - Structured query steps show the query that ran, plus failed or empty attempts - A structured query that returns no rows no longer counts as data - A valid openCypher query with no rows is not regenerated Refs: GML-2327 --- CHANGELOG.md | 2 + graphrag/app/agent/agentic_executor.py | 5 ++ graphrag/app/agent/agentic_react.py | 8 +- graphrag/app/tools/graphrag_tools.py | 80 ++++++++++++++++-- graphrag/tests/test_structural_trace.py | 104 ++++++++++++++++++++++++ 5 files changed, 189 insertions(+), 10 deletions(-) create mode 100644 graphrag/tests/test_structural_trace.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ad18e6a..9a4cb89 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,8 +4,10 @@ ### Changed - **Agentic answers reuse what the conversation already holds.** Before, the agentic engines retrieved again for a question the conversation had already answered, and the planned engine's answer step did not see the conversation at all. The planner and the reactive agent now first decide how much of the answer the conversation holds: all of it is answered without retrieval, part of it gets retrieval only for the missing parts, and questions about the current or latest state still retrieve. The answer step reads the recent conversation together with the retrieved context, and the retrieved context wins where they disagree. In planned mode, a plan that skips retrieval gets one more check, so a question the conversation does not fully answer is still looked up. Trace Logs show the check and mark a plan that answers without retrieval. +- **Trace Logs show the structured query that ran.** Each structured query step lists the query at the top of its output, followed by any attempts that returned nothing or failed, with their errors. Before, the query was buried in the result, and it was missing entirely when the query returned no rows — the case where it is most needed. ### Fixed +- **An empty structured query result no longer counts as data.** A query that returned no rows (for example `{"VS_XXX": []}`, or an empty list as text) was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. - **Agentic planning no longer gets stuck writing its plan summary.** The planner's summary and step reasons had no stated purpose or length, and the planner's output had no limit, so a model could repeat a phrase until it reached its own output maximum — costing thousands of tokens and tens of seconds, and sometimes the plan itself. These fields now ask for one short sentence, an overlong value is shortened instead of failing the plan, and the planner's output is capped. diff --git a/graphrag/app/agent/agentic_executor.py b/graphrag/app/agent/agentic_executor.py index 9878003..8760013 100644 --- a/graphrag/app/agent/agentic_executor.py +++ b/graphrag/app/agent/agentic_executor.py @@ -202,6 +202,11 @@ def _run_step(step, args, ctx, results, traces): # 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", "")} + # Structured query tools report the query they ran (and the attempts + # that returned nothing or failed). Kept out of the trace size limit. + for key in ("query", "attempts"): + if out.get(key): + trace_output[key] = out[key] if out.get("context") is not None: trace_output["result"], cut = fit_for_trace(out.get("context")) trace_output = with_trace_note(trace_output, cut) diff --git a/graphrag/app/agent/agentic_react.py b/graphrag/app/agent/agentic_react.py index 365232a..f98b2c4 100644 --- a/graphrag/app/agent/agentic_react.py +++ b/graphrag/app/agent/agentic_react.py @@ -169,11 +169,15 @@ def run_react(ctx, llm, question, conversation=None) -> GraphRAGResponse: content=json.dumps(obs, default=str), tool_call_id=tc_id or "", )) - per_call_traces.append({ + call_trace = { "tool": name, "args": cap_for_trace(args), "ok": bool(out.get("ok")), "summary": out.get("summary", ""), "duration_s": tool_dur, - }) + } + for key in ("query", "attempts"): + if out.get(key): + call_trace[key] = out[key] + per_call_traces.append(call_trace) # Record the chunk ids this tool fetched (the agent's bookkeeping, # harvested from the context the tool returned), de-duped in order. for cid in retrieved_chunk_ids(out.get("context")): diff --git a/graphrag/app/tools/graphrag_tools.py b/graphrag/app/tools/graphrag_tools.py index 6754cf1..792c56d 100644 --- a/graphrag/app/tools/graphrag_tools.py +++ b/graphrag/app/tools/graphrag_tools.py @@ -28,6 +28,7 @@ from __future__ import annotations +import json import logging from dataclasses import dataclass, field from typing import Any, Callable, Dict, Optional @@ -84,13 +85,57 @@ def _empty(summary: str) -> dict: def _result_is_empty(result: Any) -> bool: + """True when ``result`` holds no rows. + + Query results arrive as JSON text (pyTigerGraph) or nested containers + (openCypher, e.g. ``{"VS_XXX": []}``), so both are read through. A + container is empty when every value in it is empty. A scalar such as a + count of 0 is an answer, not an empty result. + """ if result is None: return True - if isinstance(result, (list, dict, str)) and len(result) == 0: - return True + if isinstance(result, str): + text = result.strip() + if not text: + return True + try: + parsed = json.loads(text) + except ValueError: + return False + return _result_is_empty(parsed) if isinstance(parsed, (list, dict)) or parsed is None else False + if isinstance(result, dict): + return all(_result_is_empty(v) for v in result.values()) + if isinstance(result, list): + return all(_result_is_empty(v) for v in result) return False +# Longest error text kept for one query attempt in the trace. +_ATTEMPT_ERROR_CHARS = 500 + + +def _attempt(query: str = "", error: str = "", outcome: str = "") -> dict: + """One query attempt as shown in the trace.""" + entry = {} + if query: + entry["query"] = query + if error: + entry["error"] = str(error)[:_ATTEMPT_ERROR_CHARS] + if outcome: + entry["outcome"] = outcome + return entry + + +def _with_queries(out: dict, query: str = "", attempts: Optional[list] = None) -> dict: + """Add the query that produced the result, and the attempts that did + not, to a tool output. The executor shows both in the trace.""" + if query: + out["query"] = query + if attempts: + out["attempts"] = attempts + return out + + # -------------------------------------------------------------------------- # Schema # -------------------------------------------------------------------------- @@ -141,6 +186,7 @@ def structural_retrieve(ctx: GraphRAGToolContext, question: str) -> dict: return _empty(f"schema mapping failed: {exc}") ctx.emit("Generating a query to answer the question") + attempts: list = [] try: step = ctx.gen_func._run( question, @@ -153,22 +199,36 @@ def structural_retrieve(ctx: GraphRAGToolContext, question: str) -> dict: except Exception as exc: logger.warning(f"structural_retrieve generate_function failed: {exc}") step = None + attempts.append(_attempt(error=exc)) result = step.get("result") if isinstance(step, dict) else None + func_call = step.get("function_call", "") if isinstance(step, dict) else "" if not _result_is_empty(result): - return _ok("structural query returned rows", step) + return _with_queries(_ok("structural query returned rows", step), func_call) + if isinstance(step, dict): + attempts.append(_attempt(func_call, step.get("error", ""), + "" if step.get("error") else "no rows")) + elif isinstance(step, Exception): + # generate_function returns (rather than raises) some errors. + attempts.append(_attempt(error=step)) # Optional cypher fallback when configured and available. if ctx.use_cypher and ctx.cypher_gen is not None: - cy = _cypher_retrieve(ctx, question) + cy = _cypher_retrieve(ctx, question, attempts) if cy["ok"]: - return cy + return _with_queries(cy, cy["context"].get("cypher", ""), attempts) + + return _with_queries(_empty("structural query returned no rows"), attempts=attempts) - return _empty("structural query returned no rows") +def _cypher_retrieve(ctx: GraphRAGToolContext, question: str, attempts: Optional[list] = None) -> dict: + """Generate and run an openCypher query, retrying on errors. -def _cypher_retrieve(ctx: GraphRAGToolContext, question: str) -> dict: - import json + Each attempt that does not produce rows is appended to ``attempts``. A + query that runs and returns no rows is a valid answer, so it is not + retried. + """ + attempts = [] if attempts is None else attempts ctx.emit("Generating a graph query") gen_history: list = [] for i in range(3): @@ -176,6 +236,7 @@ def _cypher_retrieve(ctx: GraphRAGToolContext, question: str) -> dict: cypher = ctx.cypher_gen._run(question, gen_history) except ValueError as exc: gen_history.append(f"{i}: Error: {exc}\n") + attempts.append(_attempt(error=exc)) continue response = ctx.conn.gsql(cypher) json_str = "\n".join(response.split("\n")[1:]) @@ -183,6 +244,7 @@ def _cypher_retrieve(ctx: GraphRAGToolContext, question: str) -> dict: parsed = json.loads(json_str) except Exception: gen_history.append(f"{i}: {cypher}\n\tError: {json_str}\n") + attempts.append(_attempt(cypher, json_str or response)) continue rows = parsed.get("results", [{}]) first = rows[0] if rows else None @@ -192,6 +254,8 @@ def _cypher_retrieve(ctx: GraphRAGToolContext, question: str) -> dict: {"result": first, "cypher": cypher, "reasoning": f"The following openCypher query was executed:\n{cypher}"}, ) + attempts.append(_attempt(cypher, outcome="no rows")) + return _empty("graph query returned no rows") return _empty("graph query returned no rows after retries") diff --git a/graphrag/tests/test_structural_trace.py b/graphrag/tests/test_structural_trace.py new file mode 100644 index 0000000..3e59aa8 --- /dev/null +++ b/graphrag/tests/test_structural_trace.py @@ -0,0 +1,104 @@ +"""The structured query tool reports the query it ran, and an empty query +result counts as no rows (GML-2327).""" +import json +import types +import unittest +from unittest import mock + +try: + from tools import graphrag_tools as gt + from agent import agentic_executor as ex +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +class TestResultIsEmpty(unittest.TestCase): + def test_empty_results(self): + for value in (None, "", " ", "[]", "{}", "null", '{"VS_XXX": []}', + {"VS_XXX": []}, [{"VS_XXX": []}], [], {}): + self.assertTrue(gt._result_is_empty(value), value) + + def test_results_with_rows(self): + for value in ("0", 0, '{"count": 0}', [{"v_id": "x"}], '[{"v_id": "x"}]', + {"VS_XXX": [{"v_id": "x"}]}, "not json text"): + self.assertFalse(gt._result_is_empty(value), value) + + +_FUNC = "getVertexCount('Document')" +_CYPHER = "USE GRAPH g\nINTERPRET OPENCYPHER QUERY () {\nMATCH (d:Document) RETURN d\n}" + + +def _ctx(gen_result, use_cypher=False, cypher_runs=()): + mapping = types.SimpleNamespace( + target_vertex_types=["Document"], target_vertex_attributes={}, target_vertex_ids={}, + target_edge_types=[], target_edge_attributes={}, + ) + gen_func = mock.Mock() + if isinstance(gen_result, BaseException) and not isinstance(gen_result, ValueError): + gen_func._run.side_effect = gen_result + else: + gen_func._run.return_value = gen_result + cypher_gen = mock.Mock() + cypher_gen._run.side_effect = [c for c, _ in cypher_runs] or [_CYPHER] + conn = mock.Mock() + conn.gsql.side_effect = [r for _, r in cypher_runs] or ["header\n{}"] + return types.SimpleNamespace( + emit=lambda *a, **k: None, conversation=[], use_cypher=use_cypher, + mq2s=mock.Mock(_run=mock.Mock(return_value=mapping)), + gen_func=gen_func, cypher_gen=cypher_gen, conn=conn, + ) + + +def _gsql(results): + return "header\n" + json.dumps({"results": results}) + + +class TestStructuralRetrieveQuery(unittest.TestCase): + def test_rows_report_the_query(self): + out = gt.structural_retrieve(_ctx({"function_call": _FUNC, "result": "4"}), "q") + self.assertTrue(out["ok"]) + self.assertEqual(out["query"], _FUNC) + self.assertNotIn("attempts", out) + + def test_empty_rows_are_no_rows_and_keep_the_query(self): + out = gt.structural_retrieve(_ctx({"function_call": _FUNC, "result": "[]"}), "q") + self.assertFalse(out["ok"]) + self.assertEqual(out["attempts"], [{"query": _FUNC, "outcome": "no rows"}]) + + def test_generation_error_is_recorded(self): + out = gt.structural_retrieve(_ctx(RuntimeError("bad call")), "q") + self.assertFalse(out["ok"]) + self.assertEqual(out["attempts"], [{"error": "bad call"}]) + + def test_cypher_rows_report_the_query_and_earlier_attempts(self): + ctx = _ctx({"function_call": _FUNC, "result": "[]"}, use_cypher=True, + cypher_runs=[("BAD", "header\nnot json"), (_CYPHER, _gsql([{"VS": [{"v_id": "d1"}]}]))]) + out = gt.structural_retrieve(ctx, "q") + self.assertTrue(out["ok"]) + self.assertEqual(out["query"], _CYPHER) + self.assertEqual([a.get("query") for a in out["attempts"]], [_FUNC, "BAD"]) + self.assertIn("error", out["attempts"][1]) + + def test_cypher_no_rows_is_not_retried(self): + ctx = _ctx({"function_call": _FUNC, "result": "[]"}, use_cypher=True, + cypher_runs=[(_CYPHER, _gsql([{"VS_XXX": []}])), ("UNUSED", "")]) + out = gt.structural_retrieve(ctx, "q") + self.assertFalse(out["ok"]) + self.assertEqual(ctx.cypher_gen._run.call_count, 1) + self.assertEqual(out["attempts"][-1], {"query": _CYPHER, "outcome": "no rows"}) + + +class TestTraceShowsTheQuery(unittest.TestCase): + def test_query_and_attempts_lead_the_step_output(self): + out = {"ok": True, "summary": "graph query returned rows", "context": {"result": [1]}, + "citations": [], "query": _CYPHER, "attempts": [{"query": _FUNC, "outcome": "no rows"}]} + step = mock.Mock(id="S1", tool="graphrag__structural_retrieve", kind="structural", rationale="") + traces = [] + with mock.patch.object(ex.registry, "run", return_value=out): + ex._run_step(step, {"question": "q"}, mock.Mock(), {}, traces) + self.assertEqual(list(traces[0]["output"])[:3], ["summary", "query", "attempts"]) + self.assertEqual(traces[0]["output"]["query"], _CYPHER) + + +if __name__ == "__main__": + unittest.main() From ff2124ec4ef9e2a196112026e33ad4c632d9e27b Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Wed, 14 Oct 2026 17:47:10 -0700 Subject: [PATCH 06/32] Record the error of a failed answer in Trace Logs - Trace log of a failed answer ends with an Error step that states the error - Chat error message asks to rephrase the question, not the prompt Refs: GML-2328 --- CHANGELOG.md | 1 + common/py_schemas/schemas.py | 5 +- graphrag-ui/src/pages/TraceLogs.tsx | 1 + graphrag/app/routers/ui.py | 41 +++++++++++--- graphrag/tests/test_chat_error_trace.py | 73 +++++++++++++++++++++++++ 5 files changed, 112 insertions(+), 9 deletions(-) create mode 100644 graphrag/tests/test_chat_error_trace.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9a4cb89..3036c3f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Changed - **Agentic answers reuse what the conversation already holds.** Before, the agentic engines retrieved again for a question the conversation had already answered, and the planned engine's answer step did not see the conversation at all. The planner and the reactive agent now first decide how much of the answer the conversation holds: all of it is answered without retrieval, part of it gets retrieval only for the missing parts, and questions about the current or latest state still retrieve. The answer step reads the recent conversation together with the retrieved context, and the retrieved context wins where they disagree. In planned mode, a plan that skips retrieval gets one more check, so a question the conversation does not fully answer is still looked up. Trace Logs show the check and mark a plan that answers without retrieval. +- **Failed answers show their error in Trace Logs.** When GraphRAG cannot answer, the trace log ends with an Error step that states the error, even when no other step ran. The error appears only in Trace Logs, which only the asking superuser can open; the chat message stays generic and now asks to rephrase the question instead of the prompt. - **Trace Logs show the structured query that ran.** Each structured query step lists the query at the top of its output, followed by any attempts that returned nothing or failed, with their errors. Before, the query was buried in the result, and it was missing entirely when the query returned no rows — the case where it is most needed. ### Fixed diff --git a/common/py_schemas/schemas.py b/common/py_schemas/schemas.py index 5c7bb65..590190d 100644 --- a/common/py_schemas/schemas.py +++ b/common/py_schemas/schemas.py @@ -16,7 +16,7 @@ import logging from typing import Dict, List, Optional, Union -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, PrivateAttr, field_validator class ConversationTurn(BaseModel): @@ -69,6 +69,9 @@ class GraphRAGResponse(BaseModel): answered_question: bool response_type: str query_sources: Dict = None + # The error that stopped the answer, for the trace log only. Private, so + # it is never serialized into the chat response or the chat history. + _error: Optional[str] = PrivateAttr(default=None) # --- Agentic engine (v2.0 deep-thinking mode) ------------------------------ diff --git a/graphrag-ui/src/pages/TraceLogs.tsx b/graphrag-ui/src/pages/TraceLogs.tsx index 185b86f..b838ed2 100644 --- a/graphrag-ui/src/pages/TraceLogs.tsx +++ b/graphrag-ui/src/pages/TraceLogs.tsx @@ -132,6 +132,7 @@ const NODE_LABELS: Record = { rewrite_question: "Rewrite Question", apologize: "Apologize", greet: "Greet", + error: "Error", }; function buildTraceFromMessage(message: any, userQuery?: string): TraceData { diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py index 29717e0..a44512e 100644 --- a/graphrag/app/routers/ui.py +++ b/graphrag/app/routers/ui.py @@ -128,8 +128,19 @@ def _save_trace_log(message_id: str, conversation_id: str, user_query: str, resp result = {k: v for k, v in result.items() if k != "final_retrieval"} query_sources = {**query_sources, "result": result} + # The error that stopped the answer goes into the trace log only + # (superuser- and owner-gated), never into the chat response. + error = getattr(resp, "_error", None) + if error: + query_sources["agent_steps"] = list(query_sources.get("agent_steps") or []) + [{ + "node": "error", "kind": "error", "duration_s": 0, + "input": {"question": user_query}, + "output": {"error": error}, + }] + trace_data = { "message_id": message_id, + "error": error, "conversation_id": conversation_id, "username": username, "user_query": user_query, @@ -2818,18 +2829,29 @@ def _is_superadmin(global_roles: list[str]) -> bool: return any(r in SUPERADMIN_ROLES for r in (global_roles or [])) -def _agent_error_text(e: Exception, is_superadmin: bool = False) -> str: - # asyncio.TaskGroup wraps a failing sub-task in an ExceptionGroup whose - # message is only "unhandled errors in a TaskGroup (N sub-exception(s))". - # Unwrap to the underlying cause so the admin detail shows the real error. +def _root_error(e: BaseException) -> BaseException: + """The underlying error. asyncio.TaskGroup wraps a failing sub-task in an + ExceptionGroup whose message is only "unhandled errors in a TaskGroup + (N sub-exception(s))".""" while isinstance(e, BaseExceptionGroup) and e.exceptions: e = e.exceptions[0] + return e + + +def _trace_error_text(e: BaseException) -> str: + """The error as recorded in the trace log.""" + e = _root_error(e) + return f"{type(e).__name__}: {e}" + + +def _agent_error_text(e: Exception, is_superadmin: bool = False) -> str: + e = _root_error(e) error_msg = str(e) if "does not exist" in error_msg or "not found" in error_msg.lower(): return f"Error: {error_msg}. Please check the knowledge graph name and try again." generic = ( "GraphRAG had an issue answering your question. " - "Please try again, or rephrase your prompt." + "Please try again, or rephrase your question." ) # Never return raw exception text to the client. It can carry sensitive # configuration, URLs, request fragments, or credentials, and the chat @@ -2839,7 +2861,7 @@ def _agent_error_text(e: Exception, is_superadmin: bool = False) -> str: if is_superadmin and error_msg: ref = req_id_cv.get() or "n/a" logger.error(f"agent error [ref={ref}]: {error_msg}") - return f"{generic}\n\n(Admin reference ID: {ref} — see server logs for details.)" + return f"{generic}\n\n(Admin reference ID: {ref}. The trace log and the server logs show the error.)" return generic @@ -2874,11 +2896,12 @@ async def run_agent( if ws: agent.q.clear() - except MapQuestionToSchemaException: + except MapQuestionToSchemaException as e: resp.natural_language_response = ( "A schema mapping error occurred. Please try rephrasing your question." ) resp.query_sources = {} + resp._error = _trace_error_text(e) resp.answered_question = False LogWriter.warning( f"/{graphname}/ui/chat request_id={req_id_cv.get()} agent execution failed due to MapQuestionToSchemaException" @@ -2891,9 +2914,11 @@ async def run_agent( except Exception as e: resp.natural_language_response = _agent_error_text(e, is_superadmin) # Preserve the steps the agent completed before failing so the - # (superuser-only) trace log shows how far it got (GML-2136). + # (superuser-only) trace log shows how far it got (GML-2136), and the + # error that stopped it. partial_steps = getattr(agent, "_last_agent_steps", None) resp.query_sources = {"agent_steps": partial_steps} if partial_steps else {} + resp._error = _trace_error_text(e) resp.answered_question = False LogWriter.warning( f"/{graphname}/ui/chat request_id={req_id_cv.get()} agent execution failed due to exception: {e}" diff --git a/graphrag/tests/test_chat_error_trace.py b/graphrag/tests/test_chat_error_trace.py new file mode 100644 index 0000000..9f2cce6 --- /dev/null +++ b/graphrag/tests/test_chat_error_trace.py @@ -0,0 +1,73 @@ +"""A failed answer asks to rephrase the question, and the trace log records +the error without sending it to the chat (GML-2328).""" +import json +import os +import tempfile +import unittest +from unittest import mock + +try: + from routers import ui + from common.py_schemas import GraphRAGResponse +except Exception as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +def _failed(steps=None, error="RuntimeError: quota exceeded"): + resp = GraphRAGResponse(natural_language_response=ui._agent_error_text(RuntimeError("x")), + answered_question=False, response_type="agentic", + query_sources={"agent_steps": steps} if steps else {}) + resp._error = error + return resp + + +class TestErrorMessage(unittest.TestCase): + def test_message_asks_to_rephrase_the_question(self): + text = ui._agent_error_text(RuntimeError("boom")) + self.assertIn("rephrase your question", text) + self.assertNotIn("prompt", text) + self.assertNotIn("boom", text) + + def test_admin_gets_a_reference_not_the_error(self): + text = ui._agent_error_text(RuntimeError("secret-url"), is_superadmin=True) + self.assertIn("Admin reference ID", text) + self.assertNotIn("secret-url", text) + + def test_task_group_error_is_unwrapped(self): + group = ExceptionGroup("unhandled errors in a TaskGroup", [ValueError("real cause")]) + self.assertEqual(ui._trace_error_text(group), "ValueError: real cause") + + +class TestTraceLogRecordsTheError(unittest.TestCase): + def _save(self, resp): + with tempfile.TemporaryDirectory() as d, mock.patch.object(ui, "TRACE_LOGS_DIR", d): + ui._save_trace_log("m1", "c1", "the question", resp, 1.0, "alice") + with open(os.path.join(d, "m1.json")) as f: + return json.load(f) + + def test_error_step_is_added_after_the_completed_steps(self): + data = self._save(_failed(steps=[{"node": "plan", "kind": "plan"}])) + steps = data["query_sources"]["agent_steps"] + self.assertEqual([s["node"] for s in steps], ["plan", "error"]) + self.assertEqual(steps[-1]["output"], {"error": "RuntimeError: quota exceeded"}) + self.assertEqual(data["error"], "RuntimeError: quota exceeded") + + def test_error_step_without_completed_steps(self): + data = self._save(_failed()) + self.assertEqual([s["node"] for s in data["query_sources"]["agent_steps"]], ["error"]) + + def test_error_is_not_in_the_chat_response(self): + resp = _failed() + self.assertNotIn("quota exceeded", resp.model_dump_json()) + self.assertNotIn("error", (resp.query_sources or {})) + + def test_successful_answer_has_no_error_step(self): + resp = GraphRAGResponse(natural_language_response="ok", answered_question=True, + response_type="agentic", query_sources={"agent_steps": [{"node": "plan"}]}) + data = self._save(resp) + self.assertEqual([s["node"] for s in data["query_sources"]["agent_steps"]], ["plan"]) + self.assertIsNone(data["error"]) + + +if __name__ == "__main__": + unittest.main() From 2224bf497d45dcc1becd037c41f04d57dd1cf0cc Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 12:30:24 -0700 Subject: [PATCH 07/32] Skip the pyTigerGraph function step when no functions are registered - Structured queries go straight to openCypher when no functions are registered - A question that names no relationships no longer fails the function step Refs: GML-2327 --- CHANGELOG.md | 1 + graphrag/app/tools/generate_function.py | 29 +++++++++++++---- graphrag/app/tools/validation_utils.py | 5 +++ graphrag/tests/test_structural_trace.py | 43 +++++++++++++++++++++++++ 4 files changed, 71 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3036c3f..548d36c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ - **Trace Logs show the structured query that ran.** Each structured query step lists the query at the top of its output, followed by any attempts that returned nothing or failed, with their errors. Before, the query was buried in the result, and it was missing entirely when the query returned no rows — the case where it is most needed. ### Fixed +- **Structured queries no longer start with a step that cannot succeed.** Each structured query first asked the model for a pyTigerGraph function call, but no functions are registered with the TigerGraph embedding store, so every call was rejected (or failed with `'NoneType' object is not iterable` when the question named no relationships) before the openCypher query ran. GraphRAG now detects that no functions are registered and goes straight to the openCypher query, saving one model call per structured query. - **An empty structured query result no longer counts as data.** A query that returned no rows (for example `{"VS_XXX": []}`, or an empty list as text) was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. diff --git a/graphrag/app/tools/generate_function.py b/graphrag/app/tools/generate_function.py index 073cfe3..0ffe7f0 100644 --- a/graphrag/app/tools/generate_function.py +++ b/graphrag/app/tools/generate_function.py @@ -33,6 +33,7 @@ InvalidFunctionCallException, MapQuestionToSchemaException, NoDocumentsFoundException, + NoRegisteredFunctionsException, validate_function_call, validate_schema, ) @@ -101,6 +102,27 @@ def _run( """ LogWriter.info(f"request_id={req_id_cv.get()} ENTRY GenerateFunction._run()") + # The mapping's optional fields can arrive as None (e.g. no edges in + # the question); every use below expects a list or dict. + target_vertex_types = target_vertex_types or [] + target_vertex_attributes = target_vertex_attributes or {} + target_vertex_ids = target_vertex_ids or {} + target_edge_types = target_edge_types or [] + target_edge_attributes = target_edge_attributes or {} + + # Validation accepts only registered functions. With none registered, + # no generated call can pass, so stop before any model call. + valid_function_calls = [ + x["function_header"] + for x in (self.embedding_store.list_registered_documents( + output_fields=["function_header"] + ) or []) + ] + if not valid_function_calls: + raise NoRegisteredFunctionsException( + "No pyTigerGraph functions are registered in the embedding store." + ) + if target_vertex_types == [] and target_edge_types == []: return { "error": "No vertex or edge types recognized. MapQuestionToSchema and then try again." @@ -171,13 +193,6 @@ def _run( # Prioritize pyTigerGraph docs over custom docs docs = pytg_docs + custom_docs - valid_function_calls = [ - x["function_header"] - for x in self.embedding_store.list_registered_documents( - output_fields=["function_header"] - ) - ] - if len(docs) == 0: LogWriter.warning(f"request_id={req_id_cv.get()} WARN no documents found") raise NoDocumentsFoundException diff --git a/graphrag/app/tools/validation_utils.py b/graphrag/app/tools/validation_utils.py index 6b18b35..7a5a981 100644 --- a/graphrag/app/tools/validation_utils.py +++ b/graphrag/app/tools/validation_utils.py @@ -23,6 +23,11 @@ class InvalidFunctionCallException(Exception): pass +class NoRegisteredFunctionsException(Exception): + """The embedding store has no pyTigerGraph functions registered, so no + generated function call can pass validation.""" + + def validate_schema(conn, v_types, e_types, v_attrs, e_attrs): LogWriter.info(f"request_id={req_id_cv.get()} ENTRY validate_schema()") if v_types: diff --git a/graphrag/tests/test_structural_trace.py b/graphrag/tests/test_structural_trace.py index 3e59aa8..6f1b01e 100644 --- a/graphrag/tests/test_structural_trace.py +++ b/graphrag/tests/test_structural_trace.py @@ -88,6 +88,49 @@ def test_cypher_no_rows_is_not_retried(self): self.assertEqual(out["attempts"][-1], {"query": _CYPHER, "outcome": "no rows"}) +class TestEmptyFunctionRegistry(unittest.TestCase): + """Without registered functions, no function call is generated.""" + + def test_structural_tool_records_the_stop_and_runs_cypher(self): + from tools.validation_utils import NoRegisteredFunctionsException + ctx = _ctx(NoRegisteredFunctionsException("No pyTigerGraph functions are registered."), + use_cypher=True, cypher_runs=[(_CYPHER, _gsql([{"VS": [{"v_id": "d1"}]}]))]) + out = gt.structural_retrieve(ctx, "q") + self.assertTrue(out["ok"]) + self.assertEqual(out["query"], _CYPHER) + self.assertEqual(out["attempts"], [{"error": "No pyTigerGraph functions are registered."}]) + + def _gen(self, registered, llm=None): + from tools.generate_function import GenerateFunction + store = mock.Mock() + store.list_registered_documents.return_value = registered + return GenerateFunction(mock.Mock(), llm or mock.Mock(), mock.Mock(), store), store + + def test_generator_stops_before_any_model_call(self): + from tools.validation_utils import NoRegisteredFunctionsException + llm = mock.Mock() + gen, store = self._gen([], llm) + with self.assertRaises(NoRegisteredFunctionsException): + gen._run("q", ["Document"], None, None, None, None) + llm.invoke_with_parser.assert_not_called() + store.retrieve_similar.assert_not_called() + + def test_mapping_without_edges_does_not_fail(self): + """target_edge_types=None used to raise 'NoneType' object is not iterable.""" + from langchain_core.documents import Document + llm = mock.Mock(generate_function_prompt="{question}", query_guidance_block="") + llm.invoke_with_parser.return_value = types.SimpleNamespace( + connection_func_call="getVertexCount('Document')", func_call_reasoning="count") + gen, store = self._gen([{"function_header": "getVertexCount"}], llm) + store.retrieve_similar.return_value = [Document(page_content="getVertexCount docs")] + gen.conn.getVertexTypes.return_value = ["Document"] + gen.conn.getEndpoints.return_value = [] + gen.conn.getVertexCount.return_value = 4 + out = gen._run("How many documents?", ["Document"], None, None, None, None) + self.assertEqual(out["function_call"], "getVertexCount('Document')") + self.assertEqual(json.loads(out["result"]), 4) + + class TestTraceShowsTheQuery(unittest.TestCase): def test_query_and_attempts_lead_the_step_output(self): out = {"ok": True, "summary": "graph query returned rows", "context": {"result": [1]}, From 26caf56699b289c3b4e88aafa603c2947a06a49e Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 11:58:37 -0700 Subject: [PATCH 08/32] Remove the tigergraph-mcp adapter tools from the agentic engine - Drop the agent tools that ran self-written GSQL and neighbor expansion - GraphRAG code reaches TigerGraph only through pyTigerGraph Refs: GML-2332 --- CHANGELOG.md | 3 + common/llm_services/base_llm.py | 2 +- common/requirements.txt | 7 +- graphrag/app/agent/agentic_agent.py | 10 - graphrag/app/agent/agentic_react.py | 1 - graphrag/app/tools/graphrag_tools.py | 1 - graphrag/app/tools/tg_mcp_tools.py | 193 ------------------ graphrag/app/tools/tool_registry.py | 35 ---- graphrag/tests/test_tg_mcp_sessions_closed.py | 67 ------ 9 files changed, 8 insertions(+), 311 deletions(-) delete mode 100644 graphrag/app/tools/tg_mcp_tools.py delete mode 100644 graphrag/tests/test_tg_mcp_sessions_closed.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 548d36c..612f771 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,9 @@ - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. - **Agentic planning no longer gets stuck writing its plan summary.** The planner's summary and step reasons had no stated purpose or length, and the planner's output had no limit, so a model could repeat a phrase until it reached its own output maximum — costing thousands of tokens and tens of seconds, and sometimes the plan itself. These fields now ask for one short sentence, an overlong value is shortened instead of failing the plan, and the planner's output is capped. +### Removed +- **The `tg_run_query` and `tg_get_neighbors` agent tools.** They let the agent write and run its own GSQL and expand neighbors through an in-process copy of tigergraph-mcp's functions. Structured retrieval already maps the question to the schema, writes the query, runs it and retries on errors, so the agent does not need to write queries itself, and GraphRAG code now reaches TigerGraph only through pyTigerGraph. Custom planner or agent prompts that name these tools should drop them. + ### Added - **`planner_max_output_tokens` setting.** The most output tokens the planner may use for one plan (default `4000`; reasoning models use at least `16000`). diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index ef9afbc..adbf06d 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -1097,7 +1097,7 @@ def select_retriever_prompt(self): _AGENTIC_AGENT_SYSTEM = """\ You are a GraphRAG agent answering questions over a TigerGraph knowledge graph. -You have a set of read-only tools (graph schema via graphrag__get_schema, structural query generation, several unstructured retrievers, raw GSQL via tg_run_query, neighbor expansion). The graph schema is NOT pre-loaded — fetch it with graphrag__get_schema when you need it. +You have a set of read-only tools (graph schema via graphrag__get_schema, structural query generation, several unstructured retrievers). The graph schema is NOT pre-loaded — fetch it with graphrag__get_schema when you need it. REASON, ACT, OBSERVE — repeat until you can give a complete, well-grounded answer. diff --git a/common/requirements.txt b/common/requirements.txt index 0c7f03c..64db5bb 100644 --- a/common/requirements.txt +++ b/common/requirements.txt @@ -184,9 +184,10 @@ wrapt>=1.17.2 yarl>=1.20.1 zipp>=3.23.0 -# Agentic engine (v2.0) — MCP + tigergraph-mcp for in-process, per-user -# tool execution. Requires the fastapi/starlette bump above (mcp pulls -# starlette>=0.49). mcp 2.x changes the client API that mcp_addons uses. +# Agentic engine (v2.0) — MCP client for add-on MCP servers. tigergraph-mcp +# is kept for the planned built-in in-process server (GML-2330). Requires the +# fastapi/starlette bump above (mcp pulls starlette>=0.49). mcp 2.x changes +# the client API that mcp_addons uses. mcp>=1.27.1,<2 tigergraph-mcp>=1.0.4 sse-starlette>=3.4.4 diff --git a/graphrag/app/agent/agentic_agent.py b/graphrag/app/agent/agentic_agent.py index 4dd4b9c..a4fc8f2 100644 --- a/graphrag/app/agent/agentic_agent.py +++ b/graphrag/app/agent/agentic_agent.py @@ -180,15 +180,6 @@ def question_for_agent( "citations": [], }, ) - # Per-user creds for the tigergraph-mcp tools (when available), so - # those tool calls run as the logged-in user too. - tg_cfg = None - try: - from tools.tg_mcp_tools import conn_config_from_conn - tg_cfg = conn_config_from_conn(self.conn, self.conn.graphname) - except Exception: - tg_cfg = None - # Make sure any tarball-backed stdio servers configured for this # graph are installed before we try to launch them (a server saved # since the last restart may not have been installed at startup). @@ -231,7 +222,6 @@ def question_for_agent( use_cypher=self.use_cypher, conversation=convo, progress=self.emit_progress, - tg_connection_config=tg_cfg, external_tools=external_tools, mcp_manager=mcp_manager, user=user, diff --git a/graphrag/app/agent/agentic_react.py b/graphrag/app/agent/agentic_react.py index f98b2c4..5e9a4bc 100644 --- a/graphrag/app/agent/agentic_react.py +++ b/graphrag/app/agent/agentic_react.py @@ -52,7 +52,6 @@ "graphrag__contextual_search": "Searching the documents", "graphrag__similarity_search": "Searching the documents", "graphrag__community_search": "Searching community summaries", - "tg_run_query": "Running a graph query", } diff --git a/graphrag/app/tools/graphrag_tools.py b/graphrag/app/tools/graphrag_tools.py index 792c56d..0af3d96 100644 --- a/graphrag/app/tools/graphrag_tools.py +++ b/graphrag/app/tools/graphrag_tools.py @@ -59,7 +59,6 @@ class GraphRAGToolContext: use_cypher: bool = False conversation: Optional[list] = None progress: Optional[Callable[[str], None]] = None - tg_connection_config: Optional[dict] = None # per-user creds for tg-mcp tools # External MCP-addon tools discovered for this request and the manager # that dispatches them. Populated by the agent setup; consumed by # ``tool_registry`` (catalog / run / lc_tools_spec). Empty when no diff --git a/graphrag/app/tools/tg_mcp_tools.py b/graphrag/app/tools/tg_mcp_tools.py deleted file mode 100644 index daba4ab..0000000 --- a/graphrag/app/tools/tg_mcp_tools.py +++ /dev/null @@ -1,193 +0,0 @@ -# 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 . - -"""tigergraph-mcp adapter — in-process, per-user. - -Wraps a read-only subset of ``tigergraph-mcp`` tools so the agentic engine -can leverage them for raw TigerGraph operations (interpreted-query -execution, installed-query execution, neighbor expansion). Each call runs -as the **logged-in user**: a per-request ``connection_config`` is held in a -ContextVar and injected into ``tigergraph_mcp``'s ``get_connection`` via a -process-wide patch that reads that (per-request) var — so concurrent users -never share a connection. - -Import-guarded: if ``tigergraph-mcp`` isn't installed, ``AVAILABLE`` is -False and the registry simply skips these tools (the engine falls back to -the GraphRAG-native tools). -""" - -from __future__ import annotations - -import asyncio -import contextvars -import json -import logging -from typing import Optional - -logger = logging.getLogger(__name__) - -# Per-request TigerGraph credentials for the tg-mcp tools. -_tg_conn_cfg: contextvars.ContextVar = contextvars.ContextVar( - "tg_mcp_conn_cfg", default=None -) -# Connections a tool call used, so ``_run`` can close their sessions before -# the call's event loop ends (pyTigerGraph keeps one session per loop). -_used_conns: contextvars.ContextVar = contextvars.ContextVar( - "tg_mcp_used_conns", default=None -) - -try: - from tigergraph_mcp import connection_manager as _cm - from tigergraph_mcp.tools import query_tools as _qt - from tigergraph_mcp.tools import schema_tools as _st - - _orig_get_connection = _cm.get_connection - - def _patched_get_connection(profile=None, graph_name=None, connection_config=None): - # Inject the current request's per-user creds when the caller (the - # stock tool) didn't pass an explicit connection_config. - cfg = connection_config or _tg_conn_cfg.get() - conn = _orig_get_connection( - profile=profile, graph_name=graph_name, connection_config=cfg - ) - used = _used_conns.get() - if used is not None and conn is not None: - used.append(conn) - return conn - - # Patch the name bound inside each tool module we use. - _qt.get_connection = _patched_get_connection - _st.get_connection = _patched_get_connection - AVAILABLE = True -except Exception as exc: # pragma: no cover - import guard - logger.info(f"tigergraph-mcp not available, tg-mcp tools disabled: {exc}") - AVAILABLE = False - - -def set_user_connection_config(cfg: Optional[dict]) -> None: - """Set the per-request connection_config the tg-mcp tools run under.""" - _tg_conn_cfg.set(cfg) - - -def conn_config_from_conn(conn, graphname: str) -> dict: - """Build a tigergraph-mcp connection_config from a pyTigerGraph conn.""" - return { - "host": getattr(conn, "host", None) or getattr(conn, "gsUrl", ""), - "graphname": graphname or getattr(conn, "graphname", ""), - "username": getattr(conn, "username", "") or "", - "password": getattr(conn, "password", "") or "", - "apiToken": getattr(conn, "apiToken", "") or "", - "restppPort": getattr(conn, "restppPort", "9000"), - "gsPort": getattr(conn, "gsPort", "14240"), - } - - -def _run(coro): - """Run an async tg-mcp tool from the sync executor (fresh event loop in - the current worker thread; the ContextVar value is copied into it). - - The sessions the call opened belong to that loop, so they are closed - before it ends; otherwise each call leaves one behind. - """ - async def _call(): - token = _used_conns.set([]) - try: - return await coro - finally: - for conn in _used_conns.get(): - aclose = getattr(conn, "aclose", None) - if aclose is not None: - try: - await aclose() - except Exception as exc: # noqa: BLE001 — cleanup only - logger.debug(f"tg-mcp: closing a connection failed: {exc}") - _used_conns.reset(token) - - return asyncio.run(_call()) - - -def _normalize(res, label: str) -> dict: - """Turn a tg-mcp ``List[TextContent]`` into our tool result dict.""" - item = res[0] if isinstance(res, list) and res else res - text = getattr(item, "text", None) or (item if isinstance(item, str) else str(item)) - ok = True - parsed = None - try: - body = text.strip() - if body.startswith("```"): - body = body.split("```", 2)[1] - if body.startswith("json"): - body = body[4:] - parsed = json.loads(body) - if isinstance(parsed, dict) and parsed.get("success") is False: - ok = False - except Exception: - parsed = text - return { - "ok": ok, - "summary": f"{label}: {'ok' if ok else 'failed'}", - "context": {"function_call": label, "result": parsed}, - "citations": [], - } - - -# --- read-only tool wrappers (ctx is the GraphRAGToolContext) -------------- - -def _ensure(ctx): - """Bind the per-request creds before each call (idempotent).""" - cfg = getattr(ctx, "tg_connection_config", None) - if cfg is None: - cfg = conn_config_from_conn(ctx.conn, getattr(ctx.conn, "graphname", "")) - set_user_connection_config(cfg) - - -def tg_run_query(ctx, query_text: str) -> dict: - """Run an interpreted (dynamic) GSQL query as the logged-in user.""" - if not AVAILABLE: - return {"ok": False, "summary": "tigergraph-mcp unavailable", "context": None, "citations": []} - ctx.emit("Running a graph query (tigergraph-mcp)") - _ensure(ctx) - g = getattr(ctx.conn, "graphname", "") - return _normalize(_run(_qt.run_query(query_text=query_text, graph_name=g)), "tg_run_query") - - -def tg_run_installed_query(ctx, query_name: str, params: Optional[dict] = None) -> dict: - """Run a pre-installed query by name as the logged-in user.""" - if not AVAILABLE: - return {"ok": False, "summary": "tigergraph-mcp unavailable", "context": None, "citations": []} - ctx.emit(f"Running installed query {query_name} (tigergraph-mcp)") - _ensure(ctx) - g = getattr(ctx.conn, "graphname", "") - return _normalize( - _run(_qt.run_installed_query(query_name=query_name, params=params or {}, graph_name=g)), - "tg_run_installed_query", - ) - - -def tg_get_neighbors(ctx, vertex_type: str, vertex_id: str, - edge_type: Optional[str] = None, - target_vertex_type: Optional[str] = None, - limit: Optional[int] = None) -> dict: - """Expand neighbors of a vertex as the logged-in user (no GSQL needed).""" - if not AVAILABLE: - return {"ok": False, "summary": "tigergraph-mcp unavailable", "context": None, "citations": []} - ctx.emit("Expanding neighbors (tigergraph-mcp)") - _ensure(ctx) - g = getattr(ctx.conn, "graphname", "") - return _normalize( - _run(_qt.get_neighbors( - vertex_type=vertex_type, vertex_id=vertex_id, edge_type=edge_type, - target_vertex_type=target_vertex_type, limit=limit, graph_name=g)), - "tg_get_neighbors", - ) diff --git a/graphrag/app/tools/tool_registry.py b/graphrag/app/tools/tool_registry.py index e6d503a..c24edf1 100644 --- a/graphrag/app/tools/tool_registry.py +++ b/graphrag/app/tools/tool_registry.py @@ -73,18 +73,6 @@ class CommunitySearchArgs(BaseModel): with_chunk: Optional[bool] = Field(default=None, description="Also return chunks linked to the communities.") -class TgRunQueryArgs(BaseModel): - query_text: str = Field(description="A complete interpreted GSQL query body to execute (read-only).") - - -class TgGetNeighborsArgs(BaseModel): - vertex_type: str = Field(description="Vertex type of the starting vertex.") - vertex_id: str = Field(description="Id of the starting vertex.") - edge_type: Optional[str] = Field(default=None, description="Restrict expansion to this edge type.") - target_vertex_type: Optional[str] = Field(default=None, description="Restrict neighbors to this vertex type.") - limit: Optional[int] = Field(default=None, description="Max neighbors to return.") - - @dataclass class ToolSpec: """One callable tool the planner / react loop can dispatch. @@ -199,29 +187,6 @@ def _merged_specs(ctx) -> dict: CommunitySearchArgs, gt.community_search, ) -# tigergraph-mcp read tools — registered only when the package is -# installed (import-guarded). These run as the logged-in user via the -# per-request connection shim in tg_mcp_tools. -try: - from tools import tg_mcp_tools as tgm - if tgm.AVAILABLE: - _register( - "tg_run_query", - "Execute an interpreted (dynamic) read-only GSQL query against the graph " - "and return the rows. Use when a precise graph query is needed.", - TgRunQueryArgs, tgm.tg_run_query, - ) - _register( - "tg_get_neighbors", - "Return the neighbors of a given vertex (optionally filtered by edge / target " - "type) without writing GSQL.", - TgGetNeighborsArgs, tgm.tg_get_neighbors, - ) - logger.info("tigergraph-mcp read tools registered") -except Exception as exc: # pragma: no cover - import guard - logger.info(f"tigergraph-mcp tools not registered: {exc}") - - def tool_names(ctx=None) -> list[str]: return list(_merged_specs(ctx).keys()) diff --git a/graphrag/tests/test_tg_mcp_sessions_closed.py b/graphrag/tests/test_tg_mcp_sessions_closed.py deleted file mode 100644 index 67d7417..0000000 --- a/graphrag/tests/test_tg_mcp_sessions_closed.py +++ /dev/null @@ -1,67 +0,0 @@ -import unittest - - -class TestToolCallsCloseTheirSessions(unittest.TestCase): - """Each tool call runs on its own event loop; with pyTigerGraph keeping one - session per loop, the call must close what it opened before the loop ends.""" - - @classmethod - def setUpClass(cls): - try: - from tools import tg_mcp_tools as t - except ImportError as exc: # needs the app environment - raise unittest.SkipTest(f"app environment unavailable: {exc}") - if not t.AVAILABLE: - raise unittest.SkipTest("tigergraph-mcp not installed") - cls.t = t - - def test_connections_used_by_a_call_are_closed(self): - t = self.t - closed = [] - - class FakeConn: - async def aclose(self): - closed.append(self) - - conns = [FakeConn(), FakeConn()] - - async def tool(): - # A tool asks for its connection through the patched helper. - used = t._used_conns.get() - for c in conns: - used.append(c) - return "done" - - self.assertEqual(t._run(tool()), "done") - self.assertEqual(closed, conns) - self.assertIsNone(t._used_conns.get()) - - def test_close_failure_does_not_fail_the_call(self): - t = self.t - - class Broken: - async def aclose(self): - raise RuntimeError("already closed") - - async def tool(): - t._used_conns.get().append(Broken()) - return 42 - - self.assertEqual(t._run(tool()), 42) - - def test_patched_lookup_records_the_connection(self): - t = self.t - from unittest import mock - sentinel = object() - with mock.patch.object(t, "_orig_get_connection", return_value=sentinel): - seen = [] - async def tracked(): - conn = t._patched_get_connection(graph_name="g") - seen.extend(t._used_conns.get()) - return conn - self.assertIs(t._run(tracked()), sentinel) - self.assertEqual(seen, [sentinel]) - - -if __name__ == "__main__": - unittest.main() From b49b4606176fde183996f3521093538e372bf7fe Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 11:58:43 -0700 Subject: [PATCH 09/32] Give the agentic planner the graph schema - Planner gets a compact schema summary to decide on structural retrieval - The summary is cached per graph and refreshed when the schema changes Refs: GML-2331 --- CHANGELOG.md | 1 + common/llm_services/base_llm.py | 2 +- graphrag/app/agent/agentic_graph.py | 18 ++-- graphrag/app/agent/agentic_planner.py | 91 ++++++++++++++++- graphrag/tests/test_planner_schema.py | 138 ++++++++++++++++++++++++++ 5 files changed, 238 insertions(+), 12 deletions(-) create mode 100644 graphrag/tests/test_planner_schema.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 612f771..8725105 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ - **Trace Logs show the structured query that ran.** Each structured query step lists the query at the top of its output, followed by any attempts that returned nothing or failed, with their errors. Before, the query was buried in the result, and it was missing entirely when the query returned no rows — the case where it is most needed. ### Fixed +- **The agentic planner plans from the graph schema.** Since 2.0.0 the planner made its plan without knowing the graph's vertex and edge types, so it could only guess from the wording of a question whether structural retrieval would help. It now gets a compact summary of the schema — type names, attribute names, edge endpoints and descriptions — cached per graph and refreshed when the schema changes. - **Structured queries no longer start with a step that cannot succeed.** Each structured query first asked the model for a pyTigerGraph function call, but no functions are registered with the TigerGraph embedding store, so every call was rejected (or failed with `'NoneType' object is not iterable` when the question named no relationships) before the openCypher query ran. GraphRAG now detects that no functions are registered and goes straight to the openCypher query, saving one model call per structured query. - **An empty structured query result no longer counts as data.** A query that returned no rows (for example `{"VS_XXX": []}`, or an empty list as text) was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index adbf06d..93d1c3a 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -1158,7 +1158,7 @@ def agentic_agent_prompt(self): - in what order. Express this as a small DAG of tool steps that gathers exactly the context needed, ending with one final "answer" step that consolidates all the gathered context into the response. Express ordering with depends_on and repetition with multiple steps. -The graph schema is NOT provided here — the structural and unstructured query tools load it themselves at run time, so plan retrieval steps directly. A question that needs no graph data should not include any graph-retrieval step (plan only the final answer step, or the relevant non-graph tool). +## Graph schema lists the vertex types (with their attribute names) and the edge types of the graph. Use it to decide whether the question needs STRUCTURAL retrieval: plan a structural step when the question asks for data that these types and attributes hold (counts, lookups, relationships, aggregations). The query tools load the full schema themselves, so do not plan a schema step. A question that needs no graph data should not include any graph-retrieval step (plan only the final answer step, or the relevant non-graph tool). Before you plan retrieval, list the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. Then: - Every fact is stated (for example, the same question was already asked and answered): plan only the final answer step. Do not retrieve those facts again. diff --git a/graphrag/app/agent/agentic_graph.py b/graphrag/app/agent/agentic_graph.py index f66e0c9..19e0a78 100644 --- a/graphrag/app/agent/agentic_graph.py +++ b/graphrag/app/agent/agentic_graph.py @@ -25,7 +25,7 @@ import time from agent.agentic_executor import execute_plan, fit_for_trace, with_trace_note, _run_step, _usage_since -from agent.agentic_planner import check_conversation, plan_question +from agent.agentic_planner import check_conversation, plan_question, schema_summary from agent.agentic_synthesizer import _gather, has_context, synthesize, with_conversation from common.llm_services.base_llm import get_collected_usage from common.py_schemas import GraphRAGResponse, Plan, PlanStep @@ -97,9 +97,11 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: max_replans = int(_cfg.get("agent_max_replans", MAX_REPLANS)) max_total_steps = int(_cfg.get("agent_max_total_steps", MAX_TOTAL_STEPS)) - # The schema is loaded lazily by the query tools at run time, so a question - # that needs no graph data does not trigger a schema read. emit("Planning an approach") + # The planner decides from the schema whether the question needs + # structural retrieval. Questions that need no data are answered by triage + # before this point. + schema = schema_summary(getattr(ctx, "conn", None)) results: dict = {} agent_steps: list = [] @@ -107,11 +109,12 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: # plan (timed + usage-attributed for the trace) _u0 = len(get_collected_usage() or []) _t0 = time.time() - plan = plan_question(llm, question, conversation, ctx=ctx) + plan = plan_question(llm, question, conversation, schema_rep=schema, ctx=ctx) agent_steps.append({ "node": "plan", "kind": "plan", "duration_s": round(time.time() - _t0, 3), - "input": {"question": question, "conversation": conversation or []}, + "input": {"question": question, "conversation": conversation or [], + "graph_schema": schema}, "output": _plan_output(plan), "usage": _usage_since(_u0), }) @@ -135,7 +138,7 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: emit("Refining the plan") _u0 = len(get_collected_usage() or []) _t0 = time.time() - plan = plan_question(llm, question, conversation, ctx=ctx, + plan = plan_question(llm, question, conversation, schema_rep=schema, ctx=ctx, missing_from_conversation=check.missing or "some facts the question asks for") agent_steps.append({ "node": "replan (conversation check)", "kind": "plan", @@ -187,7 +190,8 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: ] _u0 = len(get_collected_usage() or []) _t0 = time.time() - plan = plan_question(llm, question, conversation, prior_results=prior, ctx=ctx) + plan = plan_question(llm, question, conversation, schema_rep=schema, + prior_results=prior, ctx=ctx) agent_steps.append({ "node": f"replan {replans}", "kind": "plan", "duration_s": round(time.time() - _t0, 3), diff --git a/graphrag/app/agent/agentic_planner.py b/graphrag/app/agent/agentic_planner.py index a51d65e..a26c3f5 100644 --- a/graphrag/app/agent/agentic_planner.py +++ b/graphrag/app/agent/agentic_planner.py @@ -23,6 +23,8 @@ """ import logging +import threading +import time from common.llm_services.capabilities import model_capabilities from common.py_schemas import ConversationCheck, Plan @@ -31,6 +33,87 @@ logger = logging.getLogger(__name__) +# The planner decides from the schema whether a question needs structural +# retrieval (GML-2331). It gets a compact summary: type names, attribute +# names, edge endpoints and descriptions. The query tools still load the +# full schema themselves. +_SCHEMA_SUMMARY_MAX_CHARS = 12000 +# How long a cached summary is used before the schema version is read again. +_SCHEMA_RECHECK_SECONDS = 60 +_schema_cache: dict = {} +_schema_cache_lock = threading.Lock() + + +def _is_bookkeeping_attr(name: str, tg_type) -> bool: + """GraphRAG processing timestamps and embedding vectors, which no + question asks about.""" + t = str(tg_type or "").upper() + return name.startswith("epoch_") or "LIST str: + """A compact schema for the planner, from an ``AllowedSchema``.""" + lines = ["Vertex types (attributes):"] + for vt in rep.vertex_types: + attrs = [a for a, t in (rep.vertex_attributes.get(vt) or {}).items() + if not _is_bookkeeping_attr(a, t)] + desc = (rep.vertex_definitions or {}).get(vt) + lines.append(f"- {vt}({', '.join(attrs)})" + (f": {desc}" if desc else "")) + lines.append("Edge types (from -> to):") + for et in rep.edge_types: + pairs = (rep.edge_endpoints or {}).get(et) or [] + ends = ", ".join(f"{f} -> {t}" for f, t in pairs) + attrs = [a for a, t in (rep.edge_attributes.get(et) or {}).items() + if not _is_bookkeeping_attr(a, t)] + desc = (rep.edge_definitions or {}).get(et) + line = f"- {et}: {ends}" if ends else f"- {et}" + if attrs: + line += f" ({', '.join(attrs)})" + if desc: + line += f": {desc}" + lines.append(line) + text = "\n".join(lines) + if len(text) > _SCHEMA_SUMMARY_MAX_CHARS: + text = text[:_SCHEMA_SUMMARY_MAX_CHARS].rsplit("\n", 1)[0] + "\n(more types not shown)" + return text + + +def schema_summary(conn) -> str: + """The planner's schema summary for the graph of ``conn``. + + Cached per graph. After ``_SCHEMA_RECHECK_SECONDS`` the schema version + is read again, and the summary is rebuilt only when it changed. Returns + the last known summary (or "") when the schema cannot be read. + """ + if conn is None: + return "" + key = (str(getattr(conn, "host", "")), getattr(conn, "graphname", "")) + now = time.time() + with _schema_cache_lock: + entry = _schema_cache.get(key) + if entry and now - entry["checked"] < _SCHEMA_RECHECK_SECONDS: + return entry["text"] + try: + from common.db.connections import get_schema_ver + version = get_schema_ver(conn) + if entry and version is not None and version == entry["version"]: + with _schema_cache_lock: + entry["checked"] = now + return entry["text"] + from common.db.schema_utils import render_schema_rep + rep = render_schema_rep(conn) + text = render_schema_summary(rep) + with _schema_cache_lock: + _schema_cache[key] = { + "checked": now, + "version": rep.schema_version if rep.schema_version is not None else version, + "text": text, + } + return text + except Exception as exc: + logger.warning(f"planner: schema summary unavailable ({exc})") + return entry["text"] if entry else "" + def _param_type(pinfo: dict) -> str: """Render a JSON-schema property's type for the catalog. Falls back through enum / anyOf (common in external MCP tool schemas) to ``any``.""" @@ -146,10 +229,10 @@ def plan_question(llm, question, conversation=None, schema_rep="", prior_results f"## Question\n{question}", f"## Conversation\n{recent_conversation(conversation)}", ] - # Schema is normally not pre-loaded (the query tools load it themselves); - # include it only if a caller explicitly supplied one. - if schema_rep: - user_parts.append(f"## Graph schema\n{schema_rep[:6000]}") + user_parts.append( + f"## Graph schema\n{schema_rep[:_SCHEMA_SUMMARY_MAX_CHARS + 100]}" if schema_rep + else "## Graph schema\nNot available. Decide from the question." + ) user_parts.append(f"## Tools\n{catalog}") if missing_from_conversation: user_parts.append( diff --git a/graphrag/tests/test_planner_schema.py b/graphrag/tests/test_planner_schema.py new file mode 100644 index 0000000..a83a99b --- /dev/null +++ b/graphrag/tests/test_planner_schema.py @@ -0,0 +1,138 @@ +"""The planner gets a compact, cached schema summary (GML-2331).""" +import types +import unittest +from unittest import mock + +try: + import agent.agentic_planner as ap + import agent.agentic_graph as g + from common.db.schema_utils import AllowedSchema + from common.llm_services.base_llm import LLM_Model + from common.py_schemas import GraphRAGResponse, Plan, PlanStep +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +def _rep(version=1): + return AllowedSchema( + schema_version=version, + vertex_types=["Ticket", "Document"], + edge_types=["ASSIGNED_TO", "HAS_CHILD"], + vertex_attributes={ + "Ticket": {"status": "STRING", "created_at": "DATETIME", "epoch_added": "UINT", + "embedding": "LIST"}, + "Document": {"epoch_added": "UINT"}, + }, + edge_attributes={"ASSIGNED_TO": {"since": "DATETIME"}}, + vertex_definitions={"Ticket": "A Jira ticket"}, + edge_definitions={"ASSIGNED_TO": "Who works on the ticket"}, + edge_endpoints={"ASSIGNED_TO": [("Ticket", "Person")], "HAS_CHILD": [("Document", "DocumentChunk")]}, + ) + + +class TestSummary(unittest.TestCase): + def test_types_attributes_endpoints_and_descriptions(self): + text = ap.render_schema_summary(_rep()) + self.assertIn("- Ticket(status, created_at): A Jira ticket", text) + self.assertIn("- Document()", text) + self.assertIn("- ASSIGNED_TO: Ticket -> Person (since): Who works on the ticket", text) + self.assertIn("- HAS_CHILD: Document -> DocumentChunk", text) + self.assertNotIn("epoch_added", text) + self.assertNotIn("embedding", text) + + def test_long_schema_is_cut_at_a_line(self): + rep = _rep() + rep.vertex_types = [f"Type{i}" for i in range(3000)] + text = ap.render_schema_summary(rep) + self.assertLessEqual(len(text), ap._SCHEMA_SUMMARY_MAX_CHARS + 30) + self.assertTrue(text.endswith("(more types not shown)")) + + +class TestCache(unittest.TestCase): + def setUp(self): + ap._schema_cache.clear() + self.conn = types.SimpleNamespace(host="h", graphname="g") + + def _summary(self, versions, now): + with mock.patch("common.db.connections.get_schema_ver", side_effect=versions) as ver, \ + mock.patch("common.db.schema_utils.render_schema_rep", + side_effect=lambda c: _rep(version=versions_seen[-1])) as render, \ + mock.patch.object(ap.time, "time", return_value=now): + return ap.schema_summary(self.conn), ver, render + + def test_read_once_then_cached(self): + global versions_seen + versions_seen = [1] + _, ver, render = self._summary([1], 1000.0) + self.assertEqual((ver.call_count, render.call_count), (1, 1)) + _, ver, render = self._summary([1], 1030.0) # within the recheck window + self.assertEqual((ver.call_count, render.call_count), (0, 0)) + + def test_same_version_is_not_rebuilt(self): + global versions_seen + versions_seen = [1] + self._summary([1], 1000.0) + _, ver, render = self._summary([1], 1100.0) + self.assertEqual((ver.call_count, render.call_count), (1, 0)) + + def test_new_version_is_rebuilt(self): + global versions_seen + versions_seen = [1] + self._summary([1], 1000.0) + versions_seen = [2] + _, ver, render = self._summary([2], 1100.0) + self.assertEqual(render.call_count, 1) + + def test_read_error_keeps_the_last_summary(self): + global versions_seen + versions_seen = [1] + first, _, _ = self._summary([1], 1000.0) + with mock.patch("common.db.connections.get_schema_ver", side_effect=RuntimeError("down")), \ + mock.patch.object(ap.time, "time", return_value=1100.0): + self.assertEqual(ap.schema_summary(self.conn), first) + + def test_no_connection(self): + self.assertEqual(ap.schema_summary(None), "") + + +class TestPlannerGetsTheSchema(unittest.TestCase): + def _user_text(self, schema): + seen = {} + + def invoke(messages, schema_cls, **kw): + seen["user"] = messages[-1][1] + return Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) + + llm = types.SimpleNamespace(agentic_planner_prompt="sys", config={}, invoke_structured=invoke) + ap.plan_question(llm, "How many tickets?", [], schema_rep=schema) + return seen["user"] + + def test_schema_in_the_planner_input(self): + self.assertIn("## Graph schema\nVertex types", self._user_text("Vertex types (attributes):\n- Ticket(status)")) + + def test_missing_schema_is_stated(self): + self.assertIn("## Graph schema\nNot available", self._user_text("")) + + def test_engine_passes_the_schema_to_every_plan(self): + ctx = types.SimpleNamespace(graphrag_cfg={}, emit=lambda *a, **k: None, conn=object()) + answer_only = Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) + with mock.patch.object(g, "schema_summary", return_value="SCHEMA") as summ, \ + mock.patch.object(g, "plan_question", return_value=answer_only) as planner, \ + mock.patch.object(g, "check_conversation", return_value=None), \ + mock.patch.object(g, "synthesize", return_value=GraphRAGResponse( + natural_language_response="ok", answered_question=True, response_type="agentic")): + resp = g.run_agentic(ctx, llm=object(), question="q") + summ.assert_called_once() + self.assertEqual(planner.call_args.kwargs["schema_rep"], "SCHEMA") + self.assertEqual(resp.query_sources["agent_steps"][0]["input"]["graph_schema"], "SCHEMA") + + def test_prompt_uses_the_schema(self): + prompt = LLM_Model._AGENTIC_PLANNER_SYSTEM + self.assertIn("## Graph schema lists the vertex types", prompt) + self.assertNotIn("NOT provided", prompt) + + +versions_seen = [1] + +if __name__ == "__main__": + unittest.main() From 0d1f399313c704ccf99372c842dbbcd07a0fdf46 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 12:36:31 -0700 Subject: [PATCH 10/32] Use one empty-result check in both chat engines - Classic and agentic engines judge a query result empty the same way - An empty result returned as JSON text counts as no rows in the classic engine Refs: GML-2327 --- CHANGELOG.md | 2 +- common/utils/query_results.py | 50 ++++++++++++++++++++++++++++ graphrag/app/agent/agent_graph.py | 12 ++----- graphrag/app/tools/graphrag_tools.py | 27 ++------------- graphrag/tests/test_query_results.py | 34 +++++++++++++++++++ 5 files changed, 90 insertions(+), 35 deletions(-) create mode 100644 common/utils/query_results.py create mode 100644 graphrag/tests/test_query_results.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 8725105..89f770a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ ### Fixed - **The agentic planner plans from the graph schema.** Since 2.0.0 the planner made its plan without knowing the graph's vertex and edge types, so it could only guess from the wording of a question whether structural retrieval would help. It now gets a compact summary of the schema — type names, attribute names, edge endpoints and descriptions — cached per graph and refreshed when the schema changes. - **Structured queries no longer start with a step that cannot succeed.** Each structured query first asked the model for a pyTigerGraph function call, but no functions are registered with the TigerGraph embedding store, so every call was rejected (or failed with `'NoneType' object is not iterable` when the question named no relationships) before the openCypher query ran. GraphRAG now detects that no functions are registered and goes straight to the openCypher query, saving one model call per structured query. -- **An empty structured query result no longer counts as data.** A query that returned no rows (for example `{"VS_XXX": []}`, or an empty list as text) was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. +- **An empty structured query result no longer counts as data.** A query that returned no rows (for example `{"VS_XXX": []}`, or an empty list as text) was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. The classic engine uses the same check, so both engines treat an empty result the same way. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. - **Agentic planning no longer gets stuck writing its plan summary.** The planner's summary and step reasons had no stated purpose or length, and the planner's output had no limit, so a model could repeat a phrase until it reached its own output maximum — costing thousands of tokens and tens of seconds, and sometimes the plan itself. These fields now ask for one short sentence, an overlong value is shortened instead of failing the plan, and the planner's output is capped. diff --git a/common/utils/query_results.py b/common/utils/query_results.py new file mode 100644 index 0000000..deb3fbf --- /dev/null +++ b/common/utils/query_results.py @@ -0,0 +1,50 @@ +# 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 . + +"""Whether a graph query result holds any rows. + +Shared by the agentic structural tool and the classic engine, so both treat +the same result the same way. +""" + +import json +from typing import Any + + +def is_empty_result(result: Any) -> bool: + """True when ``result`` holds no rows. + + Results arrive as JSON text (pyTigerGraph function calls) or as nested + containers (openCypher, e.g. ``{"VS_XXX": []}``), so both are read + through. A container is empty when every value in it is empty. A scalar + such as a count of 0 is an answer, not an empty result. + """ + if result is None: + return True + if isinstance(result, str): + text = result.strip() + if not text: + return True + try: + parsed = json.loads(text) + except ValueError: + return False + if parsed is None or isinstance(parsed, (list, dict)): + return is_empty_result(parsed) + return False + if isinstance(result, dict): + return all(is_empty_result(v) for v in result.values()) + if isinstance(result, (list, tuple, set, frozenset, range)): + return all(is_empty_result(v) for v in result) + return False diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py index 93c7350..07ff80a 100644 --- a/graphrag/app/agent/agent_graph.py +++ b/graphrag/app/agent/agent_graph.py @@ -39,6 +39,7 @@ SiblingRetriever, CommunityRetriever) from tools import MapQuestionToSchemaException from typing_extensions import TypedDict +from common.utils.query_results import is_empty_result from common.logs.log import req_id_cv from common.metrics.prometheus_metrics import metrics as pmetrics @@ -904,16 +905,7 @@ def is_query_result_empty(self, query_result) -> bool: """ Check if the query result is empty or contains empty values. """ - if query_result in ("", [], {}, (), set(), range(0), None): - return True - - if isinstance(query_result, (list, set)): - return all(self.is_query_result_empty(item) for item in query_result) - - if isinstance(query_result, dict): - return all(self.is_query_result_empty(v) for v in query_result.values()) - - return False + return is_empty_result(query_result) # remove halucinaton check, always return grounded def check_answer_for_hallucinations(self, state): diff --git a/graphrag/app/tools/graphrag_tools.py b/graphrag/app/tools/graphrag_tools.py index 0af3d96..96f4d1c 100644 --- a/graphrag/app/tools/graphrag_tools.py +++ b/graphrag/app/tools/graphrag_tools.py @@ -33,6 +33,7 @@ from dataclasses import dataclass, field from typing import Any, Callable, Dict, Optional +from common.utils.query_results import is_empty_result from common.utils.retrieval_stats import describe_retrieval from tools import tool_guards as guards from tools.validation_utils import MapQuestionToSchemaException @@ -83,30 +84,8 @@ def _empty(summary: str) -> dict: return {"ok": False, "summary": summary, "context": None, "citations": []} -def _result_is_empty(result: Any) -> bool: - """True when ``result`` holds no rows. - - Query results arrive as JSON text (pyTigerGraph) or nested containers - (openCypher, e.g. ``{"VS_XXX": []}``), so both are read through. A - container is empty when every value in it is empty. A scalar such as a - count of 0 is an answer, not an empty result. - """ - if result is None: - return True - if isinstance(result, str): - text = result.strip() - if not text: - return True - try: - parsed = json.loads(text) - except ValueError: - return False - return _result_is_empty(parsed) if isinstance(parsed, (list, dict)) or parsed is None else False - if isinstance(result, dict): - return all(_result_is_empty(v) for v in result.values()) - if isinstance(result, list): - return all(_result_is_empty(v) for v in result) - return False +# One empty-result check for the agentic and classic engines. +_result_is_empty = is_empty_result # Longest error text kept for one query attempt in the trace. diff --git a/graphrag/tests/test_query_results.py b/graphrag/tests/test_query_results.py new file mode 100644 index 0000000..097a1c9 --- /dev/null +++ b/graphrag/tests/test_query_results.py @@ -0,0 +1,34 @@ +"""One empty-result check for the agentic and classic engines (GML-2327).""" +import unittest + +try: + from common.utils.query_results import is_empty_result + from tools import graphrag_tools as gt + from agent.agent_graph import TigerGraphAgentGraph +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + +EMPTY = [None, "", " ", "[]", "{}", "null", " [ ] ", '{"VS_XXX": []}', '[{"VS_XXX": []}]', + {"VS_XXX": []}, [{"VS_XXX": []}], [], {}, (), set(), range(0), [None, ""]] +ROWS = [0, "0", 4, "4", False, '{"count": 0}', {"count": 0}, [{"v_id": "x"}], + '[{"v_id": "x"}]', {"VS_XXX": [{"v_id": "x"}]}, "not json text", [0]] + + +class TestIsEmptyResult(unittest.TestCase): + def test_empty(self): + for value in EMPTY: + self.assertTrue(is_empty_result(value), repr(value)) + + def test_rows(self): + for value in ROWS: + self.assertFalse(is_empty_result(value), repr(value)) + + def test_both_engines_agree(self): + classic = TigerGraphAgentGraph.is_query_result_empty + for value in EMPTY + ROWS: + self.assertEqual(gt._result_is_empty(value), is_empty_result(value), repr(value)) + self.assertEqual(classic(None, value), is_empty_result(value), repr(value)) + + +if __name__ == "__main__": + unittest.main() From 39aa6b03b6afefd6c648b5b6d19f786318d0c167 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 12:57:26 -0700 Subject: [PATCH 11/32] Fix the stale unit tests so the full suite passes - Connection tests no longer leave mocked modules behind for later tests - Tests follow per-user sign-in, three-part role lookup and masked secrets - Extractor tests use the schema-based extractor setup - LLM regression tests are skipped cleanly outside a git checkout Refs: GML-2333 --- graphrag/tests/test_connections.py | 144 ++++++------------ .../test_llm_entity_relationship_extractor.py | 122 +++++++-------- graphrag/tests/test_prompt_split.py | 8 +- graphrag/tests/test_service.py | 12 +- graphrag/tests/test_ui_setup_unit.py | 44 +++--- 5 files changed, 138 insertions(+), 192 deletions(-) diff --git a/graphrag/tests/test_connections.py b/graphrag/tests/test_connections.py index cb795d6..8c6db8c 100644 --- a/graphrag/tests/test_connections.py +++ b/graphrag/tests/test_connections.py @@ -1,115 +1,63 @@ -"""Unit tests for common.db.connections apiToken support.""" +"""Unit tests for common.db.connections.elevate_db_connection_to_token. -import sys -import unittest -from unittest.mock import patch, MagicMock - -# Mock heavy dependencies before importing the module under test -_mock_modules = {} -for mod_name in [ - "common.config", - "common.logs", - "common.logs.logwriter", - "common.logs.log", - "common.metrics", - "common.metrics.tg_proxy", - "common.metrics.prometheus_metrics", - "common.embeddings", - "common.embeddings.embedding_services", - "common.embeddings.tigergraph_embedding_store", - "common.llm_services", - "common.session", - "common.status", - "langchain_core", - "langchain_core.embeddings", - "prometheus_client", -]: - if mod_name not in sys.modules: - _mock_modules[mod_name] = MagicMock() - sys.modules[mod_name] = _mock_modules[mod_name] - -# Provide the values that connections.py reads at import time -sys.modules["common.config"].security = MagicMock() +Since 1.4.1 a per-user connection authenticates as the signed-in user. The +configured ``apiToken`` is for the service's own connection only (see +``common/config.py`` and the README), so it must not reach a user's +connection. +""" -# Provide TigerGraphConnectionProxy -mock_proxy_cls = MagicMock() -sys.modules["common.metrics.tg_proxy"].TigerGraphConnectionProxy = mock_proxy_cls - -# Provide LogWriter -mock_logwriter = MagicMock() -sys.modules["common.logs.logwriter"].LogWriter = mock_logwriter +import unittest +from unittest.mock import patch -from pyTigerGraph import TigerGraphConnection, AsyncTigerGraphConnection +try: + from pyTigerGraph import AsyncTigerGraphConnection, TigerGraphConnection + import common.db.connections as connections +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") -MOCK_DB_CONFIG_BASE = { +DB_CONFIG = { "hostname": "http://test-host", "restppPort": "9000", "gsPort": "14240", - "getToken": False, + "getToken": True, + "apiToken": "service_token", "default_timeout": 300, } -class TestElevateDbConnectionWithApiToken(unittest.TestCase): - """Test that elevate_db_connection_to_token honours apiToken from db_config.""" - - def _import_and_patch(self, db_config_override): - """Import elevate_db_connection_to_token with a patched db_config.""" - sys.modules["common.config"].db_config = db_config_override - # Re-import to pick up fresh db_config reference - import importlib - import common.db.connections as conn_mod - importlib.reload(conn_mod) - return conn_mod.elevate_db_connection_to_token - - def test_static_api_token_used_directly(self): - cfg = {**MOCK_DB_CONFIG_BASE, "apiToken": "static_tok_123"} - elevate = self._import_and_patch(cfg) - - conn = elevate("http://test-host", "user", "pass", "TestGraph") - - self.assertEqual(conn.apiToken, "static_tok_123") - self.assertEqual(conn.host, "http://test-host") - self.assertEqual(conn.graphname, "TestGraph") - - def test_api_token_skips_get_token(self): - cfg = {**MOCK_DB_CONFIG_BASE, "apiToken": "tok", "getToken": True} - elevate = self._import_and_patch(cfg) - - with patch.object(TigerGraphConnection, "getToken") as mock_get: - conn = elevate("http://test-host", "user", "pass", "TestGraph") - mock_get.assert_not_called() - - self.assertEqual(conn.apiToken, "tok") - - def test_static_api_token_async_conn(self): - cfg = {**MOCK_DB_CONFIG_BASE, "apiToken": "async_tok"} - elevate = self._import_and_patch(cfg) - - conn = elevate("http://test-host", "user", "pass", "TestGraph", async_conn=True) - - self.assertIsInstance(conn, AsyncTigerGraphConnection) - self.assertEqual(conn.apiToken, "async_tok") - - def test_empty_api_token_falls_through(self): - cfg = {**MOCK_DB_CONFIG_BASE, "apiToken": ""} - elevate = self._import_and_patch(cfg) - - conn = elevate("http://test-host", "user", "pass", "TestGraph") - - # Empty token treated as not set — password auth - self.assertIsInstance(conn, TigerGraphConnection) - self.assertEqual(conn.username, "user") - - def test_no_api_token_key_falls_through(self): - cfg = {**MOCK_DB_CONFIG_BASE} - elevate = self._import_and_patch(cfg) - - conn = elevate("http://test-host", "user", "pass", "TestGraph") +class TestElevateDbConnection(unittest.TestCase): + def setUp(self): + patcher = patch.dict(connections.db_config, DB_CONFIG, clear=True) + patcher.start() + self.addCleanup(patcher.stop) + def test_user_connection_uses_the_user_credentials(self): + conn = connections.elevate_db_connection_to_token("http://test-host", "alice", "pw", "TestGraph") self.assertIsInstance(conn, TigerGraphConnection) - self.assertEqual(conn.username, "user") + self.assertEqual((conn.username, conn.password), ("alice", "pw")) + self.assertEqual((conn.host, conn.graphname), ("http://test-host", "TestGraph")) + + def test_service_token_is_not_used_for_a_user(self): + sync_conn = connections.elevate_db_connection_to_token("http://test-host", "alice", "pw", "G") + async_conn = connections.elevate_db_connection_to_token("http://test-host", "alice", "pw", "G", + async_conn=True) + self.assertIsInstance(async_conn, AsyncTigerGraphConnection) + for conn in (sync_conn, async_conn): + self.assertNotEqual(conn.apiToken, "service_token") + self.assertEqual(conn.username, "alice") + + def test_no_token_is_requested_up_front(self): + """pyTigerGraph decides when a REST++ token is needed.""" + with patch.object(TigerGraphConnection, "getToken") as get_token: + connections.elevate_db_connection_to_token("http://test-host", "alice", "pw", "G") + get_token.assert_not_called() + + def test_ports_come_from_the_config(self): + connections.db_config.update({"restppPort": "19000", "gsPort": "24240"}) + conn = connections.elevate_db_connection_to_token("http://test-host", "alice", "pw", "G") + self.assertIn("19000", str(conn.restppPort)) + self.assertIn("24240", str(conn.gsPort)) if __name__ == "__main__": diff --git a/graphrag/tests/test_llm_entity_relationship_extractor.py b/graphrag/tests/test_llm_entity_relationship_extractor.py index 12f79a1..6970992 100644 --- a/graphrag/tests/test_llm_entity_relationship_extractor.py +++ b/graphrag/tests/test_llm_entity_relationship_extractor.py @@ -12,9 +12,9 @@ # You should have received a copy of the GNU Affero General Public License # along with this program. If not, see . -"""Unit tests for the schema-aware extension of -``LLMEntityRelationshipExtractor`` — definitions are surfaced in the -prompt without breaking the legacy ``allowed_*_types`` arguments. +"""Unit tests for the schema-aware ``LLMEntityRelationshipExtractor``: +the ``allowed_schema`` it is given (types, definitions, rendered schema) +reaches the extraction prompt. We do not invoke the LLM. The tests exercise prompt assembly in ``adocument_er_extraction`` / ``document_er_extraction`` by capturing @@ -23,10 +23,40 @@ from __future__ import annotations +from common.db.schema_utils import AllowedSchema from common.extractors.LLMEntityRelationshipExtractor import ( LLMEntityRelationshipExtractor, ) +_SCHEMA_REP = ( + "Vertex Types:\n" + "Company: A corporate entity.\n" + "Fund: An investment vehicle pooling capital.\n" + "Edge Types:\n" + "PUBLISHES (FROM Company, TO Report): A company publishes a report." +) + + +def _schema(**kw): + base = dict( + schema_rep=_SCHEMA_REP, + vertex_types=["Company", "Fund"], + edge_types=["PUBLISHES"], + vertex_definitions={"Company": "A corporate entity.", + "Fund": "An investment vehicle pooling capital."}, + edge_definitions={"PUBLISHES": "A company publishes a report."}, + ) + base.update(kw) + return AllowedSchema(**base) + + +def _rendered(monkeypatch, **kw): + captured = _capture_prompt_messages(monkeypatch) + ext = LLMEntityRelationshipExtractor(_StubLLMService(), **kw) + monkeypatch.setattr(ext, "_extract_kg_from_doc", lambda *a, **kw: []) + ext.document_er_extraction("hello") + return "\n".join(m[1] for m in captured) + class _StubLLMService: entity_relationship_extraction_prompt = "(system prompt)" @@ -51,7 +81,7 @@ def _from_messages(msgs): return _FakeTemplate(msgs) monkeypatch.setattr( - "langchain.prompts.ChatPromptTemplate.from_messages", _from_messages + "langchain_core.prompts.ChatPromptTemplate.from_messages", _from_messages ) return captured @@ -74,76 +104,28 @@ def test_extractor_without_definitions_leaves_prompt_unchanged(monkeypatch): assert "(system prompt)" in rendered -def test_extractor_with_entity_definitions_appends_block(monkeypatch): - captured = _capture_prompt_messages(monkeypatch) - ext = LLMEntityRelationshipExtractor( - _StubLLMService(), - entity_type_definitions={ - "Company": "A corporate entity.", - "Fund": "An investment vehicle pooling capital.", - }, - ) - monkeypatch.setattr(ext, "_extract_kg_from_doc", lambda *a, **kw: []) - - ext.document_er_extraction("hello") - - rendered = "\n".join(m[1] for m in captured) - assert "Schema entity types with definitions" in rendered - assert "- Company: A corporate entity." in rendered - assert "- Fund: An investment vehicle pooling capital." in rendered - # Sorted ordering: Company (C) before Fund (F) - assert rendered.index("Company") < rendered.index("Fund") - # Disambiguation tip is present +def test_extractor_puts_the_schema_with_definitions_in_the_prompt(monkeypatch): + rendered = _rendered(monkeypatch, allowed_schema=_schema()) + assert "Company: A corporate entity." in rendered + assert "Fund: An investment vehicle pooling capital." in rendered + assert "PUBLISHES (FROM Company, TO Report): A company publishes a report." in rendered + # The definitions are used to choose between similar types. assert "disambiguate" in rendered.lower() -def test_extractor_renders_relationship_definitions(monkeypatch): - captured = _capture_prompt_messages(monkeypatch) - ext = LLMEntityRelationshipExtractor( - _StubLLMService(), - relationship_type_definitions={ - "PUBLISHES": "A company publishes a report.", - }, - ) - monkeypatch.setattr(ext, "_extract_kg_from_doc", lambda *a, **kw: []) - - ext.document_er_extraction("hello") - - rendered = "\n".join(m[1] for m in captured) - assert "Schema relationship types with definitions" in rendered - assert "- PUBLISHES: A company publishes a report." in rendered - - -def test_extractor_skips_empty_definitions(monkeypatch): - captured = _capture_prompt_messages(monkeypatch) - ext = LLMEntityRelationshipExtractor( - _StubLLMService(), - entity_type_definitions={"Company": "", "Fund": ""}, - ) - monkeypatch.setattr(ext, "_extract_kg_from_doc", lambda *a, **kw: []) - - ext.document_er_extraction("hello") - - rendered = "\n".join(m[1] for m in captured) - # No definitions block when every value is empty. - assert "Schema entity types with definitions" not in rendered +def test_extractor_lists_the_allowed_types(monkeypatch): + rendered = _rendered(monkeypatch, allowed_schema=_schema()) + assert "Allowed Node Types: ['Company', 'Fund']" in rendered + assert "Allowed Edge Types: ['PUBLISHES']" in rendered -def test_extractor_combines_allowed_types_and_definitions(monkeypatch): - """Legacy ``allowed_entity_types`` still renders, AND definitions - block also renders below it. They are not mutually exclusive. - """ - captured = _capture_prompt_messages(monkeypatch) - ext = LLMEntityRelationshipExtractor( - _StubLLMService(), - allowed_entity_types=["Company", "Fund"], - entity_type_definitions={"Company": "A corp."}, - ) - monkeypatch.setattr(ext, "_extract_kg_from_doc", lambda *a, **kw: []) +def test_extractor_without_schema_text_adds_no_schema_block(monkeypatch): + rendered = _rendered(monkeypatch, allowed_schema=_schema(schema_rep="")) + assert "disambiguate" not in rendered.lower() + assert "STRICT SCHEMA MODE" not in rendered - ext.document_er_extraction("hello") - rendered = "\n".join(m[1] for m in captured) - assert "Allowed Node Types" in rendered - assert "Schema entity types with definitions" in rendered - assert "- Company: A corp." in rendered +def test_strict_mode_keeps_extraction_inside_the_schema(monkeypatch): + rendered = _rendered(monkeypatch, allowed_schema=_schema(), strict_mode=True) + assert "STRICT SCHEMA MODE" in rendered + assert "Company: A corporate entity." in rendered diff --git a/graphrag/tests/test_prompt_split.py b/graphrag/tests/test_prompt_split.py index 8bac4cb..2836d6f 100644 --- a/graphrag/tests/test_prompt_split.py +++ b/graphrag/tests/test_prompt_split.py @@ -97,6 +97,10 @@ def test_repair_json_escapes(): assert m._repair_json_escapes(r'{"s":"a\nb\t\\x"}') == r'{"s":"a\nb\t\\x"}' -def test_get_user_portion_default_empty_when_no_file(): +def test_get_user_portion_default_when_no_file(): + """Without an override file, the editor shows the built-in default.""" + from common.utils.prompt_validation import sanitize_user_portion m = _model() - assert m.get_user_portion("chatbot_response.txt") == "" + expected = sanitize_user_portion(m._CHATBOT_RESPONSE_USER_DEFAULT).strip() + assert expected # the answer prompt ships with default guidance + assert m.get_user_portion("chatbot_response.txt") == expected diff --git a/graphrag/tests/test_service.py b/graphrag/tests/test_service.py index 163f3d2..177fc92 100644 --- a/graphrag/tests/test_service.py +++ b/graphrag/tests/test_service.py @@ -9,7 +9,17 @@ import time from pygit2 import Repository, Commit -branch_name = os.getenv("PR_NUMBER", Repository(".").head.shorthand) +def _git_branch() -> str: + """The current git branch, or "unknown" outside a git checkout (e.g. when + the unit suite runs from another directory). Only the LLM regression runs + report it.""" + try: + return Repository(".").head.shorthand + except Exception: + return "unknown" + + +branch_name = os.getenv("PR_NUMBER") or _git_branch() EPS = 0.001 diff --git a/graphrag/tests/test_ui_setup_unit.py b/graphrag/tests/test_ui_setup_unit.py index 14d1c35..8205873 100644 --- a/graphrag/tests/test_ui_setup_unit.py +++ b/graphrag/tests/test_ui_setup_unit.py @@ -27,7 +27,7 @@ # changes its import style, this alias must be updated accordingly. sys.modules.setdefault("app.routers.ui", sys.modules["routers.ui"]) -from app.routers.ui import _resolve_llm_config_access, _require_prompt_access +from app.routers.ui import MASKED_SECRET, _resolve_llm_config_access, _require_prompt_access def _creds(username: str = "testuser", password: str = "testpass") -> HTTPBasicCredentials: @@ -68,7 +68,7 @@ def setUp(self): # Test 1 – _resolve_llm_config_access returns "full" for superuser # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {})) + @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {}, "testuser")) def test_resolve_llm_access_superuser_returns_full(self, _mock): result = _resolve_llm_config_access(_creds(), graphname=None) self.assertEqual(result, "full") @@ -77,7 +77,7 @@ def test_resolve_llm_access_superuser_returns_full(self, _mock): # Test 2 – _resolve_llm_config_access returns "chatbot_only" for graph admin # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=([], {"mygraph": ["admin"]})) + @patch("app.routers.ui._get_user_role_details", return_value=([], {"mygraph": ["admin"]}, "testuser")) def test_resolve_llm_access_graph_admin_returns_chatbot_only(self, _mock): result = _resolve_llm_config_access(_creds(), graphname="mygraph") self.assertEqual(result, "chatbot_only") @@ -86,7 +86,7 @@ def test_resolve_llm_access_graph_admin_returns_chatbot_only(self, _mock): # Test 3 – _resolve_llm_config_access raises 403 for globalobserver # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=(["globalobserver"], {})) + @patch("app.routers.ui._get_user_role_details", return_value=(["globalobserver"], {}, "testuser")) def test_resolve_llm_access_globalobserver_raises_403(self, _mock): with self.assertRaises(HTTPException) as ctx: _resolve_llm_config_access(_creds(), graphname=None) @@ -96,7 +96,7 @@ def test_resolve_llm_access_globalobserver_raises_403(self, _mock): # Test 4 – GET /ui/config does NOT return db password # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {})) + @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {}, "testuser")) @patch( "app.routers.ui.db_config", {"hostname": "http://test-db", "username": "tigergraph", "password": "super-secret"}, @@ -107,13 +107,15 @@ def test_get_config_strips_db_password(self, _mock): response = self.client.get("/ui/config", auth=("testuser", "testpass")) self.assertEqual(response.status_code, 200) db = response.json().get("db_config", {}) - self.assertNotIn("password", db, "'password' must not be returned in db_config") + # The password is masked, never returned. + self.assertNotIn("super-secret", response.text) + self.assertEqual(db.get("password"), MASKED_SECRET) # ========================================================================= # Test 5 – GET /ui/config does NOT return LLM API keys # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {})) + @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {}, "testuser")) @patch( "app.routers.ui.llm_config", { @@ -129,19 +131,21 @@ def test_get_config_strips_llm_api_keys(self, _mock): response = self.client.get("/ui/config", auth=("testuser", "testpass")) self.assertEqual(response.status_code, 200) llm = response.json().get("llm_config", {}) - self.assertNotIn("authentication_configuration", llm) + # API keys are masked, never returned. + for secret in ("sk-top", "sk-c", "sk-e", "sk-m"): + self.assertNotIn(secret, response.text) + self.assertEqual(llm["authentication_configuration"], {"OPENAI_API_KEY": MASKED_SECRET}) for svc in ("completion_service", "embedding_service", "multimodal_service"): - self.assertNotIn( - "authentication_configuration", - llm.get(svc, {}), - f"authentication_configuration must not be returned in {svc}", + self.assertEqual( + llm[svc]["authentication_configuration"], {"OPENAI_API_KEY": MASKED_SECRET}, + f"API key must be masked in {svc}", ) # ========================================================================= # Test 8 – GET /ui/config does NOT return chat_service API keys # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {})) + @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {}, "testuser")) @patch( "app.routers.ui.llm_config", { @@ -157,17 +161,15 @@ def test_get_config_strips_chat_service_api_keys(self, _mock): self.assertEqual(response.status_code, 200) llm = response.json().get("llm_config", {}) chat_svc = llm.get("chat_service", {}) - self.assertNotIn( - "authentication_configuration", - chat_svc, - "authentication_configuration must not be returned in chat_service", - ) + # The chat service key is masked, never returned. + self.assertNotIn("gsk-secret", response.text) + self.assertEqual(chat_svc["authentication_configuration"], {"GROQ_API_KEY": MASKED_SECRET}) # ========================================================================= # Test 9 – GET /ui/config chatbot_only returns global_chat_info # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=([], {"mygraph": ["admin"]})) + @patch("app.routers.ui._get_user_role_details", return_value=([], {"mygraph": ["admin"]}, "testuser")) @patch( "app.routers.ui.llm_config", { @@ -191,7 +193,7 @@ def test_get_config_chatbot_only_returns_global_chat_info(self, _mock): # Test 6 – graph admin editing entity_extraction prompt raises 403 # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=([], {"mygraph": ["admin"]})) + @patch("app.routers.ui._get_user_role_details", return_value=([], {"mygraph": ["admin"]}, "testuser")) def test_graph_admin_entity_extraction_prompt_raises_403(self, _mock): response = self.client.post( "/ui/prompts", @@ -209,7 +211,7 @@ def test_graph_admin_entity_extraction_prompt_raises_403(self, _mock): # Test 7 – concurrent LLM save returns 409 when lock is held # ========================================================================= - @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {})) + @patch("app.routers.ui._get_user_role_details", return_value=(["superuser"], {}, "testuser")) @patch("app.routers.ui.auth", return_value=([], MagicMock())) @patch("app.routers.ui._ecc_jobs_running", return_value=False) def test_concurrent_llm_save_returns_409(self, _mock_ecc, _mock_auth, _mock_roles): From 84411dcc31187380136f528cfe46e936c9bc4b70 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 13:14:08 -0700 Subject: [PATCH 12/32] Update the 2.0.4 CHANGELOG and README - README describes schema-based planning and what Trace Logs record - CHANGELOG entries describe behavior without diagnostic detail --- CHANGELOG.md | 6 +++--- README.md | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 89f770a..78480c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,14 +9,14 @@ ### Fixed - **The agentic planner plans from the graph schema.** Since 2.0.0 the planner made its plan without knowing the graph's vertex and edge types, so it could only guess from the wording of a question whether structural retrieval would help. It now gets a compact summary of the schema — type names, attribute names, edge endpoints and descriptions — cached per graph and refreshed when the schema changes. -- **Structured queries no longer start with a step that cannot succeed.** Each structured query first asked the model for a pyTigerGraph function call, but no functions are registered with the TigerGraph embedding store, so every call was rejected (or failed with `'NoneType' object is not iterable` when the question named no relationships) before the openCypher query ran. GraphRAG now detects that no functions are registered and goes straight to the openCypher query, saving one model call per structured query. -- **An empty structured query result no longer counts as data.** A query that returned no rows (for example `{"VS_XXX": []}`, or an empty list as text) was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. The classic engine uses the same check, so both engines treat an empty result the same way. +- **Structured queries no longer start with a step that cannot succeed.** Each structured query first asked the model for a pyTigerGraph function call, but no functions are registered with the TigerGraph embedding store, so every call was rejected, or failed outright when the question named no relationships, before the openCypher query ran. GraphRAG now detects that no functions are registered and goes straight to the openCypher query, saving one model call per structured query. +- **An empty structured query result no longer counts as data.** A query that returned only empty row sets was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. The classic engine uses the same check, so both engines treat an empty result the same way. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. - **Agentic planning no longer gets stuck writing its plan summary.** The planner's summary and step reasons had no stated purpose or length, and the planner's output had no limit, so a model could repeat a phrase until it reached its own output maximum — costing thousands of tokens and tens of seconds, and sometimes the plan itself. These fields now ask for one short sentence, an overlong value is shortened instead of failing the plan, and the planner's output is capped. ### Removed -- **The `tg_run_query` and `tg_get_neighbors` agent tools.** They let the agent write and run its own GSQL and expand neighbors through an in-process copy of tigergraph-mcp's functions. Structured retrieval already maps the question to the schema, writes the query, runs it and retries on errors, so the agent does not need to write queries itself, and GraphRAG code now reaches TigerGraph only through pyTigerGraph. Custom planner or agent prompts that name these tools should drop them. +- **The `tg_run_query` and `tg_get_neighbors` agent tools.** They let the agent write and run its own GSQL and expand neighbors through tigergraph-mcp's functions. Structured retrieval already maps the question to the schema, writes the query, runs it and retries on errors, so the agent does not need to write queries itself, and GraphRAG code now reaches TigerGraph only through pyTigerGraph. Custom planner or agent prompts that name these tools should drop them. ### Added - **`planner_max_output_tokens` setting.** The most output tokens the planner may use for one plan (default `4000`; reasoning models use at least `16000`). diff --git a/README.md b/README.md index 8629b9d..192f188 100644 --- a/README.md +++ b/README.md @@ -391,10 +391,10 @@ GraphRAG offers two chat engines, chosen from the chat menu (or set as a per-gra The agent decides its **own** retrieval instead of following a fixed pipeline: it picks which methods to use (structural graph queries, vector search, community search), can call external [MCP tools](#mcp-servers-agentic-tools), answers greetings and questions about itself directly, reuses what the conversation already answers and retrieves only what is missing, and cites the chunks and queries it used. It comes in two styles: -- **Planned** *(default)* — analyzes the question up front and lays out the whole retrieval plan (which methods, how many, in what order) as a small DAG, executes it, then synthesizes one grounded answer. Predictable and efficient; a strong fit for most questions, including multi-part ones. +- **Planned** *(default)* — analyzes the question against the graph's schema up front and lays out the whole retrieval plan (which methods, how many, in what order) as a small DAG, executes it, then synthesizes one grounded answer. Predictable and efficient; a strong fit for most questions, including multi-part ones. - **Reactive** — reasons and acts one step at a time, deciding each next retrieval from what the previous step returned, and keeps going until it can answer completely and accurately. More adaptive to what it finds, at the cost of more steps (and tokens) on complex questions. -Both styles send trivial/conversational messages straight to a direct answer, and record their retrieval in the admin **Trace** — the plan/steps and reasoning, which chunks were retrieved, and which the agent selected for the answer. +Both styles send trivial/conversational messages straight to a direct answer, and record their retrieval in the admin **Trace** — the plan/steps and reasoning, the structured queries that ran, which chunks were retrieved, which the agent selected for the answer, and the error when an answer fails. **Selecting an engine.** In the chat UI, use the engine/style picker (Classic, or Agentic → Planned / Reactive). To set a per-graph default, configure `agent_style` (`auto` follows the configured default; `planned` or `reactive` force a style) — see [GraphRAG configuration](#graphrag-configuration). Depth is bounded by `agent_max_iterations` (Reactive) and `agent_max_replans` / `agent_max_total_steps` (Planned). From 5465a9c63d6be5a4ec2295868868ecfd2a889215 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 14:00:45 -0700 Subject: [PATCH 13/32] Prepare 2.0.4 release - Bump the version to 2.0.4 - Add the 2.0.4 entry to the README release list --- README.md | 1 + VERSION | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 192f188..a1b2c84 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,7 @@ --- ## Releases +* **10/15/2026**: GraphRAG v2.0.4 released. Agentic answers reuse what the conversation already holds and plan from the graph's schema, Trace Logs show the structured queries that ran and the error of a failed answer, external MCP servers load again, and fixed planning and structured-query issues. See [Release Notes](https://github.com/tigergraph/graphrag/releases/tag/v2.0.4) for details. * **10/12/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. diff --git a/VERSION b/VERSION index 50ffc5a..2165f8f 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -2.0.3 +2.0.4 From 841140a6c36c5276dc4f3c21ddba26ddb333810a Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 15:04:52 -0700 Subject: [PATCH 14/32] Cut runaway character repeats in the plan summary - A plan summary or step reason that repeats a few characters in place is cut where the repeat starts Refs: GML-2324 --- common/py_schemas/schemas.py | 17 +++++++++++++++-- graphrag/tests/test_planner_repetition_loop.py | 15 +++++++++++++++ 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/common/py_schemas/schemas.py b/common/py_schemas/schemas.py index 590190d..3c48411 100644 --- a/common/py_schemas/schemas.py +++ b/common/py_schemas/schemas.py @@ -14,6 +14,7 @@ import enum import logging +import re from typing import Dict, List, Optional, Union from pydantic import BaseModel, Field, PrivateAttr, field_validator @@ -83,9 +84,21 @@ class GraphRAGResponse(BaseModel): _PLAN_TEXT_MAX = 200 +# A unit of 1-8 characters repeated 12+ times in a row, e.g. "2026.0000…". +_CHAR_RUN = re.compile(r"(.{1,8}?)\1{11,}", re.S) + + def _collapse_repeats(text: str, min_words: int = 4, max_words: int = 60) -> str: - """Cut ``text`` at the first phrase of ``min_words``–``max_words`` words - that repeats three times in a row, keeping one copy of it.""" + """Cut ``text`` at the first runaway repetition. + + Two shapes are caught: a short run of characters repeated in place + ("2026.0000…"), cut where the run starts; and a phrase of + ``min_words``–``max_words`` words that repeats three times in a row, + keeping one copy of it. + """ + run = _CHAR_RUN.search(text) + if run: + text = text[:run.start()].rstrip(" .,;:-") words = text.split() for i in range(len(words)): for n in range(min_words, max_words + 1): diff --git a/graphrag/tests/test_planner_repetition_loop.py b/graphrag/tests/test_planner_repetition_loop.py index 5e5f668..b4b960f 100644 --- a/graphrag/tests/test_planner_repetition_loop.py +++ b/graphrag/tests/test_planner_repetition_loop.py @@ -44,6 +44,21 @@ def test_collapse_keeps_one_copy_of_the_cycle(self): out = _collapse_repeats(_LOOP) self.assertEqual(out.count("lifecycle flow initialization"), 1) + def test_character_run_is_cut_where_it_starts(self): + """A single character repeated in place (seen live on a Jira graph).""" + value = ("Query the structured graph to count Jira issues created on or after " + "August 2026." + "0" * 3000) + plan = Plan(strategy=value, steps=[PlanStep(id="A", kind="answer", tool="")]) + self.assertEqual(plan.strategy, + "Query the structured graph to count Jira issues created on or after August 2026") + + def test_repeated_short_unit_is_cut(self): + self.assertEqual(_collapse_repeats("Count the tickets " + "ab" * 50), "Count the tickets") + + def test_numbers_in_normal_text_are_kept(self): + text = "Count tickets with ids 1000 to 1000000, then answer." + self.assertEqual(_collapse_repeats(text), text) + def test_collapse_leaves_normal_text_alone(self): text = ("Search the documents for the policy, then search the documents " "for the exceptions, then answer from the documents.") From 24cbf2b4c789b825452b226e32898ae762217724 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 15:15:28 -0700 Subject: [PATCH 15/32] Plan with Gemini 3 at its recommended temperature - Gemini 3 planning, conversation check and triage use temperature 1.0 when the configured value is lower - Answers, knowledge graph building and other models keep the configured temperature Refs: GML-2335 --- CHANGELOG.md | 1 + README.md | 2 + common/llm_services/base_llm.py | 30 +++++++++ common/llm_services/capabilities.py | 15 +++++ .../test_gemini_structured_temperature.py | 63 +++++++++++++++++++ 5 files changed, 111 insertions(+) create mode 100644 graphrag/tests/test_gemini_structured_temperature.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 78480c7..77925fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ - **An empty structured query result no longer counts as data.** A query that returned only empty row sets was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. The classic engine uses the same check, so both engines treat an empty result the same way. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. +- **Gemini 3 planning no longer loops at temperature 0.** At the temperature 0 that GraphRAG's configurations set, Gemini 3 models sometimes kept writing zeros where the plan should end, until the planner's output cap, which slowed planning. Google recommends temperature 1.0 for these models, so the agentic engine's planning calls now use 1.0 for Gemini 3 when the configured value is lower. Answers, knowledge graph building and all other models keep the configured temperature. - **Agentic planning no longer gets stuck writing its plan summary.** The planner's summary and step reasons had no stated purpose or length, and the planner's output had no limit, so a model could repeat a phrase until it reached its own output maximum — costing thousands of tokens and tens of seconds, and sometimes the plan itself. These fields now ask for one short sentence, an overlong value is shortened instead of failing the plan, and the planner's output is capped. ### Removed diff --git a/README.md b/README.md index a1b2c84..59db923 100644 --- a/README.md +++ b/README.md @@ -697,6 +697,8 @@ In the `llm_config` section of `configs/server_config.json` file, copy JSON conf | `base_url` | string | No | — | Custom API endpoint URL. | | `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. | +> **Gemini 3 models:** the agentic engine plans with temperature `1.0` when `model_kwargs` sets a lower one, as Google recommends for these models. Answers and knowledge graph building use the configured temperature. + **`embedding_service` parameters:** | Parameter | Type | Required | Default | Description | diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index 93d1c3a..b5b3cf3 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -663,6 +663,35 @@ def with_output_limit(self, max_output_tokens: int): return llm.model_copy(update={name: limit}) return llm + # Models already logged as getting the structured-output temperature. + _structured_temperature_logged: set = set() + + def with_structured_temperature(self, llm): + """``llm`` at the temperature structured output needs. + + Gemini 3 models can loop below temperature 1.0, so their structured + calls (planner, conversation check, triage) use 1.0 when the + configured temperature is lower. A per-call copy: answer generation + and ECC keep the configured temperature, and other models are + returned unchanged. + """ + from common.llm_services.capabilities import ( + GEMINI_3_TEMPERATURE, gemini_loops_below_default_temperature, + ) + model = ((getattr(self, "config", None) or {}).get("llm_model") or "") + if not gemini_loops_below_default_temperature(model) or not hasattr(llm, "model_copy"): + return llm + current = getattr(llm, "temperature", None) + if current is not None and current >= GEMINI_3_TEMPERATURE: + return llm + if model not in LLM_Model._structured_temperature_logged: + LLM_Model._structured_temperature_logged.add(model) + logger.info( + f"{model}: structured output uses temperature {GEMINI_3_TEMPERATURE} " + f"instead of {current} (Gemini 3 can loop below it)" + ) + return llm.model_copy(update={"temperature": GEMINI_3_TEMPERATURE}) + def invoke_structured( self, messages: list, @@ -682,6 +711,7 @@ def invoke_structured( this call (see ``with_output_limit``). """ llm = self.with_output_limit(max_output_tokens) if max_output_tokens else self.llm + llm = LLM_Model.with_structured_temperature(self, llm) usage_data = {} with get_openai_callback() as cb: try: diff --git a/common/llm_services/capabilities.py b/common/llm_services/capabilities.py index 36140c1..fdd9c59 100644 --- a/common/llm_services/capabilities.py +++ b/common/llm_services/capabilities.py @@ -93,6 +93,21 @@ def _openai_thinking(model: str) -> bool: ) +# Gemini 3 models are tuned to sample at temperature 1.0. Below it they can +# loop, e.g. write a run of zeros where a JSON string should close (GML-2335). +GEMINI_3_TEMPERATURE = 1.0 + + +def gemini_loops_below_default_temperature(model: str) -> bool: + """Gemini 3 models, which Google says can loop below temperature 1.0 + ("Gemini 3 developer guide"). Case-insensitive; accepts a ``models/`` + prefix.""" + m = (model or "").strip().lower() + if m.startswith("models/"): + m = m[len("models/"):] + return m.startswith("gemini-3") + + def openai_rejects_temperature(model: str) -> bool: """OpenAI o-series reasoning models (o1/o3/o4) reject a custom ``temperature`` — only the default value is accepted, and sending the diff --git a/graphrag/tests/test_gemini_structured_temperature.py b/graphrag/tests/test_gemini_structured_temperature.py new file mode 100644 index 0000000..409df88 --- /dev/null +++ b/graphrag/tests/test_gemini_structured_temperature.py @@ -0,0 +1,63 @@ +"""Gemini 3 structured calls use temperature 1.0 (GML-2335).""" +import types +import unittest + +try: + from common.llm_services.base_llm import LLM_Model + from common.llm_services.capabilities import gemini_loops_below_default_temperature + from common.py_schemas import Plan +except ImportError as exc: # needs the app environment + raise unittest.SkipTest(f"app environment unavailable: {exc}") + + +class _FakeChat: + """A chat model with a temperature, as LangChain chat models have.""" + + def __init__(self, temperature): + self.temperature = temperature + self.used = None + + def model_copy(self, update): + copy = _FakeChat(update.get("temperature", self.temperature)) + return copy + + def with_structured_output(self, schema, **kw): + outer = self + return types.SimpleNamespace(invoke=lambda m: (setattr(_FakeChat, "last", outer), Plan())[1]) + + +def _svc(model, temperature=0): + svc = object.__new__(LLM_Model) + svc.llm, svc.config = _FakeChat(temperature), {"llm_model": model} + return svc + + +class TestRule(unittest.TestCase): + def test_gemini_3_models_match(self): + for m in ("gemini-3.5-flash", "gemini-3-pro", "models/gemini-3.5-flash", "Gemini-3.1-Pro"): + self.assertTrue(gemini_loops_below_default_temperature(m), m) + + def test_other_models_do_not(self): + for m in ("gemini-2.5-flash", "gpt-5.4-mini", "o3", "claude-sonnet-4", "llama3.1", "", None): + self.assertFalse(gemini_loops_below_default_temperature(m), m) + + +class TestStructuredCalls(unittest.TestCase): + def test_gemini_3_structured_call_uses_1(self): + svc = _svc("gemini-3.5-flash", 0) + svc.invoke_structured([("user", "q")], Plan, caller_name="agentic_plan") + self.assertEqual(_FakeChat.last.temperature, 1.0) + self.assertEqual(svc.llm.temperature, 0) # the service model (answers, ECC) is unchanged + + def test_higher_temperature_is_kept(self): + svc = _svc("gemini-3.5-flash", 1.2) + self.assertIs(svc.with_structured_temperature(svc.llm), svc.llm) + + def test_other_models_keep_their_temperature(self): + for model in ("gpt-5.4-mini", "gemini-2.5-flash", "claude-sonnet-4"): + svc = _svc(model, 0) + self.assertIs(svc.with_structured_temperature(svc.llm), svc.llm, model) + + +if __name__ == "__main__": + unittest.main() From 34cbaa4f63bff8260cbad6eb79bf8cb5a7c7c3c8 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 15:23:46 -0700 Subject: [PATCH 16/32] Treat a bad earlier answer as not answered - An earlier answer that says the information was not found, hedges, or is incomplete is retrieved again - Trace Logs list dropped plan steps and show the fallback retrieval - The fallback retrieval queries the graph as well as the documents Refs: GML-2325 --- CHANGELOG.md | 2 +- common/llm_services/base_llm.py | 17 ++++++---- common/py_schemas/schemas.py | 3 ++ graphrag/app/agent/agentic_graph.py | 28 ++++++++++++---- graphrag/app/agent/agentic_planner.py | 3 +- .../test_agentic_conversation_routing.py | 33 ++++++++++++++++--- 6 files changed, 65 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 77925fe..4ab75a1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,7 @@ ## [2.0.4] ### Changed -- **Agentic answers reuse what the conversation already holds.** Before, the agentic engines retrieved again for a question the conversation had already answered, and the planned engine's answer step did not see the conversation at all. The planner and the reactive agent now first decide how much of the answer the conversation holds: all of it is answered without retrieval, part of it gets retrieval only for the missing parts, and questions about the current or latest state still retrieve. The answer step reads the recent conversation together with the retrieved context, and the retrieved context wins where they disagree. In planned mode, a plan that skips retrieval gets one more check, so a question the conversation does not fully answer is still looked up. Trace Logs show the check and mark a plan that answers without retrieval. +- **Agentic answers reuse what the conversation already holds.** Before, the agentic engines retrieved again for a question the conversation had already answered, and the planned engine's answer step did not see the conversation at all. The planner and the reactive agent now first decide how much of the answer the conversation holds: all of it is answered without retrieval, part of it gets retrieval only for the missing parts, and questions about the current or latest state still retrieve. An earlier answer that says the information was not found, hedges, or is incomplete counts as not answered, so asking again retrieves again instead of repeating it. The answer step reads the recent conversation together with the retrieved context, and the retrieved context wins where they disagree. In planned mode, a plan that skips retrieval gets one more check, so a question the conversation does not fully answer is still looked up. Trace Logs show the check, mark a plan with no retrieval steps, list plan steps that named an unknown tool, and show the fallback retrieval when the planner still plans none. - **Failed answers show their error in Trace Logs.** When GraphRAG cannot answer, the trace log ends with an Error step that states the error, even when no other step ran. The error appears only in Trace Logs, which only the asking superuser can open; the chat message stays generic and now asks to rephrase the question instead of the prompt. - **Trace Logs show the structured query that ran.** Each structured query step lists the query at the top of its output, followed by any attempts that returned nothing or failed, with their errors. Before, the query was buried in the result, and it was missing entirely when the query returned no rows — the case where it is most needed. diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index b5b3cf3..16c0f90 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -1136,12 +1136,13 @@ def select_retriever_prompt(self): - If not — a part is still unanswered, a needed value or table is missing, or the results were thin — take another action to close the gap (follow a lead, widen top_k / num_hops, or switch method). Do not settle for a partial or vague answer when more retrieval could complete it. Do not commit to a full multi-step plan up front; let each next step be driven by what is still missing for a complete answer. -Before your first tool call, list the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. Then: -- Every fact is stated (for example, the same question was already asked and answered): give the final answer without calling a tool. +Before your first tool call, list the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. First judge each earlier answer that you would use. A bad answer counts as not answered, and a fact in a bad answer is not stated. An answer is bad when it says that the information was not found, not available, not explicit or not in the context; when it hedges or guesses (for example "but", "may", "it seems", "it is unclear"); when it gives only part of what was asked, or an approximate value; or when it answers a different question. Then: +- Every fact is stated (for example, the same question was already asked and answered well): give the final answer without calling a tool. - Some facts are stated: call tools only for the missing facts. Do not retrieve again what ## Conversation already states. Combine both in the final answer. - No fact is stated: retrieve as described below. This decision comes before the "Retrieval Strategy" below. Do not use an earlier answer when the question asks for the current or latest state (for example "now", "current", "latest", "again"), or when the user says that an earlier answer was wrong or incomplete. +When an earlier answer to the same question was bad, retrieve again with a different or wider approach than a single search (for example both structural and unstructured retrieval, or a higher top_k). The graph schema is required for the structural and unstructured query tools: before your first structural query or vector/unstructured retrieval, call graphrag__get_schema once to load the graph's vertex and edge types. Questions answered without graph data (e.g. by an external tool) do not need the schema. @@ -1190,12 +1191,13 @@ def agentic_agent_prompt(self): ## Graph schema lists the vertex types (with their attribute names) and the edge types of the graph. Use it to decide whether the question needs STRUCTURAL retrieval: plan a structural step when the question asks for data that these types and attributes hold (counts, lookups, relationships, aggregations). The query tools load the full schema themselves, so do not plan a schema step. A question that needs no graph data should not include any graph-retrieval step (plan only the final answer step, or the relevant non-graph tool). -Before you plan retrieval, list the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. Then: -- Every fact is stated (for example, the same question was already asked and answered): plan only the final answer step. Do not retrieve those facts again. +Before you plan retrieval, list the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. First judge each earlier answer that you would use. A bad answer counts as not answered, and a fact in a bad answer is not stated. An answer is bad when it says that the information was not found, not available, not explicit or not in the context; when it hedges or guesses (for example "but", "may", "it seems", "it is unclear"); when it gives only part of what was asked, or an approximate value; or when it answers a different question. Then: +- Every fact is stated (for example, the same question was already asked and answered well): plan only the final answer step. Do not retrieve those facts again. - Some facts are stated: plan retrieval steps only for the missing facts. Do not retrieve again what ## Conversation already states. The answer step combines both. - No fact is stated: plan the retrieval steps as described below. This decision comes before the "Retrieval Strategy" below. Do not use an earlier answer when the question asks for the current or latest state (for example "now", "current", "latest", "again"), or when the user says that an earlier answer was wrong or incomplete. +When an earlier answer to the same question was bad, retrieve again with a different or wider approach than a single search (for example both structural and unstructured retrieval, or a higher top_k). You have two kinds of retrieval: - STRUCTURAL (graphrag__structural_retrieve): generates and runs a graph query. Best for counts, lookups by attribute/id, relationships, and aggregations over typed data. It depends on the LLM generating a correct query against the live schema — it can return nothing or the wrong rows when the question doesn't map cleanly to typed graph data, so it is NOT a safe sole source of context. @@ -1222,14 +1224,15 @@ def agentic_agent_prompt(self): _AGENTIC_CONVERSATION_CHECK_SYSTEM = """\ You check one decision of a question-answering agent. The agent plans to answer the question from ## Conversation only, without new retrieval. -List the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. +List the facts the question asks for (values, names, dates). For each fact, check whether an earlier answer in ## Conversation states it explicitly. A related topic or a similar question is not enough. First judge each earlier answer that the agent would use. A bad answer counts as not answered, and a fact in a bad answer is not stated. An answer is bad when it says that the information was not found, not available, not explicit or not in the context; when it hedges or guesses (for example "but", "may", "it seems", "it is unclear"); when it gives only part of what was asked, or an approximate value; or when it answers a different question. Set "needs_retrieval" to true when one of these is correct: - ## Conversation does not state a fact that the question asks for. +- The earlier answer that the agent would use is bad. - The question asks for the current or latest state (for example "now", "current", "latest", "again"). - The user says that an earlier answer was wrong or incomplete. -Set "needs_retrieval" to false when ## Conversation states every fact, or when the question needs no data (for example a greeting or thanks). +Set "needs_retrieval" to false when a good earlier answer in ## Conversation states every fact, or when the question needs no data (for example a greeting or thanks). In "missing", name the facts that ## Conversation does not state. Return ONLY the structured result. """ @@ -1325,7 +1328,7 @@ def hyde_prompt(self): ## Rules - The contexts arrive as JSON key-context pairs. **Combine and rephrase** them to answer the question. - **Preserve** image links exactly as `![description](url)` in the final answer when used. Do NOT modify or omit them. -- The context "earlier_conversation", when present, holds earlier questions and answers from this chat. Use it to resolve references, and to answer when it holds the answer. When it conflicts with another context, the other context is correct. +- The context "earlier_conversation", when present, holds earlier questions and answers from this chat. Use it to resolve references, and to answer when it holds the answer. When it conflicts with another context, the other context is correct. An earlier answer that says information was not found is not evidence that it does not exist. ## Inputs - **Question**: {question} diff --git a/common/py_schemas/schemas.py b/common/py_schemas/schemas.py index 3c48411..8bd33a2 100644 --- a/common/py_schemas/schemas.py +++ b/common/py_schemas/schemas.py @@ -148,6 +148,9 @@ def _short_rationale(cls, v): class Plan(BaseModel): steps: List[PlanStep] = [] + # Steps the planner wrote that named an unknown tool and were dropped. + # Shown in the trace only; never part of the model's output schema. + _dropped_steps: list = PrivateAttr(default_factory=list) strategy: str = Field( "", description="One short sentence (under 25 words) on the approach. " diff --git a/graphrag/app/agent/agentic_graph.py b/graphrag/app/agent/agentic_graph.py index 19e0a78..39cdfa1 100644 --- a/graphrag/app/agent/agentic_graph.py +++ b/graphrag/app/agent/agentic_graph.py @@ -56,14 +56,19 @@ def _hybrid_fallback_step() -> PlanStep: def _search_plan(question: str) -> Plan: - """One hybrid search for the question, then the answer.""" + """A structural query and a hybrid search for the question, then the + answer. Used when the planner gives no retrieval for a question that + needs data; a structural query that finds nothing just returns no rows.""" return Plan( - strategy="Search the documents for the question, then answer.", + strategy="Query the graph and search the documents for the question, then answer.", steps=[ - PlanStep(id="S1", kind="unstructured", tool=_HYBRID_TOOL, + PlanStep(id="S1", kind="structural", tool=_STRUCTURAL_TOOL, args={"question": question}, rationale="The conversation does not hold the answer."), - PlanStep(id="A", kind="answer", tool="", depends_on=["S1"], + PlanStep(id="S2", kind="unstructured", tool=_HYBRID_TOOL, + args={"question": question}, + rationale="The conversation does not hold the answer."), + PlanStep(id="A", kind="answer", tool="", depends_on=["S1", "S2"], rationale="Answer from the retrieved context."), ], ) @@ -72,8 +77,11 @@ def _search_plan(question: str) -> Plan: def _plan_output(plan) -> dict: """A plan as shown in the trace.""" out = {"strategy": plan.strategy, "steps": [s.model_dump() for s in plan.steps]} + dropped = getattr(plan, "_dropped_steps", None) + if dropped: + out["dropped_steps"] = dropped if not _retrieves(plan): - out["routing"] = "answered without retrieval" + out["routing"] = "no retrieval steps" return out @@ -149,9 +157,15 @@ def run_agentic(ctx, llm, question, conversation=None) -> GraphRAGResponse: }) if not _retrieves(plan): # The planner skipped retrieval twice for a question the check - # says needs data: search instead of answering with nothing. - logger.info("agentic: replan still has no retrieval; using a hybrid search") + # says needs data: retrieve instead of answering with nothing. + logger.info("agentic: replan still has no retrieval; using the fallback plan") plan = _search_plan(question) + agent_steps.append({ + "node": "fallback plan", "kind": "plan", "duration_s": 0, + "input": {"reason": "The question needs data, and the plan has no retrieval steps."}, + "output": _plan_output(plan), + "usage": {}, + }) if plan.strategy: emit(plan.strategy) diff --git a/graphrag/app/agent/agentic_planner.py b/graphrag/app/agent/agentic_planner.py index a26c3f5..14a9929 100644 --- a/graphrag/app/agent/agentic_planner.py +++ b/graphrag/app/agent/agentic_planner.py @@ -165,7 +165,8 @@ def _sanitize(plan: Plan, ctx=None) -> Plan: if s.kind == "answer" or s.tool == "" or s.tool in known: steps.append(s) else: - logger.info(f"planner: dropping step with unknown tool {s.tool!r}") + logger.warning(f"planner: dropping step with unknown tool {s.tool!r}") + plan._dropped_steps.append({"id": s.id, "tool": s.tool}) if not any(s.kind == "answer" or s.tool == "" for s in steps): retrieval_ids = [s.id for s in steps] from common.py_schemas import PlanStep diff --git a/graphrag/tests/test_agentic_conversation_routing.py b/graphrag/tests/test_agentic_conversation_routing.py index c45e270..66520a6 100644 --- a/graphrag/tests/test_agentic_conversation_routing.py +++ b/graphrag/tests/test_agentic_conversation_routing.py @@ -74,7 +74,7 @@ def test_answer_only_plan_skips_retrieval_and_replanning(self): execute.assert_not_called() run_step.assert_not_called() steps = resp.query_sources["agent_steps"] - self.assertEqual(steps[0]["output"]["routing"], "answered without retrieval") + self.assertEqual(steps[0]["output"]["routing"], "no retrieval steps") self.assertIn("earlier_conversation", steps[-1]["input"]) def test_empty_retrieval_still_replans(self): @@ -86,7 +86,7 @@ def test_empty_retrieval_still_replans(self): def test_trace_marks_the_decision(self): answer_only = Plan(steps=[PlanStep(id="A", kind="answer", tool="")]) - self.assertEqual(g._plan_output(answer_only)["routing"], "answered without retrieval") + self.assertEqual(g._plan_output(answer_only)["routing"], "no retrieval steps") search = Plan(steps=[PlanStep(id="S1", kind="unstructured", tool=g._HYBRID_TOOL)]) self.assertNotIn("routing", g._plan_output(search)) @@ -120,14 +120,26 @@ def test_missing_fact_replans_with_the_gap(self): nodes = [s["node"] for s in resp.query_sources["agent_steps"]] self.assertEqual(nodes[:3], ["plan", "conversation check", "replan (conversation check)"]) - def test_second_skip_searches_anyway(self): + def test_second_skip_retrieves_anyway(self): check = ConversationCheck(missing="iPhone net sales", needs_retrieval=True) - _, _, planner, execute = self._run([_ANSWER_ONLY, _ANSWER_ONLY], check) + resp, _, planner, execute = self._run([_ANSWER_ONLY, _ANSWER_ONLY], check) self.assertEqual(planner.call_count, 2) execute.assert_called_once() plan = execute.call_args[0][0] - self.assertEqual([s.tool for s in plan.steps], [g._HYBRID_TOOL, ""]) + self.assertEqual([s.tool for s in plan.steps], [g._STRUCTURAL_TOOL, g._HYBRID_TOOL, ""]) self.assertEqual(plan.steps[0].args, {"question": "Add iPad."}) + nodes = [s["node"] for s in resp.query_sources["agent_steps"]] + self.assertEqual(nodes[:4], ["plan", "conversation check", "replan (conversation check)", "fallback plan"]) + + def test_dropped_steps_are_shown_in_the_trace(self): + from agent.agentic_planner import _sanitize + plan = Plan(steps=[PlanStep(id="S1", kind="structural", tool="structural_retrieve")]) + with mock.patch("agent.agentic_planner.registry.tool_names", return_value=[g._STRUCTURAL_TOOL]): + plan = _sanitize(plan) + out = g._plan_output(plan) + self.assertEqual(out["dropped_steps"], [{"id": "S1", "tool": "structural_retrieve"}]) + self.assertEqual(out["routing"], "no retrieval steps") + self.assertNotIn("_dropped_steps", plan.model_dump()) def test_confirmed_plan_answers_without_retrieval(self): _, _, planner, execute = self._run([_ANSWER_ONLY], ConversationCheck(needs_retrieval=False)) @@ -217,8 +229,19 @@ def test_planner_and_react_check_the_conversation_first(self): self.assertIn("A related topic or a similar question is not enough", prompt) self.assertIn('"now"', prompt) + def test_bad_earlier_answers_count_as_not_answered(self): + for prompt in (LLM_Model._AGENTIC_PLANNER_SYSTEM, LLM_Model._AGENTIC_AGENT_SYSTEM, + LLM_Model._AGENTIC_CONVERSATION_CHECK_SYSTEM): + self.assertIn("A bad answer counts as not answered", prompt) + self.assertIn("not found, not available, not explicit", prompt) + self.assertIn("The earlier answer that the agent would use is bad.", + LLM_Model._AGENTIC_CONVERSATION_CHECK_SYSTEM) + for prompt in (LLM_Model._AGENTIC_PLANNER_SYSTEM, LLM_Model._AGENTIC_AGENT_SYSTEM): + self.assertIn("retrieve again with a different or wider approach", prompt) + def test_answer_prompt_explains_the_conversation_context(self): self.assertIn('"earlier_conversation"', LLM_Model._CHATBOT_RESPONSE_SYSTEM) + self.assertIn("not evidence that it does not exist", LLM_Model._CHATBOT_RESPONSE_SYSTEM) if __name__ == "__main__": From 92bd30b4cdbf5790ffc7a40c167cba2fcbe21066 Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 15:44:09 -0700 Subject: [PATCH 17/32] Run generated openCypher through pyTigerGraph and strip code blocks correctly - A code block labelled opencypher or cypher no longer leaves the label in the query - Generated queries run in interpreted mode through pyTigerGraph, not the GSQL shell - Trace Logs show only the generated query, without the graph wrapper Refs: GML-2336 --- CHANGELOG.md | 1 + graphrag/app/agent/agent_graph.py | 13 ++++---- graphrag/app/tools/generate_cypher.py | 39 ++++++++++++++++++---- graphrag/app/tools/graphrag_tools.py | 13 +++----- graphrag/tests/test_structural_trace.py | 43 ++++++++++++++++++++++--- 5 files changed, 85 insertions(+), 24 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4ab75a1..136cdd9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ ### Fixed - **The agentic planner plans from the graph schema.** Since 2.0.0 the planner made its plan without knowing the graph's vertex and edge types, so it could only guess from the wording of a question whether structural retrieval would help. It now gets a compact summary of the schema — type names, attribute names, edge endpoints and descriptions — cached per graph and refreshed when the schema changes. - **Structured queries no longer start with a step that cannot succeed.** Each structured query first asked the model for a pyTigerGraph function call, but no functions are registered with the TigerGraph embedding store, so every call was rejected, or failed outright when the question named no relationships, before the openCypher query ran. GraphRAG now detects that no functions are registered and goes straight to the openCypher query, saving one model call per structured query. +- **Generated graph queries no longer fail when the model labels its code block.** A reply that started with an `opencypher` code block kept the word `opencypher` as the first line of the query, so TigerGraph rejected every attempt. Code blocks are now removed correctly, generated queries run through pyTigerGraph's interpreted-query call instead of the GSQL shell, and Trace Logs show only the generated query. - **An empty structured query result no longer counts as data.** A query that returned only empty row sets was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. The classic engine uses the same check, so both engines treat an empty result the same way. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py index 07ff80a..e7fcf65 100644 --- a/graphrag/app/agent/agent_graph.py +++ b/graphrag/app/agent/agent_graph.py @@ -40,6 +40,7 @@ from tools import MapQuestionToSchemaException from typing_extensions import TypedDict from common.utils.query_results import is_empty_result +from tools.generate_cypher import run_cypher from common.logs.log import req_id_cv from common.metrics.prometheus_metrics import metrics as pmetrics @@ -377,15 +378,15 @@ def generate_cypher(self, state): logger.warning(f"Cypher generation failed: {e}") gen_history.append(f"{i}: Error: {e}\n") continue - response = self.db_connection.gsql(cypher) - response_lines = response.split("\n") - json_str = "\n".join(response_lines[1:]) try: - response_json = json.loads(json_str) + response_json = {"results": run_cypher(self.db_connection, cypher)} + json_str = json.dumps(response_json) break except Exception as e: - gen_history.append(f"{i}: {cypher}\n\tError: {json_str}\n") - if response_json and not self.is_query_result_empty(response_json["results"][0]): + response = str(e) + json_str = response + gen_history.append(f"{i}: {cypher}\n\tError: {response}\n") + if response_json and response_json["results"] and not self.is_query_result_empty(response_json["results"][0]): state["context"] = { "result": response_json["results"][0], "cypher": cypher, diff --git a/graphrag/app/tools/generate_cypher.py b/graphrag/app/tools/generate_cypher.py index d9199de..f1ea84b 100644 --- a/graphrag/app/tools/generate_cypher.py +++ b/graphrag/app/tools/generate_cypher.py @@ -13,6 +13,7 @@ # along with this program. If not, see . import logging +import re from typing import Iterable from langchain_core.output_parsers import StrOutputParser from langchain_core.prompts import PromptTemplate @@ -26,6 +27,34 @@ logger = logging.getLogger(__name__) +# Language names a model may put on the first line of a code block. +_FENCE_LANGUAGES = {"cypher", "opencypher", "gsql", "sql", "graphql"} + + +def strip_code_fence(text: str) -> str: + """The query inside a model reply: without a code fence (```cypher, + ```opencypher, plain ```) and without a line that has only a language + name.""" + t = (text or "").strip() + m = re.match(r"^```[^\n]*\n(.*?)\n?```\s*$", t, re.S) + if m: + t = m.group(1).strip() + else: + t = t.strip("`").strip() + lines = t.split("\n", 1) + if lines[0].strip().lower() in _FENCE_LANGUAGES: + t = lines[1].strip() if len(lines) > 1 else "" + return t + + +def run_cypher(conn, cypher: str) -> list: + """Run an openCypher query in interpreted mode on the graph of ``conn`` + and return its result blocks. Raises the connection's exception with + TigerGraph's message when the query fails.""" + return conn.runInterpretedQuery( + "INTERPRET OPENCYPHER QUERY () FOR GRAPH $graphname {\n" + cypher + "\n}" + ) or [] + class GenerateCypher(BaseTool): """GenerateCypher Tool. @@ -98,7 +127,8 @@ def generate_cypher(self, question: str, history: Iterable[str]) -> str: PROMPT, StrOutputParser(), {"question": question, "schema": schema, "history": history}, caller_name="generate_cypher", - ).strip("```cypher").strip("```").strip() + ) + out = strip_code_fence(out) # Validate the LLM output looks like a Cypher query out_upper = out.upper() @@ -106,11 +136,8 @@ def generate_cypher(self, question: str, history: Iterable[str]) -> str: LogWriter.info(f"request_id={req_id_cv.get()} EXIT generate_cypher - LLM did not produce a valid Cypher query") raise ValueError(f"LLM did not produce a valid Cypher query: {out[:200]}") - query_header = "USE GRAPH " + self.conn.graphname + " "+ "\n" + "INTERPRET OPENCYPHER QUERY () {" + "\n" - query_footer = "\n}" - cypher = query_header + out + query_footer - LogWriter.info(f"request_id={req_id_cv.get()} EXIT generate_cypher with:\n{cypher}") - return cypher + LogWriter.info(f"request_id={req_id_cv.get()} EXIT generate_cypher with:\n{out}") + return out def _run(self, question: str, history: Iterable[str]): """Run the GenerateCypher tool. diff --git a/graphrag/app/tools/graphrag_tools.py b/graphrag/app/tools/graphrag_tools.py index 96f4d1c..b89b152 100644 --- a/graphrag/app/tools/graphrag_tools.py +++ b/graphrag/app/tools/graphrag_tools.py @@ -28,12 +28,12 @@ from __future__ import annotations -import json import logging from dataclasses import dataclass, field from typing import Any, Callable, Dict, Optional from common.utils.query_results import is_empty_result +from tools.generate_cypher import run_cypher from common.utils.retrieval_stats import describe_retrieval from tools import tool_guards as guards from tools.validation_utils import MapQuestionToSchemaException @@ -216,15 +216,12 @@ def _cypher_retrieve(ctx: GraphRAGToolContext, question: str, attempts: Optional gen_history.append(f"{i}: Error: {exc}\n") attempts.append(_attempt(error=exc)) continue - response = ctx.conn.gsql(cypher) - json_str = "\n".join(response.split("\n")[1:]) try: - parsed = json.loads(json_str) - except Exception: - gen_history.append(f"{i}: {cypher}\n\tError: {json_str}\n") - attempts.append(_attempt(cypher, json_str or response)) + rows = run_cypher(ctx.conn, cypher) + except Exception as exc: + gen_history.append(f"{i}: {cypher}\n\tError: {exc}\n") + attempts.append(_attempt(cypher, exc)) continue - rows = parsed.get("results", [{}]) first = rows[0] if rows else None if not _result_is_empty(first): return _ok( diff --git a/graphrag/tests/test_structural_trace.py b/graphrag/tests/test_structural_trace.py index 6f1b01e..ab70940 100644 --- a/graphrag/tests/test_structural_trace.py +++ b/graphrag/tests/test_structural_trace.py @@ -25,7 +25,7 @@ def test_results_with_rows(self): _FUNC = "getVertexCount('Document')" -_CYPHER = "USE GRAPH g\nINTERPRET OPENCYPHER QUERY () {\nMATCH (d:Document) RETURN d\n}" +_CYPHER = "MATCH (d:Document) RETURN d" def _ctx(gen_result, use_cypher=False, cypher_runs=()): @@ -41,7 +41,8 @@ def _ctx(gen_result, use_cypher=False, cypher_runs=()): cypher_gen = mock.Mock() cypher_gen._run.side_effect = [c for c, _ in cypher_runs] or [_CYPHER] conn = mock.Mock() - conn.gsql.side_effect = [r for _, r in cypher_runs] or ["header\n{}"] + # runInterpretedQuery returns result blocks, or raises with TigerGraph's message. + conn.runInterpretedQuery.side_effect = [r for _, r in cypher_runs] or [[]] return types.SimpleNamespace( emit=lambda *a, **k: None, conversation=[], use_cypher=use_cypher, mq2s=mock.Mock(_run=mock.Mock(return_value=mapping)), @@ -50,7 +51,8 @@ def _ctx(gen_result, use_cypher=False, cypher_runs=()): def _gsql(results): - return "header\n" + json.dumps({"results": results}) + """What runInterpretedQuery returns: the query's result blocks.""" + return results class TestStructuralRetrieveQuery(unittest.TestCase): @@ -72,7 +74,7 @@ def test_generation_error_is_recorded(self): def test_cypher_rows_report_the_query_and_earlier_attempts(self): ctx = _ctx({"function_call": _FUNC, "result": "[]"}, use_cypher=True, - cypher_runs=[("BAD", "header\nnot json"), (_CYPHER, _gsql([{"VS": [{"v_id": "d1"}]}]))]) + cypher_runs=[("BAD", RuntimeError("syntax error")), (_CYPHER, _gsql([{"VS": [{"v_id": "d1"}]}]))]) out = gt.structural_retrieve(ctx, "q") self.assertTrue(out["ok"]) self.assertEqual(out["query"], _CYPHER) @@ -131,6 +133,39 @@ def test_mapping_without_edges_does_not_fail(self): self.assertEqual(json.loads(out["result"]), 4) +class TestCypherText(unittest.TestCase): + """The generated query is the Cypher only (GML-2336).""" + + def test_fences_and_language_lines_are_removed(self): + from tools.generate_cypher import strip_code_fence + q = "MATCH (i:JiraIssue) RETURN i" + for reply in (q, f"```cypher\n{q}\n```", f"```opencypher\n{q}\n```", f"```\n{q}\n```", + f"opencypher\n{q}", f"```Cypher\n{q}```", f"`{q}`"): + self.assertEqual(strip_code_fence(reply), q, reply) + + def test_query_runs_interpreted_on_the_connection_graph(self): + from tools.generate_cypher import run_cypher + conn = mock.Mock() + conn.runInterpretedQuery.return_value = [{"T": {"n": 4}}] + self.assertEqual(run_cypher(conn, "MATCH (d:Document) RETURN count(d) AS n"), [{"T": {"n": 4}}]) + text = conn.runInterpretedQuery.call_args[0][0] + self.assertTrue(text.startswith("INTERPRET OPENCYPHER QUERY () FOR GRAPH $graphname {")) + self.assertNotIn("USE GRAPH", text) + conn.gsql.assert_not_called() + + def test_classic_engine_uses_the_same_runner(self): + from agent.agent_graph import TigerGraphAgentGraph + node = object.__new__(TigerGraphAgentGraph) + node.emit_progress = lambda *a, **k: None + node.cypher_gen = mock.Mock(_run=mock.Mock(return_value="MATCH (d:Document) RETURN count(d) AS n")) + node.db_connection = mock.Mock() + node.db_connection.runInterpretedQuery.return_value = [{"T": {"n": 4}}] + state = node.generate_cypher({"question": "How many documents?"}) + self.assertEqual(state["context"]["result"], {"T": {"n": 4}}) + self.assertEqual(state["context"]["cypher"], "MATCH (d:Document) RETURN count(d) AS n") + node.db_connection.gsql.assert_not_called() + + class TestTraceShowsTheQuery(unittest.TestCase): def test_query_and_attempts_lead_the_step_output(self): out = {"ok": True, "summary": "graph query returned rows", "context": {"result": [1]}, From f41b3a1bbcfd6c1f312d3495cf4b38aa9b39c86b Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 15:57:40 -0700 Subject: [PATCH 18/32] Check generated openCypher with TigerGraph before running it - A query that fails TigerGraph's check is not run; its errors go back to the generator - Up to 10 attempts (configurable), stopping early when the generator repeats itself - The generator uses TigerGraph's date functions, and queries that change data are rejected Refs: GML-2336 --- CHANGELOG.md | 3 +- README.md | 1 + common/llm_services/base_llm.py | 10 ++- graphrag/app/agent/agent_graph.py | 44 ++++------- graphrag/app/tools/generate_cypher.py | 101 +++++++++++++++++++++++- graphrag/app/tools/graphrag_tools.py | 44 ++++------- graphrag/tests/test_structural_trace.py | 86 +++++++++++++++++++- 7 files changed, 218 insertions(+), 71 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 136cdd9..739b316 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,7 +10,7 @@ ### Fixed - **The agentic planner plans from the graph schema.** Since 2.0.0 the planner made its plan without knowing the graph's vertex and edge types, so it could only guess from the wording of a question whether structural retrieval would help. It now gets a compact summary of the schema — type names, attribute names, edge endpoints and descriptions — cached per graph and refreshed when the schema changes. - **Structured queries no longer start with a step that cannot succeed.** Each structured query first asked the model for a pyTigerGraph function call, but no functions are registered with the TigerGraph embedding store, so every call was rejected, or failed outright when the question named no relationships, before the openCypher query ran. GraphRAG now detects that no functions are registered and goes straight to the openCypher query, saving one model call per structured query. -- **Generated graph queries no longer fail when the model labels its code block.** A reply that started with an `opencypher` code block kept the word `opencypher` as the first line of the query, so TigerGraph rejected every attempt. Code blocks are now removed correctly, generated queries run through pyTigerGraph's interpreted-query call instead of the GSQL shell, and Trace Logs show only the generated query. +- **Generated graph queries succeed more often.** On a Jira graph every attempt at a date-filtered question failed: a code-block label stayed in the query, and the query used a date function TigerGraph does not support. Code blocks are now removed correctly, and the query generator is told to use TigerGraph's `to_datetime()`. Each generated query is now checked by TigerGraph before it runs (TigerGraph 4.0 and later); a query that fails the check is not run, and its errors go back to the generator, which may try up to 10 times (`cypher_max_attempts`) and stops early when it repeats itself. A generated query that would change data is rejected. Queries run through pyTigerGraph's interpreted-query call instead of the GSQL shell, and Trace Logs show only the generated query and each failed attempt. - **An empty structured query result no longer counts as data.** A query that returned only empty row sets was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. The classic engine uses the same check, so both engines treat an empty result the same way. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. - **The agentic planner and reactive agent see the most recent conversation turns.** They received the oldest part of the conversation, so one long earlier answer could hide the latest turns that a follow-up question refers to. They now receive the most recent turns, as triage does. @@ -21,6 +21,7 @@ - **The `tg_run_query` and `tg_get_neighbors` agent tools.** They let the agent write and run its own GSQL and expand neighbors through tigergraph-mcp's functions. Structured retrieval already maps the question to the schema, writes the query, runs it and retries on errors, so the agent does not need to write queries itself, and GraphRAG code now reaches TigerGraph only through pyTigerGraph. Custom planner or agent prompts that name these tools should drop them. ### Added +- **`cypher_max_attempts` setting.** How many times a structured graph query may be generated for one question (default `10`). - **`planner_max_output_tokens` setting.** The most output tokens the planner may use for one plan (default `4000`; reasoning models use at least `16000`). ## [2.0.3] diff --git a/README.md b/README.md index 59db923..8da6525 100644 --- a/README.md +++ b/README.md @@ -547,6 +547,7 @@ Copy the below code into `configs/server_config.json`. You shouldn’t need to c | `agent_style` | string | `"planned"` | Default agentic engine style: `"planned"` (plan the whole retrieval up front) or `"reactive"` (decide each step from the last result). The chat menu can override per request. See [Chat Engines and Agents](#chat-engines-and-agents). | | `agent_max_iterations` | int | `30` | Reactive agent only: maximum reason-act-observe steps before it must answer. | | `agent_max_replans` | int | `3` | Planned agent only: how many times the planner may extend its plan when the gathered context is insufficient. | +| `cypher_max_attempts` | int | `10` | How many times a structured graph query may be generated for one question before GraphRAG gives up on it. | | `planner_max_output_tokens` | int | `4000` | Planned agent only: the most output tokens the planner may use for one plan. Reasoning models use at least `16000`. | | `agent_max_total_steps` | int | `20` | Planned agent only: hard cap on executed retrieval steps across all replans. | | `chunk_only` | bool | `true` | If true, hybrid search only retrieves document chunks, excluding entity data. | diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index 16c0f90..1180daa 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -960,10 +960,11 @@ def generate_cypher_prompt(self): - For string comparisons in `WHERE`, convert with `toLower()`. - Use multi-word, underscore-joined aliases for `ORDER BY`. Aliases / attributes used in `ORDER BY` must be in `RETURN`. Always specify `ASC` / `DESC` based on data type. - For "summarize" / "write a summary" questions, fetch all neighbour nodes and edges. -- Avoid invalid queries based on errors in the history above. +- For dates and times, use `to_datetime("YYYY-MM-DD hh:mm:ss")` (a date alone also works, e.g. `WHERE n.created >= to_datetime("2025-08-01")`). TigerGraph does not support `datetime()` or `date()`. +- Avoid invalid queries based on errors in the history above. Each error gives the line and column of the problem. ## Supported -- **Clauses**: `MATCH`, `OPTIONAL MATCH`, `MANDATORY MATCH`, `WHERE`, `RETURN`, `WITH`, `ORDER BY`, `SKIP`, `LIMIT`, `DELETE`, `DETACH DELETE` +- **Clauses**: `MATCH`, `OPTIONAL MATCH`, `MANDATORY MATCH`, `WHERE`, `RETURN`, `WITH`, `ORDER BY`, `SKIP`, `LIMIT` - **Operators**: - Math: `+`, `-`, `*`, `/`, `%`, `^` - Comparison: `=`, `<`, `<=`, `>`, `>=`, `<>`, `IS NULL`, `IS NOT NULL` @@ -974,12 +975,13 @@ def generate_cypher_prompt(self): - Math: `abs`, `sqrt`, `log`, `exp`, `sin`, `cos`, `tan`, `radians`, `degrees` - String: `left`, `right`, `substring`, `replace`, `trim`, `toLower`, `toUpper`, `split` - List: `head`, `last`, `size`, `range`, `coalesce`, `tail` + - Date / time: `to_datetime`, `datetime_format`, `datetime_to_epoch`, `year` - Other: `id`, `elementId`, `labels`, `properties`, `timestamp` - **Expressions**: `CASE` ## Unsupported -- **Clauses**: `CALL`, `CREATE`, `MERGE`, `REMOVE`, `SET`, `UNION`, `UNION ALL`, `UNWIND` -- **Functions**: `collect`, `exists`, `keys`, `nodes`, `relationships`, `length`, `percentileCont`, `percentileDisc`, `startNode`, `endNode`, `reverse` (list form) +- **Clauses**: `CALL`, `CREATE`, `MERGE`, `REMOVE`, `SET`, `DELETE`, `DETACH DELETE`, `UNION`, `UNION ALL`, `UNWIND` (queries only read data) +- **Functions**: `datetime`, `date`, `collect`, `exists`, `keys`, `nodes`, `relationships`, `length`, `percentileCont`, `percentileDisc`, `startNode`, `endNode`, `reverse` (list form) - **Syntax limits**: - `WITH` must group by exactly one vertex variable. - Path variables (`p = (...)`) not supported. diff --git a/graphrag/app/agent/agent_graph.py b/graphrag/app/agent/agent_graph.py index e7fcf65..525a2de 100644 --- a/graphrag/app/agent/agent_graph.py +++ b/graphrag/app/agent/agent_graph.py @@ -40,7 +40,7 @@ from tools import MapQuestionToSchemaException from typing_extensions import TypedDict from common.utils.query_results import is_empty_result -from tools.generate_cypher import run_cypher +from tools.generate_cypher import CYPHER_MAX_ATTEMPTS, generate_and_run_cypher from common.logs.log import req_id_cv from common.metrics.prometheus_metrics import metrics as pmetrics @@ -362,48 +362,34 @@ def generate_function(self, state): def generate_cypher(self, state): """ - Run the agent cypher generator. + Run the agent cypher generator: generate, check and run the query + (see ``generate_and_run_cypher``). """ self.emit_progress("Generating the query to answer your question") - gen_history = [] - response_json = None - cypher = None - json_str = None - response = None - - for i in range(3): - try: - cypher = self.cypher_gen._run(state["question"], gen_history) - except ValueError as e: - logger.warning(f"Cypher generation failed: {e}") - gen_history.append(f"{i}: Error: {e}\n") - continue - try: - response_json = {"results": run_cypher(self.db_connection, cypher)} - json_str = json.dumps(response_json) - break - except Exception as e: - response = str(e) - json_str = response - gen_history.append(f"{i}: {cypher}\n\tError: {response}\n") - if response_json and response_json["results"] and not self.is_query_result_empty(response_json["results"][0]): + attempts: list = [] + cfg = getattr(self, "_graphrag_cfg", None) or {} + cypher, first = generate_and_run_cypher( + self.cypher_gen, self.db_connection, state["question"], + cfg.get("cypher_max_attempts", CYPHER_MAX_ATTEMPTS), attempts, + ) + if cypher is not None and first is not None: state["context"] = { - "result": response_json["results"][0], + "result": first, "cypher": cypher, "reasoning": "The following OpenCypher query was executed to answer the question. {}".format( cypher ), } else: + last = attempts[-1] if attempts else {} state["context"] = { "error": True, - "cypher": cypher, - "result": json_str + "cypher": cypher or last.get("query"), + "result": last.get("error") or last.get("outcome") or "", } if "error_history" not in state or state["error_history"] is None: state["error_history"] = [] - - error_msg = response if response else "LLM failed to produce a valid Cypher query after 3 attempts" + error_msg = last.get("error") or last.get("outcome") or "LLM failed to produce a valid Cypher query" state["error_history"].append({"error_message": error_msg, "error_step": "generate_cypher"}) state["lookup_source"] = "cypher" diff --git a/graphrag/app/tools/generate_cypher.py b/graphrag/app/tools/generate_cypher.py index f1ea84b..17e1006 100644 --- a/graphrag/app/tools/generate_cypher.py +++ b/graphrag/app/tools/generate_cypher.py @@ -47,13 +47,108 @@ def strip_code_fence(text: str) -> str: return t +def _interpreted(conn, cypher: str) -> str: + return ("INTERPRET OPENCYPHER QUERY () FOR GRAPH " + str(conn.graphname) + + " {\n" + cypher + "\n}") + + def run_cypher(conn, cypher: str) -> list: """Run an openCypher query in interpreted mode on the graph of ``conn`` and return its result blocks. Raises the connection's exception with TigerGraph's message when the query fails.""" - return conn.runInterpretedQuery( - "INTERPRET OPENCYPHER QUERY () FOR GRAPH $graphname {\n" + cypher + "\n}" - ) or [] + return conn.runInterpretedQuery(_interpreted(conn, cypher)) or [] + + +# Clauses that change data. Retrieval only reads, whatever the model writes. +_WRITE_CLAUSE = re.compile(r"\b(CREATE|MERGE|SET|REMOVE|DELETE)\b", re.I) +_STRING_LITERAL = re.compile(r"\"(?:[^\"\\]|\\.)*\"|'(?:[^'\\]|\\.)*'") + +# TigerGraph limits worth a direct hint when the check reports them. +_HINTS = ( + (re.compile(r"\bdatetime\s*\(", re.I), + 'TigerGraph does not support datetime(). Use to_datetime("YYYY-MM-DD hh:mm:ss").'), + (re.compile(r"\bdate\s*\(", re.I), + 'TigerGraph does not support date(). Use to_datetime("YYYY-MM-DD").'), +) + + +def check_cypher(conn, cypher: str): + """Errors in ``cypher`` found without running it, as readable lines. + + Uses TigerGraph's query check (TigerGraph 4.0 and later). Returns [] for + a valid query, or None when the check is not available, so the caller + runs the query to find out. + """ + if _WRITE_CLAUSE.search(_STRING_LITERAL.sub('""', cypher)): + return ["The query changes data. Write a query that only reads (MATCH ... RETURN)."] + try: + result = conn.checkQuerySemantic(_interpreted(conn, cypher)) + except Exception as exc: + logger.info(f"generate_cypher: query check unavailable ({exc}); running the query instead") + return None + if not isinstance(result, dict): + return None + errors = [] + for err in result.get("errors") or []: + # The check counts the INTERPRET wrapper as line 1; report query lines. + line = err.get("line") + where = f"line {line - 1}, column {err.get('charpositioninline')}: " if isinstance(line, int) and line > 1 else "" + errors.append(where + str(err.get("msg", "")).strip()) + if errors: + errors.extend(hint for pattern, hint in _HINTS if pattern.search(cypher)) + return errors + + +CYPHER_MAX_ATTEMPTS = 10 + + +def generate_and_run_cypher(cypher_gen, conn, question: str, max_attempts: int = CYPHER_MAX_ATTEMPTS, + attempts=None): + """Generate an openCypher query for ``question``, check it, and run it. + + A query that fails the check is not run: the errors go back to the + generator for the next attempt. Stops at ``max_attempts``, or early when + the generator repeats a query it already wrote. Every attempt that did not + return rows is appended to ``attempts`` as ``{query, error, outcome}``. + + Returns ``(cypher, first_result_block)`` for the query that ran (the block + is None when the query returned no rows), or ``(None, None)`` when no query + ran successfully. + """ + from common.utils.query_results import is_empty_result + + attempts = [] if attempts is None else attempts + history: list = [] + seen: set = set() + for i in range(max(1, int(max_attempts or CYPHER_MAX_ATTEMPTS))): + try: + cypher = cypher_gen._run(question, history) + except ValueError as exc: + history.append(f"{i}: Error: {exc}\n") + attempts.append({"error": str(exc)[:500]}) + continue + if cypher in seen: + attempts.append({"query": cypher, "outcome": "repeated an earlier query; stopped"}) + break + seen.add(cypher) + errors = check_cypher(conn, cypher) + if errors: + message = "\n".join(errors) + history.append(f"{i}: {cypher}\n\tError: {message}\n") + attempts.append({"query": cypher, "error": message[:500], "outcome": "failed check"}) + continue + try: + rows = run_cypher(conn, cypher) + except Exception as exc: + history.append(f"{i}: {cypher}\n\tError: {exc}\n") + attempts.append({"query": cypher, "error": str(exc)[:500]}) + continue + first = rows[0] if rows else None + if is_empty_result(first): + attempts.append({"query": cypher, "outcome": "no rows"}) + return cypher, None + return cypher, first + return None, None class GenerateCypher(BaseTool): diff --git a/graphrag/app/tools/graphrag_tools.py b/graphrag/app/tools/graphrag_tools.py index b89b152..1bb73d4 100644 --- a/graphrag/app/tools/graphrag_tools.py +++ b/graphrag/app/tools/graphrag_tools.py @@ -33,7 +33,7 @@ from typing import Any, Callable, Dict, Optional from common.utils.query_results import is_empty_result -from tools.generate_cypher import run_cypher +from tools.generate_cypher import CYPHER_MAX_ATTEMPTS, generate_and_run_cypher from common.utils.retrieval_stats import describe_retrieval from tools import tool_guards as guards from tools.validation_utils import MapQuestionToSchemaException @@ -200,38 +200,22 @@ def structural_retrieve(ctx: GraphRAGToolContext, question: str) -> dict: def _cypher_retrieve(ctx: GraphRAGToolContext, question: str, attempts: Optional[list] = None) -> dict: - """Generate and run an openCypher query, retrying on errors. - - Each attempt that does not produce rows is appended to ``attempts``. A - query that runs and returns no rows is a valid answer, so it is not - retried. - """ + """Generate, check and run an openCypher query (see + ``generate_and_run_cypher``). Attempts that do not produce rows are + appended to ``attempts``.""" attempts = [] if attempts is None else attempts ctx.emit("Generating a graph query") - gen_history: list = [] - for i in range(3): - try: - cypher = ctx.cypher_gen._run(question, gen_history) - except ValueError as exc: - gen_history.append(f"{i}: Error: {exc}\n") - attempts.append(_attempt(error=exc)) - continue - try: - rows = run_cypher(ctx.conn, cypher) - except Exception as exc: - gen_history.append(f"{i}: {cypher}\n\tError: {exc}\n") - attempts.append(_attempt(cypher, exc)) - continue - first = rows[0] if rows else None - if not _result_is_empty(first): - return _ok( - "graph query returned rows", - {"result": first, "cypher": cypher, - "reasoning": f"The following openCypher query was executed:\n{cypher}"}, - ) - attempts.append(_attempt(cypher, outcome="no rows")) + max_attempts = (getattr(ctx, "graphrag_cfg", None) or {}).get("cypher_max_attempts", CYPHER_MAX_ATTEMPTS) + cypher, first = generate_and_run_cypher(ctx.cypher_gen, ctx.conn, question, max_attempts, attempts) + if cypher is None: + return _empty("graph query returned no rows after retries") + if first is None: return _empty("graph query returned no rows") - return _empty("graph query returned no rows after retries") + return _ok( + "graph query returned rows", + {"result": first, "cypher": cypher, + "reasoning": f"The following openCypher query was executed:\n{cypher}"}, + ) # -------------------------------------------------------------------------- diff --git a/graphrag/tests/test_structural_trace.py b/graphrag/tests/test_structural_trace.py index ab70940..e9ce68d 100644 --- a/graphrag/tests/test_structural_trace.py +++ b/graphrag/tests/test_structural_trace.py @@ -40,7 +40,8 @@ def _ctx(gen_result, use_cypher=False, cypher_runs=()): gen_func._run.return_value = gen_result cypher_gen = mock.Mock() cypher_gen._run.side_effect = [c for c, _ in cypher_runs] or [_CYPHER] - conn = mock.Mock() + conn = mock.Mock(graphname="G") + conn.checkQuerySemantic.return_value = {"warnings": [], "errors": []} # runInterpretedQuery returns result blocks, or raises with TigerGraph's message. conn.runInterpretedQuery.side_effect = [r for _, r in cypher_runs] or [[]] return types.SimpleNamespace( @@ -145,11 +146,11 @@ def test_fences_and_language_lines_are_removed(self): def test_query_runs_interpreted_on_the_connection_graph(self): from tools.generate_cypher import run_cypher - conn = mock.Mock() + conn = mock.Mock(graphname="G") conn.runInterpretedQuery.return_value = [{"T": {"n": 4}}] self.assertEqual(run_cypher(conn, "MATCH (d:Document) RETURN count(d) AS n"), [{"T": {"n": 4}}]) text = conn.runInterpretedQuery.call_args[0][0] - self.assertTrue(text.startswith("INTERPRET OPENCYPHER QUERY () FOR GRAPH $graphname {")) + self.assertTrue(text.startswith("INTERPRET OPENCYPHER QUERY () FOR GRAPH G {")) self.assertNotIn("USE GRAPH", text) conn.gsql.assert_not_called() @@ -158,7 +159,8 @@ def test_classic_engine_uses_the_same_runner(self): node = object.__new__(TigerGraphAgentGraph) node.emit_progress = lambda *a, **k: None node.cypher_gen = mock.Mock(_run=mock.Mock(return_value="MATCH (d:Document) RETURN count(d) AS n")) - node.db_connection = mock.Mock() + node.db_connection = mock.Mock(graphname="G") + node.db_connection.checkQuerySemantic.return_value = {"warnings": [], "errors": []} node.db_connection.runInterpretedQuery.return_value = [{"T": {"n": 4}}] state = node.generate_cypher({"question": "How many documents?"}) self.assertEqual(state["context"]["result"], {"T": {"n": 4}}) @@ -166,6 +168,82 @@ def test_classic_engine_uses_the_same_runner(self): node.db_connection.gsql.assert_not_called() +class TestCheckBeforeRun(unittest.TestCase): + """Generated queries are checked before they run (GML-2336).""" + + def _conn(self, checks, runs=None): + conn = mock.Mock(graphname="G") + conn.checkQuerySemantic.side_effect = checks + conn.runInterpretedQuery.side_effect = runs or [[{"T": {"n": 4}}]] + return conn + + def _gen(self, queries): + return mock.Mock(_run=mock.Mock(side_effect=queries)) + + def test_failed_check_feeds_the_error_back_and_skips_the_run(self): + from tools.generate_cypher import generate_and_run_cypher + bad = 'MATCH (i:JiraIssue)\nWHERE i.created >= datetime("2025-08-01")\nRETURN count(i) AS n' + good = 'MATCH (i:JiraIssue)\nWHERE i.created >= to_datetime("2025-08-01")\nRETURN count(i) AS n' + err = {"errors": [{"msg": "no viable alternative at input 'datetime'", "line": 3, "charpositioninline": 18}]} + conn = self._conn([err, {"errors": []}]) + gen = self._gen([bad, good]) + attempts = [] + cypher, first = generate_and_run_cypher(gen, conn, "q", 10, attempts) + self.assertEqual((cypher, first), (good, {"T": {"n": 4}})) + conn.runInterpretedQuery.assert_called_once() # the bad query never ran + self.assertEqual(attempts[0]["outcome"], "failed check") + self.assertIn("line 2, column 18", attempts[0]["error"]) + history = gen._run.call_args_list[1][0][1] + self.assertIn("to_datetime", history[0]) # the date hint went back to the generator + + def test_stops_when_the_generator_repeats_itself(self): + from tools.generate_cypher import generate_and_run_cypher + q = "MATCH (x:Nope) RETURN x" + conn = self._conn([{"errors": [{"msg": "unknown type"}]}] * 3) + attempts = [] + self.assertEqual(generate_and_run_cypher(self._gen([q, q, q]), conn, "q", 10, attempts), (None, None)) + self.assertEqual(attempts[-1]["outcome"], "repeated an earlier query; stopped") + self.assertEqual(len(attempts), 2) + + def test_attempt_limit(self): + from tools.generate_cypher import generate_and_run_cypher + conn = self._conn([{"errors": [{"msg": "bad"}]}] * 20) + gen = self._gen([f"MATCH (n{i}) RETURN n{i}" for i in range(20)]) + generate_and_run_cypher(gen, conn, "q", 4, []) + self.assertEqual(gen._run.call_count, 4) + + def test_write_queries_are_rejected_without_a_check(self): + from tools.generate_cypher import check_cypher + conn = mock.Mock(graphname="G") + self.assertIn("changes data", check_cypher(conn, "MATCH (n) DETACH DELETE n")[0]) + conn.checkQuerySemantic.assert_not_called() + conn.checkQuerySemantic.return_value = {"errors": []} + self.assertEqual(check_cypher(conn, 'MATCH (n {name: "set to delete"}) RETURN n'), []) + + def test_unavailable_check_runs_the_query(self): + from tools.generate_cypher import generate_and_run_cypher + conn = self._conn(RuntimeError("only on TigerGraph >= 4.0")) + cypher, first = generate_and_run_cypher(self._gen(["MATCH (d:Document) RETURN count(d) AS n"]), conn, "q") + self.assertEqual(first, {"T": {"n": 4}}) + + def test_attempt_setting_reaches_the_structural_tool(self): + ctx = _ctx({"function_call": _FUNC, "result": "[]"}, use_cypher=True, + cypher_runs=[(f"MATCH (n{i}) RETURN n{i}", []) for i in range(6)]) + ctx.conn.checkQuerySemantic.return_value = {"errors": [{"msg": "bad"}]} + ctx.graphrag_cfg = {"cypher_max_attempts": 3} + gt.structural_retrieve(ctx, "q") + self.assertEqual(ctx.cypher_gen._run.call_count, 3) + + def test_prompt_rules(self): + from common.llm_services.base_llm import LLM_Model + prompt = LLM_Model.generate_cypher_prompt.fget(types.SimpleNamespace( + _read_prompt_file=lambda p: None, prompt_path="")) + self.assertIn('to_datetime("YYYY-MM-DD hh:mm:ss")', prompt) + supported = prompt.split("## Supported")[1].split("## Unsupported")[0] + self.assertNotIn("DELETE", supported) + self.assertIn("`DETACH DELETE`", prompt.split("## Unsupported")[1]) + + class TestTraceShowsTheQuery(unittest.TestCase): def test_query_and_attempts_lead_the_step_output(self): out = {"ok": True, "summary": "graph query returned rows", "context": {"result": [1]}, From 502c8b9924c1ab746d029ceb81bb3d5d80eb6f6d Mon Sep 17 00:00:00 2001 From: Chengbiao Jin Date: Thu, 15 Oct 2026 16:12:29 -0700 Subject: [PATCH 19/32] Keep long trace content inside the Trace Logs window - Long unbroken lines in step input and output wrap inside the window - Download and Close stay visible at the bottom of the window Refs: GML-2327 --- CHANGELOG.md | 1 + graphrag-ui/src/pages/TraceLogs.tsx | 18 +++++++++++------- 2 files changed, 12 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 739b316..08d0a59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ ### Fixed - **The agentic planner plans from the graph schema.** Since 2.0.0 the planner made its plan without knowing the graph's vertex and edge types, so it could only guess from the wording of a question whether structural retrieval would help. It now gets a compact summary of the schema — type names, attribute names, edge endpoints and descriptions — cached per graph and refreshed when the schema changes. - **Structured queries no longer start with a step that cannot succeed.** Each structured query first asked the model for a pyTigerGraph function call, but no functions are registered with the TigerGraph embedding store, so every call was rejected, or failed outright when the question named no relationships, before the openCypher query ran. GraphRAG now detects that no functions are registered and goes straight to the openCypher query, saving one model call per structured query. +- **The Trace Logs window keeps long content inside it.** A long unbroken line in a step's input or output, such as a URL or a query, widened the window past the screen and pushed the Close button out of view. Long lines now wrap, and Download and Close stay at the bottom of the window. - **Generated graph queries succeed more often.** On a Jira graph every attempt at a date-filtered question failed: a code-block label stayed in the query, and the query used a date function TigerGraph does not support. Code blocks are now removed correctly, and the query generator is told to use TigerGraph's `to_datetime()`. Each generated query is now checked by TigerGraph before it runs (TigerGraph 4.0 and later); a query that fails the check is not run, and its errors go back to the generator, which may try up to 10 times (`cypher_max_attempts`) and stops early when it repeats itself. A generated query that would change data is rejected. Queries run through pyTigerGraph's interpreted-query call instead of the GSQL shell, and Trace Logs show only the generated query and each failed attempt. - **An empty structured query result no longer counts as data.** A query that returned only empty row sets was reported as "returned rows", so the agent treated it as context and did not search the documents instead. It is now reported as no rows, and the document search runs. The classic engine uses the same check, so both engines treat an empty result the same way. - **External MCP servers load again.** Images built after the MCP SDK released version 2 installed it, and the agent then skipped every external MCP server configured on the MCP Servers page, so their tools were unavailable. The SDK version is now held below 2 until GraphRAG moves to its new client API. diff --git a/graphrag-ui/src/pages/TraceLogs.tsx b/graphrag-ui/src/pages/TraceLogs.tsx index b838ed2..2a8d48b 100644 --- a/graphrag-ui/src/pages/TraceLogs.tsx +++ b/graphrag-ui/src/pages/TraceLogs.tsx @@ -360,7 +360,7 @@ const ExpandableRow: FC<{ {open && content && (
- 
{content}
+ 
{content}
)} @@ -391,7 +391,7 @@ const PlanPanel: FC<{ trace: TraceData }> = ({ trace }) => { {plan.strategy && (
Strategy -

{plan.strategy}

+

{plan.strategy}

)}
    @@ -553,7 +553,7 @@ const ToolCallExpandable: FC<{ tc: ToolCallEntry }> = ({ tc }) => {

    Input

    - 
    +            
                   {tc.input || "N/A"}
             
    
           
@@ -561,7 +561,7 @@ const ToolCallExpandable: FC<{ tc: ToolCallEntry }> = ({ tc }) => {
             
    
               Output
             
    
-            
    +            
                   {tc.output || "N/A"}
             
    
           
@@ -868,7 +868,7 @@ const TraceLogs: FC = ({ messageIdProp, onClose }) => {
       return (
          !o && onClose && onClose()}>
            e.preventDefault()}
             onOpenAutoFocus={(e) => e.preventDefault()}
           >
@@ -1076,8 +1076,12 @@ const TraceLogs: FC = ({ messageIdProp, onClose }) => {
 
         {/* Footer — Download (primary action style) + Close (outline,
             matches other dialogs). Replaces the top-right Download in
-            dialog mode so the dialog can size to content. */}
-        
    
+            dialog mode so the dialog can size to content. In the dialog it
+            stays at the bottom while the trace scrolls, so Close is always
+            visible. */}
+        
    
           {isDialog && (