🐛 fix: Handle content generation errors in GoogleClient (#5575)

This commit is contained in:
Danny Avila 2025-01-31 11:22:15 -05:00 committed by GitHub
parent 6920e23fb2
commit fdf0b41d08
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -605,91 +605,97 @@ class GoogleClient extends BaseClient {
let reply = ''; let reply = '';
if (!EXCLUDED_GENAI_MODELS.test(modelName) && !this.project_id) { try {
/** @type {GenAI} */ if (!EXCLUDED_GENAI_MODELS.test(modelName) && !this.project_id) {
const client = this.client; /** @type {GenAI} */
/** @type {GenerateContentRequest} */ const client = this.client;
const requestOptions = { /** @type {GenerateContentRequest} */
safetySettings, const requestOptions = {
contents: _payload, safetySettings,
generationConfig: googleGenConfigSchema.parse(this.modelOptions), contents: _payload,
}; generationConfig: googleGenConfigSchema.parse(this.modelOptions),
const promptPrefix = (this.options.promptPrefix ?? '').trim();
if (promptPrefix.length) {
requestOptions.systemInstruction = {
parts: [
{
text: promptPrefix,
},
],
}; };
const promptPrefix = (this.options.promptPrefix ?? '').trim();
if (promptPrefix.length) {
requestOptions.systemInstruction = {
parts: [
{
text: promptPrefix,
},
],
};
}
const delay = modelName.includes('flash') ? 8 : 15;
/** @type {GenAIUsageMetadata} */
let usageMetadata;
const result = await client.generateContentStream(requestOptions);
for await (const chunk of result.stream) {
usageMetadata = !usageMetadata
? chunk?.usageMetadata
: Object.assign(usageMetadata, chunk?.usageMetadata);
const chunkText = chunk.text();
await this.generateTextStream(chunkText, onProgress, {
delay,
});
reply += chunkText;
await sleep(streamRate);
}
if (usageMetadata) {
this.usage = {
input_tokens: usageMetadata.promptTokenCount,
output_tokens: usageMetadata.candidatesTokenCount,
};
}
return reply;
} }
const delay = modelName.includes('flash') ? 8 : 15; const { instances } = _payload;
/** @type {GenAIUsageMetadata} */ const { messages: messages, context } = instances?.[0] ?? {};
if (!this.isVisionModel && context && messages?.length > 0) {
messages.unshift(new SystemMessage(context));
}
/** @type {import('@langchain/core/messages').AIMessageChunk['usage_metadata']} */
let usageMetadata; let usageMetadata;
const result = await client.generateContentStream(requestOptions); const stream = await this.client.stream(messages, {
for await (const chunk of result.stream) { signal: abortController.signal,
streamUsage: true,
safetySettings,
});
let delay = this.options.streamRate || 8;
if (!this.options.streamRate) {
if (this.isGenerativeModel) {
delay = 15;
}
if (modelName.includes('flash')) {
delay = 5;
}
}
for await (const chunk of stream) {
usageMetadata = !usageMetadata usageMetadata = !usageMetadata
? chunk?.usageMetadata ? chunk?.usage_metadata
: Object.assign(usageMetadata, chunk?.usageMetadata); : concat(usageMetadata, chunk?.usage_metadata);
const chunkText = chunk.text(); const chunkText = chunk?.content ?? chunk;
await this.generateTextStream(chunkText, onProgress, { await this.generateTextStream(chunkText, onProgress, {
delay, delay,
}); });
reply += chunkText; reply += chunkText;
await sleep(streamRate);
} }
if (usageMetadata) { if (usageMetadata) {
this.usage = { this.usage = usageMetadata;
input_tokens: usageMetadata.promptTokenCount,
output_tokens: usageMetadata.candidatesTokenCount,
};
} }
return reply; } catch (e) {
} logger.error('[GoogleClient] There was an issue generating the completion', e);
const { instances } = _payload;
const { messages: messages, context } = instances?.[0] ?? {};
if (!this.isVisionModel && context && messages?.length > 0) {
messages.unshift(new SystemMessage(context));
}
/** @type {import('@langchain/core/messages').AIMessageChunk['usage_metadata']} */
let usageMetadata;
const stream = await this.client.stream(messages, {
signal: abortController.signal,
streamUsage: true,
safetySettings,
});
let delay = this.options.streamRate || 8;
if (!this.options.streamRate) {
if (this.isGenerativeModel) {
delay = 15;
}
if (modelName.includes('flash')) {
delay = 5;
}
}
for await (const chunk of stream) {
usageMetadata = !usageMetadata
? chunk?.usage_metadata
: concat(usageMetadata, chunk?.usage_metadata);
const chunkText = chunk?.content ?? chunk;
await this.generateTextStream(chunkText, onProgress, {
delay,
});
reply += chunkText;
}
if (usageMetadata) {
this.usage = usageMetadata;
} }
return reply; return reply;
} }