adk-typescript
Version:
TypeScript port of Google's Agent Development Kit (ADK)
189 lines (188 loc) • 6.63 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.getAgentGraph = getAgentGraph;
const BaseAgent_1 = require("../agents/BaseAgent");
const LlmAgent_1 = require("../agents/LlmAgent");
const tools_1 = require("../tools");
// Graphviz is an optional dependency - we'll try to load it but handle the case where it's not installed
let graphviz;
try {
// eslint-disable-next-line @typescript-eslint/no-var-requires
graphviz = require('graphviz');
}
catch (e) {
console.log('Graphviz module not found. Agent graph visualization will not be available.');
console.log('To install graphviz, run: npm install graphviz');
}
// Try to load retrieval tool module - handle case when it's not available
let BaseRetrievalTool;
let retrievalToolModuleLoaded = false;
try {
// eslint-disable-next-line @typescript-eslint/no-var-requires
BaseRetrievalTool = require('../tools/retrieval/BaseRetrievalTool').BaseRetrievalTool;
retrievalToolModuleLoaded = true;
}
catch (e) {
// Retrieval tool module is optional
retrievalToolModuleLoaded = false;
}
/**
* Build a graph representation of the agent and its sub-agents/tools
*
* @param graph The graphviz graph object
* @param agent The agent to build the graph for
* @param highlightPairs Optional pairs of node names to highlight
*/
function buildGraph(graph, agent, highlightPairs) {
const darkGreen = '#0F5223';
const lightGreen = '#69CB87';
const lightGray = '#cccccc';
/**
* Get the name of a node (agent or tool)
*/
function getNodeName(toolOrAgent) {
if (toolOrAgent instanceof BaseAgent_1.BaseAgent) {
return toolOrAgent.name;
}
else if (toolOrAgent instanceof tools_1.BaseTool) {
return toolOrAgent.name;
}
else {
throw new Error(`Unsupported tool type: ${toolOrAgent}`);
}
}
/**
* Get the caption for a node (with emoji)
*/
function getNodeCaption(toolOrAgent) {
if (toolOrAgent instanceof BaseAgent_1.BaseAgent) {
return '🤖 ' + toolOrAgent.name;
}
else if (retrievalToolModuleLoaded && toolOrAgent instanceof BaseRetrievalTool) {
return '🔎 ' + toolOrAgent.name;
}
else if (toolOrAgent instanceof tools_1.FunctionTool) {
return '🔧 ' + toolOrAgent.name;
}
else if (toolOrAgent instanceof tools_1.AgentTool) {
return '🤖 ' + toolOrAgent.name;
}
else if (toolOrAgent instanceof tools_1.BaseTool) {
return '🔧 ' + toolOrAgent.name;
}
else {
console.warn('Unsupported tool, type:', typeof toolOrAgent, 'obj:', toolOrAgent);
return `❓ Unsupported tool type: ${typeof toolOrAgent}`;
}
}
/**
* Get the shape for a node based on its type
*/
function getNodeShape(toolOrAgent) {
if (toolOrAgent instanceof BaseAgent_1.BaseAgent) {
return 'ellipse';
}
else if (retrievalToolModuleLoaded && toolOrAgent instanceof BaseRetrievalTool) {
return 'cylinder';
}
else if (toolOrAgent instanceof tools_1.FunctionTool) {
return 'box';
}
else if (toolOrAgent instanceof tools_1.BaseTool) {
return 'box';
}
else {
console.warn('Unsupported tool, type:', typeof toolOrAgent, 'obj:', toolOrAgent);
return 'cylinder';
}
}
/**
* Draw a node in the graph
*/
function drawNode(toolOrAgent) {
const name = getNodeName(toolOrAgent);
const shape = getNodeShape(toolOrAgent);
const caption = getNodeCaption(toolOrAgent);
if (highlightPairs) {
for (const highlightTuple of highlightPairs) {
if (highlightTuple.includes(name)) {
graph.addNode(name, {
label: caption,
style: 'filled,rounded',
fillcolor: darkGreen,
color: darkGreen,
shape: shape,
fontcolor: lightGray,
});
return;
}
}
}
// If not highlighted, draw a normal node
graph.addNode(name, {
label: caption,
shape: shape,
style: 'rounded',
color: lightGray,
fontcolor: lightGray,
});
}
/**
* Draw an edge between nodes
*/
function drawEdge(fromName, toName) {
if (highlightPairs) {
for (const [highlightFrom, highlightTo] of highlightPairs) {
if (fromName === highlightFrom && toName === highlightTo) {
graph.addEdge(fromName, toName, { color: lightGreen });
return;
}
else if (fromName === highlightTo && toName === highlightFrom) {
graph.addEdge(fromName, toName, { color: lightGreen, dir: 'back' });
return;
}
}
}
// If not highlighted, draw a normal edge
graph.addEdge(fromName, toName, { arrowhead: 'none', color: lightGray });
}
// Draw the agent node
drawNode(agent);
// Draw sub-agents
for (const subAgent of agent.subAgents) {
buildGraph(graph, subAgent, highlightPairs);
drawEdge(agent.name, subAgent.name);
}
// Draw tools if it's an LLM agent
if (agent instanceof LlmAgent_1.LlmAgent && agent.canonicalTools) {
for (const tool of agent.canonicalTools) {
drawNode(tool);
drawEdge(agent.name, getNodeName(tool));
}
}
}
/**
* Generate an agent graph visualization
*
* @param rootAgent The root agent to visualize
* @param highlightPairs Optional pairs of node names to highlight
* @param asImage Whether to return the graph as an image (PNG) or as a graphviz object
* @returns The graph as PNG binary data or as a graphviz object
*/
function getAgentGraph(rootAgent, highlightPairs, asImage = false) {
if (!graphviz) {
throw new Error('Graphviz module not found. Please install it with: npm install graphviz');
}
console.log('Building graph...');
const graph = graphviz.digraph('G');
// Set graph attributes
graph.set('rankdir', 'LR');
graph.set('bgcolor', '#333537');
buildGraph(graph, rootAgent, highlightPairs);
if (asImage) {
return graph.output('png');
}
else {
return graph;
}
}