Skip to content

Commit 3e96869

Browse files
bhoustonbhouston
authored
Merge pull request #93 from drivecore/refactor-toolagent-structure
Refactor toolAgent structure
2 parents e491f3a + ece572a commit 3e96869

File tree

8 files changed

+197
-189
lines changed

8 files changed

+197
-189
lines changed

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

Lines changed: 0 additions & 22 deletions
This file was deleted.

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

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: 16 additions & 147 deletions
Original file line numberDiff line numberDiff line change
@@ -1,157 +1,26 @@
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);
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-
};
6+
// Export the main toolAgent function
7+
export { toolAgent } from './toolAgentCore.js';
1518

1529
// Re-export everything from the module
15310
export * from './config.js';
15411
export * from './messageUtils.js';
15512
export * from './toolExecutor.js';
15613
export * from './tokenTracking.js';
15714
export * from './types.js';
15+
16+
// Export default system prompt for convenience
17+
export const getDefaultSystemPrompt = (context: Record<string, unknown>) => {
18+
return `You are an AI agent that can use tools to accomplish tasks.
19+
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+
};

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

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+
};

0 commit comments

Comments
 (0)