Setup token analysis for all functions

package.json

@@ -14,6 +14,7 @@
     "lint:fix": "eslint . --ext .ts --fix",
     "format": "prettier --write \"src/**/*.ts\" \"test/**/*.ts\"",
     "tool-summary": "tsx src/utils/analyzeTools.ts",
+    "tool-summary:langchain": "tsx src/utils/analyzeTools.ts --langchain",
     "check-tool-names:langchain": "tsx scripts/check-langchain-tool-duplicates.ts"
   },
   "engines": {
@@ -79,6 +80,7 @@
     "form-data": "^4.0.1",
     "langchain": "^0.3.8",
     "openai": "^4.77.0",
+    "tiktoken": "^1.0.18",
     "typedoc": "^0.27.6",
     "zod": "^3.24.1"
   },

src/utils/analyzeTools.ts

@@ -0,0 +1,205 @@
+import { encoding_for_model, TiktokenModel } from "tiktoken";
+import { SolanaAgentKit } from "../index";
+import * as dotenv from "dotenv";
+import { ACTIONS, Action } from "../actions/index";
+import { createSolanaTools as getLangchainTools } from "../langchain/index";
+import { createSolanaTools as getVercelTools } from "../vercel-ai/index";
+import { type CoreTool } from "ai";
+import { Tool } from "langchain/tools";
+
+dotenv.config();
+
+interface ToolAnalysis {
+  name: string;
+  actionTokens: number;
+  langchainTokens: number;
+  vercelTokens: number;
+  totalTokens: number;
+}
+
+function calculateActionTokens(action: Action, enc: any): number {
+  // You can customize how you want to weigh or encode the action
+  const description = action.description || "";
+  const schema = JSON.stringify(action.schema || {});
+  const similes = action.similes || [];
+
+  const descTokens = enc.encode(description).length;
+  const schemaTokens = enc.encode(schema).length;
+  const similesTokens = similes.reduce((acc, simile) => {
+    return acc + enc.encode(simile).length;
+  }, 0);
+
+  return descTokens + schemaTokens + similesTokens;
+}
+
+function calculateLangchainTokens(tool: Tool, enc: any): number {
+  // For Langchain, maybe you only want to encode the description
+  // or possibly there's a manifest or advanced schema. Adjust to suit your needs.
+  const description = tool.description || "";
+  const schema = JSON.stringify(tool.schema || {});
+  const descTokens = enc.encode(description).length;
+  const schemaTokens = enc.encode(schema).length;
+
+  return descTokens + schemaTokens;
+}
+
+function calculateVercelTokens(tool: CoreTool, enc: any): number {
+  // @ts-expect-error for some reason, the description is not expected to exist
+  const description = tool.description || "";
+  // @ts-expect-error for some reason, the schema is not expected to exist
+  const schemaStr = JSON.stringify(tool.schema || {});
+  const parameters = JSON.stringify(tool.parameters || {});
+
+  const descTokens = enc.encode(description).length;
+  const schemaTokens = enc.encode(schemaStr).length;
+  const parametersTokens = enc.encode(parameters).length;
+
+  return descTokens + schemaTokens + parametersTokens;
+}
+
+export async function analyzeToolTokens(
+  modelName: TiktokenModel = "gpt-4o",
+  includeLangchain: boolean = false,
+): Promise<ToolAnalysis[]> {
+  const enc = encoding_for_model(modelName);
+
+  const solanaAgentKit = new SolanaAgentKit(
+    process.env.SOLANA_PRIVATE_KEY! || "",
+    process.env.RPC_URL! || "",
+    { OPENAI_API_KEY: process.env.OPENAI_API_KEY! || "" },
+  );
+
+  // Unify tools by name
+  const toolMap = new Map<
+    string,
+    {
+      action: Action | undefined;
+      langchain: Tool | undefined;
+      vercel: CoreTool | undefined;
+    }
+  >();
+
+  // Only load Langchain tools if flag is set
+  if (includeLangchain) {
+    const langchainTools = getLangchainTools(solanaAgentKit) as Tool[];
+
+    langchainTools.forEach((tool) => {
+      toolMap.set(tool.name, {
+        action: undefined,
+        langchain: tool,
+        vercel: undefined,
+      });
+    });
+  } else {
+    const vercelTools = getVercelTools(solanaAgentKit);
+
+    Object.keys(ACTIONS).forEach((toolName) => {
+      toolMap.set(toolName, {
+        action: ACTIONS[toolName as keyof typeof ACTIONS],
+        langchain: undefined,
+        vercel: vercelTools[toolName as keyof typeof ACTIONS],
+      });
+    });
+  }
+
+  const analysis: ToolAnalysis[] = [];
+  for (const [name, implementations] of toolMap) {
+    const entry: ToolAnalysis = {
+      name,
+      actionTokens: 0,
+      langchainTokens: 0,
+      vercelTokens: 0,
+      totalTokens: 0,
+    };
+
+    // 1) Action tokens
+    if (implementations.action) {
+      entry.actionTokens = calculateActionTokens(implementations.action, enc);
+    }
+
+    // 2) Langchain tokens
+    if (implementations.langchain) {
+      entry.langchainTokens = calculateLangchainTokens(
+        implementations.langchain,
+        enc,
+      );
+    }
+
+    // 3) Vercel AI tokens
+    if (implementations.vercel) {
+      entry.vercelTokens = calculateVercelTokens(implementations.vercel, enc);
+    }
+
+    // Update total calculation to exclude Langchain if not included
+    if (includeLangchain) {
+      entry.totalTokens = entry.langchainTokens;
+    } else {
+      entry.totalTokens = entry.actionTokens + entry.vercelTokens;
+    }
+
+    analysis.push(entry);
+  }
+
+  // Sort by total tokens descending
+  const sorted = analysis.sort((a, b) => b.totalTokens - a.totalTokens);
+
+  // Print final table
+  console.log("\nCross-Implementation Tool Token Analysis:");
+  if (includeLangchain) {
+    console.table(
+      sorted.map((t) => ({
+        Name: t.name,
+        "Langchain Tokens": t.langchainTokens,
+      })),
+    );
+  } else {
+    console.table(
+      sorted.map((t) => ({
+        Name: t.name,
+        "Action Tokens": t.actionTokens,
+        "Vercel AI Tokens": t.vercelTokens,
+      })),
+    );
+  }
+
+  // Print category totals
+  const totals = sorted.reduce(
+    (acc, t) => ({
+      action: acc.action + t.actionTokens,
+      langchain: acc.langchain + t.langchainTokens,
+      vercel: acc.vercel + t.vercelTokens,
+      total: acc.total + t.totalTokens,
+    }),
+    { action: 0, langchain: 0, vercel: 0, total: 0 },
+  );
+
+  console.log("\nCategory Totals:");
+  if (includeLangchain) {
+    console.table([
+      {
+        "Langchain Total": totals.langchain,
+        "Grand Total": totals.total,
+      },
+    ]);
+  } else {
+    console.table([
+      {
+        "Action Total": totals.action,
+        "Vercel Total": totals.vercel,
+        "Grand Total": totals.total,
+      },
+    ]);
+  }
+
+  enc.free();
+  return sorted;
+}
+
+// Update standalone execution to accept CLI flag
+if (require.main === module) {
+  const includeLangchain = process.argv.includes("--langchain");
+  analyzeToolTokens("gpt-4o", includeLangchain).catch((err) => {
+    console.error(err);
+    process.exit(1);
+  });
+}