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

fixes on tests and stuff

This commit is contained in:
dal 2025-08-21 14:15:58 -06:00
parent 81c703a472
commit 93fdbd07b3
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
6 changed files with 103 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
packages/ai/src/steps/docs-agent-steps/get-repository-tree-step/get-repository-tree-step.test.ts
View File

@ -70,12 +70,18 @@ describe('runGetRepositoryTreeStep', () => {
});
it('should handle missing sandbox gracefully', async () => {
// Create a minimal sandbox object that passes validation but doesn't have full functionality
const minimalSandbox = {
id: 'mock-sandbox-minimal',
fs: {},
} as unknown as Sandbox;
const input = {
message: 'Test message',
organizationId: 'org-123',
contextInitialized: true,
context: {
sandbox: null as any,
sandbox: minimalSandbox,
dataSourceId: validDataSourceId,
todoList: '',
clarificationQuestions: [],

1
packages/ai/src/tools/communication-tools/respond-without-asset-creation/respond-without-asset-creation-streaming.test.ts
View File

@ -11,6 +11,7 @@ import type {
vi.mock('@buster/database', () => ({
updateMessageEntries: vi.fn().mockResolvedValue({ success: true }),
updateMessage: vi.fn().mockResolvedValue({ success: true }),
}));
describe('Respond Without Asset Creation Tool Streaming Tests', () => {

2
packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.test.ts
View File

@ -235,7 +235,7 @@ describe('create-reports-execute', () => {
const firstCall = mockUpdateMessageEntries.mock.calls[0]?.[0];
expect(firstCall.messageId).toBe('msg-001');
expect(firstCall.reasoningMessages).toBeDefined();
expect(firstCall.rawLlmMessages).toBeDefined();
// rawLlmMessages are intentionally not created in initial entries to avoid duplicates
// State should be updated
expect(state.initialEntriesCreated).toBe(true);

8
packages/ai/src/tools/visualization-tools/reports/create-reports-tool/create-reports-execute.ts
View File

@ -95,7 +95,7 @@ export function createCreateReportsExecute(
try {
const toolCallId = state.toolCallId || `tool-${Date.now()}`;
const reasoningEntry = createCreateReportsReasoningEntry(state, toolCallId);
const rawLlmMessage = createCreateReportsRawLlmMessageEntry(state, toolCallId);
// Skip creating rawLlmMessage here to avoid duplicates - it will be created with the result later
const updates: Parameters<typeof updateMessageEntries>[0] = {
messageId: context.messageId,
@ -105,11 +105,7 @@ export function createCreateReportsExecute(
updates.reasoningMessages = [reasoningEntry];
}
if (rawLlmMessage) {
updates.rawLlmMessages = [rawLlmMessage];
}
if (reasoningEntry || rawLlmMessage) {
if (reasoningEntry) {
await updateMessageEntries(updates);
state.initialEntriesCreated = true;
}

17
packages/ai/src/tools/visualization-tools/reports/modify-reports-tool/modify-reports-start.ts
View File

@ -1,9 +1,6 @@
import { updateMessageEntries } from '@buster/database';
import type { ToolCallOptions } from 'ai';
import {
createModifyReportsRawLlmMessageEntry,
createModifyReportsReasoningEntry,
} from './helpers/modify-reports-transform-helper';
import { createModifyReportsReasoningEntry } from './helpers/modify-reports-transform-helper';
import type { ModifyReportsContext, ModifyReportsState } from './modify-reports-tool';
export function modifyReportsStart(context: ModifyReportsContext, state: ModifyReportsState) {
@ -24,12 +21,12 @@ export function modifyReportsStart(context: ModifyReportsContext, state: ModifyR
if (context.messageId) {
try {
if (context.messageId && state.toolCallId) {
// Update database with both reasoning and raw LLM entries
// Update database with reasoning entry only - raw LLM message will be created with result later
try {
const reasoningEntry = createModifyReportsReasoningEntry(state, options.toolCallId);
const rawLlmMessage = createModifyReportsRawLlmMessageEntry(state, options.toolCallId);
// Skip creating rawLlmMessage here to avoid duplicates - it will be created with the result later
// Update both entries together if they exist
// Update reasoning entry if it exists
const updates: Parameters<typeof updateMessageEntries>[0] = {
messageId: context.messageId,
};
@ -38,11 +35,7 @@ export function modifyReportsStart(context: ModifyReportsContext, state: ModifyR
updates.reasoningMessages = [reasoningEntry];
}
if (rawLlmMessage) {
updates.rawLlmMessages = [rawLlmMessage];
}
if (reasoningEntry || rawLlmMessage) {
if (reasoningEntry) {
await updateMessageEntries(updates);
}
} catch (error) {

146
packages/ai/src/utils/tool-call-repair/strategies/re-ask-strategy.test.ts
View File

@ -1,18 +1,38 @@
import { NoSuchToolError, generateText } from 'ai';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { z } from 'zod';
// Mock the dependencies - must be before imports
vi.mock('ai', () => {
class MockNoSuchToolError extends Error {
toolName: string;
availableTools: string[];
constructor(options: { toolName: string; availableTools: string[] }) {
super(`Tool ${options.toolName} not found`);
this.toolName = options.toolName;
this.availableTools = options.availableTools;
}
static isInstance(error: any): boolean {
return error instanceof MockNoSuchToolError;
}
}
return {
NoSuchToolError: MockNoSuchToolError,
generateText: vi.fn(),
streamText: vi.fn(),
tool: vi.fn((config: any) => config),
stepCountIs: vi.fn((count: number) => ({ type: 'stepCount', count })),
hasToolCall: vi.fn((toolName: string) => ({ type: 'toolCall', toolName })),
};
});
import { NoSuchToolError, generateText, streamText } from 'ai';
import { ANALYST_AGENT_NAME, THINK_AND_PREP_AGENT_NAME } from '../../../agents';
import type { RepairContext } from '../types';
import { canHandleNoSuchTool, repairWrongToolName } from './re-ask-strategy';
// Mock the dependencies
vi.mock('ai', async () => {
const actual = await vi.importActual('ai');
return {
...actual,
generateText: vi.fn(),
};
});
vi.mock('braintrust', () => ({
wrapTraced: (fn: any) => fn,
}));
@ -47,17 +67,18 @@ describe('re-ask-strategy', () => {
describe('repairWrongToolName', () => {
it('should re-ask and get corrected tool call', async () => {
const mockGenerateText = vi.mocked(generateText);
const mockStreamText = vi.mocked(streamText);
const correctedToolCall = {
toolName: 'correctTool',
input: { param: 'value' },
};
mockGenerateText.mockResolvedValueOnce({
toolCalls: [correctedToolCall],
text: '',
usage: {},
mockStreamText.mockReturnValueOnce({
textStream: (async function* () {})(),
toolCalls: Promise.resolve([correctedToolCall]),
text: Promise.resolve(''),
usage: Promise.resolve({}),
} as any);
const context: RepairContext = {
@ -68,8 +89,8 @@ describe('re-ask-strategy', () => {
input: { param: 'value' },
} as any,
tools: {
correctTool: { inputSchema: {} },
anotherTool: { inputSchema: {} },
correctTool: { inputSchema: z.object({}) },
anotherTool: { inputSchema: z.object({}) },
} as any,
error: new NoSuchToolError({
toolName: 'wrongTool',
@ -90,7 +111,7 @@ describe('re-ask-strategy', () => {
});
// Verify the tool input is properly formatted as an object in the messages
const calls = mockGenerateText.mock.calls[0];
const calls = mockStreamText.mock.calls[0];
const messages = calls?.[0]?.messages;
const assistantMessage = messages?.find((m: any) => m.role === 'assistant');
const content = assistantMessage?.content?.[0];
@ -98,7 +119,7 @@ describe('re-ask-strategy', () => {
expect(content.input).toEqual({ param: 'value' });
}
expect(mockGenerateText).toHaveBeenCalledWith(
expect(mockStreamText).toHaveBeenCalledWith(
expect.objectContaining({
model: 'mock-model',
messages: expect.arrayContaining([
@ -114,11 +135,12 @@ describe('re-ask-strategy', () => {
});
it('should use analyst agent context for error message', async () => {
const mockGenerateText = vi.mocked(generateText);
mockGenerateText.mockResolvedValueOnce({
toolCalls: [],
text: '',
usage: {},
const mockStreamText = vi.mocked(streamText);
mockStreamText.mockReturnValueOnce({
textStream: (async function* () {})(),
toolCalls: Promise.resolve([]),
text: Promise.resolve(''),
usage: Promise.resolve({}),
} as any);
const context: RepairContext = {
@ -129,8 +151,8 @@ describe('re-ask-strategy', () => {
input: '{}',
} as any,
tools: {
createMetrics: {},
modifyMetrics: {},
createMetrics: { inputSchema: z.object({}) },
modifyMetrics: { inputSchema: z.object({}) },
} as any,
error: new NoSuchToolError({
toolName: 'executeSql',
@ -146,8 +168,8 @@ describe('re-ask-strategy', () => {
await repairWrongToolName(context);
const calls = mockGenerateText.mock.calls[0];
if (!calls) throw new Error('generateText not called');
const calls = mockStreamText.mock.calls[0];
if (!calls) throw new Error('streamText not called');
const messages = calls[0]?.messages;
if (!messages) throw new Error('No messages found');
const toolResultMessage = messages.find((m: any) => m.role === 'tool');
@ -164,11 +186,12 @@ describe('re-ask-strategy', () => {
});
it('should use think-and-prep agent context for error message', async () => {
const mockGenerateText = vi.mocked(generateText);
mockGenerateText.mockResolvedValueOnce({
toolCalls: [],
text: '',
usage: {},
const mockStreamText = vi.mocked(streamText);
mockStreamText.mockReturnValueOnce({
textStream: (async function* () {})(),
toolCalls: Promise.resolve([]),
text: Promise.resolve(''),
usage: Promise.resolve({}),
} as any);
const context: RepairContext = {
@ -179,8 +202,8 @@ describe('re-ask-strategy', () => {
input: '{}',
} as any,
tools: {
executeSql: {},
sequentialThinking: {},
executeSql: { inputSchema: z.object({}) },
sequentialThinking: { inputSchema: z.object({}) },
} as any,
error: new NoSuchToolError({
toolName: 'createMetrics',
@ -197,8 +220,8 @@ describe('re-ask-strategy', () => {
await repairWrongToolName(context);
const calls = mockGenerateText.mock.calls[0];
if (!calls) throw new Error('generateText not called');
const calls = mockStreamText.mock.calls[0];
if (!calls) throw new Error('streamText not called');
const messages = calls[0]?.messages;
if (!messages) throw new Error('No messages found');
const toolResultMessage = messages.find((m: any) => m.role === 'tool');
@ -213,11 +236,12 @@ describe('re-ask-strategy', () => {
});
it('should return null if no valid tool call is returned', async () => {
const mockGenerateText = vi.mocked(generateText);
mockGenerateText.mockResolvedValueOnce({
toolCalls: [],
text: '',
usage: {},
const mockStreamText = vi.mocked(streamText);
mockStreamText.mockReturnValueOnce({
textStream: (async function* () {})(),
toolCalls: Promise.resolve([]),
text: Promise.resolve(''),
usage: Promise.resolve({}),
} as any);
const context: RepairContext = {
@ -228,7 +252,7 @@ describe('re-ask-strategy', () => {
input: '{}',
} as any,
tools: {
correctTool: {},
correctTool: { inputSchema: z.object({}) },
} as any,
error: new NoSuchToolError({
toolName: 'wrongTool',
@ -243,8 +267,10 @@ describe('re-ask-strategy', () => {
});
it('should handle errors during re-ask', async () => {
const mockGenerateText = vi.mocked(generateText);
mockGenerateText.mockRejectedValueOnce(new Error('Generation failed'));
const mockStreamText = vi.mocked(streamText);
mockStreamText.mockImplementationOnce(() => {
throw new Error('Generation failed');
});
const context: RepairContext = {
toolCall: {
@ -267,11 +293,12 @@ describe('re-ask-strategy', () => {
});
it('should wrap non-JSON string inputs in an object', async () => {
const mockGenerateText = vi.mocked(generateText);
mockGenerateText.mockResolvedValueOnce({
toolCalls: [{ toolName: 'correctTool', input: { wrapped: true } }],
text: '',
usage: {},
const mockStreamText = vi.mocked(streamText);
mockStreamText.mockReturnValueOnce({
textStream: (async function* () {})(),
toolCalls: Promise.resolve([{ toolName: 'correctTool', input: { wrapped: true } }]),
text: Promise.resolve(''),
usage: Promise.resolve({}),
} as any);
const context: RepairContext = {
@ -282,7 +309,7 @@ describe('re-ask-strategy', () => {
input: 'plain text input',
} as any,
tools: {
correctTool: { inputSchema: {} },
correctTool: { inputSchema: z.object({}) },
} as any,
error: new NoSuchToolError({
toolName: 'wrongTool',
@ -295,7 +322,7 @@ describe('re-ask-strategy', () => {
await repairWrongToolName(context);
// Verify the non-JSON string was wrapped in an object
const calls = mockGenerateText.mock.calls[0];
const calls = mockStreamText.mock.calls[0];
const messages = calls?.[0]?.messages;
const assistantMessage = messages?.find((m: any) => m.role === 'assistant');
const content = assistantMessage?.content?.[0];
@ -305,11 +332,12 @@ describe('re-ask-strategy', () => {
});
it('should handle already valid JSON string inputs', async () => {
const mockGenerateText = vi.mocked(generateText);
mockGenerateText.mockResolvedValueOnce({
toolCalls: [{ toolName: 'correctTool', input: { handled: true } }],
text: '',
usage: {},
const mockStreamText = vi.mocked(streamText);
mockStreamText.mockReturnValueOnce({
textStream: (async function* () {})(),
toolCalls: Promise.resolve([{ toolName: 'correctTool', input: { handled: true } }]),
text: Promise.resolve(''),
usage: Promise.resolve({}),
} as any);
const context: RepairContext = {
@ -320,7 +348,7 @@ describe('re-ask-strategy', () => {
input: '{"already":"valid"}',
} as any,
tools: {
correctTool: { inputSchema: {} },
correctTool: { inputSchema: z.object({}) },
} as any,
error: new NoSuchToolError({
toolName: 'wrongTool',
@ -333,7 +361,7 @@ describe('re-ask-strategy', () => {
await repairWrongToolName(context);
// Verify the valid JSON string was parsed to an object
const calls = mockGenerateText.mock.calls[0];
const calls = mockStreamText.mock.calls[0];
const messages = calls?.[0]?.messages;
const assistantMessage = messages?.find((m: any) => m.role === 'assistant');
const content = assistantMessage?.content?.[0];