mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 14:15:52 +00:00
feat(graphrag): frontend Enable GraphRAG action + status + badge
/api/config exposes graphrag_available. 'Enable GraphRAG' source action (owner/editor + pgvector + not-already-graphrag) -> POST graphrag/enable, confirm + task-poll + 'Building graph...' + summary (EnableGraphRAGModal, unmount-guarded; pure graphragEnableUtils). GraphRAG badge for config.kind==graphrag. Mirrors the wiki convert UX. Unit G6.
This commit is contained in:
1 parent
4e0c792a6d
commit
8ccc291ef5
9 files changed
+615
No files matched your search
@@ -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"] = {
|
||||
|
||||
@@ -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) =>
|
||||
|
||||
@@ -109,6 +109,8 @@ const userService = {
|
||||
apiClient.post(endpoints.USER.CREATE_WIKI, data, token),
|
||||
convertToWiki: (sourceId: string, token: string | null): Promise<Response> =>
|
||||
apiClient.post(endpoints.USER.CONVERT_TO_WIKI(sourceId), {}, token),
|
||||
enableGraphRAG: (sourceId: string, token: string | null): Promise<Response> =>
|
||||
apiClient.post(endpoints.USER.ENABLE_GRAPHRAG(sourceId), {}, token),
|
||||
getTaskStatus: (taskId: string, token: string | null): Promise<Response> =>
|
||||
apiClient.get(endpoints.USER.TASK_STATUS(taskId), token),
|
||||
getWikiPages: (sourceId: string, token: string | null): Promise<Response> =>
|
||||
|
||||
@@ -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}}\"?",
|
||||
|
||||
@@ -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<Phase>('confirm');
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [summary, setSummary] = useState<GraphRAGSummary | null>(null);
|
||||
const pollTimer = useRef<number | null>(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 (
|
||||
<Modal
|
||||
open={modalState === 'ACTIVE'}
|
||||
onOpenChange={(o) => !o && closeModal()}
|
||||
hideTitle
|
||||
title={t('settings.sources.graphrag.enable.title')}
|
||||
size="md"
|
||||
mobileVariant="sheet"
|
||||
className="max-w-[480px]"
|
||||
isPerformingTask={phase === 'building'}
|
||||
>
|
||||
<div className="flex flex-col gap-5 px-1 py-1">
|
||||
<h2 className="text-foreground text-xl font-semibold">
|
||||
{t('settings.sources.graphrag.enable.title')}
|
||||
</h2>
|
||||
|
||||
{phase === 'confirm' && (
|
||||
<>
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{document?.name
|
||||
? t('settings.sources.graphrag.enable.intro', {
|
||||
name: document.name,
|
||||
})
|
||||
: t('settings.sources.graphrag.enable.introGeneric')}
|
||||
</p>
|
||||
<ul className="text-muted-foreground list-disc space-y-1.5 pl-5 text-sm">
|
||||
<li>{t('settings.sources.graphrag.enable.costExtraction')}</li>
|
||||
<li>{t('settings.sources.graphrag.enable.costCost')}</li>
|
||||
<li>{t('settings.sources.graphrag.enable.costPgvector')}</li>
|
||||
</ul>
|
||||
<div className="flex flex-col-reverse gap-3 sm:flex-row sm:justify-end">
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
onClick={closeModal}
|
||||
className="w-full rounded-3xl px-6 sm:w-auto"
|
||||
>
|
||||
{t('cancel')}
|
||||
</Button>
|
||||
<Button
|
||||
type="button"
|
||||
onClick={handleEnable}
|
||||
className="w-full rounded-3xl px-6 sm:w-auto"
|
||||
>
|
||||
{t('settings.sources.graphrag.enable.confirm')}
|
||||
</Button>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
||||
{phase === 'building' && (
|
||||
<div className="flex flex-col items-center gap-3 py-6">
|
||||
<Spinner size="medium" />
|
||||
<p className="text-muted-foreground text-sm">
|
||||
{t('settings.sources.graphrag.enable.inProgress')}
|
||||
</p>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{phase === 'summary' && summary && (
|
||||
<>
|
||||
<div className="rounded-xl bg-green-50 p-4 text-sm text-green-800 dark:bg-green-900/30 dark:text-green-200">
|
||||
{[
|
||||
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(' · ')}
|
||||
</div>
|
||||
<div className="flex justify-end">
|
||||
<Button
|
||||
type="button"
|
||||
onClick={closeModal}
|
||||
className="w-full rounded-3xl px-6 sm:w-auto"
|
||||
>
|
||||
{t('settings.sources.graphrag.enable.done')}
|
||||
</Button>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
|
||||
{phase === 'error' && (
|
||||
<>
|
||||
<div className="rounded-xl bg-red-50 p-4 text-sm text-red-700 dark:bg-red-900/40 dark:text-red-300">
|
||||
{error}
|
||||
</div>
|
||||
<div className="flex justify-end">
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
onClick={closeModal}
|
||||
className="w-full rounded-3xl px-6 sm:w-auto"
|
||||
>
|
||||
{t('cancel')}
|
||||
</Button>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
</Modal>
|
||||
);
|
||||
}
|
||||
@@ -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<Doc | null>(null);
|
||||
const [convertModalState, setConvertModalState] =
|
||||
useState<ActiveState>('INACTIVE');
|
||||
const [documentToGraphRAG, setDocumentToGraphRAG] = useState<Doc | null>(
|
||||
null,
|
||||
);
|
||||
const [graphRAGModalState, setGraphRAGModalState] =
|
||||
useState<ActiveState>('INACTIVE');
|
||||
const [graphRAGAvailable, setGraphRAGAvailable] = useState<boolean>(false);
|
||||
const [buildingGraphIds, setBuildingGraphIds] = useState<Set<string>>(
|
||||
() => 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 ? (
|
||||
<div className="mt-8 flex flex-col">
|
||||
{documentToView.config?.kind === 'wiki' ||
|
||||
@@ -712,6 +760,18 @@ export default function Sources({
|
||||
{t('settings.sources.ingestProcessing')}
|
||||
</span>
|
||||
)}
|
||||
{document.config?.kind === 'graphrag' && (
|
||||
<span className="bg-muted-foreground/10 text-muted-foreground flex items-center gap-1 rounded-full px-2 py-0.5 text-xs leading-[16px] font-medium">
|
||||
<Network
|
||||
size={11}
|
||||
strokeWidth={2}
|
||||
aria-hidden="true"
|
||||
/>
|
||||
{document.id && buildingGraphIds.has(document.id)
|
||||
? t('settings.sources.graphrag.building')
|
||||
: t('settings.sources.graphrag.badge')}
|
||||
</span>
|
||||
)}
|
||||
<div className="flex items-center gap-2">
|
||||
<img
|
||||
src={CalendarIcon}
|
||||
@@ -821,6 +881,32 @@ export default function Sources({
|
||||
document={documentToConvert}
|
||||
onConverted={() => refreshDocs(undefined, currentPage, rowsPerPage)}
|
||||
/>
|
||||
|
||||
<EnableGraphRAGModal
|
||||
modalState={graphRAGModalState}
|
||||
setModalState={(state) => {
|
||||
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);
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -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' });
|
||||
});
|
||||
});
|
||||
@@ -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<Response>;
|
||||
getTaskStatus: (taskId: string, token: string | null) => Promise<Response>;
|
||||
}
|
||||
|
||||
export async function startGraphRAG(
|
||||
service: Pick<GraphRAGService, 'enableGraphRAG'>,
|
||||
sourceId: string,
|
||||
token: string | null,
|
||||
): Promise<GraphRAGStart> {
|
||||
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<string, unknown>;
|
||||
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<string, unknown>)?.message as string | undefined);
|
||||
return { status: 'failed', message };
|
||||
}
|
||||
return { status: 'pending' };
|
||||
}
|
||||
|
||||
export async function pollTaskOnce(
|
||||
service: Pick<GraphRAGService, 'getTaskStatus'>,
|
||||
taskId: string,
|
||||
token: string | null,
|
||||
): Promise<GraphRAGPoll> {
|
||||
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);
|
||||
}
|
||||
@@ -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:
|
||||
|
||||
Reference in new issue
Block a user