pit-manager/scripts/visualize-chain-stats.ts

Centralized prompt management system for Human Behavior AI agents

#!/usr/bin/env tsx
/**
 * Visualize aggregate chain tracking statistics in ASCII tree format
 * Shows hierarchical agent relationships and average metrics per chain execution
 */

import { createClient } from '@supabase/supabase-js';
import * as dotenv from 'dotenv';
import * as path from 'path';

// Load environment variables
dotenv.config({ path: path.resolve(process.cwd(), '.env') });

interface ChainExecutionGroup {
  id: string;
  group_id: string;
  chain_id: string;
  repo_id: string;
  session_id?: string;
  process_id?: string;
  start_time: string;
  end_time?: string;
  status: string;
  total_executions: number;
  total_tokens: number;
  metadata: any;
}

interface Execution {
  id: string;
  tag: string;
  tokens: any;
  chain_group_id: string;
  parent_execution_id?: string;
}

interface AgentStats {
  callsPerChain: number[];
  tokensPerChain: number[];
}

interface AgentNode {
  name: string;
  calls: number;
  totalTokens: number;
  stats: AgentStats;
  children: Map<string, AgentNode>;
}

class ChainVisualizationAnalyzer {
  private supabaseClient: any;
  private repoId: string;

  constructor() {
    const url = process.env.PIT_SUPABASE_URL;
    const serviceKey = process.env.PIT_SUPABASE_SERVICE_ROLE_KEY;
    const repoKey = process.env.PIT_REPO_KEY;

    if (!url || !serviceKey || !repoKey) {
      throw new Error('Missing required environment variables: PIT_SUPABASE_URL, PIT_SUPABASE_SERVICE_ROLE_KEY, PIT_REPO_KEY');
    }

    this.supabaseClient = createClient(url, serviceKey);
    
    // Extract repo_id from repo_key by validating it
    this.repoId = '';
  }

  async initialize() {
    // Validate repo key to get repo_id
    const repoKey = process.env.PIT_REPO_KEY!;
    const keyHash = this.hashRepoKey(repoKey);
    
    const { data, error } = await this.supabaseClient
      .from('repo_keys')
      .select('repo_id')
      .eq('key_hash', keyHash)
      .single();

    if (error || !data) {
      throw new Error('Invalid repository key');
    }

    this.repoId = data.repo_id;
    console.log('Repository ID:', this.repoId);
  }

  private hashRepoKey(repoKey: string): string {
    const crypto = require('crypto');
    return crypto.createHash('sha256').update(repoKey).digest('hex');
  }

  async analyzeChain(chainId: string) {
    console.log(`\nAnalyzing chain: ${chainId}\n`);
    console.log('=' .repeat(80));

    // Get all chain execution groups for this chain
    const { data: chainGroups, error: groupError } = await this.supabaseClient
      .from('chain_execution_groups')
      .select('*')
      .eq('repo_id', this.repoId)
      .eq('chain_id', chainId);
      // Removed status filter to include all chains for demonstration

    if (groupError || !chainGroups || chainGroups.length === 0) {
      console.log('No chain execution groups found for this chain.');
      return;
    }

    const completedGroups = chainGroups.filter(g => g.status === 'completed');
    const runningGroups = chainGroups.filter(g => g.status === 'running');
    const failedGroups = chainGroups.filter(g => g.status === 'failed');

    console.log(`Found ${chainGroups.length} chain execution(s):`);
    console.log(`  - Completed: ${completedGroups.length}`);
    console.log(`  - Running: ${runningGroups.length}`);
    console.log(`  - Failed: ${failedGroups.length}\n`);

    // Build agent hierarchy
    const rootNode: AgentNode = {
      name: chainId,
      calls: chainGroups.length,
      totalTokens: 0,
      stats: { callsPerChain: [], tokensPerChain: [] },
      children: new Map()
    };

    // Process each chain group separately to track per-chain stats
    const chainResults: Map<string, Map<string, { calls: number, tokens: number }>> = new Map();
    
    for (const group of chainGroups) {
      const chainStats = await this.processChainGroupForStats(group);
      chainResults.set(group.group_id, chainStats);
    }

    // Build logical hierarchy based on code analysis since parent_execution_id is not working yet
    this.buildLogicalHierarchy(rootNode, chainResults);

    // Debug: Show what we collected
    console.log('Agents captured:', Array.from(rootNode.children.keys()).join(', '));
    console.log();


    // Generate and print ASCII tree
    this.printTree(rootNode, chainGroups.length, '', true, true);
  }

  private async processChainGroupForStats(group: ChainExecutionGroup): Promise<Map<string, { calls: number, tokens: number }>> {
    // Get all executions for this chain group
    const { data: executions, error } = await this.supabaseClient
      .from('executions')
      .select('*')
      .eq('chain_group_id', group.id);

    if (error || !executions) {
      console.warn(`Failed to get executions for group ${group.group_id}`);
      return new Map();
    }

    // Count calls and tokens per agent for this chain
    const agentStats = new Map<string, { calls: number, tokens: number }>();

    for (const exec of executions) {
      const agentName = exec.tag || 'unknown-agent';
      
      if (!agentStats.has(agentName)) {
        agentStats.set(agentName, { calls: 0, tokens: 0 });
      }

      const stats = agentStats.get(agentName)!;
      stats.calls++;

      // Add tokens from execution
      const tokens = exec.tokens || {};
      const totalTokens = typeof tokens === 'object' ? (tokens.total || 0) : 0;
      stats.tokens += totalTokens;
    }

    return agentStats;
  }

  private buildAggregatedStats(rootNode: AgentNode, chainResults: Map<string, Map<string, { calls: number, tokens: number }>>) {
    // Get all unique agent names
    const allAgents = new Set<string>();
    chainResults.forEach(chainStats => {
      chainStats.forEach((_, agentName) => allAgents.add(agentName));
    });

    // Build nodes for each agent
    for (const agentName of allAgents) {
      const node: AgentNode = {
        name: agentName,
        calls: 0,
        totalTokens: 0,
        stats: { callsPerChain: [], tokensPerChain: [] },
        children: new Map()
      };

      // Collect stats from each chain
      chainResults.forEach(chainStats => {
        const agentStats = chainStats.get(agentName);
        if (agentStats) {
          node.stats.callsPerChain.push(agentStats.calls);
          node.stats.tokensPerChain.push(agentStats.tokens);
          node.calls += agentStats.calls;
          node.totalTokens += agentStats.tokens;
        } else {
          // Agent didn't run in this chain
          node.stats.callsPerChain.push(0);
          node.stats.tokensPerChain.push(0);
        }
      });

      rootNode.children.set(agentName, node);
    }
  }

  private buildLogicalHierarchy(rootNode: AgentNode, chainResults: Map<string, Map<string, { calls: number, tokens: number }>>) {
    // Based on code analysis, the logical hierarchy should be:
    // line-analysis-agent (multiple chunks) - ROOT level agents
    //   ├── summary-agent (child of last line-analysis chunk)
    //   │   └── title-agent (child of summary-agent)
    //   │       └── product-features-agent (child of title-agent)
    //   ├── create-raw-tickets-agent (child of last line-analysis chunk, multiple instances)
    //   │   └── check-duplicate-agent (child of create-raw-tickets-agent)
    //   ├── notification-filter-agent (child of last line-analysis chunk)
    //   └── slack-summary-agent (child of last line-analysis chunk)

    // Get all unique agent names
    const allAgents = new Set<string>();
    chainResults.forEach(chainStats => {
      chainStats.forEach((_, agentName) => allAgents.add(agentName));
    });

    // Build stats for all agents first
    const agentStats = new Map<string, AgentNode>();
    for (const agentName of allAgents) {
      const node: AgentNode = {
        name: agentName,
        calls: 0,
        totalTokens: 0,
        stats: { callsPerChain: [], tokensPerChain: [] },
        children: new Map()
      };

      // Collect stats from each chain
      chainResults.forEach(chainStats => {
        const agentStat = chainStats.get(agentName);
        if (agentStat) {
          node.stats.callsPerChain.push(agentStat.calls);
          node.stats.tokensPerChain.push(agentStat.tokens);
          node.calls += agentStat.calls;
          node.totalTokens += agentStat.tokens;
        } else {
          // Agent didn't run in this chain
          node.stats.callsPerChain.push(0);
          node.stats.tokensPerChain.push(0);
        }
      });

      agentStats.set(agentName, node);
    }

    // Build logical hierarchy based on code analysis
    const getOrCreateNode = (name: string) => agentStats.get(name) || {
      name,
      calls: 0,
      totalTokens: 0,
      stats: { callsPerChain: [], tokensPerChain: [] },
      children: new Map()
    };

    // Start with line-analysis-agent as root level (multiple chunks)
    if (agentStats.has('line-analysis-agent')) {
      const lineAnalysisNode = getOrCreateNode('line-analysis-agent');
      rootNode.children.set('line-analysis-agent', lineAnalysisNode);
      
      // summary-agent is child of line-analysis-agent
      if (agentStats.has('summary-agent')) {
        const summaryNode = getOrCreateNode('summary-agent');
        lineAnalysisNode.children.set('summary-agent', summaryNode);
        
        // title-agent is child of summary-agent
        if (agentStats.has('title-agent')) {
          const titleNode = getOrCreateNode('title-agent');
          summaryNode.children.set('title-agent', titleNode);
          
          // product-features-agent is child of title-agent
          if (agentStats.has('product-features-agent')) {
            const productFeaturesNode = getOrCreateNode('product-features-agent');
            titleNode.children.set('product-features-agent', productFeaturesNode);
          }
        }
      }
      
      // create-raw-tickets-agent is also child of line-analysis-agent
      if (agentStats.has('create-raw-tickets-agent')) {
        const createTicketsNode = getOrCreateNode('create-raw-tickets-agent');
        lineAnalysisNode.children.set('create-raw-tickets-agent', createTicketsNode);
        
        // check-duplicate-agent is child of create-raw-tickets-agent
        if (agentStats.has('check-duplicate-agent')) {
          const checkDuplicateNode = getOrCreateNode('check-duplicate-agent');
          createTicketsNode.children.set('check-duplicate-agent', checkDuplicateNode);
        }
      }
      
      // notification-filter-agent is child of line-analysis-agent
      if (agentStats.has('notification-filter-agent')) {
        const notificationFilterNode = getOrCreateNode('notification-filter-agent');
        lineAnalysisNode.children.set('notification-filter-agent', notificationFilterNode);
      }
      
      // slack-summary-agent is child of line-analysis-agent
      if (agentStats.has('slack-summary-agent')) {
        const slackSummaryNode = getOrCreateNode('slack-summary-agent');
        lineAnalysisNode.children.set('slack-summary-agent', slackSummaryNode);
      }
      
      // workflow-assignment-agent is child of line-analysis-agent
      if (agentStats.has('workflow-assignment-agent')) {
        const workflowAssignmentNode = getOrCreateNode('workflow-assignment-agent');
        lineAnalysisNode.children.set('workflow-assignment-agent', workflowAssignmentNode);
      }
    }

    // Add any remaining agents that weren't part of the logical hierarchy as root level
    for (const [agentName, node] of agentStats) {
      if (!this.isAgentInHierarchy(rootNode, agentName)) {
        rootNode.children.set(agentName, node);
      }
    }
  }

  private isAgentInHierarchy(rootNode: AgentNode, agentName: string): boolean {
    // Check if agent is already placed in the hierarchy
    const checkNode = (node: AgentNode): boolean => {
      if (node.name === agentName) return true;
      for (const child of node.children.values()) {
        if (checkNode(child)) return true;
      }
      return false;
    };
    
    return checkNode(rootNode);
  }

  private async processChainGroup(group: ChainExecutionGroup, rootNode: AgentNode) {
    // Get all executions for this chain group
    const { data: executions, error } = await this.supabaseClient
      .from('executions')
      .select('*')
      .eq('chain_group_id', group.id);

    if (error || !executions) {
      console.warn(`Failed to get executions for group ${group.group_id}`);
      return;
    }

    // Build parent-child relationships
    const executionMap = new Map<string, Execution>();
    const childrenMap = new Map<string, string[]>();

    for (const exec of executions) {
      executionMap.set(exec.id, exec);
      
      if (exec.parent_execution_id) {
        if (!childrenMap.has(exec.parent_execution_id)) {
          childrenMap.set(exec.parent_execution_id, []);
        }
        childrenMap.get(exec.parent_execution_id)!.push(exec.id);
      }
    }

    // Find root executions (no parent)
    const rootExecutions = executions.filter(e => !e.parent_execution_id);

    // Process each root execution recursively
    for (const rootExec of rootExecutions) {
      this.addToHierarchy(rootNode, rootExec, executionMap, childrenMap, []);
    }
  }

  private addToHierarchy(
    parentNode: AgentNode,
    execution: Execution,
    executionMap: Map<string, Execution>,
    childrenMap: Map<string, string[]>,
    path: string[]
  ) {
    const agentName = execution.tag || 'unknown-agent';
    
    // Get or create child node
    if (!parentNode.children.has(agentName)) {
      parentNode.children.set(agentName, {
        name: agentName,
        calls: 0,
        totalTokens: 0,
        children: new Map()
      });
    }

    const node = parentNode.children.get(agentName)!;
    node.calls++;
    
    // Add tokens from execution
    const tokens = execution.tokens || {};
    const totalTokens = typeof tokens === 'object' ? (tokens.total || 0) : 0;
    node.totalTokens += totalTokens;

    // Process children
    const children = childrenMap.get(execution.id) || [];
    for (const childId of children) {
      const childExec = executionMap.get(childId);
      if (childExec && !path.includes(childId)) { // Avoid cycles
        this.addToHierarchy(node, childExec, executionMap, childrenMap, [...path, childId]);
      }
    }
  }

  private calculateVariance(values: number[]): number {
    if (values.length <= 1) return 0;
    const mean = values.reduce((a, b) => a + b, 0) / values.length;
    const variance = values.reduce((acc, val) => acc + Math.pow(val - mean, 2), 0) / values.length;
    return variance;
  }

  private printTree(node: AgentNode, totalChainExecutions: number, prefix = '', isLast = true, isRoot = false) {
    // Print current node
    if (isRoot) {
      // Root node
      console.log(`${node.name} (executed ${totalChainExecutions} separate times)`);
      console.log('');
    } else {
      const connector = isLast ? '└── ' : '├── ';
      
      // Calculate statistics
      const callsPerChain = node.stats.callsPerChain;
      const tokensPerChain = node.stats.tokensPerChain;
      
      const avgCalls = callsPerChain.length > 0 ? 
        (callsPerChain.reduce((a, b) => a + b, 0) / callsPerChain.length) : 0;
      
      const avgTokensPerCall = node.calls > 0 ? (node.totalTokens / node.calls) : 0;
      const avgTokensPerChain = tokensPerChain.length > 0 ? 
        (tokensPerChain.reduce((a, b) => a + b, 0) / tokensPerChain.length) : 0;
      
      const callsVariance = this.calculateVariance(callsPerChain);
      const tokensVariance = this.calculateVariance(tokensPerChain);
      
      // Main line with averages
      console.log(
        `${prefix}${connector}${node.name} ` +
        `(avg ${avgCalls.toFixed(1)} calls/chain, σ²=${callsVariance.toFixed(1)}) ` +
        `(avg ${Math.round(avgTokensPerCall).toLocaleString()} tokens/call) ` +
        `(avg ${Math.round(avgTokensPerChain).toLocaleString()} tokens/chain)`
      );
      
      // Show per-chain breakdown indented
      if (callsPerChain.length > 0) {
        const detailPrefix = prefix + (isLast ? '    ' : '│   ') + '    ';
        console.log(`${detailPrefix}Per-chain calls: [${callsPerChain.join(', ')}]`);
        console.log(`${detailPrefix}Per-chain tokens: [${tokensPerChain.map(t => Math.round(t).toLocaleString()).join(', ')}]`);
      }
    }

    // Print children
    const children = Array.from(node.children.values())
      .sort((a, b) => {
        // Sort by average calls per chain (descending)
        const aAvg = a.stats.callsPerChain.length > 0 ? 
          (a.stats.callsPerChain.reduce((x, y) => x + y, 0) / a.stats.callsPerChain.length) : 0;
        const bAvg = b.stats.callsPerChain.length > 0 ? 
          (b.stats.callsPerChain.reduce((x, y) => x + y, 0) / b.stats.callsPerChain.length) : 0;
        return bAvg - aAvg;
      });
    
    children.forEach((child, index) => {
      const isLastChild = index === children.length - 1;
      const extension = isRoot ? '' : (isLast ? '    ' : '│   ');
      this.printTree(child, totalChainExecutions, prefix + extension, isLastChild, false);
    });
  }

  async listAvailableChains() {
    const { data: chains, error } = await this.supabaseClient
      .from('chain_execution_groups')
      .select('chain_id, status')
      .eq('repo_id', this.repoId);

    if (error || !chains) {
      console.log('No chains found.');
      return [];
    }

    // Get unique chain IDs
    const uniqueChains = [...new Set(chains.map(c => c.chain_id))];
    
    console.log('\nAvailable chains:');
    console.log('-----------------');
    uniqueChains.forEach(chain => {
      console.log(`  - ${chain}`);
    });
    console.log();

    return uniqueChains;
  }
}

// Main execution
async function main() {
  try {
    const analyzer = new ChainVisualizationAnalyzer();
    await analyzer.initialize();

    // Get chain ID from command line or use default
    const chainId = process.argv[2] || 'session-analysis-pipeline';

    if (chainId === '--list') {
      await analyzer.listAvailableChains();
    } else {
      await analyzer.analyzeChain(chainId);
    }
  } catch (error) {
    console.error('Error:', error);
    process.exit(1);
  }
}

main();