From daab5ea0bc65ae72e66d6e9bf4e0f0be8a892a9e Mon Sep 17 00:00:00 2001 From: Enrico Ros Date: Wed, 24 Jul 2024 04:15:26 -0700 Subject: [PATCH] FF: fix crashes on undefined tokens --- src/modules/aix/server/api/aix.wiretypes.ts | 14 +++++++++----- .../chatGenerate/ChatGenerateTransmitter.ts | 11 +++++++++-- .../dispatch/chatGenerate/parsers/gemini.parser.ts | 5 ++++- .../dispatch/chatGenerate/parsers/openai.parser.ts | 10 ++++++++-- .../server/dispatch/wiretypes/gemini.wiretypes.ts | 2 +- 5 files changed, 31 insertions(+), 11 deletions(-) diff --git a/src/modules/aix/server/api/aix.wiretypes.ts b/src/modules/aix/server/api/aix.wiretypes.ts index 527f6033f..04882ff2b 100644 --- a/src/modules/aix/server/api/aix.wiretypes.ts +++ b/src/modules/aix/server/api/aix.wiretypes.ts @@ -412,6 +412,10 @@ export namespace AixWire_API_ChatGenerate { /** * This is the protocol for both the control objects sent by the tRPC streaming procedures, * and the thePartTransmitter/PartReassembler. + * + * VITAL: when transmitting anything that's "undefined", leave it out of the + * object rather than setting it as 'undefined' as 'superjson' will mess it up + * and tRPC decoding will be broken (very important!) */ export namespace AixAPI_Particles { @@ -438,19 +442,19 @@ export namespace AixAPI_Particles { export type ParticleOp = | { p: 't_', t: string /* incremental text, despite not having the 'i_' prefix for brevity */ } - | { p: 'inline-image', mimeType: string, i_b64?: string } + | { p: 'inline-image', mimeType: string, i_b64?: string /* never undefined */ } | { p: 'ii_', i_b64: string } - | { p: 'inline-doc', type: string, ref: string, l1Title: string, i_text?: string } + | { p: 'inline-doc', type: string, ref: string, l1Title: string, i_text?: string /* never undefined */ } | { p: 'id_', i_text: string } - | { p: 'function-call', id: string, name: string, i_args?: string } + | { p: 'function-call', id: string, name: string, i_args?: string /* never undefined */ } | { p: 'fc_', i_args: string } | { p: 'code-call', id: string, language: string, code: string } - | { p: 'code-response', id: string, output: string, error?: string } + | { p: 'code-response', id: string, output: string, error?: string /* never undefined */ } + // NOTE: see the Vital notice export type ChatGenerateCounts = { chatIn?: number, chatOut?: number, - chatTotal?: number, chatOutRate?: number, chatTimeInner?: number, }; diff --git a/src/modules/aix/server/dispatch/chatGenerate/ChatGenerateTransmitter.ts b/src/modules/aix/server/dispatch/chatGenerate/ChatGenerateTransmitter.ts index e5c30e5f8..047646548 100644 --- a/src/modules/aix/server/dispatch/chatGenerate/ChatGenerateTransmitter.ts +++ b/src/modules/aix/server/dispatch/chatGenerate/ChatGenerateTransmitter.ts @@ -239,8 +239,15 @@ export class ChatGenerateTransmitter implements IPartTransmitter { /** Update the counters, sent twice (after the first call, and then at the end of the transmission) */ setCounters(counts: AixAPI_Particles.ChatGenerateCounts) { if (!this.accCounts) - this.accCounts = {}; - Object.assign(this.accCounts, counts); + this.accCounts = {} as AixAPI_Particles.ChatGenerateCounts; + + // similar to Object.assign, but takes care of removing the "undefined" entries + for (const key in counts) { + const value = (counts as any)[key] as number | undefined; + if (value !== undefined) + (this.accCounts as any)[key] = value; + } + this.freshCounts = true; } diff --git a/src/modules/aix/server/dispatch/chatGenerate/parsers/gemini.parser.ts b/src/modules/aix/server/dispatch/chatGenerate/parsers/gemini.parser.ts index bbe705cae..e5e4c7ac1 100644 --- a/src/modules/aix/server/dispatch/chatGenerate/parsers/gemini.parser.ts +++ b/src/modules/aix/server/dispatch/chatGenerate/parsers/gemini.parser.ts @@ -118,7 +118,10 @@ export function createGeminiGenerateContentResponseParser(modelId: string): Chat // -> Stats if (generationChunk.usageMetadata) - pt.setCounters({ chatIn: generationChunk.usageMetadata.promptTokenCount, chatOut: generationChunk.usageMetadata.candidatesTokenCount }); + pt.setCounters({ + chatIn: generationChunk.usageMetadata.promptTokenCount, + chatOut: generationChunk.usageMetadata.candidatesTokenCount, + }); }; } diff --git a/src/modules/aix/server/dispatch/chatGenerate/parsers/openai.parser.ts b/src/modules/aix/server/dispatch/chatGenerate/parsers/openai.parser.ts index b06a3f6bd..854d77b2f 100644 --- a/src/modules/aix/server/dispatch/chatGenerate/parsers/openai.parser.ts +++ b/src/modules/aix/server/dispatch/chatGenerate/parsers/openai.parser.ts @@ -82,7 +82,10 @@ export function createOpenAIChatCompletionsChunkParser(): ChatGenerateParseFunct // -> Stats if (json.usage) { if (json.usage.completion_tokens !== undefined) - pt.setCounters({ chatIn: json.usage.prompt_tokens || -1, chatOut: json.usage.completion_tokens }); + pt.setCounters({ + chatIn: json.usage.prompt_tokens || -1, + chatOut: json.usage.completion_tokens, + }); // [OpenAI] Expected correct case: the last object has usage, but an empty choices array if (!json.choices.length) @@ -196,7 +199,10 @@ export function createOpenAIChatCompletionsParserNS(): ChatGenerateParseFunction // -> Stats if (json.usage) - pt.setCounters({ chatIn: json.usage.prompt_tokens, chatOut: json.usage.completion_tokens, chatTotal: json.usage.total_tokens }); + pt.setCounters({ + chatIn: json.usage.prompt_tokens, + chatOut: json.usage.completion_tokens, + }); // Assumption/validate: expect 1 completion, or stop if (json.choices.length !== 1) diff --git a/src/modules/aix/server/dispatch/wiretypes/gemini.wiretypes.ts b/src/modules/aix/server/dispatch/wiretypes/gemini.wiretypes.ts index 66bb2d657..afd848fdc 100644 --- a/src/modules/aix/server/dispatch/wiretypes/gemini.wiretypes.ts +++ b/src/modules/aix/server/dispatch/wiretypes/gemini.wiretypes.ts @@ -469,7 +469,7 @@ export namespace GeminiWire_API_Generate_Content { const UsageMetadata_schema = z.object({ promptTokenCount: z.number(), candidatesTokenCount: z.number().optional(), // .optional: in case the first message is 'RECITATION' there could be no output token count - totalTokenCount: z.number(), + // totalTokenCount: z.number(), // cachedContentTokenCount: z.number().optional(), // Not supported for now, hence disabled });