diff --git a/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.ts b/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.ts index 885939375..14fe6c1c2 100644 --- a/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.ts +++ b/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.ts @@ -4,9 +4,20 @@ import { wrapTraced } from 'braintrust'; import z from 'zod'; import { Sonnet4 } from '../../llm'; import { createExecuteSqlTool, createSequentialThinkingTool } from '../../tools'; -import { createMessageUserClarifyingQuestionTool, MESSAGE_USER_CLARIFYING_QUESTION_TOOL_NAME } from '../../tools/communication-tools/message-user-clarifying-question/message-user-clarifying-question'; -import { createRespondWithoutAssetCreationTool, RESPOND_WITHOUT_ASSET_CREATION_TOOL_NAME } from '../../tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-tool'; -import { createSubmitThoughtsTool, SUBMIT_THOUGHTS_TOOL_NAME } from '../../tools/communication-tools/submit-thoughts-tool/submit-thoughts-tool'; +import { + MESSAGE_USER_CLARIFYING_QUESTION_TOOL_NAME, + createMessageUserClarifyingQuestionTool, +} from '../../tools/communication-tools/message-user-clarifying-question/message-user-clarifying-question'; +import { + RESPOND_WITHOUT_ASSET_CREATION_TOOL_NAME, + createRespondWithoutAssetCreationTool, +} from '../../tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-tool'; +import { + SUBMIT_THOUGHTS_TOOL_NAME, + createSubmitThoughtsTool, +} from '../../tools/communication-tools/submit-thoughts-tool/submit-thoughts-tool'; +import { EXECUTE_SQL_TOOL_NAME } from '../../tools/database-tools/execute-sql/execute-sql'; +import { SEQUENTIAL_THINKING_TOOL_NAME } from '../../tools/planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool'; import { healToolWithLlm } from '../../utils'; import { type AnalysisMode, @@ -106,11 +117,11 @@ export function createThinkAndPrepAgent(thinkAndPrepAgentSchema: ThinkAndPrepAge streamText({ model: Sonnet4, tools: { - sequentialThinking, - executeSql: executeSqlTool, - respondWithoutAssetCreation, - submitThoughts, - messageUserClarifyingQuestion, + [SEQUENTIAL_THINKING_TOOL_NAME]: sequentialThinking, + [EXECUTE_SQL_TOOL_NAME]: executeSqlTool, + [RESPOND_WITHOUT_ASSET_CREATION_TOOL_NAME]: respondWithoutAssetCreation, + [SUBMIT_THOUGHTS_TOOL_NAME]: submitThoughts, + [MESSAGE_USER_CLARIFYING_QUESTION_TOOL_NAME]: messageUserClarifyingQuestion, }, messages: [systemMessage, datasetsSystemMessage, ...currentMessages], stopWhen: STOP_CONDITIONS, diff --git a/packages/ai/src/tools/database-tools/execute-sql/execute-sql-execute.ts b/packages/ai/src/tools/database-tools/execute-sql/execute-sql-execute.ts index 8c8bd70ca..05b10bc36 100644 --- a/packages/ai/src/tools/database-tools/execute-sql/execute-sql-execute.ts +++ b/packages/ai/src/tools/database-tools/execute-sql/execute-sql-execute.ts @@ -6,11 +6,12 @@ import { createPermissionErrorMessage, validateSqlPermissions, } from '../../../utils/sql-permissions'; -import type { - ExecuteSqlContext, - ExecuteSqlInput, - ExecuteSqlOutput, - ExecuteSqlState, +import { + EXECUTE_SQL_TOOL_NAME, + type ExecuteSqlContext, + type ExecuteSqlInput, + type ExecuteSqlOutput, + type ExecuteSqlState, } from './execute-sql'; import { createExecuteSqlRawLlmMessageEntry, @@ -366,6 +367,6 @@ export function createExecuteSqlExecute(state: ExecuteSqlState, context: Execute } } }, - { name: 'execute-sql' } + { name: EXECUTE_SQL_TOOL_NAME } ); } diff --git a/packages/ai/src/tools/database-tools/execute-sql/execute-sql.ts b/packages/ai/src/tools/database-tools/execute-sql/execute-sql.ts index ba62e1f36..69d9c1734 100644 --- a/packages/ai/src/tools/database-tools/execute-sql/execute-sql.ts +++ b/packages/ai/src/tools/database-tools/execute-sql/execute-sql.ts @@ -6,6 +6,8 @@ import { createExecuteSqlExecute } from './execute-sql-execute'; import { createExecuteSqlFinish } from './execute-sql-finish'; import { createExecuteSqlStart } from './execute-sql-start'; +export const EXECUTE_SQL_TOOL_NAME = 'executeSql'; + export const ExecuteSqlInputSchema = z.object({ statements: z.array(z.string()).describe( `Array of lightweight, optimized SQL statements to execute. diff --git a/packages/ai/src/tools/planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool.ts b/packages/ai/src/tools/planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool.ts index 1ec8dfa28..314dbaf80 100644 --- a/packages/ai/src/tools/planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool.ts +++ b/packages/ai/src/tools/planning-thinking-tools/sequential-thinking-tool/sequential-thinking-tool.ts @@ -5,6 +5,8 @@ import { createSequentialThinkingExecute } from './sequential-thinking-tool-exec import { createSequentialThinkingFinish } from './sequential-thinking-tool-finish'; import { createSequentialThinkingStart } from './sequential-thinking-tool-start'; +export const SEQUENTIAL_THINKING_TOOL_NAME = 'sequentialThinking'; + export const SequentialThinkingInputSchema = z.object({ thought: z .string()