Include Schema Error Handling for Vertex and Google Auth methods (#5780)
Co-authored-by: Jacob Richman <jacob314@gmail.com>
This commit is contained in:
parent
8e6a565adb
commit
3a3b138195
|
@ -33,7 +33,7 @@ import {
|
|||
} from '../telemetry/types.js';
|
||||
import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
|
||||
import { hasCycleInSchema } from '../tools/tools.js';
|
||||
import { isStructuredError } from '../utils/quotaErrorDetection.js';
|
||||
import { StructuredError } from './turn.js';
|
||||
|
||||
/**
|
||||
* Returns true if the response is valid, false otherwise.
|
||||
|
@ -352,7 +352,6 @@ export class GeminiChat {
|
|||
} catch (error) {
|
||||
const durationMs = Date.now() - startTime;
|
||||
this._logApiError(durationMs, error, prompt_id);
|
||||
await this.maybeIncludeSchemaDepthContext(error);
|
||||
this.sendPromise = Promise.resolve();
|
||||
throw error;
|
||||
}
|
||||
|
@ -452,7 +451,6 @@ export class GeminiChat {
|
|||
const durationMs = Date.now() - startTime;
|
||||
this._logApiError(durationMs, error, prompt_id);
|
||||
this.sendPromise = Promise.resolve();
|
||||
await this.maybeIncludeSchemaDepthContext(error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
@ -523,6 +521,34 @@ export class GeminiChat {
|
|||
return lastChunkWithMetadata?.usageMetadata;
|
||||
}
|
||||
|
||||
async maybeIncludeSchemaDepthContext(error: StructuredError): Promise<void> {
|
||||
// Check for potentially problematic cyclic tools with cyclic schemas
|
||||
// and include a recommendation to remove potentially problematic tools.
|
||||
if (
|
||||
isSchemaDepthError(error.message) ||
|
||||
isInvalidArgumentError(error.message)
|
||||
) {
|
||||
const tools = (await this.config.getToolRegistry()).getAllTools();
|
||||
const cyclicSchemaTools: string[] = [];
|
||||
for (const tool of tools) {
|
||||
if (
|
||||
(tool.schema.parametersJsonSchema &&
|
||||
hasCycleInSchema(tool.schema.parametersJsonSchema)) ||
|
||||
(tool.schema.parameters && hasCycleInSchema(tool.schema.parameters))
|
||||
) {
|
||||
cyclicSchemaTools.push(tool.displayName);
|
||||
}
|
||||
}
|
||||
if (cyclicSchemaTools.length > 0) {
|
||||
const extraDetails =
|
||||
`\n\nThis error was probably caused by cyclic schema references in one of the following tools, try disabling them with excludeTools:\n\n - ` +
|
||||
cyclicSchemaTools.join(`\n - `) +
|
||||
`\n`;
|
||||
error.message += extraDetails;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private async *processStreamResponse(
|
||||
streamResponse: AsyncGenerator<GenerateContentResponse>,
|
||||
inputContent: Content,
|
||||
|
@ -684,34 +710,13 @@ export class GeminiChat {
|
|||
content.parts[0].thought === true
|
||||
);
|
||||
}
|
||||
|
||||
private async maybeIncludeSchemaDepthContext(error: unknown): Promise<void> {
|
||||
// Check for potentially problematic cyclic tools with cyclic schemas
|
||||
// and include a recommendation to remove potentially problematic tools.
|
||||
if (isStructuredError(error) && isSchemaDepthError(error.message)) {
|
||||
const tools = (await this.config.getToolRegistry()).getAllTools();
|
||||
const cyclicSchemaTools: string[] = [];
|
||||
for (const tool of tools) {
|
||||
if (
|
||||
(tool.schema.parametersJsonSchema &&
|
||||
hasCycleInSchema(tool.schema.parametersJsonSchema)) ||
|
||||
(tool.schema.parameters && hasCycleInSchema(tool.schema.parameters))
|
||||
) {
|
||||
cyclicSchemaTools.push(tool.displayName);
|
||||
}
|
||||
}
|
||||
if (cyclicSchemaTools.length > 0) {
|
||||
const extraDetails =
|
||||
`\n\nThis error was probably caused by cyclic schema references in one of the following tools, try disabling them:\n\n - ` +
|
||||
cyclicSchemaTools.join(`\n - `) +
|
||||
`\n`;
|
||||
error.message += extraDetails;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** Visible for Testing */
|
||||
export function isSchemaDepthError(errorMessage: string): boolean {
|
||||
return errorMessage.includes('maximum schema depth exceeded');
|
||||
}
|
||||
|
||||
export function isInvalidArgumentError(errorMessage: string): boolean {
|
||||
return errorMessage.includes('Request contains an invalid argument');
|
||||
}
|
||||
|
|
|
@ -17,12 +17,14 @@ import { GeminiChat } from './geminiChat.js';
|
|||
|
||||
const mockSendMessageStream = vi.fn();
|
||||
const mockGetHistory = vi.fn();
|
||||
const mockMaybeIncludeSchemaDepthContext = vi.fn();
|
||||
|
||||
vi.mock('@google/genai', async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import('@google/genai')>();
|
||||
const MockChat = vi.fn().mockImplementation(() => ({
|
||||
sendMessageStream: mockSendMessageStream,
|
||||
getHistory: mockGetHistory,
|
||||
maybeIncludeSchemaDepthContext: mockMaybeIncludeSchemaDepthContext,
|
||||
}));
|
||||
return {
|
||||
...actual,
|
||||
|
@ -46,6 +48,7 @@ describe('Turn', () => {
|
|||
type MockedChatInstance = {
|
||||
sendMessageStream: typeof mockSendMessageStream;
|
||||
getHistory: typeof mockGetHistory;
|
||||
maybeIncludeSchemaDepthContext: typeof mockMaybeIncludeSchemaDepthContext;
|
||||
};
|
||||
let mockChatInstance: MockedChatInstance;
|
||||
|
||||
|
@ -54,6 +57,7 @@ describe('Turn', () => {
|
|||
mockChatInstance = {
|
||||
sendMessageStream: mockSendMessageStream,
|
||||
getHistory: mockGetHistory,
|
||||
maybeIncludeSchemaDepthContext: mockMaybeIncludeSchemaDepthContext,
|
||||
};
|
||||
turn = new Turn(mockChatInstance as unknown as GeminiChat, 'prompt-id-1');
|
||||
mockGetHistory.mockReturnValue([]);
|
||||
|
@ -200,7 +204,7 @@ describe('Turn', () => {
|
|||
{ role: 'model', parts: [{ text: 'Previous history' }] },
|
||||
];
|
||||
mockGetHistory.mockReturnValue(historyContent);
|
||||
|
||||
mockMaybeIncludeSchemaDepthContext.mockResolvedValue(undefined);
|
||||
const events = [];
|
||||
for await (const event of turn.run(
|
||||
reqParts,
|
||||
|
|
|
@ -275,6 +275,7 @@ export class Turn {
|
|||
message: getErrorMessage(error),
|
||||
status,
|
||||
};
|
||||
await this.chat.maybeIncludeSchemaDepthContext(structuredError);
|
||||
yield { type: GeminiEventType.Error, value: { error: structuredError } };
|
||||
return;
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue