Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

add in the todos.

This commit is contained in:
dal 2025-10-06 12:31:54 -06:00
parent 6c77a4f7a3
commit d776ef6fba
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
10 changed files with 216 additions and 184 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

24
packages/ai/src/agents/analyst-agent/analyst-agent.ts
View File

@ -123,10 +123,10 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
const docsSystemMessage = docsContent
? ({
role: 'system',
content: `<data_catalog_docs>\n${docsContent}\n</data_catalog_docs>`,
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
} as ModelMessage)
role: 'system',
content: `<data_catalog_docs>\n${docsContent}\n</data_catalog_docs>`,
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
} as ModelMessage)
: null;
async function stream({ messages }: AnalystStreamOptions) {
@ -183,19 +183,19 @@ export function createAnalystAgent(analystAgentOptions: AnalystAgentOptions) {
// Create analyst instructions system message with proper escaping
const analystInstructionsMessage = analystInstructions
? ({
role: 'system',
content: `<organization_instructions>\n${analystInstructions}\n</organization_instructions>`,
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
} as ModelMessage)
role: 'system',
content: `<organization_instructions>\n${analystInstructions}\n</organization_instructions>`,
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
} as ModelMessage)
: null;
// Create user personalization system message
const userPersonalizationSystemMessage = userPersonalizationMessageContent
? ({
role: 'system',
content: userPersonalizationMessageContent,
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
} as ModelMessage)
role: 'system',
content: userPersonalizationMessageContent,
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
} as ModelMessage)
: null;
return wrapTraced(

76
packages/ai/src/agents/analytics-engineer-agent/analytics-engineer-agent.ts
View File

@ -1,84 +1,19 @@
import type { LanguageModelV2 } from '@ai-sdk/provider';
import type { Sandbox } from '@buster/sandbox';
import { type ModelMessage, hasToolCall, stepCountIs, streamText } from 'ai';
import { wrapTraced } from 'braintrust';
import z from 'zod';
import { DEFAULT_ANTHROPIC_OPTIONS } from '../../llm/providers/gateway';
import { Sonnet4 } from '../../llm/sonnet-4';
import { createIdleTool } from '../../tools';
import { IDLE_TOOL_NAME } from '../../tools/communication-tools/idle-tool/idle-tool';
import {
createEditFileTool,
createLsTool,
createMultiEditFileTool,
createWriteFileTool,
} from '../../tools/file-tools';
import { createBashTool } from '../../tools/file-tools/bash-tool/bash-tool';
import { BASH_TOOL_NAME } from '../../tools/file-tools/bash-tool/bash-tool';
import { EDIT_FILE_TOOL_NAME } from '../../tools/file-tools/edit-file-tool/edit-file-tool';
import { GREP_TOOL_NAME, createGrepTool } from '../../tools/file-tools/grep-tool/grep-tool';
import { LS_TOOL_NAME } from '../../tools/file-tools/ls-tool/ls-tool';
import { MULTI_EDIT_FILE_TOOL_NAME } from '../../tools/file-tools/multi-edit-file-tool/multi-edit-file-tool';
import { READ_FILE_TOOL_NAME, createReadFileTool } from '../../tools/file-tools/read-file-tool/read-file-tool';
import { WRITE_FILE_TOOL_NAME } from '../../tools/file-tools/write-file-tool/write-file-tool';
import { createTaskTool } from '../../tools/task-tools/task-tool/task-tool';
import { type AgentContext, repairToolCall } from '../../utils/tool-call-repair';
import { createAnalyticsEngineerToolset } from './create-analytics-engineer-toolset';
import { getDocsAgentSystemPrompt as getAnalyticsEngineerAgentSystemPrompt } from './get-analytics-engineer-agent-system-prompt';
import type {
AnalyticsEngineerAgentOptions,
AnalyticsEngineerAgentStreamOptions,
TodoItem,
} from './types';
export const ANALYST_ENGINEER_AGENT_NAME = 'analyticsEngineerAgent';
const STOP_CONDITIONS = [stepCountIs(100), hasToolCall(IDLE_TOOL_NAME)];
export const TodoItemSchema = z.object({
id: z.string().describe('Unique identifier for the todo item. Use existing ID to update, or generate new ID for new items'),
content: z.string().describe('The content/description of the todo'),
status: z.enum(['pending', 'in_progress', 'completed']).describe('Current status of the todo'),
createdAt: z.string().datetime().optional().describe('ISO timestamp when todo was created (optional, will be set automatically for new items)'),
completedAt: z.string().datetime().optional().describe('ISO timestamp when todo was completed (optional)'),
});
export type TodoItem = z.infer<typeof TodoItemSchema>;
const AnalyticsEngineerAgentOptionsSchema = z.object({
folder_structure: z.string().describe('The file structure of the dbt repository'),
userId: z.string(),
chatId: z.string(),
dataSourceId: z.string(),
organizationId: z.string(),
messageId: z.string(),
todosList:
z
.array(TodoItemSchema)
.optional()
.describe('Array of todo items to write/update. Include all todos with their current state.'),
model: z
.custom<LanguageModelV2>()
.optional()
.describe('Custom language model to use (defaults to Sonnet4)'),
isSubagent: z
.boolean()
.optional()
.describe('Flag indicating this is a subagent (prevents infinite recursion)'),
abortSignal: z
.custom<AbortSignal>()
.optional()
.describe('Optional abort signal to cancel agent execution'),
});
const AnalyticsEngineerAgentStreamOptionsSchema = z.object({
messages: z.array(z.custom<ModelMessage>()).describe('The messages to send to the docs agent'),
});
export type AnalyticsEngineerAgentStreamOptions = z.infer<
typeof AnalyticsEngineerAgentStreamOptionsSchema
>;
export type AnalyticsEngineerAgentOptions = z.infer<
typeof AnalyticsEngineerAgentOptionsSchema
>;
export function createAnalyticsEngineerAgent(
analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions
) {
@ -89,7 +24,6 @@ export function createAnalyticsEngineerAgent(
} as ModelMessage;
async function stream({ messages }: AnalyticsEngineerAgentStreamOptions) {
const toolSet = await createAnalyticsEngineerToolset(analyticsEngineerAgentOptions);
const streamFn = () =>

65
packages/ai/src/agents/analytics-engineer-agent/create-analytics-engineer-toolset.ts
View File

@ -1,11 +1,32 @@
import type { AnalyticsEngineerAgentOptions } from "..";
import { createAnalyticsEngineerAgent } from "..";
import { BASH_TOOL_NAME, EDIT_FILE_TOOL_NAME, GREP_TOOL_NAME, IDLE_TOOL_NAME, LS_TOOL_NAME, MULTI_EDIT_FILE_TOOL_NAME, READ_FILE_TOOL_NAME, WRITE_FILE_TOOL_NAME, createBashTool, createEditFileTool, createGrepTool, createIdleTool, createLsTool, createMultiEditFileTool, createReadFileTool, createTaskTool, createWriteFileTool } from "../../tools";
import type { AgentFactory } from "../../tools/task-tools/task-tool/task-tool";
import {
BASH_TOOL_NAME,
EDIT_FILE_TOOL_NAME,
GREP_TOOL_NAME,
IDLE_TOOL_NAME,
LS_TOOL_NAME,
MULTI_EDIT_FILE_TOOL_NAME,
READ_FILE_TOOL_NAME,
TODO_WRITE_TOOL_NAME,
WRITE_FILE_TOOL_NAME,
createBashTool,
createEditFileTool,
createGrepTool,
createIdleTool,
createLsTool,
createMultiEditFileTool,
createReadFileTool,
createTaskTool,
createTodoWriteTool,
createWriteFileTool,
} from '../../tools';
import type { AgentFactory } from '../../tools/task-tools/task-tool/task-tool';
import { createAnalyticsEngineerAgent } from './analytics-engineer-agent';
import type { AnalyticsEngineerAgentOptions } from './types';
export async function createAnalyticsEngineerToolset(analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions) {
const idleTool = createIdleTool({
});
export async function createAnalyticsEngineerToolset(
analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions
) {
const idleTool = createIdleTool({});
const writeFileTool = createWriteFileTool({
messageId: analyticsEngineerAgentOptions.messageId,
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
@ -34,21 +55,26 @@ export async function createAnalyticsEngineerToolset(analyticsEngineerAgentOptio
messageId: analyticsEngineerAgentOptions.messageId,
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
});
const todosTool = createTodoWriteTool({
chatId: analyticsEngineerAgentOptions.chatId,
workingDirectory: analyticsEngineerAgentOptions.folder_structure,
todosList: analyticsEngineerAgentOptions.todosList,
});
// Conditionally create task tool (only for main agent, not for subagents)
const taskTool = !analyticsEngineerAgentOptions.isSubagent
? createTaskTool({
messageId: analyticsEngineerAgentOptions.messageId,
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
// Pass the agent factory function to enable task agent creation
// This needs to match the AgentFactory type signature
createAgent: ((options: AnalyticsEngineerAgentOptions) => {
return createAnalyticsEngineerAgent({
...options,
// Inherit model from parent agent if provided
model: analyticsEngineerAgentOptions.model,
});
}) as unknown as AgentFactory,
})
messageId: analyticsEngineerAgentOptions.messageId,
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
// Pass the agent factory function to enable task agent creation
// This needs to match the AgentFactory type signature
createAgent: ((options: AnalyticsEngineerAgentOptions) => {
return createAnalyticsEngineerAgent({
...options,
// Inherit model from parent agent if provided
model: analyticsEngineerAgentOptions.model,
});
}) as unknown as AgentFactory,
})
: null;
return {
@ -60,6 +86,7 @@ export async function createAnalyticsEngineerToolset(analyticsEngineerAgentOptio
[EDIT_FILE_TOOL_NAME]: editFileTool,
[MULTI_EDIT_FILE_TOOL_NAME]: multiEditFileTool,
[LS_TOOL_NAME]: lsTool,
[TODO_WRITE_TOOL_NAME]: todosTool,
...(taskTool ? { taskTool } : {}),
};
}

62
packages/ai/src/agents/analytics-engineer-agent/types.ts Normal file
View File

@ -0,0 +1,62 @@
import type { LanguageModelV2 } from '@ai-sdk/provider';
import type { ModelMessage } from 'ai';
import z from 'zod';
export const TodoItemSchema = z.object({
id: z
.string()
.describe(
'Unique identifier for the todo item. Use existing ID to update, or generate new ID for new items'
),
content: z.string().describe('The content/description of the todo'),
status: z.enum(['pending', 'in_progress', 'completed']).describe('Current status of the todo'),
createdAt: z
.string()
.datetime()
.optional()
.describe(
'ISO timestamp when todo was created (optional, will be set automatically for new items)'
),
completedAt: z
.string()
.datetime()
.optional()
.describe('ISO timestamp when todo was completed (optional)'),
});
export type TodoItem = z.infer<typeof TodoItemSchema>;
export const AnalyticsEngineerAgentOptionsSchema = z.object({
folder_structure: z.string().describe('The file structure of the dbt repository'),
userId: z.string(),
chatId: z.string(),
dataSourceId: z.string(),
organizationId: z.string(),
messageId: z.string(),
todosList: z
.array(TodoItemSchema)
.default([])
.describe('Array of todo items to write/update. Include all todos with their current state.'),
model: z
.custom<LanguageModelV2>()
.optional()
.describe('Custom language model to use (defaults to Sonnet4)'),
isSubagent: z
.boolean()
.optional()
.describe('Flag indicating this is a subagent (prevents infinite recursion)'),
abortSignal: z
.custom<AbortSignal>()
.optional()
.describe('Optional abort signal to cancel agent execution'),
});
export const AnalyticsEngineerAgentStreamOptionsSchema = z.object({
messages: z.array(z.custom<ModelMessage>()).describe('The messages to send to the docs agent'),
});
export type AnalyticsEngineerAgentStreamOptions = z.infer<
typeof AnalyticsEngineerAgentStreamOptionsSchema
>;
export type AnalyticsEngineerAgentOptions = z.infer<typeof AnalyticsEngineerAgentOptionsSchema>;

3
packages/ai/src/index.ts
View File

@ -34,6 +34,3 @@ export type {
LsToolInput,
LsToolOutput,
} from './tools/file-tools/ls-tool/ls-tool';
// Export typed tool events for type-safe tool callbacks
export type { ToolEvent, ToolEventCallback } from './agents/analytics-engineer-agent/tool-events';

2
packages/ai/src/tools/database-tools/super-execute-sql/super-execute-sql.ts
View File

@ -1,6 +1,6 @@
import { tool } from 'ai';
import { z } from 'zod';
import type { AnalyticsEngineerAgentOptions } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent';
import type { AnalyticsEngineerAgentOptions } from '../../../agents/analytics-engineer-agent/types';
import { createSuperExecuteSqlExecute } from './super-execute-sql-execute';
export const SuperExecuteSqlInputSchema = z.object({

70
packages/ai/src/tools/index.ts
View File

@ -1,32 +1,71 @@
// Communication tools
export { createDoneTool, DONE_TOOL_NAME } from './communication-tools/done-tool/done-tool';
export { createIdleTool, IDLE_TOOL_NAME } from './communication-tools/idle-tool/idle-tool';
export { createSubmitThoughtsTool, SUBMIT_THOUGHTS_TOOL_NAME } from './communication-tools/submit-thoughts-tool/submit-thoughts-tool';
export {
createSubmitThoughtsTool,
SUBMIT_THOUGHTS_TOOL_NAME,
} from './communication-tools/submit-thoughts-tool/submit-thoughts-tool';
// Planning/thinking tools
export { createSequentialThinkingTool, SEQUENTIAL_THINKING_TOOL_NAME } from './planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool';
export {
createSequentialThinkingTool,
SEQUENTIAL_THINKING_TOOL_NAME,
} from './planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool';
// Task tools
export { createTaskTool } from './task-tools/task-tool/task-tool';
// Visualization tools
export { createCreateMetricsTool, CREATE_METRICS_TOOL_NAME } from './visualization-tools/metrics/create-metrics-tool/create-metrics-tool';
export { createModifyMetricsTool, MODIFY_METRICS_TOOL_NAME } from './visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool';
export { createCreateDashboardsTool, CREATE_DASHBOARDS_TOOL_NAME } from './visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool';
export { createModifyDashboardsTool, MODIFY_DASHBOARDS_TOOL_NAME } from './visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool';
export { createCreateReportsTool, CREATE_REPORTS_TOOL_NAME } from './visualization-tools/reports/create-reports-tool/create-reports-tool';
export { createModifyReportsTool, MODIFY_REPORTS_TOOL_NAME } from './visualization-tools/reports/modify-reports-tool/modify-reports-tool';
export {
createCreateMetricsTool,
CREATE_METRICS_TOOL_NAME,
} from './visualization-tools/metrics/create-metrics-tool/create-metrics-tool';
export {
createModifyMetricsTool,
MODIFY_METRICS_TOOL_NAME,
} from './visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool';
export {
createCreateDashboardsTool,
CREATE_DASHBOARDS_TOOL_NAME,
} from './visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool';
export {
createModifyDashboardsTool,
MODIFY_DASHBOARDS_TOOL_NAME,
} from './visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool';
export {
createCreateReportsTool,
CREATE_REPORTS_TOOL_NAME,
} from './visualization-tools/reports/create-reports-tool/create-reports-tool';
export {
createModifyReportsTool,
MODIFY_REPORTS_TOOL_NAME,
} from './visualization-tools/reports/modify-reports-tool/modify-reports-tool';
// Database tools
export { createExecuteSqlTool, EXECUTE_SQL_TOOL_NAME } from './database-tools/execute-sql/execute-sql';
export {
createExecuteSqlTool,
EXECUTE_SQL_TOOL_NAME,
} from './database-tools/execute-sql/execute-sql';
export { executeSqlDocsAgent } from './database-tools/super-execute-sql/super-execute-sql';
// File tools
export { createLsTool, LS_TOOL_NAME } from './file-tools/ls-tool/ls-tool';
export { createReadFileTool, READ_FILE_TOOL_NAME } from './file-tools/read-file-tool/read-file-tool';
export { createWriteFileTool, WRITE_FILE_TOOL_NAME } from './file-tools/write-file-tool/write-file-tool';
export { createEditFileTool, EDIT_FILE_TOOL_NAME } from './file-tools/edit-file-tool/edit-file-tool';
export { createMultiEditFileTool, MULTI_EDIT_FILE_TOOL_NAME } from './file-tools/multi-edit-file-tool/multi-edit-file-tool';
export {
createReadFileTool,
READ_FILE_TOOL_NAME,
} from './file-tools/read-file-tool/read-file-tool';
export {
createWriteFileTool,
WRITE_FILE_TOOL_NAME,
} from './file-tools/write-file-tool/write-file-tool';
export {
createEditFileTool,
EDIT_FILE_TOOL_NAME,
} from './file-tools/edit-file-tool/edit-file-tool';
export {
createMultiEditFileTool,
MULTI_EDIT_FILE_TOOL_NAME,
} from './file-tools/multi-edit-file-tool/multi-edit-file-tool';
export { createBashTool, BASH_TOOL_NAME } from './file-tools/bash-tool/bash-tool';
export { createGrepTool, GREP_TOOL_NAME } from './file-tools/grep-tool/grep-tool';
@ -36,7 +75,10 @@ export { createWebSearchTool } from './web-tools/web-search-tool';
// More planning/thinking tools
export { createCheckOffTodoListTool } from './planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool';
export { createUpdateClarificationsFileTool } from './planning-thinking-tools/update-clarifications-file-tool/update-clarifications-file-tool';
export { createTodoWriteTool, TODO_WRITE_TOOL_NAME } from './planning-thinking-tools/todo-write-tool/todo-write-tool';
export {
createTodoWriteTool,
TODO_WRITE_TOOL_NAME,
} from './planning-thinking-tools/todo-write-tool/todo-write-tool';
// Legacy exports for backward compatibility (to be deprecated)
export { checkOffTodoList } from './planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool';

45
packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool-execute.ts
View File

@ -1,14 +1,18 @@
import type { TodoWriteToolContext, TodoWriteToolInput, TodoWriteToolOutput } from './todo-write-tool';
import type { TodoItem } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent';
import type { TodoItem } from '../../../agents/analytics-engineer-agent/types';
import type {
TodoWriteToolContext,
TodoWriteToolInput,
TodoWriteToolOutput,
} from './todo-write-tool';
/**
* Processes todos by setting timestamps and handling status changes
*/
function processTodos(inputTodos: TodoItem[], existingTodos: TodoItem[]): TodoItem[] {
const existingById = new Map(existingTodos.map(todo => [todo.id, todo]));
function processTodos(inputTodos: TodoItem[], existingTodos: TodoItem[] = []): TodoItem[] {
const existingById = new Map(existingTodos.map((todo) => [todo.id, todo]));
const now = new Date().toISOString();
return inputTodos.map(todo => {
return inputTodos.map((todo) => {
const existing = existingById.get(todo.id);
const processed = { ...todo };
@ -39,38 +43,19 @@ function processTodos(inputTodos: TodoItem[], existingTodos: TodoItem[]): TodoIt
*/
export function createTodoWriteToolExecute(context: TodoWriteToolContext) {
return async function execute(input: TodoWriteToolInput): Promise<TodoWriteToolOutput> {
const { chatId, workingDirectory } = context;
const { chatId, todosList = [] } = context;
const { todos: inputTodos } = input;
console.info(`Writing ${inputTodos.length} todo(s) for chat ${chatId}`);
try {
// Load existing todos from disk
let existingTodos: TodoItem[] = [];
try {
const loaded = await loadTodos(chatId, workingDirectory);
existingTodos = loaded?.todos || [];
} catch (error) {
console.warn('Failed to load existing todos:', error);
}
// Process the todos (set timestamps, handle status changes)
const processedTodos = processTodos(inputTodos, existingTodos);
const processedTodos = processTodos(inputTodos, todosList);
// Save to disk
try {
await saveTodos(chatId, workingDirectory, processedTodos);
} catch (error) {
console.error('Failed to save todos to disk:', error);
return {
success: false,
todos: processedTodos,
message: `Failed to save todos: ${error instanceof Error ? error.message : 'Unknown error'}`,
};
}
// Update the in-memory todosList by clearing and repopulating
todosList.splice(0, todosList.length, ...processedTodos);
console.info(`Successfully saved ${processedTodos.length} todo(s)`);
console.info(`Successfully updated ${processedTodos.length} todo(s) in memory`);
return {
success: true,
@ -83,7 +68,7 @@ export function createTodoWriteToolExecute(context: TodoWriteToolContext) {
return {
success: false,
todos: [],
todos: todosList,
message: `Error: ${errorMessage}`,
};
}

38
packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.test.ts
View File

@ -1,5 +1,5 @@
import { describe, expect, it, vi } from 'vitest';
import type { TodoItem } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent';
import type { TodoItem } from '../../../agents/analytics-engineer-agent/types';
import type { TodoWriteToolInput } from './todo-write-tool';
import { createTodoWriteToolExecute } from './todo-write-tool-execute';
@ -8,13 +8,9 @@ describe('createTodoWriteToolExecute', () => {
const workingDirectory = '/test/directory';
it('should create new todos with timestamps', async () => {
// Mock the conversation-history module
vi.doMock('../../../../../apps/cli/src/utils/conversation-history', () => ({
loadTodos: vi.fn().mockResolvedValue(null),
saveTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: [], updatedAt: new Date().toISOString() }),
}));
const todosList: TodoItem[] = [];
const execute = createTodoWriteToolExecute({ chatId, workingDirectory });
const execute = createTodoWriteToolExecute({ chatId, workingDirectory, todosList });
const input: TodoWriteToolInput = {
todos: [
{
@ -36,11 +32,12 @@ describe('createTodoWriteToolExecute', () => {
expect(result.todos).toHaveLength(2);
expect(result.todos[0]?.createdAt).toBeDefined();
expect(result.todos[1]?.createdAt).toBeDefined();
expect(todosList).toHaveLength(2);
});
it('should preserve createdAt for existing todos', async () => {
const existingCreatedAt = '2024-01-01T00:00:00.000Z';
const existingTodos: TodoItem[] = [
const todosList: TodoItem[] = [
{
id: '1',
content: 'Existing todo',
@ -49,12 +46,7 @@ describe('createTodoWriteToolExecute', () => {
},
];
vi.doMock('../../../../../apps/cli/src/utils/conversation-history', () => ({
loadTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }),
saveTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }),
}));
const execute = createTodoWriteToolExecute({ chatId, workingDirectory });
const execute = createTodoWriteToolExecute({ chatId, workingDirectory, todosList });
const input: TodoWriteToolInput = {
todos: [
{
@ -72,7 +64,7 @@ describe('createTodoWriteToolExecute', () => {
});
it('should set completedAt when status changes to completed', async () => {
const existingTodos: TodoItem[] = [
const todosList: TodoItem[] = [
{
id: '1',
content: 'Todo to complete',
@ -81,12 +73,7 @@ describe('createTodoWriteToolExecute', () => {
},
];
vi.doMock('../../../../../apps/cli/src/utils/conversation-history', () => ({
loadTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }),
saveTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }),
}));
const execute = createTodoWriteToolExecute({ chatId, workingDirectory });
const execute = createTodoWriteToolExecute({ chatId, workingDirectory, todosList });
const input: TodoWriteToolInput = {
todos: [
{
@ -104,7 +91,7 @@ describe('createTodoWriteToolExecute', () => {
});
it('should clear completedAt when status changes from completed', async () => {
const existingTodos: TodoItem[] = [
const todosList: TodoItem[] = [
{
id: '1',
content: 'Completed todo',
@ -114,12 +101,7 @@ describe('createTodoWriteToolExecute', () => {
},
];
vi.doMock('../../../../../apps/cli/src/utils/conversation-history', () => ({
loadTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }),
saveTodos: vi.fn().mockResolvedValue({ chatId, workingDirectory, todos: existingTodos, updatedAt: new Date().toISOString() }),
}));
const execute = createTodoWriteToolExecute({ chatId, workingDirectory });
const execute = createTodoWriteToolExecute({ chatId, workingDirectory, todosList });
const input: TodoWriteToolInput = {
todos: [
{

13
packages/ai/src/tools/planning-thinking-tools/todo-write-tool/todo-write-tool.ts
View File

@ -1,6 +1,6 @@
import { tool } from 'ai';
import { z } from 'zod';
import { TodoItemSchema } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent';
import { TodoItemSchema } from '../../../agents/analytics-engineer-agent/types';
import { createTodoWriteToolExecute } from './todo-write-tool-execute';
export const TODO_WRITE_TOOL_NAME = 'todoWrite';
@ -12,21 +12,24 @@ const TodoWriteToolOutputSchema = z.object({
});
const TodoWriteToolInputSchema = z.object({
todos: z.array(TodoItemSchema).describe('Array of todo items to write/update. Include all todos with their current state.'),
todos: z
.array(TodoItemSchema)
.describe('Array of todo items to write/update. Include all todos with their current state.'),
});
const TodoWriteToolContextSchema = z.object({
chatId: z.string().describe('The chat/conversation ID to associate todos with'),
workingDirectory: z.string().describe('The working directory for the chat'),
todosList: z.array(TodoItemSchema).default([]).describe('In-memory array of todo items to manipulate'),
});
export type TodoWriteToolInput = z.infer<typeof TodoWriteToolInputSchema>;
export type TodoWriteToolOutput = z.infer<typeof TodoWriteToolOutputSchema>;
export type TodoWriteToolContext = z.infer<typeof TodoWriteToolContextSchema>;
export function createTodoWriteTool<TAgentContext extends TodoWriteToolContext = TodoWriteToolContext>(
context: TAgentContext
) {
export function createTodoWriteTool<
TAgentContext extends TodoWriteToolContext = TodoWriteToolContext,
>(context: TAgentContext) {
const execute = createTodoWriteToolExecute(context);
return tool({