diff --git a/backend/app/api/admin_routes/knowledge_base/graph/routes.py b/backend/app/api/admin_routes/knowledge_base/graph/routes.py index 95a2bbf4f..d9c6e1daa 100644 --- a/backend/app/api/admin_routes/knowledge_base/graph/routes.py +++ b/backend/app/api/admin_routes/knowledge_base/graph/routes.py @@ -1,7 +1,10 @@ import logging from typing import List +import json from fastapi import APIRouter, HTTPException, status +from fastapi.responses import StreamingResponse +from fastapi.encoders import jsonable_encoder from app.api.admin_routes.knowledge_base.graph.models import ( SynopsisEntityCreate, @@ -259,4 +262,31 @@ def get_entire_knowledge_graph(session: SessionDep, kb_id: int): raise e except Exception as e: # TODO: throw InternalServerError - raise e \ No newline at end of file + raise e + +@router.get("/admin/knowledge_bases/{kb_id}/graph/entire_graph/stream") +def stream_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) + + def generate(): + for chunk in graph_store.stream_entire_knowledge_graph(chunk_size=5000): + yield f"data: {json.dumps(jsonable_encoder(chunk))}\n\n" + yield f"data: {json.dumps({'type': 'complete'})}\n\n" + + return StreamingResponse( + generate(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "Access-Control-Allow-Origin": "*", + } + ) + + except KBNotFound as e: + raise e + except Exception as e: + logger.exception(e) + raise InternalServerError() \ No newline at end of file diff --git a/backend/app/rag/indices/knowledge_graph/graph_store/tidb_graph_store.py b/backend/app/rag/indices/knowledge_graph/graph_store/tidb_graph_store.py index ef7bcda3b..450e70e5a 100644 --- a/backend/app/rag/indices/knowledge_graph/graph_store/tidb_graph_store.py +++ b/backend/app/rag/indices/knowledge_graph/graph_store/tidb_graph_store.py @@ -11,7 +11,7 @@ from llama_index.embeddings.openai import OpenAIEmbedding, OpenAIEmbeddingModelType import sqlalchemy from sqlmodel import Session, asc, func, select, text, SQLModel -from sqlalchemy.orm import aliased, defer, joinedload +from sqlalchemy.orm import aliased, defer, joinedload, noload from tidb_vector.sqlalchemy import VectorAdaptor from sqlalchemy import or_, desc @@ -1182,3 +1182,91 @@ def get_entire_knowledge_graph(self) -> RetrievedKnowledgeGraph: entities=entities, relationships=relationships, ) + + def stream_entire_knowledge_graph(self, chunk_size: int = 5000): + """Stream entire knowledge graph in chunks + + Args: + chunk_size: Number of entities/relationships per chunk + + Yields: + Dict containing chunk type and data + """ + # Stream entities + entity_query = ( + select(self._entity_model) + .options( + defer(self._entity_model.description_vec), + defer(self._entity_model.meta_vec), + ) + .order_by(self._entity_model.id) + ) + last_entity_id = 0 + + while True: + chunk_query = entity_query.where( + self._entity_model.id > last_entity_id + ).limit(chunk_size) + db_entities = self._session.exec(chunk_query).all() + + if not db_entities: + break + + 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, + ) + ) + + last_entity_id = db_entities[-1].id + yield {"type": "entities", "data": entities} + + # Stream relationships + relationship_query = ( + select(self._relationship_model) + .options( + defer(self._relationship_model.description_vec), + defer(self._relationship_model.chunk_id), + noload(self._relationship_model.source_entity), + noload(self._relationship_model.target_entity), + ) + .order_by(self._relationship_model.id) + ) + logger.info(f"Relationship query: {relationship_query}") + last_relationship_id = 0 + + while True: + chunk_query = relationship_query.where( + self._relationship_model.id > last_relationship_id + ).limit(chunk_size) + logger.info(f"Executing relationship chunk query: {chunk_query}") + db_relationships = self._session.exec(chunk_query).all() + + if not db_relationships: + break + + 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=None, # Skip rag_description for streaming performance + meta=rel.meta, + weight=rel.weight, + last_modified_at=rel.last_modified_at, + ) + ) + + last_relationship_id = db_relationships[-1].id + yield {"type": "relationships", "data": relationships} diff --git a/frontend/app/src/api/graph.ts b/frontend/app/src/api/graph.ts index d36bd9030..c53228d30 100644 --- a/frontend/app/src/api/graph.ts +++ b/frontend/app/src/api/graph.ts @@ -1,5 +1,6 @@ import { authenticationHeaders, handleResponse, requestUrl } from '@/lib/request'; import { zodJsonDate } from '@/lib/zod'; +import { bufferedReadableStreamTransformer } from '@/lib/buffered-readable-stream'; import { z, type ZodType } from 'zod'; export interface KnowledgeGraph { @@ -181,6 +182,60 @@ export async function getEntireKnowledgeGraph (kbId: number, params: KBRetrieveK .then(handleResponse(knowledgeGraphSchema)); } +export async function streamEntireKnowledgeGraph (kbId: number): Promise { + const entities: KnowledgeGraphEntity[] = []; + const relationships: KnowledgeGraphRelationship[] = []; + + const response = await fetch(requestUrl(`/api/v1/admin/knowledge_bases/${kbId}/graph/entire_graph/stream`), { + method: 'GET', + headers: await authenticationHeaders(), + credentials: 'include', + }); + + if (!response.ok) { + throw new Error(`${response.status} ${response.statusText}`); + } + + if (!response.body) { + throw new Error('Empty response body'); + } + + const reader = response.body.pipeThrough(bufferedReadableStreamTransformer()).getReader(); + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + + if (value.trim() && value.startsWith('data: ')) { + const dataStr = value.substring(6).trim(); + if (dataStr) { + try { + const data = JSON.parse(dataStr); + + if (data.type === 'entities') { + entities.push(...data.data); + // console.log(`Received ${data.data.length} entities, total: ${entities.length}`); + } else if (data.type === 'relationships') { + relationships.push(...data.data); + // console.log(`Received ${data.data.length} relationships, total: ${relationships.length}`); + } else if (data.type === 'complete') { + // console.log(`Streaming complete. Final counts - entities: ${entities.length}, relationships: ${relationships.length}`); + return { entities, relationships }; + } + } catch (error) { + console.warn('Failed to parse streaming data:', error, 'Data:', dataStr); + } + } + } + } + } finally { + reader.releaseLock(); + } + + return { entities, relationships }; +} + export async function getRelationship (kbId: number, id: number) { return await fetch(requestUrl(`/api/v1/admin/knowledge_bases/${kbId}/graph/relationships/${id}`), { headers: { diff --git a/frontend/app/src/components/graph/GraphEditor.tsx b/frontend/app/src/components/graph/GraphEditor.tsx index 225aec66b..2649528cb 100644 --- a/frontend/app/src/components/graph/GraphEditor.tsx +++ b/frontend/app/src/components/graph/GraphEditor.tsx @@ -1,7 +1,7 @@ 'use client'; import { getChatMessageSubgraph } from '@/api/chats'; -import { getEntitySubgraph, getEntireKnowledgeGraph, type KnowledgeGraph, search } from '@/api/graph'; +import { getEntitySubgraph, streamEntireKnowledgeGraph, 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'; @@ -93,6 +93,9 @@ function SubgraphSelector ({ knowledgeBaseId, query, onQueryChange }: { knowledg - setInput(event.target.value)} - onKeyDown={event => { - if (isHotkey('Enter', event)) { - onQueryChange(`${type}:${input}`); - } - }} - /> + {type !== 'entire-knowledge-graph' && ( + <> + setInput(event.target.value)} + onKeyDown={event => { + if (isHotkey('Enter', event)) { + onQueryChange(`${type}:${input}`); + } + }} + /> + + )} Create Synopsis Entity @@ -164,6 +172,7 @@ function getFetchInfo (kbId: number, query: string | null): [string | false, () const param = parsedQuery[1]; + switch (parsedQuery[0]) { // case 'trace': // return ['get', `/api/v1/traces/${parsedQuery[1]}/knowledge-graph-retrieval`]; @@ -175,6 +184,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`, () => streamEntireKnowledgeGraph(kbId)]; } return [false, () => Promise.reject()]; diff --git a/frontend/app/src/components/graph/network/CanvasNetworkRenderer.ts b/frontend/app/src/components/graph/network/CanvasNetworkRenderer.ts index 93f00a388..58abb2fb2 100644 --- a/frontend/app/src/components/graph/network/CanvasNetworkRenderer.ts +++ b/frontend/app/src/components/graph/network/CanvasNetworkRenderer.ts @@ -8,7 +8,6 @@ export class CanvasNetworkRenderer void) | undefined; private _onUpdateNode: ((id: IdType) => void) | undefined; @@ -21,34 +20,37 @@ export class CanvasNetworkRenderer(); private highlightedLinks = new Set(); - private readonly linkDefaultDistance = 30; private readonly chargeDefaultStrength = -80; private readonly linkHighlightDistance = 120; private readonly chargeHighlightStrength = -200; private readonly linkDefaultWidth = 1; - // Clustering - private clusterMode = 'enabled'; private clustersCalculated = false; + + private adjacencyMap = new Map, connectedLinks: Set }>(); + private adjacencyCalculated = false; + + scale = 1; + private initialLayoutComplete = false; + private viewportBounds = { x0: -Infinity, y0: -Infinity, x1: Infinity, y1: Infinity }; - // Colors private colors = { textColor: '#000000', - nodeColor: '#1f77b4', nodeHighlighted: '#18a0b1', nodeSelected: '#72fefb', - linkColor: '#999999', + linkDefaultColor: '#999999', linkHighlighted: '#18a0b1', - linkSelected: '#72fefb', - clusterColors: [ - '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', - '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf' - ] + linkSelected: '#72fefb' }; - - scale = 1; + private zoomLevels = { + one: 0.1, + two: 0.2, + three: 0.3, + four: 0.4, + five: 0.8, + } constructor( private network: ReadonlyNetwork, @@ -60,14 +62,22 @@ export class CanvasNetworkRenderer) { const nodeMap = new Map(); this.nodes = this.network.nodes().map((node, index) => { + const nodeRadius = 8; + const fontSize = Math.max(8, nodeRadius * 0.3); + const label = options.getNodeLabel?.(node) ?? (node as any).name ?? node.id; + const labelColor = options.getNodeLabelColor?.(node) ?? this.colors.textColor; + nodeMap.set(node.id, index); return { id: node.id, index, - radius: 8, - label: options.getNodeLabel?.(node), + radius: nodeRadius, + label, details: options.getNodeDetails?.(node), meta: options.getNodeMeta?.(node), + fontSize, + fontString: `${fontSize}px Sans-Serif`, + labelColor, ...options.getNodeInitialAttrs?.(node, index), }; }); @@ -82,8 +92,51 @@ export class CanvasNetworkRenderer 50 || this.links.length > 50; + private updateViewportBounds() { + if (!this._graph || !this._el) return; + + const canvas = this._el.querySelector('canvas'); + if (!canvas) return; + + const width = canvas.width; + const height = canvas.height; + + const topLeft = this._graph.screen2GraphCoords(0, 0); + const bottomRight = this._graph.screen2GraphCoords(width, height); + + const padding = 100 / this.scale; + + this.viewportBounds = { + x0: topLeft.x - padding, + y0: topLeft.y - padding, + x1: bottomRight.x + padding, + y1: bottomRight.y + padding + }; + } + + private isNodeInViewport(node: any): boolean { + const x = node.x ?? 0; + const y = node.y ?? 0; + return x >= this.viewportBounds.x0 && + x <= this.viewportBounds.x1 && + y >= this.viewportBounds.y0 && + y <= this.viewportBounds.y1; + } + + private isLinkInViewport(link: any): boolean { + const sourceX = link.source.x ?? 0; + const sourceY = link.source.y ?? 0; + const targetX = link.target.x ?? 0; + const targetY = link.target.y ?? 0; + + if ((sourceX < this.viewportBounds.x0 && targetX < this.viewportBounds.x0) || + (sourceX > this.viewportBounds.x1 && targetX > this.viewportBounds.x1) || + (sourceY < this.viewportBounds.y0 && targetY < this.viewportBounds.y0) || + (sourceY > this.viewportBounds.y1 && targetY > this.viewportBounds.y1)) { + return false; + } + + return true; } mount(container: HTMLElement) { @@ -94,27 +147,25 @@ export class CanvasNetworkRenderer { - this.drawNodeWithLabel(node, ctx); + .autoPauseRedraw(false) + .warmupTicks(50) + .nodeAutoColorBy('clusterId') + .nodeCanvasObject((node: any, ctx: CanvasRenderingContext2D, globalScale: number) => { + if (this.isNodeInViewport(node)) { + this.drawNodeWithLabel(node, ctx, globalScale); + } }) - .linkWidth(this.linkDefaultWidth) - .linkColor((link: any) => { - if (this.selectedLink && this.selectedLink.id === link.id) { - return this.colors.linkSelected; - } else if (this.highlightedLinks.has(link.id)) { - return this.colors.linkHighlighted; - } else { - return this.colors.linkColor; + .linkCanvasObject((link: any, ctx: CanvasRenderingContext2D, globalScale: number) => { + this.scale = globalScale; + if (this.scale > this.zoomLevels.three && this.isLinkInViewport(link)) { + this.drawLink(link, ctx); } }) - .linkDirectionalArrowLength(6) - .linkDirectionalArrowRelPos(1) - .linkCurvature(0.1) + .linkCanvasObjectMode(() => 'replace') .onNodeClick((node: any, event: MouseEvent) => { this.onNodeClick(node, event); }) @@ -127,7 +178,31 @@ export class CanvasNetworkRenderer d.id).distance(this.linkDefaultDistance)) - .d3Force("charge", d3.forceManyBody().strength(this.chargeDefaultStrength)); + .d3Force("charge", d3.forceManyBody() + .strength(this.chargeDefaultStrength) + .theta(1.2) + ) + .onZoom((transform: any) => { + this.scale = transform.k; + }) + .onRenderFramePre(() => { + this.updateViewportBounds(); + }); + + this._graph = graph; + + setTimeout(() => { + this.initialLayoutComplete = true; + graph.d3Force('x', null); + graph.d3Force('y', null); + graph.d3Force("charge")?.distanceMax(300).strength(0); + + const data = graph.graphData(); + data.nodes.forEach((node: any) => { + node.fx = node.x; + node.fy = node.y; + }); + }, 2000); container.style.overflow = 'hidden'; @@ -139,7 +214,6 @@ export class CanvasNetworkRenderer { const canvas = container.querySelector('canvas'); if (canvas) { @@ -162,8 +236,7 @@ export class CanvasNetworkRenderer= this.zoomLevels.five) { + ctx.font = node.fontString; + ctx.textAlign = 'center'; + ctx.textBaseline = 'middle'; + ctx.fillStyle = node.labelColor; + ctx.fillText(node.label, node.x, node.y + nodeRadius + node.fontSize * 0.7); + } + } + + private drawLink(link: any, ctx: CanvasRenderingContext2D) { + const source = link.source; + const target = link.target; + + // Determine link color + let color = this.colors.linkDefaultColor; + if (this.selectedLink && this.selectedLink.id === link.id) { + color = this.colors.linkSelected; + } else if (this.highlightedLinks.has(link.id)) { + color = this.colors.linkHighlighted; + } + + ctx.strokeStyle = color; + ctx.lineWidth = Math.min(this.linkDefaultWidth * this.scale, 1); + + ctx.beginPath(); + ctx.moveTo(source.x, source.y); + ctx.lineTo(target.x, target.y); + ctx.stroke(); + + if (this.scale > this.zoomLevels.four) { + this.drawArrow(ctx, source, target, color); + } + } + + private drawArrow( + ctx: CanvasRenderingContext2D, + source: any, + target: any, + color: string + ) { + const arrowLength = 6; + const arrowAngle = Math.PI / 6; + + const dx = target.x - source.x; + const dy = target.y - source.y; + const angle = Math.atan2(dy, dx); + + const targetRadius = target.radius || 8; + const arrowX = target.x - Math.cos(angle) * targetRadius; + const arrowY = target.y - Math.sin(angle) * targetRadius; + + ctx.fillStyle = color; + ctx.beginPath(); + ctx.moveTo(arrowX, arrowY); + ctx.lineTo( + arrowX - arrowLength * Math.cos(angle - arrowAngle), + arrowY - arrowLength * Math.sin(angle - arrowAngle) + ); + ctx.lineTo( + arrowX - arrowLength * Math.cos(angle + arrowAngle), + arrowY - arrowLength * Math.sin(angle + arrowAngle) + ); + ctx.closePath(); + ctx.fill(); } private onNodeClick(node: any, event: MouseEvent) { @@ -249,24 +378,27 @@ export class CanvasNetworkRenderer(); - const connectedLinkIds = new Set(); - - this._graph.graphData().links.forEach((link: any) => { - if (link.source.id === node.id) { - connectedNodeIds.add(link.target.id); - connectedLinkIds.add(link.id); - } else if (link.target.id === node.id) { - connectedNodeIds.add(link.source.id); - connectedLinkIds.add(link.id); - } - }); + const adjacency = this.adjacencyMap.get(node.id); + if (!adjacency) { + return; + } + + const connectedNodeIds = adjacency.connectedNodes; + const connectedLinkIds = adjacency.connectedLinks; this.highlightedNodes.clear(); connectedNodeIds.forEach(nodeId => this.highlightedNodes.add(nodeId)); this.highlightedLinks.clear(); connectedLinkIds.forEach(linkId => this.highlightedLinks.add(linkId)); + + const data = this._graph.graphData(); + data.nodes.forEach((n: any) => { + if (connectedNodeIds.has(n.id)) { + n.fx = null; + n.fy = null; + } + }); this._graph.d3Force("link").distance((link: any) => { if (connectedLinkIds.has(link.id)) { @@ -279,7 +411,7 @@ export class CanvasNetworkRenderer n.id === id); - if (node) { - this.selectedNode = node; - this.selectedLink = null; - this.highlightConnections(node); - } - } - - blurNode(): void { - this.clearHighlight(); - } - - focusLink(id: IdType): void { - const link = this.links.find(l => l.id === id); - if (link) { - this.selectedLink = link; - this.selectedNode = null; - this.highlightLink(link); - } - } - - blurLink(): void { - this.clearHighlight(); + if (!this.initialLayoutComplete) return; + this._graph.d3Force("charge").strength(0); + + setTimeout(() => { + const data = this._graph.graphData(); + data.nodes.forEach((n: any) => { + n.fx = n.x; + n.fy = n.y; + }); + }, 500); + + this._graph.d3ReheatSimulation(); } private calculateAndCacheClusters() { @@ -327,11 +445,45 @@ export class CanvasNetworkRenderer { - (node as any).clusterId = clusters.get(node.id) || 0; + const clusterId = clusters.get(node.id) || 0; + (node as any).clusterId = clusterId; }); this.clustersCalculated = true; } + + private calculateAndCacheAdjacency() { + if (!this.nodes || !this.links) { + return; + } + + this.adjacencyMap.clear(); + + this.nodes.forEach(node => { + this.adjacencyMap.set(node.id, { + connectedNodes: new Set(), + connectedLinks: new Set() + }); + }); + + this.links.forEach(link => { + const sourceId = typeof link.source === 'object' ? link.source.id : link.source; + const targetId = typeof link.target === 'object' ? link.target.id : link.target; + + const sourceAdjacency = this.adjacencyMap.get(sourceId); + const targetAdjacency = this.adjacencyMap.get(targetId); + + if (sourceAdjacency && targetAdjacency) { + sourceAdjacency.connectedNodes.add(targetId); + targetAdjacency.connectedNodes.add(sourceId); + + sourceAdjacency.connectedLinks.add(link.id); + targetAdjacency.connectedLinks.add(link.id); + } + }); + + this.adjacencyCalculated = true; + } private findClusters(): Map { const clusters = new Map(); @@ -360,11 +512,11 @@ export class CanvasNetworkRenderer { + for (const neighborId of neighbors) { if (!visited.has(neighborId)) { dfs(neighborId, currentClusterId); } - }); + } }; this.nodes.forEach(node => { @@ -377,14 +529,15 @@ export class CanvasNetworkRenderer