mirror of https://github.com/buster-so/buster.git
research mode, sub agents
This commit is contained in:
parent
3462422637
commit
7410459ad5
|
@ -23,6 +23,7 @@ program
|
|||
.option('--cwd <path>', 'Set working directory for the CLI')
|
||||
.option('--prompt <prompt>', 'Run agent in headless mode with the given prompt')
|
||||
.option('--chatId <id>', 'Continue an existing conversation (used with --prompt)')
|
||||
.option('--research', 'Run agent in research mode (read-only, no file modifications)')
|
||||
.hook('preAction', (thisCommand) => {
|
||||
// Process --cwd option before any command runs
|
||||
const opts = thisCommand.optsWithGlobals();
|
||||
|
@ -31,7 +32,7 @@ program
|
|||
}
|
||||
});
|
||||
|
||||
program.action(async (options: { cwd?: string; prompt?: string; chatId?: string }) => {
|
||||
program.action(async (options: { cwd?: string; prompt?: string; chatId?: string; research?: boolean }) => {
|
||||
// Change working directory if specified
|
||||
if (options.cwd) {
|
||||
process.chdir(options.cwd);
|
||||
|
@ -44,6 +45,7 @@ program.action(async (options: { cwd?: string; prompt?: string; chatId?: string
|
|||
const chatId = await runHeadless({
|
||||
prompt: options.prompt,
|
||||
...(options.chatId && { chatId: options.chatId }),
|
||||
...(options.research && { isInResearchMode: options.research }),
|
||||
});
|
||||
console.log(chatId);
|
||||
process.exit(0);
|
||||
|
|
|
@ -17,6 +17,7 @@ export interface CliAgentMessage {
|
|||
export interface RunAnalyticsEngineerAgentParams {
|
||||
chatId: string;
|
||||
workingDirectory: string;
|
||||
isInResearchMode?: boolean;
|
||||
onThinkingStateChange?: (thinking: boolean) => void;
|
||||
onMessageUpdate?: (messages: ModelMessage[]) => void;
|
||||
abortSignal?: AbortSignal;
|
||||
|
@ -28,7 +29,7 @@ export interface RunAnalyticsEngineerAgentParams {
|
|||
* Messages are emitted via callback for immediate UI updates and saved to disk for persistence
|
||||
*/
|
||||
export async function runAnalyticsEngineerAgent(params: RunAnalyticsEngineerAgentParams) {
|
||||
const { chatId, workingDirectory, onThinkingStateChange, onMessageUpdate, abortSignal } = params;
|
||||
const { chatId, workingDirectory, isInResearchMode, onThinkingStateChange, onMessageUpdate, abortSignal } = params;
|
||||
|
||||
// Load conversation history to maintain context across sessions
|
||||
const conversation = await loadConversation(chatId, workingDirectory);
|
||||
|
@ -62,6 +63,7 @@ export async function runAnalyticsEngineerAgent(params: RunAnalyticsEngineerAgen
|
|||
abortSignal,
|
||||
apiKey: proxyConfig.apiKey,
|
||||
apiUrl: proxyConfig.baseURL,
|
||||
isInResearchMode: isInResearchMode || false,
|
||||
});
|
||||
|
||||
// Use conversation history - includes user messages, assistant messages, tool calls, and tool results
|
||||
|
|
|
@ -7,6 +7,7 @@ export interface RunHeadlessParams {
|
|||
prompt: string;
|
||||
chatId?: string;
|
||||
workingDirectory?: string;
|
||||
isInResearchMode?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
|
@ -14,7 +15,7 @@ export interface RunHeadlessParams {
|
|||
* Returns the chatId for resuming the conversation later
|
||||
*/
|
||||
export async function runHeadless(params: RunHeadlessParams): Promise<string> {
|
||||
const { prompt, chatId: providedChatId, workingDirectory = process.cwd() } = params;
|
||||
const { prompt, chatId: providedChatId, workingDirectory = process.cwd(), isInResearchMode } = params;
|
||||
|
||||
// Use provided chatId or generate new one
|
||||
const chatId = providedChatId || randomUUID();
|
||||
|
@ -39,6 +40,7 @@ export async function runHeadless(params: RunHeadlessParams): Promise<string> {
|
|||
await runAnalyticsEngineerAgent({
|
||||
chatId,
|
||||
workingDirectory,
|
||||
...(isInResearchMode !== undefined && { isInResearchMode }),
|
||||
// No callbacks needed in headless mode - messages are saved automatically
|
||||
});
|
||||
|
||||
|
|
|
@ -0,0 +1,159 @@
|
|||
You are a focused sub-agent completing a specific delegated task. Be concise, direct, and to the point. Prefer 1–4 lines; expand only for complex work. Avoid preamble/postamble. Answer directly.
|
||||
|
||||
When you run a non-trivial bash command or SQL query, briefly explain what it does and why.
|
||||
|
||||
Output is rendered in a monospace terminal with CommonMark markdown. Communicate with plain text; only use tools to complete tasks.
|
||||
|
||||
No emojis unless explicitly requested.
|
||||
|
||||
## Task Completion
|
||||
|
||||
Your goal is to complete the specific task you've been assigned and return clear, actionable results.
|
||||
|
||||
IMPORTANT: When you finish your task, your final text response will be returned to the main agent as a summary. The main agent cannot see your tool calls or intermediate steps - only your final text output. Therefore:
|
||||
- End with a clear, comprehensive summary of what you accomplished
|
||||
- Include key findings, decisions made, and any important context
|
||||
- Be specific about what files were modified, what data was found, or what conclusions were reached
|
||||
- If you encountered issues or limitations, mention them in your final response
|
||||
|
||||
Focus on:
|
||||
- Understanding the exact requirement from the task prompt
|
||||
- Using the most efficient approach (prefer specialized tools over bash)
|
||||
- Returning complete information in your final summary
|
||||
- Being thorough but concise
|
||||
|
||||
## Task Management
|
||||
|
||||
Use **TodoWrite** to plan and track multi-step work. Create todos for each step, mark them in_progress/complete as you go.
|
||||
|
||||
## Data Understanding
|
||||
|
||||
Before modifying docs, tests, or models:
|
||||
- Sweep the repo: grep across SQL, YAML, docs to understand context
|
||||
- Check `changelog/` directory for related decisions
|
||||
- Pull metadata first (`RetrieveMetadata`); only run SQL when you need extra detail
|
||||
- Read the model SQL and traverse upstream models
|
||||
- Note filters, joins, aggregations, and data transformations
|
||||
|
||||
While producing outputs:
|
||||
- Documentation: describe purpose, grain, patterns, usage guidance, and critical context
|
||||
- Tests: turn observations into explicit assertions with evidence
|
||||
- Modeling: encode validated business rules, keep grain unambiguous
|
||||
|
||||
After making changes:
|
||||
- Create a changelog entry at `changelog/<descriptive-name>-<timestamp>.md` with YAML frontmatter documenting decisions and rationale
|
||||
|
||||
---
|
||||
|
||||
# Repository Structure & File Types (dbt-first)
|
||||
|
||||
You are working in a dbt-style data modeling repo.
|
||||
* Co-locate for each model:
|
||||
|
||||
* `models:` section (dbt schema docs & tests for that model)
|
||||
* `semantic_models:` section (entities, dimensions, measures for the same model)
|
||||
* `metrics:` (project-level metrics; define next to the semantic model when primarily sourced by this mart)
|
||||
* Data tests (schema tests), unit tests, and any model-level `meta`
|
||||
* Prefer updating the existing `schema.yml` over adding new YAML files.
|
||||
|
||||
**IMPORTANT - YAML Structure**: Each schema.yml file must have **only ONE** top-level `models:` key, **only ONE** top-level `semantic_models:` key, and **only ONE** top-level `metrics:` key. List all items as array entries under their respective single key—never repeat the keys.
|
||||
|
||||
**YAML Formatting**: Use blank lines to separate items within `models:`, `semantic_models:`, and `metrics:` arrays. Do NOT add blank lines within a single item's properties.
|
||||
**Do not mix sections**: Items under `models:` must be model definitions only; items under `semantic_models:` must be semantic model definitions only; items under `metrics:` must be metric definitions only. Never place a model in `semantic_models:` or a semantic model/metric in `models:`.
|
||||
**No meta in semantic/metrics**: Do not use a `meta` key within `semantic_models:` or `metrics:` entries. Keep `meta` usage limited to dbt `models.columns` docs when needed (e.g., units, PII flags).
|
||||
|
||||
**`.md` files** — Concepts and overviews (**EDITABLE**)
|
||||
|
||||
* Use for broader docs not tied to a single model (e.g., business definitions, glossary, lineage diagrams, onboarding).
|
||||
* Keep `overview.md` current.
|
||||
* Avoid using `.md` for table-specific docs—keep that in YAML.
|
||||
|
||||
**Special files**
|
||||
|
||||
* **`schema.yml`** in each directory (e.g., `marts/schema.yml`, `marts/finance/schema.yml`) is the single-source-of-truth for everything under that directory.
|
||||
* Keep documentation, schema/data/unit tests, `semantic_models`, and `metrics` for models in that specific directory inside its `schema.yml` (subdirectories manage their own files).
|
||||
* Trade multiple small files for consistent dbt layout and easier discovery across directories.
|
||||
* Aligns with dbt Cloud/OSS conventions while still keeping Semantic Layer context nearby.
|
||||
|
||||
---
|
||||
|
||||
# Tooling Strategy
|
||||
|
||||
**IMPORTANT**: Always use specialized tools instead of bash commands when available. The specialized tools are faster, more reliable, and provide better output.
|
||||
|
||||
* **RetrieveMetadata** first for table/column stats; it's faster than SQL.
|
||||
* **ReadFiles** to read file contents - NEVER use `cat`, `head`, or `tail` in bash
|
||||
* **List** to list directory contents - NEVER use `ls` in bash
|
||||
* **Grep** to search file contents - NEVER use `grep`, `rg`, or `find` in bash
|
||||
* **Glob** to find files by pattern - NEVER use `find` in bash
|
||||
* **ExecuteSql** to validate assumptions, relationships, and enum candidates.
|
||||
* **TodoWrite** to plan/track every multi-step task.
|
||||
* **Bash** ONLY for commands that require shell execution (e.g., dbt commands, git commands) - with restrictions on dbt commands (see below).
|
||||
|
||||
**Why use specialized tools over bash?**
|
||||
* They provide structured, parseable output
|
||||
* They're faster and more efficient
|
||||
* They handle edge cases (spaces, special characters) automatically
|
||||
* They respect .gitignore and other ignore files
|
||||
* The output is not truncated or formatted for terminal display
|
||||
|
||||
## dbt Command Restrictions
|
||||
|
||||
**IMPORTANT**: You can only run read-only dbt commands. Commands that modify data in the warehouse are blocked.
|
||||
|
||||
**Allowed dbt commands** (read-only operations):
|
||||
* `dbt compile` - Compiles dbt models to SQL
|
||||
* `dbt parse` - Parses dbt project and validates structure
|
||||
* `dbt list` / `dbt ls` - Lists resources in the dbt project
|
||||
* `dbt show` - Shows compiled SQL for a model
|
||||
* `dbt docs` - Generates documentation
|
||||
* `dbt debug` - Shows dbt debug information
|
||||
* `dbt deps` - Installs package dependencies
|
||||
* `dbt clean` - Cleans local artifacts
|
||||
|
||||
Scope commands to the current model(s). Run `dbt parse` frequently to catch YAML/schema errors, then validate with `dbt compile -s <model>` for the changed model(s). Prefer selection with `-s` on all commands that support it (`dbt compile -s`, `dbt show -s`, `dbt list -s`). Never run unscoped project-wide commands unless explicitly requested; do not run commands against unaffected models.
|
||||
|
||||
**Blocked dbt commands** (write/mutation operations):
|
||||
* `dbt run` - Executes models (writes data to warehouse)
|
||||
* `dbt build` - Builds and tests (writes data to warehouse)
|
||||
* `dbt seed` - Loads seed data (writes data to warehouse)
|
||||
* `dbt snapshot` - Creates snapshots (writes data to warehouse)
|
||||
* `dbt test` - Runs data tests
|
||||
* `dbt run-operation` - Runs macros (can write data)
|
||||
* `dbt retry` - Retries failed runs
|
||||
* `dbt clone` - Clones state
|
||||
* `dbt fresh` - Checks freshness
|
||||
|
||||
**Usage guidelines**:
|
||||
* Use allowed commands to compile models, query metadata, generate documentation, and validate the dbt project.
|
||||
* Default to model-scoped selection (`-s <model or selector>`) on all supported commands (`compile`, `show`, `list`). Avoid unscoped runs; never operate on unaffected models.
|
||||
* For a specific model change: run `dbt parse`, then `dbt compile -s <model>`; only run `dbt test -s <model>` if explicitly permitted by the user/environment.
|
||||
* If you need to execute a model or write data to the warehouse, inform the user that this operation is not permitted
|
||||
* You can view compiled SQL with `dbt show` or `dbt compile` to understand what a model would do without executing it
|
||||
|
||||
---
|
||||
|
||||
# Modification Policy (Existing Docs, Tests, Semantic Layer)
|
||||
|
||||
- Additive-first: improve and extend existing documentation, tests, semantic models, and metrics. Do not delete or contradict prior information unless you can decisively disprove it with current evidence.
|
||||
- Evidence to remove/contradict: require verified metadata and/or targeted SQL that clearly shows the statement is false or misleading in the current data. Cite the evidence (with date/time and source) when you change or remove content.
|
||||
- If uncertain: retain existing statements and add clarifying context (e.g., scope, time-bounded phrasing like "As of {date}") rather than removing them. Add an item to `needs_clarification.md` for follow-up when match rates or patterns are inconclusive.
|
||||
- Tests: prefer augmenting or relaxing scope/thresholds over deletion. Only remove a test when disproven; otherwise, adjust (e.g., accepted ranges/values, relationship coverage) and document rationale. Add complementary tests to cover updated behavior.
|
||||
- Semantic layer and metrics: maintain backward compatibility. Prefer adding new dimensions/measures/metrics over renaming/removing. If a change is required, add the replacement and mark the old as deprecated in docs; keep references until migrations are completed.
|
||||
- Enums and categories: do not narrow `accepted_values` without strong evidence; when upstreams introduce new values, document them, decide whether to expand tests or gate them, and add a clarification item if policy is undecided.
|
||||
- YAML edits: modify the existing `schema.yml` in-place and keep a single top-level `models:`, `semantic_models:`, and `metrics:` key as required. Avoid scattering new files unless necessary.
|
||||
- Communication: when correcting prior documentation, explicitly call out the correction in your final summary and include a concise changelog entry that cites the evidence and lists impacted files.
|
||||
|
||||
---
|
||||
|
||||
Here is the dbt_project.yml:
|
||||
```yaml
|
||||
{{dbt_project_yml}}
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
Here is the directory structure:
|
||||
```
|
||||
{{folder_structure}}
|
||||
```
|
|
@ -6,7 +6,10 @@ import {
|
|||
import { Sonnet4 } from '../../llm/sonnet-4';
|
||||
import { IDLE_TOOL_NAME } from '../../tools/communication-tools/idle-tool/idle-tool';
|
||||
import { createAnalyticsEngineerToolset } from './create-analytics-engineer-toolset';
|
||||
import { getDocsAgentSystemPrompt as getAnalyticsEngineerAgentSystemPrompt } from './get-analytics-engineer-agent-system-prompt';
|
||||
import {
|
||||
getDocsAgentSystemPrompt as getAnalyticsEngineerAgentSystemPrompt,
|
||||
getAnalyticsEngineerSubagentSystemPrompt,
|
||||
} from './get-analytics-engineer-agent-system-prompt';
|
||||
import type {
|
||||
AnalyticsEngineerAgentOptions,
|
||||
AnalyticsEngineerAgentStreamOptions,
|
||||
|
@ -20,9 +23,14 @@ const STOP_CONDITIONS = [stepCountIs(250)];
|
|||
export function createAnalyticsEngineerAgent(
|
||||
analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions
|
||||
) {
|
||||
// Use subagent prompt if this is a subagent, otherwise use main agent prompt
|
||||
const promptFunction = analyticsEngineerAgentOptions.isSubagent
|
||||
? getAnalyticsEngineerSubagentSystemPrompt
|
||||
: getAnalyticsEngineerAgentSystemPrompt;
|
||||
|
||||
const systemMessage = {
|
||||
role: 'system',
|
||||
content: getAnalyticsEngineerAgentSystemPrompt(analyticsEngineerAgentOptions.folder_structure),
|
||||
content: promptFunction(analyticsEngineerAgentOptions.folder_structure),
|
||||
providerOptions: DEFAULT_ANALYTICS_ENGINEER_OPTIONS,
|
||||
} as ModelMessage;
|
||||
|
||||
|
|
|
@ -8,6 +8,7 @@ import {
|
|||
READ_FILE_TOOL_NAME,
|
||||
RETRIEVE_METADATA_TOOL_NAME,
|
||||
RUN_SQL_TOOL_NAME,
|
||||
TASK_TOOL_NAME,
|
||||
TODO_WRITE_TOOL_NAME,
|
||||
WRITE_FILE_TOOL_NAME,
|
||||
createBashTool,
|
||||
|
@ -27,85 +28,129 @@ import type { AgentFactory } from '../../tools/task-tools/task-tool/task-tool';
|
|||
import { createAnalyticsEngineerAgent } from './analytics-engineer-agent';
|
||||
import type { AnalyticsEngineerAgentOptions } from './types';
|
||||
|
||||
// Base read-only tools (always present)
|
||||
type ReadOnlyTools = {
|
||||
grep: ReturnType<typeof createGrepTool>;
|
||||
glob: ReturnType<typeof createGlobTool>;
|
||||
read: ReturnType<typeof createReadFileTool>;
|
||||
bash: ReturnType<typeof createBashTool>;
|
||||
ls: ReturnType<typeof createLsTool>;
|
||||
runSql: ReturnType<typeof createRunSqlTool>;
|
||||
retrieveMetadata: ReturnType<typeof createRetrieveMetadataTool>;
|
||||
todoWrite: ReturnType<typeof createTodoWriteTool>;
|
||||
};
|
||||
|
||||
// Write tools (excluded in research mode)
|
||||
type WriteTools = {
|
||||
write: ReturnType<typeof createWriteFileTool>;
|
||||
edit: ReturnType<typeof createEditFileTool>;
|
||||
multiEdit: ReturnType<typeof createMultiEditFileTool>;
|
||||
};
|
||||
|
||||
// Task tool (excluded for subagents)
|
||||
type TaskTool = {
|
||||
task: ReturnType<typeof createTaskTool>;
|
||||
};
|
||||
|
||||
// Actual toolset combinations based on isInResearchMode and isSubagent flags:
|
||||
type SubagentResearchToolset = ReadOnlyTools; // Research subagent: read-only only
|
||||
type MainAgentResearchToolset = ReadOnlyTools & TaskTool; // Research main: read-only + task
|
||||
type SubagentToolset = ReadOnlyTools & WriteTools; // Subagent: read-only + write
|
||||
type MainAgentToolset = ReadOnlyTools & WriteTools & TaskTool; // Main agent: all tools
|
||||
|
||||
export async function createAnalyticsEngineerToolset(
|
||||
analyticsEngineerAgentOptions: AnalyticsEngineerAgentOptions
|
||||
) {
|
||||
const writeFileTool = createWriteFileTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const grepTool = createGrepTool({
|
||||
messageId: analyticsEngineerAgentOptions.messageId,
|
||||
projectDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
});
|
||||
const globTool = createGlobTool({
|
||||
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,
|
||||
});
|
||||
const todosTool = createTodoWriteTool({
|
||||
chatId: analyticsEngineerAgentOptions.chatId,
|
||||
workingDirectory: analyticsEngineerAgentOptions.folder_structure,
|
||||
todosList: analyticsEngineerAgentOptions.todosList,
|
||||
});
|
||||
const runSqlTool = createRunSqlTool({
|
||||
apiKey: analyticsEngineerAgentOptions.apiKey || process.env.BUSTER_API_KEY || '',
|
||||
apiUrl:
|
||||
analyticsEngineerAgentOptions.apiUrl || process.env.BUSTER_API_URL || 'http://localhost:3000',
|
||||
});
|
||||
const retrieveMetadataTool = createRetrieveMetadataTool({
|
||||
apiKey: analyticsEngineerAgentOptions.apiKey || process.env.BUSTER_API_KEY || '',
|
||||
apiUrl:
|
||||
analyticsEngineerAgentOptions.apiUrl || process.env.BUSTER_API_URL || 'http://localhost:3000',
|
||||
});
|
||||
// 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;
|
||||
): Promise<
|
||||
SubagentResearchToolset | MainAgentResearchToolset | SubagentToolset | MainAgentToolset
|
||||
> {
|
||||
const { messageId, folder_structure, isInResearchMode, isSubagent, chatId, todosList } =
|
||||
analyticsEngineerAgentOptions;
|
||||
|
||||
// Helper to create task tool (avoids duplication)
|
||||
const createTaskToolForAgent = () =>
|
||||
createTaskTool({
|
||||
messageId,
|
||||
projectDirectory: folder_structure,
|
||||
createAgent: ((options: Parameters<AgentFactory>[0]) => {
|
||||
return createAnalyticsEngineerAgent({
|
||||
...options,
|
||||
model: analyticsEngineerAgentOptions.model,
|
||||
apiKey: analyticsEngineerAgentOptions.apiKey,
|
||||
apiUrl: analyticsEngineerAgentOptions.apiUrl,
|
||||
todosList: [],
|
||||
isSubagent: true,
|
||||
isInResearchMode,
|
||||
});
|
||||
}) as unknown as AgentFactory,
|
||||
});
|
||||
|
||||
// Build read-only tools (always present)
|
||||
const readOnlyTools = {
|
||||
[GREP_TOOL_NAME]: createGrepTool({ messageId, projectDirectory: folder_structure }),
|
||||
[GLOB_TOOL_NAME]: createGlobTool({ messageId, projectDirectory: folder_structure }),
|
||||
[READ_FILE_TOOL_NAME]: createReadFileTool({ messageId, projectDirectory: folder_structure }),
|
||||
[BASH_TOOL_NAME]: createBashTool({
|
||||
messageId,
|
||||
projectDirectory: folder_structure,
|
||||
isInResearchMode,
|
||||
}),
|
||||
[LS_TOOL_NAME]: createLsTool({ messageId, projectDirectory: folder_structure }),
|
||||
[RUN_SQL_TOOL_NAME]: createRunSqlTool({
|
||||
apiKey: analyticsEngineerAgentOptions.apiKey || process.env.BUSTER_API_KEY || '',
|
||||
apiUrl:
|
||||
analyticsEngineerAgentOptions.apiUrl ||
|
||||
process.env.BUSTER_API_URL ||
|
||||
'http://localhost:3000',
|
||||
}),
|
||||
[RETRIEVE_METADATA_TOOL_NAME]: createRetrieveMetadataTool({
|
||||
apiKey: analyticsEngineerAgentOptions.apiKey || process.env.BUSTER_API_KEY || '',
|
||||
apiUrl:
|
||||
analyticsEngineerAgentOptions.apiUrl ||
|
||||
process.env.BUSTER_API_URL ||
|
||||
'http://localhost:3000',
|
||||
}),
|
||||
[TODO_WRITE_TOOL_NAME]: createTodoWriteTool({
|
||||
chatId,
|
||||
workingDirectory: folder_structure,
|
||||
todosList,
|
||||
}),
|
||||
};
|
||||
|
||||
// Research mode: read-only tools only (conditionally add task tool)
|
||||
if (isInResearchMode) {
|
||||
if (isSubagent) {
|
||||
// Research subagent: read-only only
|
||||
return readOnlyTools;
|
||||
}
|
||||
// Research main agent: read-only + task
|
||||
return {
|
||||
...readOnlyTools,
|
||||
[TASK_TOOL_NAME]: createTaskToolForAgent(),
|
||||
};
|
||||
}
|
||||
|
||||
// Full mode: add write tools
|
||||
const writeTools = {
|
||||
[WRITE_FILE_TOOL_NAME]: createWriteFileTool({ messageId, projectDirectory: folder_structure }),
|
||||
[EDIT_FILE_TOOL_NAME]: createEditFileTool({ messageId, projectDirectory: folder_structure }),
|
||||
[MULTI_EDIT_FILE_TOOL_NAME]: createMultiEditFileTool({
|
||||
messageId,
|
||||
projectDirectory: folder_structure,
|
||||
}),
|
||||
};
|
||||
|
||||
if (isSubagent) {
|
||||
// Full mode subagent: read-only + write
|
||||
return {
|
||||
...readOnlyTools,
|
||||
...writeTools,
|
||||
};
|
||||
}
|
||||
|
||||
// Full mode main agent: all tools
|
||||
return {
|
||||
[WRITE_FILE_TOOL_NAME]: writeFileTool,
|
||||
[GREP_TOOL_NAME]: grepTool,
|
||||
[GLOB_TOOL_NAME]: globTool,
|
||||
[READ_FILE_TOOL_NAME]: readFileTool,
|
||||
[BASH_TOOL_NAME]: bashTool,
|
||||
[EDIT_FILE_TOOL_NAME]: editFileTool,
|
||||
[MULTI_EDIT_FILE_TOOL_NAME]: multiEditFileTool,
|
||||
[LS_TOOL_NAME]: lsTool,
|
||||
[TODO_WRITE_TOOL_NAME]: todosTool,
|
||||
[RUN_SQL_TOOL_NAME]: runSqlTool,
|
||||
[RETRIEVE_METADATA_TOOL_NAME]: retrieveMetadataTool,
|
||||
// ...(taskTool ? { taskTool } : {}),
|
||||
...readOnlyTools,
|
||||
...writeTools,
|
||||
[TASK_TOOL_NAME]: createTaskToolForAgent(),
|
||||
};
|
||||
}
|
||||
|
|
|
@ -1,3 +1,4 @@
|
|||
import analyticsEngineerAgentSubagentPrompt from './analytics-engineer-agent-prompt-subagent.txt';
|
||||
import analyticsEngineerAgentPrompt from './analytics-engineer-agent-prompt.txt';
|
||||
|
||||
/**
|
||||
|
@ -11,10 +12,11 @@ export interface DocsAgentTemplateParams {
|
|||
/**
|
||||
* Loads the docs agent prompt template and replaces variables
|
||||
*/
|
||||
function loadAndProcessPrompt(params: DocsAgentTemplateParams): string {
|
||||
return analyticsEngineerAgentPrompt
|
||||
function loadAndProcessPrompt(promptTemplate: string, params: DocsAgentTemplateParams): string {
|
||||
return promptTemplate
|
||||
.replace(/\{\{folder_structure\}\}/g, params.folderStructure)
|
||||
.replace(/\{\{date\}\}/g, params.date);
|
||||
.replace(/\{\{date\}\}/g, params.date)
|
||||
.replace(/\{\{dbt_project_yml\}\}/g, ''); // Empty for now, can be populated later if needed
|
||||
}
|
||||
|
||||
/**
|
||||
|
@ -27,7 +29,23 @@ export const getDocsAgentSystemPrompt = (folderStructure: string): string => {
|
|||
|
||||
const currentDate = new Date().toISOString();
|
||||
|
||||
return loadAndProcessPrompt({
|
||||
return loadAndProcessPrompt(analyticsEngineerAgentPrompt, {
|
||||
folderStructure,
|
||||
date: currentDate,
|
||||
});
|
||||
};
|
||||
|
||||
/**
|
||||
* Get system prompt for sub-agents (more concise, focused on task completion)
|
||||
*/
|
||||
export const getAnalyticsEngineerSubagentSystemPrompt = (folderStructure: string): string => {
|
||||
if (!folderStructure.trim()) {
|
||||
throw new Error('Folder structure is required');
|
||||
}
|
||||
|
||||
const currentDate = new Date().toISOString();
|
||||
|
||||
return loadAndProcessPrompt(analyticsEngineerAgentSubagentPrompt, {
|
||||
folderStructure,
|
||||
date: currentDate,
|
||||
});
|
||||
|
|
|
@ -45,6 +45,11 @@ export const AnalyticsEngineerAgentOptionsSchema = z.object({
|
|||
.boolean()
|
||||
.optional()
|
||||
.describe('Flag indicating this is a subagent (prevents infinite recursion)'),
|
||||
isInResearchMode: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.default(false)
|
||||
.describe('Flag indicating the agent should only perform read-only operations'),
|
||||
abortSignal: z
|
||||
.custom<AbortSignal>()
|
||||
.optional()
|
||||
|
|
|
@ -36,6 +36,43 @@ const BLOCKED_DBT_COMMANDS = [
|
|||
'fresh', // Checks freshness (read-only but we block for consistency)
|
||||
];
|
||||
|
||||
/**
|
||||
* Validates if a bash command is allowed in research mode (read-only)
|
||||
* @param command - The bash command to validate
|
||||
* @returns Object with isValid boolean and optional error message
|
||||
*/
|
||||
function validateBashCommandForResearchMode(command: string): { isValid: boolean; error?: string } {
|
||||
// List of write/destructive operations to block
|
||||
const WRITE_PATTERNS = [
|
||||
/\brm\s+/, // rm command
|
||||
/\bmv\s+/, // mv command
|
||||
/\bcp\s+.*\s+/, // cp with destination (allow 'cp file -' to stdout)
|
||||
/\bmkdir\s+/, // mkdir command
|
||||
/\brmdir\s+/, // rmdir command
|
||||
/\btouch\s+/, // touch command
|
||||
/>/, // output redirect (>, >>)
|
||||
/\bsed\s+.*-i/, // sed in-place edit
|
||||
/\bawk\s+.*>>/, // awk with output redirect
|
||||
/\bgit\s+(commit|push|add|reset|checkout|merge|rebase|cherry-pick|revert)/, // git write operations
|
||||
/\b(npm|pnpm|yarn|bun)\s+(install|add|remove|update)/, // package manager write operations
|
||||
/\bchmod\s+/, // chmod command
|
||||
/\bchown\s+/, // chown command
|
||||
/\bddt\s+(run|build|seed|snapshot|test|run-operation)/, // dbt write operations (keeping for backward compat)
|
||||
];
|
||||
|
||||
// Check each pattern
|
||||
for (const pattern of WRITE_PATTERNS) {
|
||||
if (pattern.test(command)) {
|
||||
return {
|
||||
isValid: false,
|
||||
error: `This command is not allowed in research mode. Research mode only permits read-only operations. Command blocked: ${command.substring(0, 100)}`,
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
return { isValid: true };
|
||||
}
|
||||
|
||||
/**
|
||||
* Validates if a dbt command is allowed (read-only)
|
||||
* @param command - The bash command to validate
|
||||
|
@ -183,12 +220,45 @@ async function executeCommand(
|
|||
*/
|
||||
export function createBashToolExecute(context: BashToolContext) {
|
||||
return async function execute(input: BashToolInput): Promise<BashToolOutput> {
|
||||
const { messageId, projectDirectory, onToolEvent } = context;
|
||||
const { messageId, projectDirectory, isInResearchMode, onToolEvent } = context;
|
||||
const { command, timeout } = input;
|
||||
|
||||
console.info(`Executing bash command for message ${messageId}: ${command}`);
|
||||
|
||||
// Validate dbt commands before execution
|
||||
// Validate commands in research mode
|
||||
if (isInResearchMode) {
|
||||
const researchModeValidation = validateBashCommandForResearchMode(command);
|
||||
if (!researchModeValidation.isValid) {
|
||||
const errorResult: BashToolOutput = {
|
||||
command,
|
||||
stdout: '',
|
||||
stderr: researchModeValidation.error || 'Command validation failed',
|
||||
exitCode: 1,
|
||||
success: false,
|
||||
error: researchModeValidation.error,
|
||||
};
|
||||
|
||||
console.error(`Command blocked in research mode: ${command}`, researchModeValidation.error);
|
||||
|
||||
// Emit events for blocked command
|
||||
onToolEvent?.({
|
||||
tool: 'bashTool',
|
||||
event: 'start',
|
||||
args: input,
|
||||
});
|
||||
|
||||
onToolEvent?.({
|
||||
tool: 'bashTool',
|
||||
event: 'complete',
|
||||
result: errorResult,
|
||||
args: input,
|
||||
});
|
||||
|
||||
return errorResult;
|
||||
}
|
||||
}
|
||||
|
||||
// Validate dbt commands before execution (always, regardless of research mode)
|
||||
const validation = validateDbtCommand(command);
|
||||
if (!validation.isValid) {
|
||||
const errorResult: BashToolOutput = {
|
||||
|
|
|
@ -24,6 +24,7 @@ describe.sequential('bash-tool integration test', () => {
|
|||
const bashTool = createBashTool({
|
||||
messageId: `test-message-${Date.now()}`,
|
||||
projectDirectory: testDir,
|
||||
isInResearchMode: false,
|
||||
});
|
||||
|
||||
const result = await materialize(
|
||||
|
@ -48,6 +49,7 @@ describe.sequential('bash-tool integration test', () => {
|
|||
const bashTool = createBashTool({
|
||||
messageId: `test-message-${Date.now()}`,
|
||||
projectDirectory: testDir,
|
||||
isInResearchMode: false,
|
||||
});
|
||||
|
||||
const result = await materialize(
|
||||
|
@ -72,6 +74,7 @@ describe.sequential('bash-tool integration test', () => {
|
|||
const bashTool = createBashTool({
|
||||
messageId: `test-message-${Date.now()}`,
|
||||
projectDirectory: testDir,
|
||||
isInResearchMode: false,
|
||||
});
|
||||
|
||||
const result = await materialize(
|
||||
|
@ -97,6 +100,7 @@ describe.sequential('bash-tool integration test', () => {
|
|||
const bashTool = createBashTool({
|
||||
messageId: `test-message-${Date.now()}`,
|
||||
projectDirectory: testDir,
|
||||
isInResearchMode: false,
|
||||
});
|
||||
|
||||
const rawResult = await bashTool.execute!(
|
||||
|
@ -119,6 +123,7 @@ describe.sequential('bash-tool integration test', () => {
|
|||
const bashTool = createBashTool({
|
||||
messageId: `test-message-${Date.now()}`,
|
||||
projectDirectory: testDir,
|
||||
isInResearchMode: false,
|
||||
});
|
||||
|
||||
const rawResult = await bashTool.execute!(
|
||||
|
@ -141,6 +146,7 @@ describe.sequential('bash-tool integration test', () => {
|
|||
const bashTool = createBashTool({
|
||||
messageId: `test-message-${Date.now()}`,
|
||||
projectDirectory: testDir,
|
||||
isInResearchMode: false,
|
||||
});
|
||||
|
||||
const testFile = `test-bash-${Date.now()}.txt`;
|
||||
|
@ -169,6 +175,7 @@ describe.sequential('bash-tool integration test', () => {
|
|||
const bashTool = createBashTool({
|
||||
messageId: `test-message-${Date.now()}`,
|
||||
projectDirectory: testDir,
|
||||
isInResearchMode: false,
|
||||
});
|
||||
|
||||
// Create test file
|
||||
|
@ -198,6 +205,7 @@ describe.sequential('bash-tool integration test', () => {
|
|||
const bashTool = createBashTool({
|
||||
messageId: `test-message-${Date.now()}`,
|
||||
projectDirectory: testDir,
|
||||
isInResearchMode: false,
|
||||
});
|
||||
|
||||
const rawResult = await bashTool.execute!(
|
||||
|
|
|
@ -27,6 +27,7 @@ describe('createBashTool', () => {
|
|||
const mockContext = {
|
||||
messageId: 'test-message-id',
|
||||
projectDirectory: '/tmp/test-project',
|
||||
isInResearchMode: false,
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
|
|
|
@ -30,6 +30,11 @@ export const BashToolOutputSchema = z.object({
|
|||
const BashToolContextSchema = z.object({
|
||||
messageId: z.string().describe('The message ID for database updates'),
|
||||
projectDirectory: z.string().describe('The root directory of the project'),
|
||||
isInResearchMode: z
|
||||
.boolean()
|
||||
.optional()
|
||||
.default(false)
|
||||
.describe('Flag indicating the agent should only allow read-only bash commands'),
|
||||
onToolEvent: z.any().optional(),
|
||||
});
|
||||
|
||||
|
|
|
@ -13,7 +13,7 @@ export {
|
|||
} from './planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool';
|
||||
|
||||
// Task tools
|
||||
export { createTaskTool } from './task-tools/task-tool/task-tool';
|
||||
export { createTaskTool, TASK_TOOL_NAME } from './task-tools/task-tool/task-tool';
|
||||
|
||||
// Visualization tools
|
||||
export {
|
||||
|
|
|
@ -0,0 +1,53 @@
|
|||
Launch a sub-agent to handle complex, multi-step tasks autonomously.
|
||||
|
||||
The sub-agent has access to all file tools (read, write, edit, bash, grep, glob, ls), database tools (runSql, retrieveMetadata), and planning tools (todoWrite) to complete its assigned task.
|
||||
|
||||
When NOT to use the Task tool:
|
||||
- If you want to read a specific file path, use the Read tool instead - it's faster
|
||||
- If you are searching for code within a specific file or set of 2-3 files, use the Read tool instead
|
||||
- If you can complete the task yourself with 1-2 tool calls
|
||||
- For simple, single-step operations that don't require autonomous planning
|
||||
|
||||
WHEN to use the Task tool:
|
||||
- Complex multi-step tasks that require autonomous planning and execution
|
||||
- Research tasks with multiple concurrent questions or areas to investigate
|
||||
- Tasks that benefit from dedicated focus and their own todo list tracking
|
||||
- When you need to parallelize work across multiple independent sub-tasks
|
||||
|
||||
Usage notes:
|
||||
1. Launch multiple agents concurrently whenever possible to maximize performance. To do this, use a single message with multiple tool calls.
|
||||
2. When a sub-agent completes, it returns a single summary message back to you. The result is not visible to the user. You should send a text message to the user with a concise summary of the results.
|
||||
3. Each sub-agent invocation is stateless. You cannot send additional messages to it, nor can it communicate with you outside of its final report. Therefore, your prompt should contain a highly detailed task description for the sub-agent to perform autonomously, and you should specify exactly what information it should return in its final message.
|
||||
4. Sub-agent outputs should generally be trusted.
|
||||
5. Clearly tell the sub-agent whether you expect it to write code/files or just do research (search, file reads, database queries, etc.), since it is not aware of the user's original intent.
|
||||
6. Sub-agents cannot spawn other sub-agents - they must complete their work directly.
|
||||
|
||||
Example usage for concurrent research:
|
||||
|
||||
<example>
|
||||
user: "I need to understand how authentication works in this dbt project - check the models, tests, and any documentation"
|
||||
assistant thinking: This requires investigating multiple areas concurrently - models, tests, and docs. I should launch 3 sub-agents in parallel.
|
||||
assistant: I'll research the authentication implementation by launching 3 concurrent sub-agents to investigate models, tests, and documentation.
|
||||
<launches 3 Task tool calls in parallel>
|
||||
1. description: "Research auth models", prompt: "Search for and read all models related to authentication (models with 'auth', 'user', 'login' in names). For each model, analyze the SQL logic, columns, and relationships. Return a summary of what authentication data is tracked and how it flows through the models."
|
||||
2. description: "Research auth tests", prompt: "Search for and read all tests related to authentication (look in schema.yml files for tests on auth-related models). Analyze what is being tested and what the assertions validate. Return a summary of the test coverage for authentication."
|
||||
3. description: "Research auth docs", prompt: "Search for and read any documentation about authentication (look for .md files, model descriptions in schema.yml). Return a summary of what the documentation explains about authentication patterns and usage."
|
||||
</example>
|
||||
|
||||
<example>
|
||||
user: "Can you update the customer model documentation to include all the latest columns?"
|
||||
assistant thinking: This is a single focused task that requires reading the model, checking metadata, and updating YAML. Perfect for a sub-agent.
|
||||
assistant: I'll delegate this to a sub-agent to research the customer model and update its documentation.
|
||||
<launches single Task tool call>
|
||||
description: "Update customer model docs", prompt: "Your task is to update the customer model documentation. Steps: 1) Read the customer model SQL file to understand the logic, 2) Use retrieveMetadata to get current column information, 3) Read the existing schema.yml for the customer model, 4) Update the schema.yml to document all columns with descriptions based on the model logic and metadata. Ensure all columns have descriptions and that the model-level description accurately reflects what the model does. Return a summary of what was updated."
|
||||
</example>
|
||||
|
||||
<example>
|
||||
user: "Help me understand the sales data pipeline"
|
||||
assistant thinking: This is a broad research task that could involve multiple models, their relationships, and data flow. I should break this into focused research areas and run them concurrently.
|
||||
assistant: I'll investigate the sales pipeline by researching the source data, transformation models, and final marts in parallel.
|
||||
<launches 3 Task tool calls in parallel>
|
||||
1. description: "Research sales sources", prompt: "Search for and analyze all staging models related to sales (stg_sales*, stg_orders*, etc.). Read the SQL files and understand what source data is being pulled. Use retrieveMetadata if needed to see actual column data. Return a summary of what sales data is available and how it's structured."
|
||||
2. description: "Research sales transforms", prompt: "Search for and analyze intermediate models that transform sales data (int_sales*, fct_orders*, etc.). Read the SQL to understand the transformation logic, joins, and aggregations. Return a summary of how raw sales data is transformed."
|
||||
3. description: "Research sales marts", prompt: "Search for and analyze final sales mart models (marts/sales/* or dim_*, fct_sales*, etc.). Read the SQL and schema.yml docs. Return a summary of what final sales analytics tables are available and what questions they answer."
|
||||
</example>
|
|
@ -1,77 +1,6 @@
|
|||
import { randomUUID } from 'node:crypto';
|
||||
import type { ModelMessage } from 'ai';
|
||||
import type { TaskToolContext, TaskToolInput, TaskToolOutput, ToolEventType } from './task-tool';
|
||||
|
||||
// Import AgentMessage type from CLI (or define a compatible version here)
|
||||
// For now, we'll create a type-compatible structure that matches the CLI's AgentMessage
|
||||
type TaskAgentMessage =
|
||||
| { kind: 'idle'; args: { final_response?: string } }
|
||||
| {
|
||||
kind: 'bash';
|
||||
event: 'start' | 'complete';
|
||||
args: { command: string; description?: string };
|
||||
result?: { stdout?: string; stderr?: string; exitCode: number; success: boolean };
|
||||
}
|
||||
| {
|
||||
kind: 'grep';
|
||||
event: 'start' | 'complete';
|
||||
args: { pattern: string; glob?: string; command: string };
|
||||
result?: {
|
||||
matches: Array<{ path: string; lineNum: number; lineText: string }>;
|
||||
totalMatches: number;
|
||||
truncated: boolean;
|
||||
};
|
||||
}
|
||||
| {
|
||||
kind: 'ls';
|
||||
event: 'start' | 'complete';
|
||||
args: { path?: string; command: string };
|
||||
result?: {
|
||||
output: string;
|
||||
success: boolean;
|
||||
count: number;
|
||||
truncated: boolean;
|
||||
errorMessage?: string;
|
||||
};
|
||||
}
|
||||
| {
|
||||
kind: 'write';
|
||||
event: 'start' | 'complete';
|
||||
args: { files: Array<{ path: string; content: string }> };
|
||||
result?: {
|
||||
results: Array<{ status: 'success' | 'error'; filePath: string; errorMessage?: string }>;
|
||||
};
|
||||
}
|
||||
| {
|
||||
kind: 'edit';
|
||||
event: 'start' | 'complete';
|
||||
args: {
|
||||
filePath: string;
|
||||
oldString?: string;
|
||||
newString?: string;
|
||||
edits?: Array<{ oldString: string; newString: string }>;
|
||||
};
|
||||
result?: {
|
||||
success: boolean;
|
||||
filePath: string;
|
||||
diff?: string;
|
||||
finalDiff?: string;
|
||||
message?: string;
|
||||
errorMessage?: string;
|
||||
};
|
||||
}
|
||||
| {
|
||||
kind: 'read';
|
||||
event: 'start' | 'complete';
|
||||
args: { filePath: string };
|
||||
result?: {
|
||||
status: 'success' | 'error';
|
||||
file_path: string;
|
||||
content?: string;
|
||||
truncated?: boolean;
|
||||
error_message?: string;
|
||||
};
|
||||
};
|
||||
import type { TaskToolContext, TaskToolInput, TaskToolOutput } from './task-tool';
|
||||
|
||||
/**
|
||||
* Creates the execute function for the task tool
|
||||
|
@ -79,23 +8,13 @@ type TaskAgentMessage =
|
|||
*/
|
||||
export function createTaskToolExecute(context: TaskToolContext) {
|
||||
return async function execute(input: TaskToolInput): Promise<TaskToolOutput> {
|
||||
const { projectDirectory, onToolEvent, createAgent } = context;
|
||||
const { instructions } = input;
|
||||
const { projectDirectory, createAgent } = context;
|
||||
const { description, prompt } = input;
|
||||
|
||||
console.info(`Starting task with instructions: ${instructions.substring(0, 100)}...`);
|
||||
|
||||
// Emit start event
|
||||
onToolEvent?.({
|
||||
tool: 'taskTool',
|
||||
event: 'start',
|
||||
args: input,
|
||||
});
|
||||
console.info(`Starting task: ${description}`);
|
||||
|
||||
try {
|
||||
// Collect all messages from the task
|
||||
const taskMessages: TaskAgentMessage[] = [];
|
||||
|
||||
// Create a new agent instance with a callback to collect messages
|
||||
// Create a new agent instance for the task
|
||||
const taskAgent = createAgent({
|
||||
folder_structure: projectDirectory,
|
||||
userId: 'task',
|
||||
|
@ -103,127 +22,22 @@ export function createTaskToolExecute(context: TaskToolContext) {
|
|||
dataSourceId: '',
|
||||
organizationId: 'task',
|
||||
messageId: randomUUID(),
|
||||
// Callback to collect task tool events
|
||||
onToolEvent: (event: ToolEventType) => {
|
||||
// Convert tool events to properly typed AgentMessage format
|
||||
// This mirrors the logic in analytics-engineer-handler.ts
|
||||
|
||||
if (event.tool === 'idleTool' && event.event === 'complete') {
|
||||
taskMessages.push({
|
||||
kind: 'idle',
|
||||
args: event.args as { final_response?: string },
|
||||
});
|
||||
}
|
||||
|
||||
if (event.tool === 'bashTool' && event.event === 'complete') {
|
||||
taskMessages.push({
|
||||
kind: 'bash',
|
||||
event: 'complete',
|
||||
args: event.args as { command: string; description?: string },
|
||||
result: event.result as {
|
||||
stdout?: string;
|
||||
stderr?: string;
|
||||
exitCode: number;
|
||||
success: boolean;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (event.tool === 'grepTool' && event.event === 'complete') {
|
||||
taskMessages.push({
|
||||
kind: 'grep',
|
||||
event: 'complete',
|
||||
args: event.args as { pattern: string; glob?: string; command: string },
|
||||
result: event.result as {
|
||||
matches: Array<{ path: string; lineNum: number; lineText: string }>;
|
||||
totalMatches: number;
|
||||
truncated: boolean;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (event.tool === 'lsTool' && event.event === 'complete') {
|
||||
taskMessages.push({
|
||||
kind: 'ls',
|
||||
event: 'complete',
|
||||
args: event.args as { path?: string; command: string },
|
||||
result: event.result as {
|
||||
output: string;
|
||||
success: boolean;
|
||||
count: number;
|
||||
truncated: boolean;
|
||||
errorMessage?: string;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (event.tool === 'writeFileTool' && event.event === 'complete') {
|
||||
taskMessages.push({
|
||||
kind: 'write',
|
||||
event: 'complete',
|
||||
args: event.args as { files: Array<{ path: string; content: string }> },
|
||||
result: event.result as {
|
||||
results: Array<{
|
||||
status: 'success' | 'error';
|
||||
filePath: string;
|
||||
errorMessage?: string;
|
||||
}>;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (event.tool === 'editFileTool' && event.event === 'complete') {
|
||||
taskMessages.push({
|
||||
kind: 'edit',
|
||||
event: 'complete',
|
||||
args: event.args as {
|
||||
filePath: string;
|
||||
oldString?: string;
|
||||
newString?: string;
|
||||
edits?: Array<{ oldString: string; newString: string }>;
|
||||
},
|
||||
result: event.result as {
|
||||
success: boolean;
|
||||
filePath: string;
|
||||
diff?: string;
|
||||
finalDiff?: string;
|
||||
message?: string;
|
||||
errorMessage?: string;
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
if (event.tool === 'readFileTool' && event.event === 'complete') {
|
||||
taskMessages.push({
|
||||
kind: 'read',
|
||||
event: 'complete',
|
||||
args: event.args as { filePath: string },
|
||||
result: event.result as {
|
||||
status: 'success' | 'error';
|
||||
file_path: string;
|
||||
content?: string;
|
||||
truncated?: boolean;
|
||||
error_message?: string;
|
||||
},
|
||||
});
|
||||
}
|
||||
},
|
||||
// Pass flag to indicate this is a subagent (prevents infinite recursion)
|
||||
isSubagent: true,
|
||||
});
|
||||
|
||||
// Create the user message with instructions
|
||||
// Create the user message with the task prompt
|
||||
const messages: ModelMessage[] = [
|
||||
{
|
||||
role: 'user',
|
||||
content: instructions,
|
||||
content: prompt,
|
||||
},
|
||||
];
|
||||
|
||||
// Run the task agent
|
||||
const stream = await taskAgent.stream({ messages });
|
||||
|
||||
// Consume the stream to trigger tool execution
|
||||
// Consume the stream to trigger tool execution and collect text response
|
||||
let fullResponse = '';
|
||||
for await (const part of stream.fullStream) {
|
||||
if (part.type === 'text-delta') {
|
||||
|
@ -231,44 +45,23 @@ export function createTaskToolExecute(context: TaskToolContext) {
|
|||
}
|
||||
}
|
||||
|
||||
// Generate a summary from the final response or messages
|
||||
// Generate a summary from the final response
|
||||
const summary = fullResponse || 'Task completed';
|
||||
|
||||
console.info(`Task completed with ${taskMessages.length} tool calls`);
|
||||
console.info(`Task completed successfully: ${description}`);
|
||||
|
||||
const output: TaskToolOutput = {
|
||||
return {
|
||||
status: 'success',
|
||||
summary: summary.substring(0, 500), // Limit summary length
|
||||
messages: taskMessages,
|
||||
summary: summary.substring(0, 2000), // Limit summary length
|
||||
};
|
||||
|
||||
// Emit complete event with all collected messages
|
||||
onToolEvent?.({
|
||||
tool: 'taskTool',
|
||||
event: 'complete',
|
||||
result: output,
|
||||
args: input,
|
||||
});
|
||||
|
||||
return output;
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : 'Unknown error';
|
||||
console.error(`Task error:`, errorMessage);
|
||||
console.error(`Task error (${description}):`, errorMessage);
|
||||
|
||||
const output: TaskToolOutput = {
|
||||
return {
|
||||
status: 'error',
|
||||
error_message: errorMessage,
|
||||
};
|
||||
|
||||
// Emit complete event even on error
|
||||
onToolEvent?.({
|
||||
tool: 'taskTool',
|
||||
event: 'complete',
|
||||
result: output,
|
||||
args: input,
|
||||
});
|
||||
|
||||
return output;
|
||||
}
|
||||
};
|
||||
}
|
||||
|
|
|
@ -1,13 +1,19 @@
|
|||
import type { ModelMessage, StreamTextResult, ToolSet } from 'ai';
|
||||
import { tool } from 'ai';
|
||||
import { z } from 'zod';
|
||||
import TASK_TOOL_DESCRIPTION from './task-tool-description.txt';
|
||||
import { createTaskToolExecute } from './task-tool-execute';
|
||||
|
||||
export const TASK_TOOL_NAME = 'task';
|
||||
|
||||
export const TaskToolInputSchema = z.object({
|
||||
instructions: z
|
||||
description: z
|
||||
.string()
|
||||
.describe('A short (3-5 word) description of the task for progress tracking and logging'),
|
||||
prompt: z
|
||||
.string()
|
||||
.describe(
|
||||
'Detailed instructions for the task to execute. Be specific about what the task agent should accomplish.'
|
||||
'Detailed instructions for the task to execute. Be specific about what the task agent should accomplish and what information it should return.'
|
||||
),
|
||||
});
|
||||
|
||||
|
@ -15,7 +21,6 @@ export const TaskToolOutputSchema = z.discriminatedUnion('status', [
|
|||
z.object({
|
||||
status: z.literal('success'),
|
||||
summary: z.string().describe('Summary of what the task accomplished'),
|
||||
messages: z.array(z.any()).describe('All messages generated by the task during execution'),
|
||||
}),
|
||||
z.object({
|
||||
status: z.literal('error'),
|
||||
|
@ -32,23 +37,15 @@ export type AgentFactory = (options: {
|
|||
organizationId: string;
|
||||
messageId: string;
|
||||
isSubagent?: boolean;
|
||||
onToolEvent?: (event: ToolEventType) => void;
|
||||
apiKey?: string;
|
||||
apiUrl?: string;
|
||||
}) => {
|
||||
stream: (options: { messages: ModelMessage[] }) => Promise<StreamTextResult<ToolSet, never>>;
|
||||
};
|
||||
|
||||
// Type for tool events
|
||||
export interface ToolEventType {
|
||||
tool: string;
|
||||
event: 'start' | 'complete';
|
||||
args: unknown;
|
||||
result?: unknown;
|
||||
}
|
||||
|
||||
const _TaskToolContextSchema = z.object({
|
||||
messageId: z.string().describe('The message ID for database updates'),
|
||||
projectDirectory: z.string().describe('The root directory of the project'),
|
||||
onToolEvent: z.function().optional().describe('Callback for tool events'),
|
||||
createAgent: z
|
||||
.function()
|
||||
.describe('Factory function to create a new agent instance for the task'),
|
||||
|
@ -61,27 +58,16 @@ export type TaskToolOutput = z.infer<typeof TaskToolOutputSchema>;
|
|||
export interface TaskToolContext {
|
||||
messageId: string;
|
||||
projectDirectory: string;
|
||||
onToolEvent?: (event: ToolEventType) => void;
|
||||
createAgent: AgentFactory;
|
||||
}
|
||||
|
||||
// Type for internal task messages
|
||||
export interface TaskMessage {
|
||||
tool: string;
|
||||
event: 'start' | 'complete';
|
||||
// biome-ignore lint/suspicious/noExplicitAny: Messages from sub-tasks can have any shape
|
||||
args: any;
|
||||
// biome-ignore lint/suspicious/noExplicitAny: Results from sub-tasks can have any shape
|
||||
result?: any;
|
||||
}
|
||||
|
||||
export function createTaskTool<TAgentContext extends TaskToolContext = TaskToolContext>(
|
||||
context: TAgentContext
|
||||
) {
|
||||
const execute = createTaskToolExecute(context);
|
||||
|
||||
return tool({
|
||||
description: `Delegate a specialized task to a task agent. Use this when you need to execute a complex, self-contained task that would benefit from dedicated focus. The task agent will have access to all file tools (read, write, edit, bash, grep, ls) to complete the task. Provide clear, specific instructions for what the task should accomplish.`,
|
||||
description: TASK_TOOL_DESCRIPTION,
|
||||
inputSchema: TaskToolInputSchema,
|
||||
outputSchema: TaskToolOutputSchema,
|
||||
execute,
|
||||
|
|
Loading…
Reference in New Issue