feat: support document's subgraph

src/app/api/v1/indexes/[name]/chunks/[uri]/import_state/route.ts

@@ -0,0 +1,27 @@
+import { getIndexByName } from '@/core/repositories/index_';
+import { KnowledgeGraphClient } from '@/lib/knowledge-graph/client';
+import { defineHandler } from '@/lib/next/handler';
+import { notFound } from 'next/navigation';
+import z from 'zod';
+
+export const GET = defineHandler({
+  params: z.object({
+    name: z.string(),
+    uri: z.string(),
+  }),
+  auth: 'admin',
+}, async ({ params }) => {
+  const index = await getIndexByName(params.name);
+
+  if (!index) {
+    notFound();
+  }
+
+  if (index.config.provider !== 'knowledge-graph') {
+    notFound();
+  }
+
+  return new KnowledgeGraphClient().getChunkImportState(params.uri);
+});
+
+export const dynamic = 'force-dynamic';

src/app/api/v1/indexes/[name]/chunks/[uri]/subgraph/route.ts

@@ -0,0 +1,27 @@
+import { getIndexByName } from '@/core/repositories/index_';
+import { KnowledgeGraphClient } from '@/lib/knowledge-graph/client';
+import { defineHandler } from '@/lib/next/handler';
+import { notFound } from 'next/navigation';
+import z from 'zod';
+
+export const GET = defineHandler({
+  params: z.object({
+    name: z.string(),
+    uri: z.string(),
+  }),
+  auth: 'admin',
+}, async ({ params }) => {
+  const index = await getIndexByName(params.name);
+
+  if (!index) {
+    notFound();
+  }
+
+  if (index.config.provider !== 'knowledge-graph') {
+    notFound();
+  }
+
+  return new KnowledgeGraphClient().getChunkSubgraph(params.uri);
+});
+
+export const dynamic = 'force-dynamic';

src/components/graph/GraphEditor.tsx

@@ -2,6 +2,7 @@ import { type ServerGraphData } from '@/components/graph/api';
 import { LinkDetails } from '@/components/graph/components/LinkDetails';
 import { NetworkViewer, type NetworkViewerDetailsProps } from '@/components/graph/components/NetworkViewer';
 import { NodeDetails } from '@/components/graph/components/NodeDetails';
+import type { IdType } from '@/components/graph/network/Network';
 import { useNetwork } from '@/components/graph/useNetwork';
 import { Alert, AlertDescription, AlertTitle } from '@/components/ui/alert';
 import { Input } from '@/components/ui/input';
@@ -10,7 +11,7 @@ import { useSearchParam } from '@/components/use-search-param';
 import { getErrorMessage } from '@/lib/errors';
 import { fetcher } from '@/lib/fetch';
 import isHotkey from 'is-hotkey';
-import { useRef, useState } from 'react';
+import { useEffect, useRef, useState } from 'react';
 import { createPortal } from 'react-dom';
 import useSWR from 'swr';
 
@@ -39,14 +40,27 @@ export function GraphEditor ({}: {}) {
             loadingTitle={'Loading knowledge graph...'}
             network={network}
             Details={(props) => (
-              ref.current && createPortal(<Editor {...props} onEnterEntitySubgraph={id => {
-                props.onTargetChange(undefined);
-                setQuery(`entity:${id}`);
-              }} />, ref.current!)
+              ref.current && createPortal(
+                <Editor
+                  {...props}
+                  onEnterSubgraph={(type, id) => {
+                    props.onTargetChange(undefined);
+                    switch (type) {
+                      case `entity`:
+                        setQuery(`entity:${id}`);
+                        break;
+                      case 'document':
+                        setQuery(`document:${id}`);
+                        break;
+                    }
+                  }}
+                />,
+                ref.current!,
+              )
             )}
           />
         </div>
-        <div className="w-96 flex-shrink-0 relative" ref={ref} />
+        <div className="w-96 flex-shrink-0 relative" style={{ padding: '0.1px' }} ref={ref} />
       </div>
     </div>
   );
@@ -58,6 +72,12 @@ function SubgraphSelector ({ query, onQueryChange }: { query: string | null, onQ
   const [type, setType] = useState<string>(initialType);
   const [input, setInput] = useState<string>(initialInput);
 
+  useEffect(() => {
+    const [type = 'sample-question', input = 'What is TiDB?'] = parseQuery(query) ?? [];
+    setType(type);
+    setInput(input);
+  }, [query]);
+
   return (
     <div className="flex gap-2">
       <Select value={type} onValueChange={type => {
@@ -70,6 +90,7 @@ function SubgraphSelector ({ query, onQueryChange }: { query: string | null, onQ
         <SelectContent>
           <SelectItem value="trace">Langfuse Trace ID (UUID)</SelectItem>
           <SelectItem value="entity">Entity ID</SelectItem>
+          <SelectItem value="document">Document URI</SelectItem>
           <SelectItem value="sample-question">Sample Question</SelectItem>
         </SelectContent>
         <Input
@@ -87,16 +108,16 @@ function SubgraphSelector ({ query, onQueryChange }: { query: string | null, onQ
   );
 }
 
-function Editor ({ network, target, onTargetChange, onEnterEntitySubgraph }: NetworkViewerDetailsProps & { onEnterEntitySubgraph: (entityId: number) => void }) {
+function Editor ({ network, target, onTargetChange, onEnterSubgraph }: NetworkViewerDetailsProps & { onEnterSubgraph: (type: string, entityId: IdType) => void }) {
   if (target) {
     if (target.type === 'link') {
-      return <LinkDetails relationship={network.link(target.id)!} onClickTarget={onTargetChange} />;
+      return <LinkDetails relationship={network.link(target.id)!} onClickTarget={onTargetChange} onEnterSubgraph={onEnterSubgraph} />;
     } else if (target.type === 'node') {
-      return <NodeDetails entity={network.node(target.id)!} onClickTarget={onTargetChange} onEnterEntitySubgraph={onEnterEntitySubgraph} />;
+      return <NodeDetails entity={network.node(target.id)!} onClickTarget={onTargetChange} onEnterSubgraph={onEnterSubgraph} />;
     }
   }
 
-  return <div className='flex items-center justify-center h-40 text-sm text-muted-foreground font-bold'>
+  return <div className="flex items-center justify-center h-40 text-sm text-muted-foreground font-bold">
     Select an entity or relationship
   </div>;
 }
@@ -114,6 +135,8 @@ function getFetchUrl (query: string | null): ['get', string] | null {
   switch (parsedQuery[0]) {
     case 'trace':
       return ['get', `/api/v1/traces/${parsedQuery[1]}/knowledge-graph-retrieval`];
+    case 'document':
+      return ['get', `/api/v1/indexes/graph/chunks/${encodeURIComponent(parsedQuery[1])}/subgraph`];
     case 'entity':
       return ['get', `/api/v1/indexes/graph/entities/${parsedQuery[1]}/subgraph`];
     case 'sample-question':

src/components/graph/components/EditingButton.tsx

@@ -3,14 +3,15 @@ import { ExternalLinkIcon, Loader2Icon, PencilIcon } from 'lucide-react';
 
 export interface EditingButtonProps {
   editing: boolean;
-  onEnterEntitySubgraph?: () => void;
+  onEnterSubgraph?: () => void;
+  subGraphTitle?: string;
   onStartEdit: () => void;
   onSave: () => void;
   onReset: () => void;
   busy: boolean;
 }
 
-export function EditingButton ({ onEnterEntitySubgraph, editing, onStartEdit, onReset, onSave, busy }: EditingButtonProps) {
+export function EditingButton ({ onEnterSubgraph, subGraphTitle = 'Subgraph', editing, onStartEdit, onReset, onSave, busy }: EditingButtonProps) {
   return editing
     ? (
       <div className="flex gap-2 items-center">
@@ -36,9 +37,9 @@ export function EditingButton ({ onEnterEntitySubgraph, editing, onStartEdit, on
           <PencilIcon className="w-3 h-3 mr-2" />
           Edit
         </Button>
-        {onEnterEntitySubgraph && <Button size="sm" variant="secondary" onClick={onEnterEntitySubgraph}>
+        {onEnterSubgraph && <Button size="sm" variant="secondary" onClick={onEnterSubgraph}>
           <ExternalLinkIcon className="w-3 h-3 mr-2" />
-          Subgraph
+          {subGraphTitle}
         </Button>}
       </div>
     );

src/components/graph/components/LinkDetails.tsx

@@ -17,10 +17,12 @@ export function LinkDetails ({
   relationship,
   onClickTarget,
   onUpdate,
+  onEnterSubgraph,
 }: {
   relationship: Relationship,
   onClickTarget?: (target: { type: string, id: IdType }) => void;
   onUpdate?: (newRelationship: Relationship) => void;
+  onEnterSubgraph: (type: string, entityId: IdType) => void
 }) {
   const network = useContext(NetworkContext);
 
@@ -69,8 +71,14 @@ export function LinkDetails ({
         <span className="text-sm text-muted-foreground font-normal ">
           <b>#{relationship.id}</b> relationship
         </span>
-        <EditingButton editing={editing} onStartEdit={() => setEditing(true)} onSave={handleSave} onReset={handleReset} busy={busy} />
+        <EditingButton editing={editing} onStartEdit={() => setEditing(true)} onSave={handleSave} onReset={handleReset} busy={busy} onEnterSubgraph={() => onEnterSubgraph('document', relationship.meta.doc_id)} subGraphTitle='Document subgraph' />
       </div>
+      <section>
+        <h6 className="text-xs font-bold text-accent-foreground mb-1">Document URI</h6>
+        <p className="block w-full text-xs text-accent-foreground">
+          <a className="underline" href={relationship.meta.doc_id} target="_blank">{relationship.meta.doc_id}</a>
+        </p>
+      </section>
       <TextareaField label="Description" ref={dirtyRelationship.descriptionRef} defaultValue={relationship.description} disabled={controlsDisabled} />
       <InputField label="Weight" ref={dirtyRelationship.weightRef} defaultValue={relationship.weight} disabled={controlsDisabled} min={0} step={1} type="number" />
       <JsonField label="meta" ref={dirtyRelationship.metaRef} defaultValue={relationship.meta} disabled={controlsDisabled} />

src/components/graph/components/NodeDetails.tsx

@@ -17,12 +17,12 @@ export function NodeDetails ({
   entity,
   onClickTarget,
   onUpdate,
-  onEnterEntitySubgraph,
+  onEnterSubgraph,
 }: {
   entity: Entity,
   onClickTarget?: (target: { type: string, id: IdType }) => void;
   onUpdate?: (newData: Entity) => void
-  onEnterEntitySubgraph: (id: number) => void
+  onEnterSubgraph: (type: string, entityId: IdType) => void
 }) {
   const [editing, setEditing] = useState(false);
   const network = useContext(NetworkContext);
@@ -69,7 +69,7 @@ export function NodeDetails ({
         <span className="text-sm text-muted-foreground font-normal ">
           <b>#{entity.id}</b> {entity.entity_type} entity
         </span>
-        <EditingButton onEnterEntitySubgraph={() => onEnterEntitySubgraph(entity.id as number)} editing={editing} onStartEdit={() => setEditing(true)} onSave={handleSave} onReset={handleReset} busy={busy} />
+        <EditingButton onEnterSubgraph={() => onEnterSubgraph('entity', entity.id)} editing={editing} onStartEdit={() => setEditing(true)} onSave={handleSave} onReset={handleReset} busy={busy} />
       </div>
       <InputField label="Name" ref={dirtyEntity.nameRef} defaultValue={entity.name} disabled={controlsDisabled} />
       <TextareaField label="Description" ref={dirtyEntity.descriptionRef} defaultValue={entity.description} disabled={controlsDisabled} />

src/lib/knowledge-graph/client.ts

@@ -174,4 +174,33 @@ export class KnowledgeGraphClient {
     };
   }
 
+  async getChunkSubgraph (uri: string) {
+    const url = `${this.baseURL}/api/graph/chunks/${encodeURIComponent(uri)}/subgraph`;
+    const res = await fetch(url, {
+      method: 'GET',
+      headers: {
+        'Accept': 'application/json',
+      },
+    }).then(handleErrors).then(res => res.json());
+
+    return res as {
+      entities: Entity[]
+      relationships: Relationship[]
+    };
+  }
+
+  async getChunkImportState (uri: string) {
+    const url = `${this.baseURL}/api/import/check`;
+    const res = await fetch(url, {
+      method: 'POST',
+      body: JSON.stringify({
+        uri,
+      }),
+      headers: {
+        'Content-Type': 'application/json',
+      },
+    }).then(handleErrors).then(res => res.json());
+
+    return res;
+  }
 }