New graph visualizer

backend/app/api/admin_routes/knowledge_base/graph/routes.py

@@ -243,3 +243,20 @@ def legacy_search_graph(session: SessionDep, kb_id: int, request: GraphSearchReq
     except Exception as e:
         # TODO: throw InternalServerError
         raise e
+
+@router.post("/admin/knowledge_bases/{kb_id}/graph/entire_graph")
+def get_entire_knowledge_graph(session: SessionDep, kb_id: int):
+    try:
+        kb = knowledge_base_repo.must_get(session, kb_id)
+        graph_store = get_kb_tidb_graph_store(session, kb)
+        retrieved_kg = graph_store.get_entire_knowledge_graph()
+        return {
+            "entities": retrieved_kg.entities,
+            "relationships": retrieved_kg.relationships,
+        }
+
+    except KBNotFound as e:
+        raise e
+    except Exception as e:
+        # TODO: throw InternalServerError
+        raise e

backend/app/rag/indices/knowledge_graph/graph_store/tidb_graph_store.py

@@ -312,7 +312,7 @@ def create_relationship(
             self._session.flush()
 
     def get_subgraph_by_relationship_ids(
-        self, ids: list[int], **kwargs
+            self, ids: list[int], **kwargs
     ) -> RetrievedKnowledgeGraph:
         stmt = (
             select(self._relationship_model)
@@ -1103,3 +1103,60 @@ def get_chunks_by_relationships(
             }
             for chunk in chunks
         ]
+
+    def get_entire_knowledge_graph(self) -> RetrievedKnowledgeGraph:
+        """Retrieve all entities and relationships from the knowledge graph store.
+
+        Returns:
+            RetrievedKnowledgeGraph containing all entities and relationships
+        """
+        # Query all entities
+        entity_query = select(self._entity_model).order_by(self._entity_model.id)
+        db_entities = self._session.exec(entity_query).all()
+
+        # Query all relationships with their related entities
+        relationship_query = (
+            select(self._relationship_model)
+            .options(
+                joinedload(self._relationship_model.source_entity),
+                joinedload(self._relationship_model.target_entity),
+            )
+            .order_by(self._relationship_model.id)
+        )
+        db_relationships = self._session.exec(relationship_query).all()
+
+        # Convert entities to RetrievedEntity objects
+        entities = []
+        for entity in db_entities:
+            entities.append(
+                RetrievedEntity(
+                    id=entity.id,
+                    knowledge_base_id=self.knowledge_base.id,
+                    name=entity.name,
+                    description=entity.description,
+                    meta=entity.meta,
+                    entity_type=entity.entity_type,
+                )
+            )
+
+        # Convert relationships to RetrievedRelationship objects
+        relationships = []
+        for rel in db_relationships:
+            relationships.append(
+                RetrievedRelationship(
+                    id=rel.id,
+                    knowledge_base_id=self.knowledge_base.id,
+                    source_entity_id=rel.source_entity_id,
+                    target_entity_id=rel.target_entity_id,
+                    description=rel.description,
+                    rag_description=f"{rel.source_entity.name} -> {rel.description} -> {rel.target_entity.name}",
+                    meta=rel.meta,
+                    weight=rel.weight,
+                    last_modified_at=rel.last_modified_at,
+                )
+            )
+        return RetrievedKnowledgeGraph(
+            knowledge_base=self.knowledge_base.to_descriptor(),
+            entities=entities,
+            relationships=relationships,
+        )

docker-compose.dev.yml

@@ -64,4 +64,4 @@ services:
     image: redis:6.0.16 
     volumes:
       - ./redis-data:/data
-    command: ["redis-server", "--loglevel", "warning"]
+    command: ["redis-server", "--loglevel", "warning"]

frontend/app/package.json

@@ -49,6 +49,7 @@
     "cmdk": "^1.0.4",
     "date-fns": "^3.6.0",
     "embla-carousel-react": "^8.5.1",
+    "force-graph": "^1.49.6",
     "input-otp": "^1.4.1",
     "js-cookie": "^3.0.5",
     "lucide-react": "^0.400.0",

frontend/app/src/api/graph.ts

@@ -91,6 +91,24 @@ export interface GraphSearchParams {
   with_degree?: boolean;
 }
 
+export interface KBRetrieveKnowledgeGraphParams {
+  query: string;
+  llm_id: number;
+  retrieval_config: {
+    knowledge_graph: KnowledgeGraphRetrievalConfig;
+  }
+}
+
+export interface KnowledgeGraphRetrievalConfig {
+  depth?: number;
+  include_meta?: boolean;
+  with_degree?: boolean;
+  metadata_filter?: {
+    enabled?: boolean;
+    filters?: Record<string, any>
+  }
+}
+
 export async function search (kbId: number, params: GraphSearchParams) {
   return await fetch(requestUrl(`/api/v1/admin/knowledge_bases/${kbId}/graph/search`), {
     method: 'post',
@@ -151,6 +169,18 @@ export async function getEntitySubgraph (kbId: number, id: number) {
     .then(handleResponse(knowledgeGraphSchema));
 }
 
+export async function getEntireKnowledgeGraph (kbId: number, params: KBRetrieveKnowledgeGraphParams) {
+  return await fetch(requestUrl(`/api/v1/admin/knowledge_bases/${kbId}/graph/entire_graph`), {
+    method: 'post',
+    headers: {
+      ...await authenticationHeaders(),
+      'Content-Type': 'application/json',
+    },
+    body: JSON.stringify(params),
+  })
+  .then(handleResponse(knowledgeGraphSchema));
+}
+
 export async function getRelationship (kbId: number, id: number) {
   return await fetch(requestUrl(`/api/v1/admin/knowledge_bases/${kbId}/graph/relationships/${id}`), {
     headers: {

frontend/app/src/components/graph/GraphEditor.tsx

@@ -1,7 +1,7 @@
 'use client';
 
 import { getChatMessageSubgraph } from '@/api/chats';
-import { getEntitySubgraph, type KnowledgeGraph, search } from '@/api/graph';
+import { getEntitySubgraph, getEntireKnowledgeGraph, type KnowledgeGraph, search } from '@/api/graph';
 import { LinkDetails } from '@/components/graph/components/LinkDetails';
 import { NetworkViewer, type NetworkViewerDetailsProps } from '@/components/graph/components/NetworkViewer';
 import { NodeDetails } from '@/components/graph/components/NodeDetails';
@@ -21,6 +21,7 @@ import useSWR from 'swr';
 
 export function GraphEditor ({ knowledgeBaseId }: { knowledgeBaseId: number }) {
   const [query, setQuery] = useSearchParam('query', 'sample-question:What is TiDB?');
+  const [graphStyle, setGraphStyle] = useState<'new' | 'legacy'>('new');
 
   const [key, fetcher] = getFetchInfo(knowledgeBaseId, query);
 
@@ -33,18 +34,20 @@ export function GraphEditor ({ knowledgeBaseId }: { knowledgeBaseId: number }) {
   return (
     <div className="p-4 space-y-4">
       <SubgraphSelector knowledgeBaseId={knowledgeBaseId} query={query} onQueryChange={setQuery} />
+      <GraphStyleSelector style={graphStyle} onStyleChange={setGraphStyle} />
       {(error != null) && <Alert variant="destructive">
         <AlertTitle>Failed to fetch subgraph</AlertTitle>
         <AlertDescription>{getErrorMessage(error)}</AlertDescription>
       </Alert>}
       <div className="w-full flex gap-4">
         <div className="flex-1">
           <NetworkViewer
-            key={query}
+            key={`${query}-${graphStyle}`}
             className="border rounded h-auto aspect-square"
             loading={isLoading}
             loadingTitle={'Loading knowledge graph...'}
             network={network}
+            useCanvasRenderer={graphStyle === 'new'}
             Details={(props) => (
               ref.current && createPortal(
                 <Editor
@@ -90,6 +93,9 @@ function SubgraphSelector ({ knowledgeBaseId, query, onQueryChange }: { knowledg
       <Select value={type} onValueChange={type => {
         setType(type);
         setInput('');
+        if (type === 'entire-knowledge-graph') {
+          onQueryChange(`${type}:`);
+        }
       }}>
         <SelectTrigger className="w-max">
           <SelectValue />
@@ -100,25 +106,46 @@ function SubgraphSelector ({ knowledgeBaseId, query, onQueryChange }: { knowledg
           <SelectItem value="message-subgraph">Message Subgraph</SelectItem>
           <SelectItem value="trace" disabled>Langfuse Trace ID (UUID)</SelectItem>
           <SelectItem value="document" disabled>Document URI</SelectItem>
+          <SelectItem value="entire-knowledge-graph">Entire Knowledge Graph</SelectItem>
         </SelectContent>
       </Select>
-      <Input
-        className="flex-1"
-        value={input}
-        onChange={event => setInput(event.target.value)}
-        onKeyDown={event => {
-          if (isHotkey('Enter', event)) {
-            onQueryChange(`${type}:${input}`);
-          }
-        }}
-      />
+      {type !== 'entire-knowledge-graph' && (
+        <>
+          <Input
+            className="flex-1"
+            value={input}
+            onChange={event => setInput(event.target.value)}
+            onKeyDown={event => {
+              if (isHotkey('Enter', event)) {
+                onQueryChange(`${type}:${input}`);
+              }
+            }}
+          />
+        </>
+      )}
       <Link className={buttonVariants({})} href={`/knowledge-bases/${knowledgeBaseId}/knowledge-graph-explorer/create-synopsis-entity`}>
         Create Synopsis Entity
       </Link>
     </div>
   );
 }
 
+function GraphStyleSelector ({ style, onStyleChange }: { style: 'new' | 'legacy', onStyleChange: (style: 'new' | 'legacy') => void }) {
+  return (
+    <div>
+      <Select value={style} onValueChange={(value) => onStyleChange(value as 'new' | 'legacy')}>
+        <SelectTrigger className="w-max">
+          <SelectValue />
+        </SelectTrigger>
+        <SelectContent>
+          <SelectItem value="legacy">Legacy (SVG)</SelectItem>
+          <SelectItem value="new">New (Canvas)</SelectItem>
+        </SelectContent>
+      </Select>
+    </div>
+  )
+}
+
 function Editor ({ knowledgeBaseId, network, target, onTargetChange, onEnterSubgraph }: NetworkViewerDetailsProps & { knowledgeBaseId: number, onEnterSubgraph: (type: string, entityId: IdType) => void }) {
   if (target) {
     if (target.type === 'link') {
@@ -145,6 +172,18 @@ function getFetchInfo (kbId: number, query: string | null): [string | false, ()
 
   const param = parsedQuery[1];
 
+  const entireKnowledgeGraphParams = {
+    query: "",
+    llm_id: 1,
+    retrieval_config: {
+      knowledge_graph: {
+        depth: 20,
+        include_meta: true,
+        with_degree: true
+      }
+    }
+  }
+
   switch (parsedQuery[0]) {
     // case 'trace':
     //   return ['get', `/api/v1/traces/${parsedQuery[1]}/knowledge-graph-retrieval`];
@@ -156,6 +195,8 @@ function getFetchInfo (kbId: number, query: string | null): [string | false, ()
       return [`api.knowledge-bases.${kbId}.graph.search?query=${param}`, () => search(kbId, { query: param })];
     case 'message-subgraph':
       return [`api.chats.get-message-subgraph?id=${param}`, () => getChatMessageSubgraph(parseInt(param))];
+    case 'entire-knowledge-graph':
+      return [`api.knowledge-bases.${kbId}.graph.entire-knowledge-graph`, () => getEntireKnowledgeGraph(kbId, entireKnowledgeGraphParams)];
   }
 
   return [false, () => Promise.reject()];

frontend/app/src/components/graph/components/NetworkCanvas.tsx

@@ -1,28 +1,39 @@
 import { useEffect, useRef, useState } from 'react';
 import type { IdType, NetworkLink, NetworkNode, ReadonlyNetwork } from '../network/Network';
 import { NetworkRenderer, type NetworkRendererOptions } from '../network/NetworkRenderer';
+import { CanvasNetworkRenderer } from '../network/CanvasNetworkRenderer';
 
 export interface NetworkCanvasProps<Node extends NetworkNode, Link extends NetworkLink> extends NetworkRendererOptions<Node, Link> {
   network: ReadonlyNetwork<Node, Link>;
   target: { type: string, id: IdType } | undefined;
   className?: string;
+  useCanvasRenderer?: boolean;
 }
 
-export function NetworkCanvas<Node extends NetworkNode, Link extends NetworkLink> ({ className, network, target, ...options }: NetworkCanvasProps<Node, Link>) {
+export function NetworkCanvas<Node extends NetworkNode, Link extends NetworkLink> ({ className, network, target, useCanvasRenderer = false, ...options }: NetworkCanvasProps<Node, Link>) {
   const ref = useRef<HTMLDivElement>(null);
-  const [renderer, setRenderer] = useState<NetworkRenderer<Node, Link>>();
+  const [renderer, setRenderer] = useState<NetworkRenderer<Node, Link> | CanvasNetworkRenderer<Node, Link>>();
 
   useEffect(() => {
-    const renderer = new NetworkRenderer(network, options);
+    // Cleanup previous renderer if it exists (needed for renderer switching)
+    if (renderer) {
+      renderer.unmount();
+    }
+
+    const newRenderer = useCanvasRenderer 
+      ? new CanvasNetworkRenderer(network, options)
+      : new NetworkRenderer(network, options);
+
     if (ref.current) {
-      renderer.mount(ref.current);
+      newRenderer.mount(ref.current);
     }
-    setRenderer(renderer);
+    setRenderer(newRenderer);
 
     return () => {
-      renderer.unmount();
+      newRenderer.unmount();
+      setRenderer(undefined);
     };
-  }, [network]);
+  }, [network, useCanvasRenderer]);
 
   useEffect(() => {
     if (!renderer) {
@@ -41,7 +52,7 @@ export function NetworkCanvas<Node extends NetworkNode, Link extends NetworkLink
         renderer.focusLink(target.id);
         return () => renderer.blurLink();
     }
-  }, [target]);
+  }, [target, renderer]);
 
   return (
     <div className={className} ref={ref} />

frontend/app/src/components/graph/components/NetworkViewer.tsx

@@ -13,6 +13,7 @@ export interface NetworkViewerProps {
   loading: boolean;
   loadingTitle: ReactNode;
   Details: FC<NetworkViewerDetailsProps>;
+  useCanvasRenderer?: boolean;
 }
 
 export interface NetworkViewerDetailsProps {
@@ -31,7 +32,7 @@ function randomPosition (radius: number, kbSpacing: number, kbIndex: number, kbC
   };
 }
 
-export function NetworkViewer ({ network, loading, loadingTitle, className, Details }: NetworkViewerProps) {
+export function NetworkViewer ({ network, loading, loadingTitle, className, Details, useCanvasRenderer = false }: NetworkViewerProps) {
   const [target, setTarget] = useState<{ type: string, id: IdType }>();
 
   const knowledgeGraphIndexMap = useMemo(() => {
@@ -145,6 +146,7 @@ export function NetworkViewer ({ network, loading, loadingTitle, className, Deta
           className={cn('w-full h-full overflow-hidden')}
           network={network}
           target={target}
+          useCanvasRenderer={useCanvasRenderer}
           {...networkOptions}
         />
         <Details