Skip to content

Commit 5f3bb0f

Browse files
committed
refactor: Move toolAgent implementation to toolAgentCore.ts and update imports
- Move main toolAgent implementation from index.ts to toolAgentCore.ts\n- Update index.ts to re-export from toolAgentCore.ts\n- Remove deprecated toolAgent.ts file\n- Update all direct imports to use the new structure\n- Update documentation in README.md\n\nCloses #92
1 parent 7ffcd57 commit 5f3bb0f

File tree

5 files changed

+190
-155
lines changed

5 files changed

+190
-155
lines changed
Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
# Tool Agent Module
22

3-
This directory contains the refactored Tool Agent implementation, split into smaller, focused modules for improved maintainability and testability.
3+
This directory contains the Tool Agent implementation, split into smaller, focused modules for improved maintainability and testability.
44

55
## Module Structure
66

7-
- **index.ts**: Main entry point and orchestration of the tool agent functionality
7+
- **index.ts**: Re-exports from toolAgentCore.ts and other modules
8+
- **toolAgentCore.ts**: Main implementation of the tool agent functionality
89
- **config.ts**: Configuration-related code and default settings
910
- **messageUtils.ts**: Utilities for handling and formatting messages
1011
- **toolExecutor.ts**: Logic for executing tool calls
@@ -14,10 +15,10 @@ This directory contains the refactored Tool Agent implementation, split into sma
1415
## Usage
1516

1617
```typescript
17-
import { toolAgent } from './toolAgent/index.js';
18-
import { Tool, ToolContext } from './toolAgent/types.js';
18+
import { toolAgent } from '../../core/toolAgent/index.js';
19+
import { Tool, ToolContext } from '../../core/types.js';
1920

20-
// Use the toolAgent function as before
21+
// Use the toolAgent function
2122
const result = await toolAgent(prompt, tools, config, context);
2223
```
2324

@@ -28,7 +29,3 @@ const result = await toolAgent(prompt, tools, config, context);
2829
- **Clearer responsibilities**: Each module has a single purpose
2930
- **Easier onboarding**: New developers can understand the system more quickly
3031
- **Simpler future extensions**: Modular design makes it easier to extend functionality
31-
32-
## Migration
33-
34-
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

Lines changed: 29 additions & 143 deletions
Original file line numberDiff line numberDiff line change
@@ -1,152 +1,38 @@
1-
import { CoreMessage, ToolSet, generateText, tool as makeTool } from 'ai';
2-
3-
import { getAnthropicApiKeyError } from '../../utils/errors.js';
4-
5-
import { DEFAULT_CONFIG } from './config.js';
6-
import {
7-
addCacheControlToMessages,
8-
createCacheControlMessageFromSystemPrompt,
9-
createToolCallParts,
10-
formatToolCalls,
11-
} from './messageUtils.js';
12-
import { logTokenUsage } from './tokenTracking.js';
13-
import { executeTools } from './toolExecutor.js';
14-
import { Tool, ToolAgentResult, ToolContext } from './types.js';
15-
161
/**
17-
* Main tool agent function that orchestrates the conversation with the AI
18-
* and handles tool execution
2+
* Main entry point for the toolAgent module
3+
* Re-exports all functionality from the modular structure
194
*/
20-
export const toolAgent = async (
21-
initialPrompt: string,
22-
tools: Tool[],
23-
config = DEFAULT_CONFIG,
24-
context: ToolContext,
25-
): Promise<ToolAgentResult> => {
26-
const { logger, tokenTracker } = context;
27-
28-
logger.verbose('Starting agent execution');
29-
logger.verbose('Initial prompt:', initialPrompt);
30-
31-
let interactions = 0;
32-
33-
const apiKey = process.env.ANTHROPIC_API_KEY;
34-
if (!apiKey) throw new Error(getAnthropicApiKeyError());
35-
36-
const messages: CoreMessage[] = [
37-
{
38-
role: 'user',
39-
content: [{ type: 'text', text: initialPrompt }],
40-
},
41-
];
42-
43-
logger.debug('User message:', initialPrompt);
44-
45-
// Get the system prompt once at the start
46-
const systemPrompt = config.getSystemPrompt(context);
47-
48-
for (let i = 0; i < config.maxIterations; i++) {
49-
logger.verbose(
50-
`Requesting completion ${i + 1} with ${messages.length} messages with ${
51-
JSON.stringify(messages).length
52-
} bytes`,
53-
);
54-
55-
interactions++;
56-
57-
const toolSet: ToolSet = {};
58-
tools.forEach((tool) => {
59-
toolSet[tool.name] = makeTool({
60-
description: tool.description,
61-
parameters: tool.parameters,
62-
});
63-
});
64-
65-
// Apply cache control to messages for token caching
66-
const messagesWithCacheControl = [
67-
createCacheControlMessageFromSystemPrompt(systemPrompt),
68-
...addCacheControlToMessages(messages),
69-
];
70-
71-
const generateTextProps = {
72-
model: config.model,
73-
temperature: config.temperature,
74-
messages: messagesWithCacheControl,
75-
tools: toolSet,
76-
};
77-
const { text, toolCalls } = await generateText(generateTextProps);
785

79-
const localToolCalls = formatToolCalls(toolCalls);
6+
// Export the main toolAgent function
7+
export { toolAgent } from './toolAgentCore.js';
808

81-
if (!text.length) {
82-
// Instead of treating empty response as completion, remind the agent
83-
logger.verbose('Received empty response from agent, sending reminder');
84-
messages.push({
85-
role: 'user',
86-
content: [
87-
{
88-
type: 'text',
89-
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.',
90-
},
91-
],
92-
});
93-
continue;
94-
}
95-
96-
messages.push({
97-
role: 'assistant',
98-
content: [{ type: 'text', text: text }],
99-
});
100-
101-
if (text) {
102-
logger.info(text);
103-
}
104-
105-
if (toolCalls.length > 0) {
106-
const toolCallParts = createToolCallParts(toolCalls);
107-
108-
messages.push({
109-
role: 'assistant',
110-
content: toolCallParts,
111-
});
112-
}
113-
114-
const { sequenceCompleted, completionResult, respawn } = await executeTools(
115-
localToolCalls,
116-
tools,
117-
messages,
118-
context,
119-
);
120-
121-
if (respawn) {
122-
logger.info('Respawning agent with new context');
123-
// Reset messages to just the new context
124-
messages.length = 0;
125-
messages.push({
126-
role: 'user',
127-
content: [{ type: 'text', text: respawn.context }],
128-
});
129-
continue;
130-
}
131-
132-
if (sequenceCompleted) {
133-
const result: ToolAgentResult = {
134-
result: completionResult ?? 'Sequence explicitly completed',
135-
interactions,
136-
};
137-
logTokenUsage(tokenTracker);
138-
return result;
139-
}
140-
}
9+
// Re-export everything from the module
10+
export * from './config.js';
11+
export * from './messageUtils.js';
12+
export * from './toolExecutor.js';
13+
export * from './tokenTracking.js';
14+
export * from './types.js';
14115

142-
logger.warn('Maximum iterations reached');
143-
const result = {
144-
result: 'Maximum sub-agent iterations reach without successful completion',
145-
interactions,
146-
};
16+
// Export default system prompt for convenience
17+
export const getDefaultSystemPrompt = (context: any) => {
18+
return `You are an AI agent that can use tools to accomplish tasks.
14719
148-
logTokenUsage(tokenTracker);
149-
return result;
20+
Current Context:
21+
Directory: ${context.workingDirectory}
22+
Files:
23+
${context.directoryListing ?? 'No directory listing available'}
24+
System: ${context.systemInfo ?? 'No system info available'}
25+
DateTime: ${new Date().toString()}`;
26+
};
27+
export const getDefaultSystemPrompt = (context: any) => {
28+
return `You are an AI agent that can use tools to accomplish tasks.
29+
30+
Current Context:
31+
Directory: ${context.workingDirectory}
32+
Files:
33+
${context.directoryListing ?? 'No directory listing available'}
34+
System: ${context.systemInfo ?? 'No system info available'}
35+
DateTime: ${new Date().toString()}`;
15036
};
15137

15238
// Re-export everything from the module
Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
import { CoreMessage, ToolSet, generateText, tool as makeTool } from 'ai';
2+
3+
import { getAnthropicApiKeyError } from '../../utils/errors.js';
4+
5+
import { DEFAULT_CONFIG } from './config.js';
6+
import {
7+
addCacheControlToMessages,
8+
createCacheControlMessageFromSystemPrompt,
9+
createToolCallParts,
10+
formatToolCalls,
11+
} from './messageUtils.js';
12+
import { logTokenUsage } from './tokenTracking.js';
13+
import { executeTools } from './toolExecutor.js';
14+
import { Tool, ToolAgentResult, ToolContext } from './types.js';
15+
16+
/**
17+
* Main tool agent function that orchestrates the conversation with the AI
18+
* and handles tool execution
19+
*/
20+
export const toolAgent = async (
21+
initialPrompt: string,
22+
tools: Tool[],
23+
config = DEFAULT_CONFIG,
24+
context: ToolContext,
25+
): Promise<ToolAgentResult> => {
26+
const { logger, tokenTracker } = context;
27+
28+
logger.verbose('Starting agent execution');
29+
logger.verbose('Initial prompt:', initialPrompt);
30+
31+
let interactions = 0;
32+
33+
const apiKey = process.env.ANTHROPIC_API_KEY;
34+
if (!apiKey) throw new Error(getAnthropicApiKeyError());
35+
36+
const messages: CoreMessage[] = [
37+
{
38+
role: 'user',
39+
content: [{ type: 'text', text: initialPrompt }],
40+
},
41+
];
42+
43+
logger.debug('User message:', initialPrompt);
44+
45+
// Get the system prompt once at the start
46+
const systemPrompt = config.getSystemPrompt(context);
47+
48+
for (let i = 0; i < config.maxIterations; i++) {
49+
logger.verbose(
50+
`Requesting completion ${i + 1} with ${messages.length} messages with ${
51+
JSON.stringify(messages).length
52+
} bytes`,
53+
);
54+
55+
interactions++;
56+
57+
const toolSet: ToolSet = {};
58+
tools.forEach((tool) => {
59+
toolSet[tool.name] = makeTool({
60+
description: tool.description,
61+
parameters: tool.parameters,
62+
});
63+
});
64+
65+
// Apply cache control to messages for token caching
66+
const messagesWithCacheControl = [
67+
createCacheControlMessageFromSystemPrompt(systemPrompt),
68+
...addCacheControlToMessages(messages),
69+
];
70+
71+
const generateTextProps = {
72+
model: config.model,
73+
temperature: config.temperature,
74+
messages: messagesWithCacheControl,
75+
tools: toolSet,
76+
};
77+
const { text, toolCalls } = await generateText(generateTextProps);
78+
79+
const localToolCalls = formatToolCalls(toolCalls);
80+
81+
if (!text.length) {
82+
// Instead of treating empty response as completion, remind the agent
83+
logger.verbose('Received empty response from agent, sending reminder');
84+
messages.push({
85+
role: 'user',
86+
content: [
87+
{
88+
type: 'text',
89+
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.',
90+
},
91+
],
92+
});
93+
continue;
94+
}
95+
96+
messages.push({
97+
role: 'assistant',
98+
content: [{ type: 'text', text: text }],
99+
});
100+
101+
if (text) {
102+
logger.info(text);
103+
}
104+
105+
if (toolCalls.length > 0) {
106+
const toolCallParts = createToolCallParts(toolCalls);
107+
108+
messages.push({
109+
role: 'assistant',
110+
content: toolCallParts,
111+
});
112+
}
113+
114+
const { sequenceCompleted, completionResult, respawn } = await executeTools(
115+
localToolCalls,
116+
tools,
117+
messages,
118+
context,
119+
);
120+
121+
if (respawn) {
122+
logger.info('Respawning agent with new context');
123+
// Reset messages to just the new context
124+
messages.length = 0;
125+
messages.push({
126+
role: 'user',
127+
content: [{ type: 'text', text: respawn.context }],
128+
});
129+
continue;
130+
}
131+
132+
if (sequenceCompleted) {
133+
const result: ToolAgentResult = {
134+
result: completionResult ?? 'Sequence explicitly completed',
135+
interactions,
136+
};
137+
logTokenUsage(tokenTracker);
138+
return result;
139+
}
140+
}
141+
142+
logger.warn('Maximum iterations reached');
143+
const result = {
144+
result: 'Maximum sub-agent iterations reach without successful completion',
145+
interactions,
146+
};
147+
148+
logTokenUsage(tokenTracker);
149+
return result;
150+
};

packages/agent/src/index.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ export * from './tools/interaction/userPrompt.js';
2525
// Core
2626
export * from './core/executeToolCall.js';
2727
export * from './core/types.js';
28-
export * from './core/toolAgent.js';
28+
export * from './core/toolAgent/index.js';
2929
export * from './core/toolAgent/config.js';
3030

3131
// Utils

packages/agent/src/tools/interaction/subAgent.ts

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,8 +2,10 @@ import { z } from 'zod';
22
import { zodToJsonSchema } from 'zod-to-json-schema';
33

44
import { getModel } from '../../core/toolAgent/config.js';
5-
import { getDefaultSystemPrompt } from '../../core/toolAgent/index.js';
6-
import { toolAgent } from '../../core/toolAgent.js';
5+
import {
6+
getDefaultSystemPrompt,
7+
toolAgent,
8+
} from '../../core/toolAgent/index.js';
79
import { Tool, ToolContext } from '../../core/types.js';
810
import { getTools } from '../getTools.js';
911

0 commit comments

Comments
 (0)