Merge pull request #47 from drivecore/token-caching-and-usage

Token caching and usage

Author: bhouston

package.json

@@ -2,6 +2,7 @@
   "name": "mycoder-monorepo",
   "version": "0.0.1",
   "type": "module",
+  "private": true,
   "packageManager": "[email protected]",
   "engines": {
     "node": ">=18.0.0"

packages/agent/README.md

@@ -36,6 +36,7 @@ Get an API key from https://www.anthropic.com/api
 - Categories: Interaction, I/O, System, Data Management
 - Parallel execution capability
 - Type-safe definitions
+- Input token caching to reduce API costs
 
 ### Agent System
 

packages/agent/package.json

@@ -1,6 +1,6 @@
 {
   "name": "mycoder-agent",
-  "version": "0.1.3",
+  "version": "0.1.4",
   "description": "Agent module for mycoder - an AI-powered software development assistant",
   "type": "module",
   "main": "dist/index.js",

packages/agent/src/core/executeToolCall.ts

@@ -1,6 +1,6 @@
 import { Logger } from '../utils/logger.js';
 
-import { Tool, ToolCall } from './types.js';
+import { Tool, ToolCall, ToolContext } from './types.js';
 
 const OUTPUT_LIMIT = 12 * 1024; // 10KB limit
 
@@ -10,20 +10,22 @@ const OUTPUT_LIMIT = 12 * 1024; // 10KB limit
 export const executeToolCall = async (
   toolCall: ToolCall,
   tools: Tool[],
-  parentLogger: Logger,
-  options?: { workingDirectory?: string },
+  context: ToolContext,
 ): Promise<string> => {
   const logger = new Logger({
     name: `Tool:${toolCall.name}`,
-    parent: parentLogger,
+    parent: context.logger,
   });
 
   const tool = tools.find((t) => t.name === toolCall.name);
   if (!tool) {
     throw new Error(`No tool with the name '${toolCall.name}' exists.`);
   }
 
-  const toolContext = { logger };
+  const toolContext = {
+    ...context,
+    logger,
+  };
 
   // for each parameter log it and its name
   if (tool.logParameters) {
@@ -36,10 +38,7 @@ export const executeToolCall = async (
   }
 
   // TODO: validate JSON schema for input
-  const output = await tool.execute(toolCall.input, {
-    logger,
-    workingDirectory: options?.workingDirectory,
-  });
+  const output = await tool.execute(toolCall.input, toolContext);
 
   // for each result log it and its name
   if (tool.logReturns) {

packages/agent/src/core/tokens.ts

@@ -0,0 +1,53 @@
+import Anthropic from '@anthropic-ai/sdk';
+
+export type TokenUsage = {
+  input: number;
+  inputCacheWrites: number;
+  inputCacheReads: number;
+  output: number;
+};
+
+export const getTokenUsage = (response: Anthropic.Message): TokenUsage => {
+  return {
+    input: response.usage.input_tokens,
+    inputCacheWrites: response.usage.cache_creation_input_tokens ?? 0,
+    inputCacheReads: response.usage.cache_read_input_tokens ?? 0,
+    output: response.usage.output_tokens,
+  };
+};
+
+export const addTokenUsage = (a: TokenUsage, b: TokenUsage): TokenUsage => {
+  return {
+    input: a.input + b.input,
+    inputCacheWrites: a.inputCacheWrites + b.inputCacheWrites,
+    inputCacheReads: a.inputCacheReads + b.inputCacheReads,
+    output: a.output + b.output,
+  };
+};
+
+const PER_MILLION = 1 / 1000000;
+const TOKEN_COST: TokenUsage = {
+  input: 3 * PER_MILLION,
+  inputCacheWrites: 3.75 * PER_MILLION,
+  inputCacheReads: 0.3 * PER_MILLION,
+  output: 15 * PER_MILLION,
+};
+
+const formatter = new Intl.NumberFormat('en-US', {
+  style: 'currency',
+  currency: 'USD',
+  minimumFractionDigits: 2,
+});
+
+export const getTokenCost = (usage: TokenUsage): string => {
+  return formatter.format(
+    usage.input * TOKEN_COST.input +
+      usage.inputCacheWrites * TOKEN_COST.inputCacheWrites +
+      usage.inputCacheReads * TOKEN_COST.inputCacheReads +
+      usage.output * TOKEN_COST.output,
+  );
+};
+
+export const getTokenString = (usage: TokenUsage): string => {
+  return `input: ${usage.input} input-cache-writes: ${usage.inputCacheWrites} input-cache-reads: ${usage.inputCacheReads} output: ${usage.output} COST: ${getTokenCost(usage)}`;
+};

packages/agent/src/core/toolAgent.cache.test.ts

@@ -0,0 +1,102 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { MockLogger } from '../utils/mockLogger.js';
+
+import { toolAgent } from './toolAgent.js';
+
+const logger = new MockLogger();
+
+process.env.ANTHROPIC_API_KEY = 'sk-ant-api03-1234567890';
+
+// Mock Anthropic client
+vi.mock('@anthropic-ai/sdk', () => {
+  return {
+    default: class MockAnthropic {
+      messages = {
+        create: vi.fn().mockImplementation(() => {
+          return {
+            id: 'msg_123',
+            model: 'claude-3-7-sonnet-latest',
+            type: 'message',
+            role: 'assistant',
+            content: [
+              {
+                type: 'text',
+                text: 'I will help with that.',
+              },
+              {
+                type: 'tool_use',
+                id: 'tu_123',
+                name: 'sequenceComplete',
+                input: {
+                  result: 'Test complete',
+                },
+              },
+            ],
+            usage: {
+              input_tokens: 100,
+              output_tokens: 50,
+              // Simulating cached tokens
+              cache_read_input_tokens: 30,
+              cache_creation_input_tokens: 70,
+            },
+          };
+        }),
+      };
+      constructor() {}
+    },
+  };
+});
+
+// Mock tool
+const mockTool = {
+  name: 'sequenceComplete',
+  description: 'Completes the sequence',
+  parameters: {
+    type: 'object' as const,
+    properties: {
+      result: {
+        type: 'string' as const,
+      },
+    },
+    additionalProperties: false,
+    required: ['result'],
+  },
+  returns: {
+    type: 'string' as const,
+  },
+  execute: vi.fn().mockImplementation(async (params) => {
+    console.log('   Parameters:');
+    Object.entries(params).forEach(([key, value]) => {
+      console.log(`     - ${key}: ${JSON.stringify(value)}`);
+    });
+    console.log();
+    console.log('   Results:');
+    console.log(`     - ${params.result}`);
+    console.log();
+    return params.result;
+  }),
+};
+
+describe('toolAgent input token caching', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  it('should track cached tokens in the result', async () => {
+    const result = await toolAgent('test prompt', [mockTool], undefined, {
+      logger,
+      headless: true,
+      workingDirectory: '.',
+      tokenLevel: 'debug',
+    });
+
+    // Verify that cached tokens are tracked
+    expect(result.tokens.inputCacheReads).toBeDefined();
+    expect(result.tokens.inputCacheReads).toBe(30);
+
+    // Verify total token counts
+    expect(result.tokens.input).toBe(100);
+    expect(result.tokens.output).toBe(50);
+  });
+});

packages/agent/src/core/toolAgent.respawn.test.ts

@@ -2,7 +2,9 @@ import { describe, it, expect, vi, beforeEach } from 'vitest';
 
 import { toolAgent } from '../../src/core/toolAgent.js';
 import { getTools } from '../../src/tools/getTools.js';
-import { Logger } from '../../src/utils/logger.js';
+import { MockLogger } from '../utils/mockLogger.js';
+
+const logger = new MockLogger();
 
 // Mock Anthropic SDK
 vi.mock('@anthropic-ai/sdk', () => {
@@ -32,7 +34,6 @@ vi.mock('@anthropic-ai/sdk', () => {
 });
 
 describe('toolAgent respawn functionality', () => {
-  const mockLogger = new Logger({ name: 'test' });
   const tools = getTools();
 
   beforeEach(() => {
@@ -41,13 +42,23 @@ describe('toolAgent respawn functionality', () => {
   });
 
   it('should handle respawn tool calls', async () => {
-    const result = await toolAgent('initial prompt', tools, mockLogger, {
-      maxIterations: 2, // Need at least 2 iterations for respawn + empty response
-      model: 'test-model',
-      maxTokens: 100,
-      temperature: 0,
-      getSystemPrompt: () => 'test system prompt',
-    });
+    const result = await toolAgent(
+      'initial prompt',
+      tools,
+      {
+        maxIterations: 2, // Need at least 2 iterations for respawn + empty response
+        model: 'test-model',
+        maxTokens: 100,
+        temperature: 0,
+        getSystemPrompt: () => 'test system prompt',
+      },
+      {
+        logger,
+        headless: true,
+        workingDirectory: '.',
+        tokenLevel: 'debug',
+      },
+    );
 
     expect(result.result).toBe(
       'Maximum sub-agent iterations reach without successful completion',

packages/agent/src/core/toolAgent.test.ts

@@ -101,7 +101,12 @@ describe('toolAgent', () => {
         input: { input: 'test' },
       },
       [mockTool],
-      logger,
+      {
+        logger,
+        headless: true,
+        workingDirectory: '.',
+        tokenLevel: 'debug',
+      },
     );
 
     expect(result.includes('Processed: test')).toBeTruthy();
@@ -116,7 +121,12 @@ describe('toolAgent', () => {
           input: {},
         },
         [mockTool],
-        logger,
+        {
+          logger,
+          headless: true,
+          workingDirectory: '.',
+          tokenLevel: 'debug',
+        },
       ),
     ).rejects.toThrow("No tool with the name 'nonexistentTool' exists.");
   });
@@ -147,7 +157,12 @@ describe('toolAgent', () => {
           input: {},
         },
         [errorTool],
-        logger,
+        {
+          logger,
+          headless: true,
+          workingDirectory: '.',
+          tokenLevel: 'debug',
+        },
       ),
     ).rejects.toThrow('Deliberate failure');
   });
@@ -166,8 +181,13 @@ describe('toolAgent', () => {
     const result = await toolAgent(
       'Test prompt',
       [sequenceCompleteTool],
-      logger,
       testConfig,
+      {
+        logger,
+        headless: true,
+        workingDirectory: '.',
+        tokenLevel: 'debug',
+      },
     );
 
     // Verify that create was called twice (once for empty response, once for completion)
@@ -184,8 +204,13 @@ describe('toolAgent', () => {
     const result = await toolAgent(
       'Test prompt',
       [sequenceCompleteTool],
-      logger,
       testConfig,
+      {
+        logger,
+        headless: true,
+        workingDirectory: '.',
+        tokenLevel: 'debug',
+      },
     );
 
     expect(result.result).toBe('Test complete');