mirror of https://github.com/buster-so/buster.git
changes and clean up
This commit is contained in:
parent
ce9e51e481
commit
6c77a4f7a3
|
@ -59,7 +59,6 @@ export async function runAnalyticsEngineerAgent(params: RunAnalyticsEngineerAgen
|
|||
messageId: randomUUID(),
|
||||
model: proxyModel,
|
||||
abortSignal,
|
||||
skipTracing: true, // Skip Braintrust tracing in CLI mode
|
||||
});
|
||||
|
||||
// Use conversation history - includes user messages, assistant messages, tool calls, and tool results
|
||||
|
|
|
@ -15,6 +15,26 @@ const ConversationSchema = z.object({
|
|||
|
||||
export type Conversation = z.infer<typeof ConversationSchema>;
|
||||
|
||||
// Schema for a todo item
|
||||
const TodoItemSchema = z.object({
|
||||
id: z.string().describe('Unique identifier for the todo item'),
|
||||
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().describe('ISO timestamp when todo was created'),
|
||||
completedAt: z.string().datetime().optional().describe('ISO timestamp when todo was completed'),
|
||||
});
|
||||
|
||||
// Schema for todos associated with a chat
|
||||
const TodoListSchema = z.object({
|
||||
chatId: z.string().uuid().describe('Unique chat/conversation ID'),
|
||||
workingDirectory: z.string().describe('Absolute path of the working directory'),
|
||||
updatedAt: z.string().datetime().describe('ISO timestamp when todos were last updated'),
|
||||
todos: z.array(TodoItemSchema).describe('Array of todo items'),
|
||||
});
|
||||
|
||||
export type TodoItem = z.infer<typeof TodoItemSchema>;
|
||||
export type TodoList = z.infer<typeof TodoListSchema>;
|
||||
|
||||
// Base directory for all history
|
||||
const HISTORY_DIR = join(homedir(), '.buster', 'history');
|
||||
|
||||
|
@ -188,3 +208,52 @@ export async function deleteConversation(chatId: string, workingDirectory: strin
|
|||
const { unlink } = await import('node:fs/promises');
|
||||
await unlink(filePath);
|
||||
}
|
||||
|
||||
/**
|
||||
* Gets the file path for todos associated with a chat
|
||||
*/
|
||||
function getTodoFilePath(chatId: string, workingDirectory: string): string {
|
||||
return join(getHistoryDir(workingDirectory), `${chatId}.todos.json`);
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads todos for a specific chat
|
||||
* Returns null if no todos exist for this chat
|
||||
*/
|
||||
export async function loadTodos(
|
||||
chatId: string,
|
||||
workingDirectory: string
|
||||
): Promise<TodoList | null> {
|
||||
try {
|
||||
const filePath = getTodoFilePath(chatId, workingDirectory);
|
||||
const data = await readFile(filePath, 'utf-8');
|
||||
const parsed = JSON.parse(data);
|
||||
return TodoListSchema.parse(parsed);
|
||||
} catch (error) {
|
||||
// File doesn't exist or is invalid
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Saves todos for a specific chat
|
||||
*/
|
||||
export async function saveTodos(
|
||||
chatId: string,
|
||||
workingDirectory: string,
|
||||
todos: TodoItem[]
|
||||
): Promise<TodoList> {
|
||||
await ensureHistoryDir(workingDirectory);
|
||||
|
||||
const todoList: TodoList = {
|
||||
chatId,
|
||||
workingDirectory,
|
||||
updatedAt: new Date().toISOString(),
|
||||
todos,
|
||||
};
|
||||
|
||||
const filePath = getTodoFilePath(chatId, workingDirectory);
|
||||
await writeFile(filePath, JSON.stringify(todoList, null, 2), { mode: 0o600 });
|
||||
|
||||
return todoList;
|
||||
}
|
||||
|
|
|
@ -108,7 +108,7 @@ export async function deployHandler(
|
|||
};
|
||||
|
||||
// Upsert the dataset
|
||||
const { updated } = await upsertDataset(tx, datasetParams);
|
||||
const { updated } = await upsertDataset(datasetParams);
|
||||
|
||||
if (updated) {
|
||||
modelResult.updated.push({
|
||||
|
|
|
@ -1,29 +0,0 @@
|
|||
import type { Sandbox } from '@buster/sandbox';
|
||||
import { z } from 'zod';
|
||||
|
||||
// Best practice: Use const object for keys
|
||||
export const AnalyticsEngineerAgentContextKeys = {
|
||||
TodoList: 'todoList',
|
||||
ClarificationQuestions: 'clarificationQuestions',
|
||||
DataSourceId: 'dataSourceId',
|
||||
} as const;
|
||||
|
||||
// Extract type from const object
|
||||
export type AnalyticsEngineerAgentContextKey = keyof typeof AnalyticsEngineerAgentContextKeys;
|
||||
|
||||
export const ClarifyingQuestionSchema = z.object({
|
||||
issue: z.string(),
|
||||
context: z.string(),
|
||||
clarificationQuestion: z.string(),
|
||||
});
|
||||
|
||||
export type MessageUserClarifyingQuestion = z.infer<typeof ClarifyingQuestionSchema>;
|
||||
|
||||
// Use the const keys in your schema
|
||||
export const AnalyticsEngineerAgentContextSchema = z.object({
|
||||
[AnalyticsEngineerAgentContextKeys.TodoList]: z.string(),
|
||||
[AnalyticsEngineerAgentContextKeys.ClarificationQuestions]: z.array(ClarifyingQuestionSchema),
|
||||
[AnalyticsEngineerAgentContextKeys.DataSourceId]: z.string().uuid(),
|
||||
});
|
||||
|
||||
export type AnalyticsEngineerAgentContext = z.infer<typeof AnalyticsEngineerAgentContextSchema>;
|
|
@ -23,13 +23,23 @@ import { READ_FILE_TOOL_NAME, createReadFileTool } from '../../tools/file-tools/
|
|||
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 { ToolEventCallback } from './tool-events';
|
||||
|
||||
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(),
|
||||
|
@ -37,14 +47,11 @@ const AnalyticsEngineerAgentOptionsSchema = z.object({
|
|||
dataSourceId: z.string(),
|
||||
organizationId: z.string(),
|
||||
messageId: z.string(),
|
||||
sandbox: z
|
||||
.custom<Sandbox>(
|
||||
(val) => {
|
||||
return val && typeof val === 'object' && 'id' in val && 'fs' in val;
|
||||
},
|
||||
{ message: 'Invalid Sandbox instance' }
|
||||
)
|
||||
.optional(),
|
||||
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()
|
||||
|
@ -57,22 +64,20 @@ const AnalyticsEngineerAgentOptionsSchema = z.object({
|
|||
.custom<AbortSignal>()
|
||||
.optional()
|
||||
.describe('Optional abort signal to cancel agent execution'),
|
||||
skipTracing: z.boolean().optional().describe('Skip Braintrust tracing (useful for CLI mode)'),
|
||||
});
|
||||
|
||||
const AnalyticsEngineerAgentStreamOptionsSchema = z.object({
|
||||
messages: z.array(z.custom<ModelMessage>()).describe('The messages to send to the docs agent'),
|
||||
});
|
||||
|
||||
export type AnalyticsEngineerAgentOptions = z.infer<typeof AnalyticsEngineerAgentOptionsSchema> & {
|
||||
onToolEvent?: ToolEventCallback;
|
||||
};
|
||||
export type AnalyticsEngineerAgentStreamOptions = z.infer<
|
||||
typeof AnalyticsEngineerAgentStreamOptionsSchema
|
||||
>;
|
||||
|
||||
// Extended type for passing to tools (includes sandbox)
|
||||
export type DocsAgentContextWithSandbox = AnalyticsEngineerAgentOptions & { sandbox: Sandbox };
|
||||
export type AnalyticsEngineerAgentOptions = z.infer<
|
||||
typeof AnalyticsEngineerAgentOptionsSchema
|
||||
>;
|
||||
|
||||
|
||||
export function createAnalyticsEngineerAgent(
|
||||
analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions
|
||||
|
@ -83,120 +88,23 @@ export function createAnalyticsEngineerAgent(
|
|||
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
||||
} as ModelMessage;
|
||||
|
||||
const idleTool = createIdleTool({
|
||||
});
|
||||
const writeFileTool = createWriteFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const grepTool = createGrepTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const readFileTool = createReadFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const bashTool = createBashTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const editFileTool = createEditFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const multiEditFileTool = createMultiEditFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const lsTool = createLsTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
|
||||
// Conditionally create task tool (only for main agent, not for subagents)
|
||||
const taskTool = !analyticsEngineerAgentOptions.isSubagent
|
||||
? createTaskTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
// Wrap onToolEvent to ensure type compatibility - no optional chaining
|
||||
...(analyticsEngineerAgentOptions.onToolEvent && {
|
||||
// biome-ignore lint/suspicious/noExplicitAny: Bridging between ToolEventType and parent's ToolEvent type
|
||||
onToolEvent: (event) => analyticsEngineerAgentOptions.onToolEvent?.(event as any),
|
||||
}),
|
||||
// Pass the agent factory function to enable task agent creation
|
||||
// This needs to match the AgentFactory type signature
|
||||
createAgent: (options) => {
|
||||
return createAnalyticsEngineerAgent({
|
||||
...options,
|
||||
// Inherit model from parent agent if provided
|
||||
model: analyticsEngineerAgentOptions.model,
|
||||
});
|
||||
},
|
||||
})
|
||||
: null;
|
||||
|
||||
// Create planning tools with simple context
|
||||
async function stream({ messages }: AnalyticsEngineerAgentStreamOptions) {
|
||||
const agentContext: AgentContext = {
|
||||
agentName: ANALYST_ENGINEER_AGENT_NAME,
|
||||
availableTools: [IDLE_TOOL_NAME, GREP_TOOL_NAME, WRITE_FILE_TOOL_NAME, READ_FILE_TOOL_NAME, BASH_TOOL_NAME, EDIT_FILE_TOOL_NAME, MULTI_EDIT_FILE_TOOL_NAME, LS_TOOL_NAME],
|
||||
};
|
||||
|
||||
// Build tools object conditionally including task tool
|
||||
// biome-ignore lint/suspicious/noExplicitAny: tools object contains various tool types
|
||||
const tools: Record<string, any> = {
|
||||
[IDLE_TOOL_NAME]: idleTool,
|
||||
[GREP_TOOL_NAME]: grepTool,
|
||||
[WRITE_FILE_TOOL_NAME]: writeFileTool,
|
||||
[READ_FILE_TOOL_NAME]: readFileTool,
|
||||
[BASH_TOOL_NAME]: bashTool,
|
||||
[EDIT_FILE_TOOL_NAME]: editFileTool,
|
||||
[MULTI_EDIT_FILE_TOOL_NAME]: multiEditFileTool,
|
||||
[LS_TOOL_NAME]: lsTool,
|
||||
};
|
||||
|
||||
// Add task tool only if not a subagent (prevent recursion)
|
||||
if (taskTool) {
|
||||
tools.taskTool = taskTool;
|
||||
}
|
||||
const toolSet = await createAnalyticsEngineerToolset(analyticsEngineerAgentOptions);
|
||||
|
||||
const streamFn = () =>
|
||||
streamText({
|
||||
model: analyticsEngineerAgentOptions.model || Sonnet4,
|
||||
providerOptions: DEFAULT_ANTHROPIC_OPTIONS,
|
||||
tools,
|
||||
tools: toolSet,
|
||||
messages: [systemMessage, ...messages],
|
||||
stopWhen: STOP_CONDITIONS,
|
||||
toolChoice: 'required',
|
||||
maxOutputTokens: 10000,
|
||||
temperature: 0,
|
||||
experimental_context: analyticsEngineerAgentOptions,
|
||||
...(analyticsEngineerAgentOptions.abortSignal && {
|
||||
abortSignal: analyticsEngineerAgentOptions.abortSignal,
|
||||
}),
|
||||
experimental_repairToolCall: async (repairContext) => {
|
||||
return repairToolCall({
|
||||
toolCall: repairContext.toolCall,
|
||||
tools: repairContext.tools,
|
||||
error: repairContext.error,
|
||||
messages: repairContext.messages,
|
||||
...(repairContext.system && { system: repairContext.system }),
|
||||
...(repairContext.inputSchema && { inputSchema: repairContext.inputSchema }),
|
||||
agentContext,
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
// Skip tracing for CLI mode
|
||||
if (analyticsEngineerAgentOptions.skipTracing) {
|
||||
return streamFn();
|
||||
}
|
||||
|
||||
// Use Braintrust tracing for production
|
||||
return wrapTraced(streamFn, {
|
||||
name: 'Docs Agent',
|
||||
})();
|
||||
return streamFn();
|
||||
}
|
||||
|
||||
return {
|
||||
|
|
|
@ -0,0 +1,65 @@
|
|||
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";
|
||||
|
||||
export async function createAnalyticsEngineerToolset(analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions) {
|
||||
const idleTool = createIdleTool({
|
||||
});
|
||||
const writeFileTool = createWriteFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const grepTool = createGrepTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const readFileTool = createReadFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const bashTool = createBashTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const editFileTool = createEditFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const multiEditFileTool = createMultiEditFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const lsTool = createLsTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
// 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,
|
||||
})
|
||||
: null;
|
||||
|
||||
return {
|
||||
[IDLE_TOOL_NAME]: idleTool,
|
||||
[WRITE_FILE_TOOL_NAME]: writeFileTool,
|
||||
[GREP_TOOL_NAME]: grepTool,
|
||||
[READ_FILE_TOOL_NAME]: readFileTool,
|
||||
[BASH_TOOL_NAME]: bashTool,
|
||||
[EDIT_FILE_TOOL_NAME]: editFileTool,
|
||||
[MULTI_EDIT_FILE_TOOL_NAME]: multiEditFileTool,
|
||||
[LS_TOOL_NAME]: lsTool,
|
||||
...(taskTool ? { taskTool } : {}),
|
||||
};
|
||||
}
|
|
@ -1,53 +0,0 @@
|
|||
import type { IdleInput, IdleOutput } from '../../tools/communication-tools/idle-tool/idle-tool';
|
||||
import type { BashToolInput, BashToolOutput } from '../../tools/file-tools/bash-tool/bash-tool';
|
||||
import type {
|
||||
EditFileToolInput,
|
||||
EditFileToolOutput,
|
||||
} from '../../tools/file-tools/edit-file-tool/edit-file-tool';
|
||||
import type { GrepToolInput, GrepToolOutput } from '../../tools/file-tools/grep-tool/grep-tool';
|
||||
import type { LsToolInput, LsToolOutput } from '../../tools/file-tools/ls-tool/ls-tool';
|
||||
import type {
|
||||
ReadFileToolInput,
|
||||
ReadFileToolOutput,
|
||||
} from '../../tools/file-tools/read-file-tool/read-file-tool';
|
||||
import type {
|
||||
WriteFileToolInput,
|
||||
WriteFileToolOutput,
|
||||
} from '../../tools/file-tools/write-file-tool/write-file-tool';
|
||||
|
||||
/**
|
||||
* Discriminated union of all possible tool events in the analytics engineer agent
|
||||
* This provides full type safety from tools -> CLI display
|
||||
*/
|
||||
export type ToolEvent =
|
||||
// Idle tool events
|
||||
| { tool: 'idleTool'; event: 'start'; args: IdleInput }
|
||||
| { tool: 'idleTool'; event: 'complete'; result: IdleOutput; args: IdleInput }
|
||||
// Bash tool events
|
||||
| { tool: 'bashTool'; event: 'start'; args: BashToolInput }
|
||||
| { tool: 'bashTool'; event: 'complete'; result: BashToolOutput; args: BashToolInput }
|
||||
// Grep tool events
|
||||
| { tool: 'grepTool'; event: 'start'; args: GrepToolInput }
|
||||
| { tool: 'grepTool'; event: 'complete'; result: GrepToolOutput; args: GrepToolInput }
|
||||
// Read file tool events
|
||||
| { tool: 'readFileTool'; event: 'start'; args: ReadFileToolInput }
|
||||
| { tool: 'readFileTool'; event: 'complete'; result: ReadFileToolOutput; args: ReadFileToolInput }
|
||||
// Write file tool events
|
||||
| { tool: 'writeFileTool'; event: 'start'; args: WriteFileToolInput }
|
||||
| {
|
||||
tool: 'writeFileTool';
|
||||
event: 'complete';
|
||||
result: WriteFileToolOutput;
|
||||
args: WriteFileToolInput;
|
||||
}
|
||||
// Edit file tool events
|
||||
| { tool: 'editFileTool'; event: 'start'; args: EditFileToolInput }
|
||||
| { tool: 'editFileTool'; event: 'complete'; result: EditFileToolOutput; args: EditFileToolInput }
|
||||
// Ls tool events
|
||||
| { tool: 'lsTool'; event: 'start'; args: LsToolInput }
|
||||
| { tool: 'lsTool'; event: 'complete'; result: LsToolOutput; args: LsToolInput };
|
||||
|
||||
/**
|
||||
* Callback type for tool events - single typed callback for all tools
|
||||
*/
|
||||
export type ToolEventCallback = (event: ToolEvent) => void;
|
|
@ -1,32 +1,42 @@
|
|||
export { createDoneTool } from './communication-tools/done-tool/done-tool';
|
||||
export { createIdleTool } from './communication-tools/idle-tool/idle-tool';
|
||||
export { createSubmitThoughtsTool } from './communication-tools/submit-thoughts-tool/submit-thoughts-tool';
|
||||
export { createSequentialThinkingTool } from './planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool';
|
||||
// Task tools - factory functions
|
||||
// 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';
|
||||
|
||||
// Planning/thinking tools
|
||||
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 - factory functions
|
||||
export { createCreateMetricsTool } from './visualization-tools/metrics/create-metrics-tool/create-metrics-tool';
|
||||
export { createModifyMetricsTool } from './visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool';
|
||||
export { createCreateDashboardsTool } from './visualization-tools/dashboards/create-dashboards-tool/create-dashboards-tool';
|
||||
export { createModifyDashboardsTool } from './visualization-tools/dashboards/modify-dashboards-tool/modify-dashboards-tool';
|
||||
export { createCreateReportsTool } from './visualization-tools/reports/create-reports-tool/create-reports-tool';
|
||||
export { createModifyReportsTool } from './visualization-tools/reports/modify-reports-tool/modify-reports-tool';
|
||||
export { createExecuteSqlTool } from './database-tools/execute-sql/execute-sql';
|
||||
|
||||
// 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';
|
||||
|
||||
// Database tools
|
||||
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 - factory functions
|
||||
export { createLsTool } from './file-tools/ls-tool/ls-tool';
|
||||
export { createReadFileTool } from './file-tools/read-file-tool/read-file-tool';
|
||||
export { createWriteFileTool } from './file-tools/write-file-tool/write-file-tool';
|
||||
export { createEditFileTool } from './file-tools/edit-file-tool/edit-file-tool';
|
||||
export { createMultiEditFileTool } from './file-tools/multi-edit-file-tool/multi-edit-file-tool';
|
||||
export { createBashTool } from './file-tools/bash-tool/bash-tool';
|
||||
export { createGrepTool } from './file-tools/grep-tool/grep-tool';
|
||||
// Web tools - factory functions
|
||||
|
||||
// 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 { createBashTool, BASH_TOOL_NAME } from './file-tools/bash-tool/bash-tool';
|
||||
export { createGrepTool, GREP_TOOL_NAME } from './file-tools/grep-tool/grep-tool';
|
||||
|
||||
// Web tools
|
||||
export { createWebSearchTool } from './web-tools/web-search-tool';
|
||||
|
||||
// Planning/thinking tools - factory functions
|
||||
// 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';
|
||||
|
||||
// Legacy exports for backward compatibility (to be deprecated)
|
||||
export { checkOffTodoList } from './planning-thinking-tools/check-off-todo-list-tool/check-off-todo-list-tool';
|
||||
|
|
|
@ -0,0 +1,91 @@
|
|||
import type { TodoWriteToolContext, TodoWriteToolInput, TodoWriteToolOutput } from './todo-write-tool';
|
||||
import type { TodoItem } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent';
|
||||
|
||||
/**
|
||||
* 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]));
|
||||
const now = new Date().toISOString();
|
||||
|
||||
return inputTodos.map(todo => {
|
||||
const existing = existingById.get(todo.id);
|
||||
const processed = { ...todo };
|
||||
|
||||
// Set createdAt for new todos
|
||||
if (!existing) {
|
||||
processed.createdAt = processed.createdAt || now;
|
||||
} else {
|
||||
// Keep existing createdAt
|
||||
processed.createdAt = existing.createdAt;
|
||||
}
|
||||
|
||||
// Set completedAt when status changes to completed
|
||||
if (processed.status === 'completed' && (!existing || existing.status !== 'completed')) {
|
||||
processed.completedAt = now;
|
||||
}
|
||||
|
||||
// Clear completedAt if status changes from completed to something else
|
||||
if (processed.status !== 'completed' && existing?.status === 'completed') {
|
||||
processed.completedAt = undefined;
|
||||
}
|
||||
|
||||
return processed;
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates the execute function for the todo write tool
|
||||
*/
|
||||
export function createTodoWriteToolExecute(context: TodoWriteToolContext) {
|
||||
return async function execute(input: TodoWriteToolInput): Promise<TodoWriteToolOutput> {
|
||||
const { chatId, workingDirectory } = 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);
|
||||
|
||||
// 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'}`,
|
||||
};
|
||||
}
|
||||
|
||||
console.info(`Successfully saved ${processedTodos.length} todo(s)`);
|
||||
|
||||
return {
|
||||
success: true,
|
||||
todos: processedTodos,
|
||||
message: `Successfully updated ${processedTodos.length} todo(s)`,
|
||||
};
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
|
||||
console.error('Error in todo write tool:', errorMessage);
|
||||
|
||||
return {
|
||||
success: false,
|
||||
todos: [],
|
||||
message: `Error: ${errorMessage}`,
|
||||
};
|
||||
}
|
||||
};
|
||||
}
|
|
@ -0,0 +1,138 @@
|
|||
import { describe, expect, it, vi } from 'vitest';
|
||||
import type { TodoItem } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent';
|
||||
import type { TodoWriteToolInput } from './todo-write-tool';
|
||||
import { createTodoWriteToolExecute } from './todo-write-tool-execute';
|
||||
|
||||
describe('createTodoWriteToolExecute', () => {
|
||||
const chatId = 'test-chat-id';
|
||||
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 execute = createTodoWriteToolExecute({ chatId, workingDirectory });
|
||||
const input: TodoWriteToolInput = {
|
||||
todos: [
|
||||
{
|
||||
id: '1',
|
||||
content: 'First todo',
|
||||
status: 'pending',
|
||||
},
|
||||
{
|
||||
id: '2',
|
||||
content: 'Second todo',
|
||||
status: 'in_progress',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const result = await execute(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.todos).toHaveLength(2);
|
||||
expect(result.todos[0]?.createdAt).toBeDefined();
|
||||
expect(result.todos[1]?.createdAt).toBeDefined();
|
||||
});
|
||||
|
||||
it('should preserve createdAt for existing todos', async () => {
|
||||
const existingCreatedAt = '2024-01-01T00:00:00.000Z';
|
||||
const existingTodos: TodoItem[] = [
|
||||
{
|
||||
id: '1',
|
||||
content: 'Existing todo',
|
||||
status: 'pending',
|
||||
createdAt: existingCreatedAt,
|
||||
},
|
||||
];
|
||||
|
||||
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 input: TodoWriteToolInput = {
|
||||
todos: [
|
||||
{
|
||||
id: '1',
|
||||
content: 'Updated todo',
|
||||
status: 'completed',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const result = await execute(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.todos[0]?.createdAt).toBe(existingCreatedAt);
|
||||
});
|
||||
|
||||
it('should set completedAt when status changes to completed', async () => {
|
||||
const existingTodos: TodoItem[] = [
|
||||
{
|
||||
id: '1',
|
||||
content: 'Todo to complete',
|
||||
status: 'in_progress',
|
||||
createdAt: '2024-01-01T00:00:00.000Z',
|
||||
},
|
||||
];
|
||||
|
||||
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 input: TodoWriteToolInput = {
|
||||
todos: [
|
||||
{
|
||||
id: '1',
|
||||
content: 'Todo to complete',
|
||||
status: 'completed',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const result = await execute(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.todos[0]?.completedAt).toBeDefined();
|
||||
});
|
||||
|
||||
it('should clear completedAt when status changes from completed', async () => {
|
||||
const existingTodos: TodoItem[] = [
|
||||
{
|
||||
id: '1',
|
||||
content: 'Completed todo',
|
||||
status: 'completed',
|
||||
createdAt: '2024-01-01T00:00:00.000Z',
|
||||
completedAt: '2024-01-02T00:00:00.000Z',
|
||||
},
|
||||
];
|
||||
|
||||
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 input: TodoWriteToolInput = {
|
||||
todos: [
|
||||
{
|
||||
id: '1',
|
||||
content: 'Completed todo',
|
||||
status: 'in_progress',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const result = await execute(input);
|
||||
|
||||
expect(result.success).toBe(true);
|
||||
expect(result.todos[0]?.completedAt).toBeUndefined();
|
||||
});
|
||||
});
|
|
@ -0,0 +1,38 @@
|
|||
import { tool } from 'ai';
|
||||
import { z } from 'zod';
|
||||
import { TodoItemSchema } from '../../../agents/analytics-engineer-agent/analytics-engineer-agent';
|
||||
import { createTodoWriteToolExecute } from './todo-write-tool-execute';
|
||||
|
||||
export const TODO_WRITE_TOOL_NAME = 'todoWrite';
|
||||
|
||||
const TodoWriteToolOutputSchema = z.object({
|
||||
success: z.boolean().describe('Whether the operation was successful'),
|
||||
todos: z.array(TodoItemSchema).describe('The full current state of all todos'),
|
||||
message: z.string().optional().describe('Optional message about the operation'),
|
||||
});
|
||||
|
||||
const TodoWriteToolInputSchema = z.object({
|
||||
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'),
|
||||
});
|
||||
|
||||
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
|
||||
) {
|
||||
const execute = createTodoWriteToolExecute(context);
|
||||
|
||||
return tool({
|
||||
description: `Write and manage todo items for the current chat session. Accepts an array of todo items with their current state (pending, in_progress, or completed). All todos are persisted to disk and associated with the current chat. Returns the full current state of all todos.`,
|
||||
inputSchema: TodoWriteToolInputSchema,
|
||||
outputSchema: TodoWriteToolOutputSchema,
|
||||
execute,
|
||||
});
|
||||
}
|
Loading…
Reference in New Issue