diff --git a/src/providers/fireworks-ai/chatComplete.test.ts b/src/providers/fireworks-ai/chatComplete.test.ts new file mode 100644 index 000000000..8882c6f47 --- /dev/null +++ b/src/providers/fireworks-ai/chatComplete.test.ts @@ -0,0 +1,125 @@ +import { FireworksAIChatCompleteStreamChunkTransform } from './chatComplete'; + +describe('FireworksAIChatCompleteStreamChunkTransform', () => { + it('transforms a normal stream chunk with choices', () => { + const input = JSON.stringify({ + id: 'chatcmpl-123', + object: 'chat.completion.chunk', + created: 1700000000, + model: 'accounts/fireworks/models/llama-v3p1-405b-instruct', + choices: [ + { + index: 0, + delta: { role: 'assistant', content: 'Hello' }, + finish_reason: null, + logprobs: null, + }, + ], + usage: null, + }); + + const result = FireworksAIChatCompleteStreamChunkTransform( + `data: ${input}` + ); + const parsed = JSON.parse(result.replace('data: ', '').trim()); + + expect(parsed.id).toBe('chatcmpl-123'); + expect(parsed.provider).toBe('fireworks-ai'); + expect(parsed.choices).toHaveLength(1); + expect(parsed.choices[0].delta.content).toBe('Hello'); + expect(parsed.choices[0].finish_reason).toBeNull(); + }); + + it('handles empty choices array without crashing', () => { + const input = JSON.stringify({ + id: 'chatcmpl-123', + object: 'chat.completion.chunk', + created: 1700000000, + model: 'accounts/fireworks/models/llama-v3p1-405b-instruct', + choices: [], + usage: { + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + }, + }); + + const result = FireworksAIChatCompleteStreamChunkTransform( + `data: ${input}` + ); + const parsed = JSON.parse(result.replace('data: ', '').trim()); + + expect(parsed.id).toBe('chatcmpl-123'); + expect(parsed.provider).toBe('fireworks-ai'); + expect(parsed.choices).toEqual([]); + expect(parsed.usage).toEqual({ + prompt_tokens: 10, + completion_tokens: 5, + total_tokens: 15, + }); + }); + + it('handles [DONE] chunk', () => { + const result = FireworksAIChatCompleteStreamChunkTransform('data: [DONE]'); + expect(result).toBe('data: [DONE]\n\n'); + }); + + it('includes usage when present in chunk with choices', () => { + const input = JSON.stringify({ + id: 'chatcmpl-456', + object: 'chat.completion.chunk', + created: 1700000000, + model: 'accounts/fireworks/models/llama-v3p1-405b-instruct', + choices: [ + { + index: 0, + delta: {}, + finish_reason: 'stop', + logprobs: null, + }, + ], + usage: { + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + }, + }); + + const result = FireworksAIChatCompleteStreamChunkTransform( + `data: ${input}` + ); + const parsed = JSON.parse(result.replace('data: ', '').trim()); + + expect(parsed.usage).toEqual({ + prompt_tokens: 20, + completion_tokens: 10, + total_tokens: 30, + }); + expect(parsed.choices[0].finish_reason).toBe('stop'); + }); + + it('omits usage when null', () => { + const input = JSON.stringify({ + id: 'chatcmpl-789', + object: 'chat.completion.chunk', + created: 1700000000, + model: 'accounts/fireworks/models/llama-v3p1-405b-instruct', + choices: [ + { + index: 0, + delta: { content: 'Hi' }, + finish_reason: null, + logprobs: null, + }, + ], + usage: null, + }); + + const result = FireworksAIChatCompleteStreamChunkTransform( + `data: ${input}` + ); + const parsed = JSON.parse(result.replace('data: ', '').trim()); + + expect(parsed.usage).toBeUndefined(); + }); +}); diff --git a/src/providers/fireworks-ai/chatComplete.ts b/src/providers/fireworks-ai/chatComplete.ts index 5c7fcfe66..5bcb65cf7 100644 --- a/src/providers/fireworks-ai/chatComplete.ts +++ b/src/providers/fireworks-ai/chatComplete.ts @@ -222,6 +222,7 @@ export const FireworksAIChatCompleteStreamChunkTransform: ( return `data: ${chunk}\n\n`; } const parsedChunk: FireworksAIStreamChunk = JSON.parse(chunk); + const choice = parsedChunk.choices?.[0]; return ( `data: ${JSON.stringify({ id: parsedChunk.id, @@ -229,14 +230,16 @@ export const FireworksAIChatCompleteStreamChunkTransform: ( created: parsedChunk.created, model: parsedChunk.model, provider: FIREWORKS_AI, - choices: [ - { - index: parsedChunk.choices[0].index, - delta: parsedChunk.choices[0].delta, - finish_reason: parsedChunk.choices[0].finish_reason, - logprobs: parsedChunk.choices[0].logprobs, - }, - ], + choices: choice + ? [ + { + index: choice.index, + delta: choice.delta, + finish_reason: choice.finish_reason, + logprobs: choice.logprobs, + }, + ] + : [], ...(parsedChunk.usage ? { usage: parsedChunk.usage } : {}), })}` + '\n\n' );