From a9e818a144d04d5ae0ba07c6de284a1bc5ad0702 Mon Sep 17 00:00:00 2001 From: Prins Kumar Date: Mon, 31 Aug 2026 15:43:54 +0530 Subject: [PATCH] Register installed GSQL as GraphRAG tools from KG Admin. Tag queries with a native GSQL description so the agent catalog can pick them by match, without treating every installed query as a tool. Co-authored-by: Cursor --- common/llm_services/base_llm.py | 12 +- graphrag-ui/src/pages/setup/KGAdmin.tsx | 450 +++++++++++++++++++++++- graphrag/app/agent/agentic_agent.py | 14 + graphrag/app/agent/agentic_react.py | 6 +- graphrag/app/routers/ui.py | 228 ++++++++++++ graphrag/app/tools/gsql_query_tools.py | 308 ++++++++++++++++ 6 files changed, 1012 insertions(+), 6 deletions(-) create mode 100644 graphrag/app/tools/gsql_query_tools.py diff --git a/common/llm_services/base_llm.py b/common/llm_services/base_llm.py index fe77ac1..17d10e8 100644 --- a/common/llm_services/base_llm.py +++ b/common/llm_services/base_llm.py @@ -1038,7 +1038,8 @@ def select_retriever_prompt(self): # Operator-customizable retrieval strategy for the react agent: the first # action, then each next action driven by what the previous result returned. _AGENTIC_AGENT_USER_DEFAULT = """\ -- For most questions, make your FIRST action a vector search (graphrag__hybrid_search or graphrag__contextual_search) — it gives the broadest grounding. Skip it only when you are highly confident the question is a pure structured-data request (an exact count, an attribute/id lookup, a relationship traversal, or an aggregation over typed graph data) that a generated graph query fully answers on its own. +- If a graphrag__gsql__* tool is available and its description matches the question, you may call it. If none match, ignore those tools. Do not call a list/register tool first, and do not call a gsql tool first unless its description matches. +- For most other questions, make your FIRST action a vector search (graphrag__hybrid_search or graphrag__contextual_search) — it gives the broadest grounding. Skip it only when you are highly confident the question is a pure structured-data request (an exact count, an attribute/id lookup, a relationship traversal, an aggregation over typed graph data, or a matching graphrag__gsql__* tool) that a graph query fully answers on its own. - Let each observation drive the next action: if the passages you got back name specific entities or relationships you still need hard facts about, follow up with a structural query; if a result is thin, empty, or off-target, widen its parameters (top_k, num_hops) or switch method rather than repeating the same call. - Before answering, check that every part of the question is covered with the specific facts and figures it asks for; if a required value, table, or entity is still missing, retrieve again (widen top_k / num_hops or switch method) rather than answering vaguely or partially. - For a specific value, row, total, ranking, or year-over-year comparison, use graphrag__hybrid_search or graphrag__contextual_search with top_k >= 10 (they return atomic table chunks that keep full row/column structure), and quote the exact label, column, year, or unit from the question so the retriever can match it.""" @@ -1067,7 +1068,8 @@ 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). -You have two kinds of retrieval: +You have three kinds of retrieval: +- INSTALLED (graphrag__gsql__*): a user-registered installed GSQL query. Use it only when that tool's description matches the question. Do not call one just because it is listed, and do not call a lookup/list tool first. - 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. @@ -1088,8 +1090,10 @@ def agentic_agent_prompt(self): # Strategy (operator-customizable) — moved out of the fixed rules so it can # be tuned without touching the role / act model / plan mechanics. _AGENTIC_PLANNER_USER_DEFAULT = """\ -- Prioritize including at least one vector search step (graphrag__hybrid_search or graphrag__contextual_search) unless you are highly confident the question is a pure structured-data request — an exact count, an attribute/id lookup, a relationship traversal, or an aggregation over typed graph data — that a generated graph query fully answers on its own. Whenever the answer could plausibly live in document text (what/why/how/describe/summarize, definitions, explanations, figures), include a vector search step. When unsure, include vector search. -- Use BOTH kinds when a question needs facts from the graph AND supporting text; you may run several of each, in any order. When you use STRUCTURAL, pair it with a vector search step unless the question is a pure structured-data request. +- If a graphrag__gsql__* tool is in the catalog and its description matches the question, include that tool. If none match, ignore them and plan hybrid/community/structural exactly as today. Do not call a list/register tool; do not call a gsql tool first unless its description matches. +- You may use a graphrag__gsql__* tool and a vector search together when the question needs both the dedicated query result and supporting passages, in any order. +- Prioritize including at least one vector search step (graphrag__hybrid_search or graphrag__contextual_search) unless you are highly confident the question is a pure structured-data request — an exact count, an attribute/id lookup, a relationship traversal, an aggregation over typed graph data, or a question fully answered by a matching graphrag__gsql__* tool — that a generated or installed graph query fully answers on its own. Whenever the answer could plausibly live in document text (what/why/how/describe/summarize, definitions, explanations, figures), include a vector search step. When unsure, include vector search. +- Use BOTH structural and unstructured kinds when a question needs facts from the graph AND supporting text; you may run several of each, in any order. When you use STRUCTURAL, pair it with a vector search step unless the question is a pure structured-data request. - Prefer the smallest plan that will work. Trivial/greeting questions need only the final answer step. - Tabular / numeric questions (a specific value, a row, a column total, a ranking, or a year-over-year comparison from a table or chart): prefer graphrag__contextual_search or graphrag__hybrid_search with top_k>=10 (these return atomic table chunks that preserve full row/column structure); avoid graphrag__similarity_search alone; quote any specific table label, column header, year, or unit from the question (e.g. "ROE 2023"); for "compare X across years/regions/categories" set top_k>=15.""" diff --git a/graphrag-ui/src/pages/setup/KGAdmin.tsx b/graphrag-ui/src/pages/setup/KGAdmin.tsx index 2cf23b5..c899fcb 100644 --- a/graphrag-ui/src/pages/setup/KGAdmin.tsx +++ b/graphrag-ui/src/pages/setup/KGAdmin.tsx @@ -2,7 +2,7 @@ import React, { useState, useEffect, useRef } from "react"; import { Button } from "@/components/ui/button"; import { Input } from "@/components/ui/input"; import { TagInput, TypeHint } from "@/components/ui/tag-input"; -import { Database, Loader2, RefreshCw, Upload, Wrench } from "lucide-react"; +import { Database, Loader2, RefreshCw, Upload, Wrench, FileCode, List } from "lucide-react"; import { pauseIdleTimer, resumeIdleTimer, pingIdleTimer } from "@/hooks/useIdleTimeout"; import { Dialog, @@ -19,6 +19,7 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { useConfirm } from "@/hooks/useConfirm"; import { useAlert } from "@/hooks/useAlert"; import { resolveUploadConflicts } from "@/utils/uploadConflicts"; @@ -35,6 +36,22 @@ const INPUT_CLIP_FIX: React.CSSProperties = { lineHeight: "1.5", }; +type QueryDraft = { name: string; description: string; gsql: string }; +type ListedQuery = { function_header: string; description: string; docstring?: string }; + +const DESCRIPTION_TEMPLATE = `When to use: +- + +Do not use when: +- + +Example questions: +- `; + +function emptyQueryDraft(): QueryDraft { + return { name: "", description: DESCRIPTION_TEMPLATE, gsql: "" }; +} + /** * Returns a human-readable error string when a graph name violates naming rules, * or null when the name is valid. @@ -68,6 +85,7 @@ const KGAdmin = () => { const [refreshDialogOpen, setRefreshDialogOpen] = useState(false); const [ingestDialogOpen, setIngestDialogOpen] = useState(false); const [migrationDialogOpen, setMigrationDialogOpen] = useState(false); + const [registerDialogOpen, setRegisterDialogOpen] = useState(false); // Migration Assistant state const [migrationGraph, setMigrationGraph] = useState(""); @@ -92,6 +110,19 @@ const KGAdmin = () => { // "" | "regenerate_embeddings" | "regenerate_summaries" — which regen is running const [migrationRegenerating, setMigrationRegenerating] = useState(""); const [migrationMessage, setMigrationMessage] = useState(""); + + // Register Queries state + const [registerGraph, setRegisterGraph] = useState(""); + const [registerMode, setRegisterMode] = useState<"single" | "multiple">("single"); + const [registeredQueries, setRegisteredQueries] = useState([]); + const [installedQueries, setInstalledQueries] = useState([]); + const [queryDrafts, setQueryDrafts] = useState([emptyQueryDraft()]); + const [registerLoading, setRegisterLoading] = useState(false); + const [registerSaving, setRegisterSaving] = useState(false); + const [registerMessage, setRegisterMessage] = useState(""); + const [queryListFilter, setQueryListFilter] = useState(""); + const [registerPage, setRegisterPage] = useState<"registered" | "original">("registered"); + const registerStatusRef = useRef(null); // Reset states when dialogs close const handleInitializeDialogChange = (open: boolean) => { if (!open && isConfirmDialogOpen) { @@ -1478,6 +1509,30 @@ const KGAdmin = () => { + {/* Register Queries Card */} +
+
+
+ +
+

+ Register Queries +

+

+ Paste GSQL to create and install a query, then register it as a GraphRAG tool. Leave GSQL empty to tag an already-installed query. +

+
+
+ +
+
+ {/* Initialize Dialog */} @@ -3095,6 +3150,399 @@ const KGAdmin = () => { + + {/* Register Queries Dialog */} + + e.preventDefault()} + > + + + Register Queries + + + Register installed GSQL as GraphRAG tools, or pick an original query already on the graph. + + + +
+
+ + +
+ + { + if (registerSaving) return; + setQueryListFilter(""); + setRegisterPage(value as "registered" | "original"); + }} + className="w-full" + > + + + + Registered Queries + + + + Original Queries + + + + {registerLoading && ( +
+ + Loading queries… +
+ )} + + +

+ Paste GSQL to create and install a query, then register it as a GraphRAG tool. Leave GSQL empty to tag an already-installed query. +

+ {!registerLoading && registerGraph && ( +
+ setQueryListFilter(e.target.value)} + placeholder="Filter registered queries by name…" + disabled={registerSaving} + className="dark:border-[#3D3D3D] dark:bg-shadeA dark:text-white" + /> + +
+
+ Registered queries ({registeredQueries.length}) +
+

+ Tagged with [GRAPHRAG_TOOL]. The agent can pick these as tools. +

+ {registeredQueries.length === 0 ? ( +

+ None registered yet. Paste GSQL below, or open Original Queries and click Use. +

+ ) : visibleRegisteredQueries.length === 0 ? ( +

+ No registered queries match the filter. +

+ ) : ( +
+ {visibleRegisteredQueries.map((q) => ( +
+
+
+
+ {q.function_header} +
+ + Registered + +
+
+ {q.description || "No description"} +
+
+ +
+ ))} +
+ )} +
+
+ )} + +
+ +
+ + +
+
+ +
+ {visibleDrafts.map((draft, index) => ( +
+ {registerMode === "multiple" && ( +
+
+ Query {index + 1} +
+ {visibleDrafts.length > 1 && ( + + )} +
+ )} +
+ + updateDraft(index, { name: e.target.value })} + placeholder="e.g. my_custom_query" + disabled={registerSaving} + className="dark:border-[#3D3D3D] dark:bg-shadeA dark:text-white" + /> +
+
+ +