Refactor toolAgent structure

packages/agent/src/core/toolAgent/README.md

@@ -1,10 +1,11 @@
 # Tool Agent Module
 
-This directory contains the refactored Tool Agent implementation, split into smaller, focused modules for improved maintainability and testability.
+This directory contains the Tool Agent implementation, split into smaller, focused modules for improved maintainability and testability.
 
 ## Module Structure
 
-- **index.ts**: Main entry point and orchestration of the tool agent functionality
+- **index.ts**: Re-exports from toolAgentCore.ts and other modules
+- **toolAgentCore.ts**: Main implementation of the tool agent functionality
 - **config.ts**: Configuration-related code and default settings
 - **messageUtils.ts**: Utilities for handling and formatting messages
 - **toolExecutor.ts**: Logic for executing tool calls
@@ -14,10 +15,10 @@ This directory contains the refactored Tool Agent implementation, split into sma
 ## Usage
 
 ```typescript
-import { toolAgent } from './toolAgent/index.js';
-import { Tool, ToolContext } from './toolAgent/types.js';
+import { toolAgent } from '../../core/toolAgent/index.js';
+import { Tool, ToolContext } from '../../core/types.js';
 
-// Use the toolAgent function as before
+// Use the toolAgent function
 const result = await toolAgent(prompt, tools, config, context);
 ```
 
@@ -28,7 +29,3 @@ const result = await toolAgent(prompt, tools, config, context);
 - **Clearer responsibilities**: Each module has a single purpose
 - **Easier onboarding**: New developers can understand the system more quickly
 - **Simpler future extensions**: Modular design makes it easier to extend functionality
-
-## Migration
-
-The original `toolAgent.ts` file now re-exports from this directory for backward compatibility, but it will display a deprecation warning. New code should import directly from the toolAgent directory.

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

@@ -1,157 +1,26 @@
-import { CoreMessage, ToolSet, generateText, tool as makeTool } from 'ai';
-
-import { getAnthropicApiKeyError } from '../../utils/errors.js';
-
-import { DEFAULT_CONFIG } from './config.js';
-import {
-  addCacheControlToMessages,
-  createCacheControlMessageFromSystemPrompt,
-  createToolCallParts,
-  formatToolCalls,
-} from './messageUtils.js';
-import { logTokenUsage } from './tokenTracking.js';
-import { executeTools } from './toolExecutor.js';
-import { Tool, ToolAgentResult, ToolContext } from './types.js';
-
 /**
- * Main tool agent function that orchestrates the conversation with the AI
- * and handles tool execution
+ * Main entry point for the toolAgent module
+ * Re-exports all functionality from the modular structure
  */
-export const toolAgent = async (
-  initialPrompt: string,
-  tools: Tool[],
-  config = DEFAULT_CONFIG,
-  context: ToolContext,
-): Promise<ToolAgentResult> => {
-  const { logger, tokenTracker } = context;
-
-  logger.verbose('Starting agent execution');
-  logger.verbose('Initial prompt:', initialPrompt);
-
-  let interactions = 0;
-
-  const apiKey = process.env.ANTHROPIC_API_KEY;
-  if (!apiKey) throw new Error(getAnthropicApiKeyError());
-
-  const messages: CoreMessage[] = [
-    {
-      role: 'user',
-      content: [{ type: 'text', text: initialPrompt }],
-    },
-  ];
-
-  logger.debug('User message:', initialPrompt);
-
-  // Get the system prompt once at the start
-  const systemPrompt = config.getSystemPrompt(context);
-
-  for (let i = 0; i < config.maxIterations; i++) {
-    logger.verbose(
-      `Requesting completion ${i + 1} with ${messages.length} messages with ${
-        JSON.stringify(messages).length
-      } bytes`,
-    );
-
-    interactions++;
-
-    const toolSet: ToolSet = {};
-    tools.forEach((tool) => {
-      toolSet[tool.name] = makeTool({
-        description: tool.description,
-        parameters: tool.parameters,
-      });
-    });
-
-    // Apply cache control to messages for token caching
-    const messagesWithCacheControl = [
-      createCacheControlMessageFromSystemPrompt(systemPrompt),
-      ...addCacheControlToMessages(messages),
-    ];
-
-    const generateTextProps = {
-      model: config.model,
-      temperature: config.temperature,
-      messages: messagesWithCacheControl,
-      tools: toolSet,
-    };
-    const { text, toolCalls } = await generateText(generateTextProps);
 
-    const localToolCalls = formatToolCalls(toolCalls);
-
-    if (!text.length) {
-      // Instead of treating empty response as completion, remind the agent
-      logger.verbose('Received empty response from agent, sending reminder');
-      messages.push({
-        role: 'user',
-        content: [
-          {
-            type: 'text',
-            text: 'I notice you sent an empty response. If you are done with your tasks, please call the sequenceComplete tool with your results. If you are waiting for other tools to complete, you can use the sleep tool to wait before checking again.',
-          },
-        ],
-      });
-      continue;
-    }
-
-    messages.push({
-      role: 'assistant',
-      content: [{ type: 'text', text: text }],
-    });
-
-    if (text) {
-      logger.info(text);
-    }
-
-    if (toolCalls.length > 0) {
-      const toolCallParts = createToolCallParts(toolCalls);
-
-      messages.push({
-        role: 'assistant',
-        content: toolCallParts,
-      });
-    }
-
-    const { sequenceCompleted, completionResult, respawn } = await executeTools(
-      localToolCalls,
-      tools,
-      messages,
-      context,
-    );
-
-    if (respawn) {
-      logger.info('Respawning agent with new context');
-      // Reset messages to just the new context
-      messages.length = 0;
-      messages.push({
-        role: 'user',
-        content: [{ type: 'text', text: respawn.context }],
-      });
-      continue;
-    }
-
-    if (sequenceCompleted) {
-      const result: ToolAgentResult = {
-        result: completionResult ?? 'Sequence explicitly completed',
-        interactions,
-      };
-      logTokenUsage(tokenTracker);
-      return result;
-    }
-  }
-
-  logger.warn('Maximum iterations reached');
-  const result = {
-    result: 'Maximum sub-agent iterations reach without successful completion',
-    interactions,
-  };
-
-  logTokenUsage(tokenTracker);
-  return result;
-};
+// Export the main toolAgent function
+export { toolAgent } from './toolAgentCore.js';
 
 // Re-export everything from the module
 export * from './config.js';
 export * from './messageUtils.js';
 export * from './toolExecutor.js';
 export * from './tokenTracking.js';
 export * from './types.js';
+
+// Export default system prompt for convenience
+export const getDefaultSystemPrompt = (context: Record<string, unknown>) => {
+  return `You are an AI agent that can use tools to accomplish tasks.
+
+Current Context:
+Directory: ${context.workingDirectory}
+Files:
+${context.directoryListing ?? 'No directory listing available'}
+System: ${context.systemInfo ?? 'No system info available'}
+DateTime: ${new Date().toString()}`;
+};

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

@@ -0,0 +1,150 @@
+import { CoreMessage, ToolSet, generateText, tool as makeTool } from 'ai';
+
+import { getAnthropicApiKeyError } from '../../utils/errors.js';
+
+import { DEFAULT_CONFIG } from './config.js';
+import {
+  addCacheControlToMessages,
+  createCacheControlMessageFromSystemPrompt,
+  createToolCallParts,
+  formatToolCalls,
+} from './messageUtils.js';
+import { logTokenUsage } from './tokenTracking.js';
+import { executeTools } from './toolExecutor.js';
+import { Tool, ToolAgentResult, ToolContext } from './types.js';
+
+/**
+ * Main tool agent function that orchestrates the conversation with the AI
+ * and handles tool execution
+ */
+export const toolAgent = async (
+  initialPrompt: string,
+  tools: Tool[],
+  config = DEFAULT_CONFIG,
+  context: ToolContext,
+): Promise<ToolAgentResult> => {
+  const { logger, tokenTracker } = context;
+
+  logger.verbose('Starting agent execution');
+  logger.verbose('Initial prompt:', initialPrompt);
+
+  let interactions = 0;
+
+  const apiKey = process.env.ANTHROPIC_API_KEY;
+  if (!apiKey) throw new Error(getAnthropicApiKeyError());
+
+  const messages: CoreMessage[] = [
+    {
+      role: 'user',
+      content: [{ type: 'text', text: initialPrompt }],
+    },
+  ];
+
+  logger.debug('User message:', initialPrompt);
+
+  // Get the system prompt once at the start
+  const systemPrompt = config.getSystemPrompt(context);
+
+  for (let i = 0; i < config.maxIterations; i++) {
+    logger.verbose(
+      `Requesting completion ${i + 1} with ${messages.length} messages with ${
+        JSON.stringify(messages).length
+      } bytes`,
+    );
+
+    interactions++;
+
+    const toolSet: ToolSet = {};
+    tools.forEach((tool) => {
+      toolSet[tool.name] = makeTool({
+        description: tool.description,
+        parameters: tool.parameters,
+      });
+    });
+
+    // Apply cache control to messages for token caching
+    const messagesWithCacheControl = [
+      createCacheControlMessageFromSystemPrompt(systemPrompt),
+      ...addCacheControlToMessages(messages),
+    ];
+
+    const generateTextProps = {
+      model: config.model,
+      temperature: config.temperature,
+      messages: messagesWithCacheControl,
+      tools: toolSet,
+    };
+    const { text, toolCalls } = await generateText(generateTextProps);
+
+    const localToolCalls = formatToolCalls(toolCalls);
+
+    if (!text.length) {
+      // Instead of treating empty response as completion, remind the agent
+      logger.verbose('Received empty response from agent, sending reminder');
+      messages.push({
+        role: 'user',
+        content: [
+          {
+            type: 'text',
+            text: 'I notice you sent an empty response. If you are done with your tasks, please call the sequenceComplete tool with your results. If you are waiting for other tools to complete, you can use the sleep tool to wait before checking again.',
+          },
+        ],
+      });
+      continue;
+    }
+
+    messages.push({
+      role: 'assistant',
+      content: [{ type: 'text', text: text }],
+    });
+
+    if (text) {
+      logger.info(text);
+    }
+
+    if (toolCalls.length > 0) {
+      const toolCallParts = createToolCallParts(toolCalls);
+
+      messages.push({
+        role: 'assistant',
+        content: toolCallParts,
+      });
+    }
+
+    const { sequenceCompleted, completionResult, respawn } = await executeTools(
+      localToolCalls,
+      tools,
+      messages,
+      context,
+    );
+
+    if (respawn) {
+      logger.info('Respawning agent with new context');
+      // Reset messages to just the new context
+      messages.length = 0;
+      messages.push({
+        role: 'user',
+        content: [{ type: 'text', text: respawn.context }],
+      });
+      continue;
+    }
+
+    if (sequenceCompleted) {
+      const result: ToolAgentResult = {
+        result: completionResult ?? 'Sequence explicitly completed',
+        interactions,
+      };
+      logTokenUsage(tokenTracker);
+      return result;
+    }
+  }
+
+  logger.warn('Maximum iterations reached');
+  const result = {
+    result: 'Maximum sub-agent iterations reach without successful completion',
+    interactions,
+  };
+
+  logTokenUsage(tokenTracker);
+  return result;
+};