diff --git a/application/app.py b/application/app.py index ad9f0140..a1f5f4e9 100644 --- a/application/app.py +++ b/application/app.py @@ -126,9 +126,12 @@ def health(): @app.route("/api/config") def get_config(): + from application.graphrag import graphrag_available + response = { "auth_type": settings.AUTH_TYPE, "requires_auth": settings.AUTH_TYPE in ["simple_jwt", "session_jwt", "oidc"], + "graphrag_available": graphrag_available(), } if settings.AUTH_TYPE == "oidc": response["oidc"] = { diff --git a/frontend/src/api/endpoints.ts b/frontend/src/api/endpoints.ts index 141584ff..7cbbd07a 100644 --- a/frontend/src/api/endpoints.ts +++ b/frontend/src/api/endpoints.ts @@ -58,6 +58,7 @@ const endpoints = { SOURCE_CONFIG: (id: string) => `/api/sources/${id}/config`, CREATE_WIKI: '/api/sources/wiki', CONVERT_TO_WIKI: (id: string) => `/api/sources/${id}/wiki/convert`, + ENABLE_GRAPHRAG: (id: string) => `/api/sources/${id}/graphrag/enable`, TASK_STATUS: (taskId: string) => `/api/task_status?task_id=${taskId}`, WIKI_PAGES: (id: string) => `/api/sources/${id}/wiki/pages`, WIKI_PAGE: (id: string, path: string) => diff --git a/frontend/src/api/services/userService.ts b/frontend/src/api/services/userService.ts index f2ec6dd6..121669a3 100644 --- a/frontend/src/api/services/userService.ts +++ b/frontend/src/api/services/userService.ts @@ -109,6 +109,8 @@ const userService = { apiClient.post(endpoints.USER.CREATE_WIKI, data, token), convertToWiki: (sourceId: string, token: string | null): Promise => apiClient.post(endpoints.USER.CONVERT_TO_WIKI(sourceId), {}, token), + enableGraphRAG: (sourceId: string, token: string | null): Promise => + apiClient.post(endpoints.USER.ENABLE_GRAPHRAG(sourceId), {}, token), getTaskStatus: (taskId: string, token: string | null): Promise => apiClient.get(endpoints.USER.TASK_STATUS(taskId), token), getWikiPages: (sourceId: string, token: string | null): Promise => diff --git a/frontend/src/locale/en.json b/frontend/src/locale/en.json index bae7bbd8..eb489da6 100644 --- a/frontend/src/locale/en.json +++ b/frontend/src/locale/en.json @@ -149,6 +149,33 @@ } } }, + "graphrag": { + "badge": "GraphRAG", + "building": "Building graph…", + "enable": { + "action": "Enable GraphRAG", + "title": "Enable GraphRAG", + "intro": "Build a knowledge graph over \"{{name}}\" so agents can reason across linked entities and relationships.", + "introGeneric": "Build a knowledge graph over this source so agents can reason across linked entities and relationships.", + "costExtraction": "Entities and relationships are extracted by running an LLM over this source's chunks.", + "costCost": "This makes one LLM call per chunk, so it consumes tokens and may take a while on large sources.", + "costPgvector": "GraphRAG is available only with the pgvector store.", + "confirm": "Enable GraphRAG", + "inProgress": "Building graph…", + "summaryNodes_one": "{{count}} entity", + "summaryNodes_other": "{{count}} entities", + "summaryEdges_one": "{{count}} relationship", + "summaryEdges_other": "{{count}} relationships", + "summaryChunks_one": "{{count}} chunk processed", + "summaryChunks_other": "{{count}} chunks processed", + "done": "Done", + "errors": { + "generic": "Enabling GraphRAG failed. Please try again.", + "forbidden": "You don't have permission to enable GraphRAG on this source.", + "conflict": "This source is still ingesting. Wait for it to finish before enabling GraphRAG." + } + } + }, "editConfig": "Source settings", "shareWithTeam": "Share with team", "deleteWarning": "Are you sure you want to delete \"{{name}}\"?", diff --git a/frontend/src/settings/EnableGraphRAGModal.tsx b/frontend/src/settings/EnableGraphRAGModal.tsx new file mode 100644 index 00000000..d7694d0c --- /dev/null +++ b/frontend/src/settings/EnableGraphRAGModal.tsx @@ -0,0 +1,242 @@ +import { useEffect, useRef, useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { useSelector } from 'react-redux'; + +import userService from '../api/services/userService'; +import Spinner from '../components/Spinner'; +import { Button } from '../components/ui/button'; +import { Modal } from '../components/ui/modal'; +import { ActiveState, Doc } from '../models/misc'; +import { selectToken } from '../preferences/preferenceSlice'; + +import { + GraphRAGSummary, + pollTaskOnce, + startGraphRAG, +} from './graphragEnableUtils'; + +const POLL_INTERVAL_MS = 2000; + +type Phase = 'confirm' | 'building' | 'summary' | 'error'; + +interface EnableGraphRAGModalProps { + modalState: ActiveState; + setModalState: (state: ActiveState) => void; + document: Doc | null; + onEnabled: () => void; +} + +export default function EnableGraphRAGModal({ + modalState, + setModalState, + document, + onEnabled, +}: EnableGraphRAGModalProps) { + const { t } = useTranslation(); + const token = useSelector(selectToken); + + const [phase, setPhase] = useState('confirm'); + const [error, setError] = useState(null); + const [summary, setSummary] = useState(null); + const pollTimer = useRef(null); + const isActiveRef = useRef(true); + + const clearPoll = () => { + if (pollTimer.current !== null) { + window.clearTimeout(pollTimer.current); + pollTimer.current = null; + } + }; + + const stopPolling = () => { + isActiveRef.current = false; + clearPoll(); + }; + + useEffect(() => { + if (modalState === 'ACTIVE') { + isActiveRef.current = true; + setPhase('confirm'); + setError(null); + setSummary(null); + } + return stopPolling; + }, [modalState, document]); + + const closeModal = () => { + stopPolling(); + setModalState('INACTIVE'); + }; + + const fail = (message?: string) => { + setError(message ?? t('settings.sources.graphrag.enable.errors.generic')); + setPhase('error'); + }; + + const poll = (taskId: string) => { + pollTaskOnce(userService, taskId, token) + .then((result) => { + if (!isActiveRef.current) return; + if (result.status === 'pending') { + pollTimer.current = window.setTimeout( + () => poll(taskId), + POLL_INTERVAL_MS, + ); + return; + } + if (result.status === 'done') { + setSummary(result.summary); + setPhase('summary'); + onEnabled(); + return; + } + fail(result.message); + }) + .catch(() => { + if (!isActiveRef.current) return; + fail(); + }); + }; + + const handleEnable = async () => { + if (!document?.id) return; + setPhase('building'); + setError(null); + const start = await startGraphRAG(userService, document.id, token); + if (start.status === 'enabled') { + setSummary({ + nodes: 0, + edges: 0, + chunksProcessed: 0, + skippedOverCap: 0, + failedChunks: 0, + }); + setPhase('summary'); + onEnabled(); + return; + } + if (start.status === 'task') { + onEnabled(); + poll(start.taskId); + return; + } + if (start.status === 'forbidden') { + fail(t('settings.sources.graphrag.enable.errors.forbidden')); + return; + } + if (start.status === 'conflict') { + fail( + start.message ?? t('settings.sources.graphrag.enable.errors.conflict'), + ); + return; + } + fail(start.message); + }; + + return ( + !o && closeModal()} + hideTitle + title={t('settings.sources.graphrag.enable.title')} + size="md" + mobileVariant="sheet" + className="max-w-[480px]" + isPerformingTask={phase === 'building'} + > +
+

+ {t('settings.sources.graphrag.enable.title')} +

+ + {phase === 'confirm' && ( + <> +

+ {document?.name + ? t('settings.sources.graphrag.enable.intro', { + name: document.name, + }) + : t('settings.sources.graphrag.enable.introGeneric')} +

+
    +
  • {t('settings.sources.graphrag.enable.costExtraction')}
    • +
  • {t('settings.sources.graphrag.enable.costCost')}
    • +
  • {t('settings.sources.graphrag.enable.costPgvector')}
    • +
+
+ + +
+ + )} + + {phase === 'building' && ( +
+ +

+ {t('settings.sources.graphrag.enable.inProgress')} +

+
+ )} + + {phase === 'summary' && summary && ( + <> +
+ {[ + t('settings.sources.graphrag.enable.summaryNodes', { + count: summary.nodes, + }), + t('settings.sources.graphrag.enable.summaryEdges', { + count: summary.edges, + }), + t('settings.sources.graphrag.enable.summaryChunks', { + count: summary.chunksProcessed, + }), + ].join(' · ')} +
+
+ +
+ + )} + + {phase === 'error' && ( + <> +
+ {error} +
+
+ +
+ + )} +
+ + ); +} diff --git a/frontend/src/settings/Sources.tsx b/frontend/src/settings/Sources.tsx index 8cb86932..cdd01e8b 100644 --- a/frontend/src/settings/Sources.tsx +++ b/frontend/src/settings/Sources.tsx @@ -1,5 +1,6 @@ import { BookOpen, + Network, Search as SearchIcon, SlidersHorizontal, Users, @@ -52,6 +53,7 @@ import ConnectorTree from '../components/ConnectorTree'; import Chunks from '../components/Chunks'; import WikiViewer from '../components/WikiViewer'; import ConvertToWikiModal from './ConvertToWikiModal'; +import EnableGraphRAGModal from './EnableGraphRAGModal'; import SourceConfigModal from './SourceConfigModal'; type SourceMenuOption = { @@ -120,6 +122,15 @@ export default function Sources({ const [documentToConvert, setDocumentToConvert] = useState(null); const [convertModalState, setConvertModalState] = useState('INACTIVE'); + const [documentToGraphRAG, setDocumentToGraphRAG] = useState( + null, + ); + const [graphRAGModalState, setGraphRAGModalState] = + useState('INACTIVE'); + const [graphRAGAvailable, setGraphRAGAvailable] = useState(false); + const [buildingGraphIds, setBuildingGraphIds] = useState>( + () => new Set(), + ); const [syncMenuState, setSyncMenuState] = useState<{ isOpen: boolean; docId: string | null; @@ -338,6 +349,7 @@ export default function Sources({ document: Doc, ): SourceMenuOption[] => { const isWiki = document.config?.kind === 'wiki' || document.type === 'wiki'; + const isGraphRAG = document.config?.kind === 'graphrag'; // 'team' viewers cannot write; convert is owner/editor only. const canEdit = document.ownership !== 'team' || document.team_access === 'editor'; @@ -430,6 +442,28 @@ export default function Sources({ }); } + if ( + document.id && + !isWiki && + !isGraphRAG && + canEdit && + graphRAGAvailable && + document.ingestStatus !== 'processing' && + document.ingestStatus !== 'failed' + ) { + actions.push({ + icon: Network, + label: t('settings.sources.graphrag.enable.action'), + onClick: () => { + setDocumentToGraphRAG(document); + setGraphRAGModalState('ACTIVE'); + }, + iconWidth: 16, + iconHeight: 16, + variant: 'default', + }); + } + // Sharing is an owner-only action: hide it for sources shared into the // user's workspace by a team. if (document.ownership !== 'team' && document.id) { @@ -462,6 +496,20 @@ export default function Sources({ refreshDocs(undefined, 1, rowsPerPage); }, [debouncedSearchTerm]); + useEffect(() => { + let cancelled = false; + userService + .getConfig() + .then((response) => response.json()) + .then((config) => { + if (!cancelled) setGraphRAGAvailable(!!config?.graphrag_available); + }) + .catch(() => undefined); + return () => { + cancelled = true; + }; + }, []); + return documentToView ? (
{documentToView.config?.kind === 'wiki' || @@ -712,6 +760,18 @@ export default function Sources({ {t('settings.sources.ingestProcessing')} )} + {document.config?.kind === 'graphrag' && ( + + + )}
refreshDocs(undefined, currentPage, rowsPerPage)} /> + + { + setGraphRAGModalState(state); + if (state === 'INACTIVE') { + const closedId = documentToGraphRAG?.id; + if (closedId) { + setBuildingGraphIds((prev) => { + const next = new Set(prev); + next.delete(closedId); + return next; + }); + } + setDocumentToGraphRAG(null); + } + }} + document={documentToGraphRAG} + onEnabled={() => { + const enabledId = documentToGraphRAG?.id; + if (enabledId) { + setBuildingGraphIds((prev) => new Set(prev).add(enabledId)); + } + refreshDocs(undefined, currentPage, rowsPerPage); + }} + />
); } diff --git a/frontend/src/settings/graphragEnableUtils.test.ts b/frontend/src/settings/graphragEnableUtils.test.ts new file mode 100644 index 00000000..f597d23e --- /dev/null +++ b/frontend/src/settings/graphragEnableUtils.test.ts @@ -0,0 +1,143 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { + interpretTaskStatus, + pollTaskOnce, + startGraphRAG, +} from './graphragEnableUtils'; + +const jsonResponse = (status: number, body: unknown): Response => + ({ + ok: status >= 200 && status < 300, + status, + json: () => Promise.resolve(body), + }) as unknown as Response; + +describe('startGraphRAG', () => { + it('returns the task id when extraction is enqueued', async () => { + const enableGraphRAG = vi + .fn() + .mockResolvedValue(jsonResponse(200, { success: true, task_id: 't-9' })); + + const result = await startGraphRAG({ enableGraphRAG }, 'src-1', 'tok'); + + expect(enableGraphRAG).toHaveBeenCalledWith('src-1', 'tok'); + expect(result).toEqual({ status: 'task', taskId: 't-9' }); + }); + + it('returns enabled when success has no task id', async () => { + const enableGraphRAG = vi + .fn() + .mockResolvedValue(jsonResponse(200, { success: true })); + + const result = await startGraphRAG({ enableGraphRAG }, 'src-1', 'tok'); + + expect(result).toEqual({ status: 'enabled' }); + }); + + it('maps 403 to forbidden', async () => { + const enableGraphRAG = vi.fn().mockResolvedValue(jsonResponse(403, {})); + const result = await startGraphRAG({ enableGraphRAG }, 'src-1', 'tok'); + expect(result).toEqual({ status: 'forbidden' }); + }); + + it('maps 409 to conflict with the server message', async () => { + const enableGraphRAG = vi + .fn() + .mockResolvedValue(jsonResponse(409, { message: 'still ingesting' })); + const result = await startGraphRAG({ enableGraphRAG }, 'src-1', 'tok'); + expect(result).toEqual({ status: 'conflict', message: 'still ingesting' }); + }); + + it('maps a 400 unavailable to an error with the message', async () => { + const enableGraphRAG = vi + .fn() + .mockResolvedValue( + jsonResponse(400, { success: false, message: 'pgvector required' }), + ); + const result = await startGraphRAG({ enableGraphRAG }, 'src-1', 'tok'); + expect(result).toEqual({ status: 'error', message: 'pgvector required' }); + }); + + it('maps a network throw to a generic error', async () => { + const enableGraphRAG = vi.fn().mockRejectedValue(new Error('offline')); + const result = await startGraphRAG({ enableGraphRAG }, 'src-1', 'tok'); + expect(result).toEqual({ status: 'error' }); + }); +}); + +describe('interpretTaskStatus', () => { + it('parses a SUCCESS summary (nodes + edges + chunks)', () => { + const result = interpretTaskStatus('SUCCESS', { + nodes: 12, + edges: 7, + chunks_processed: 30, + skipped_over_cap: 2, + failed_chunks: 1, + }); + expect(result).toEqual({ + status: 'done', + summary: { + nodes: 12, + edges: 7, + chunksProcessed: 30, + skippedOverCap: 2, + failedChunks: 1, + }, + }); + }); + + it('defaults missing summary fields', () => { + const result = interpretTaskStatus('SUCCESS', {}); + expect(result).toEqual({ + status: 'done', + summary: { + nodes: 0, + edges: 0, + chunksProcessed: 0, + skippedOverCap: 0, + failedChunks: 0, + }, + }); + }); + + it('maps FAILURE to failed', () => { + expect(interpretTaskStatus('FAILURE', 'boom')).toEqual({ + status: 'failed', + message: 'boom', + }); + }); + + it('treats other statuses as pending', () => { + expect(interpretTaskStatus('PENDING', null)).toEqual({ status: 'pending' }); + expect(interpretTaskStatus('STARTED', null)).toEqual({ status: 'pending' }); + }); +}); + +describe('pollTaskOnce', () => { + it('returns the parsed summary on success', async () => { + const getTaskStatus = vi.fn().mockResolvedValue( + jsonResponse(200, { + status: 'SUCCESS', + result: { nodes: 3, edges: 1, chunks_processed: 5 }, + }), + ); + const result = await pollTaskOnce({ getTaskStatus }, 't-1', 'tok'); + expect(result).toEqual({ + status: 'done', + summary: { + nodes: 3, + edges: 1, + chunksProcessed: 5, + skippedOverCap: 0, + failedChunks: 0, + }, + }); + }); + + it('returns pending on a non-ok response', async () => { + const getTaskStatus = vi.fn().mockResolvedValue(jsonResponse(503, {})); + const result = await pollTaskOnce({ getTaskStatus }, 't-1', 'tok'); + expect(result).toEqual({ status: 'pending' }); + }); +}); diff --git a/frontend/src/settings/graphragEnableUtils.ts b/frontend/src/settings/graphragEnableUtils.ts new file mode 100644 index 00000000..f0daab44 --- /dev/null +++ b/frontend/src/settings/graphragEnableUtils.ts @@ -0,0 +1,89 @@ +export interface GraphRAGSummary { + nodes: number; + edges: number; + chunksProcessed: number; + skippedOverCap: number; + failedChunks: number; +} + +export type GraphRAGStart = + | { status: 'enabled' } + | { status: 'task'; taskId: string } + | { status: 'forbidden' } + | { status: 'conflict'; message?: string } + | { status: 'error'; message?: string }; + +export type GraphRAGPoll = + | { status: 'pending' } + | { status: 'done'; summary: GraphRAGSummary } + | { status: 'failed'; message?: string }; + +interface GraphRAGService { + enableGraphRAG: (sourceId: string, token: string | null) => Promise; + getTaskStatus: (taskId: string, token: string | null) => Promise; +} + +export async function startGraphRAG( + service: Pick, + sourceId: string, + token: string | null, +): Promise { + let response: Response; + try { + response = await service.enableGraphRAG(sourceId, token); + } catch { + return { status: 'error' }; + } + const data = await response.json().catch(() => ({})); + if (response.ok && data?.success) { + if (data.task_id) return { status: 'task', taskId: data.task_id }; + return { status: 'enabled' }; + } + if (response.status === 403) return { status: 'forbidden' }; + if (response.status === 409) + return { status: 'conflict', message: data?.message }; + return { status: 'error', message: data?.message }; +} + +function parseSummary(result: unknown): GraphRAGSummary { + const r = (result || {}) as Record; + return { + nodes: Number(r.nodes) || 0, + edges: Number(r.edges) || 0, + chunksProcessed: Number(r.chunks_processed) || 0, + skippedOverCap: Number(r.skipped_over_cap) || 0, + failedChunks: Number(r.failed_chunks) || 0, + }; +} + +export function interpretTaskStatus( + status: string | undefined, + result: unknown, +): GraphRAGPoll { + if (status === 'SUCCESS') + return { status: 'done', summary: parseSummary(result) }; + if (status === 'FAILURE') { + const message = + typeof result === 'string' + ? result + : ((result as Record)?.message as string | undefined); + return { status: 'failed', message }; + } + return { status: 'pending' }; +} + +export async function pollTaskOnce( + service: Pick, + taskId: string, + token: string | null, +): Promise { + let response: Response; + try { + response = await service.getTaskStatus(taskId, token); + } catch { + return { status: 'pending' }; + } + if (!response.ok) return { status: 'pending' }; + const data = await response.json().catch(() => ({})); + return interpretTaskStatus(data?.status, data?.result); +} diff --git a/tests/test_app_routes.py b/tests/test_app_routes.py index bef3300f..7c2fb0ab 100644 --- a/tests/test_app_routes.py +++ b/tests/test_app_routes.py @@ -53,6 +53,28 @@ class TestConfigRoute: assert "requires_auth" in data assert "oidc" not in data + @pytest.mark.unit + def test_exposes_graphrag_available(self, client): + with patch("application.app.settings") as mock_settings, patch( + "application.graphrag.graphrag_available", return_value=True + ): + mock_settings.AUTH_TYPE = None + response = client.get("/api/config") + assert response.status_code == 200 + data = json.loads(response.data) + assert data["graphrag_available"] is True + + @pytest.mark.unit + def test_graphrag_unavailable_when_flag_off(self, client): + with patch("application.app.settings") as mock_settings, patch( + "application.graphrag.graphrag_available", return_value=False + ): + mock_settings.AUTH_TYPE = None + response = client.get("/api/config") + assert response.status_code == 200 + data = json.loads(response.data) + assert data["graphrag_available"] is False + @pytest.mark.unit def test_oidc_config_exposes_login_paths(self, client): with patch("application.app.settings") as mock_settings: