fix: chat completion streaming when enabling logprobs

Author: dgellow

examples/chat-completions.ts

@@ -8,17 +8,18 @@ async function main() {
     .stream({
       model: 'mistralai/Mixtral-8x7B-Instruct-v0.1',
       messages: [{ role: 'user', content: 'Say this is a test' }],
+      logprobs: 1,
     })
     .on('message', (msg) => console.log(msg))
     .on('content', (diff) => process.stdout.write(diff));
 
   for await (const chunk of runner) {
     // Note: comment out the next line to print chunks as they are streamed from the API
-    // console.log('chunk', chunk);
+    console.log('chunk', chunk);
   }
 
   const result = await runner.finalMessage();
-  console.log(result);
+  console.log({ result });
 }
 
 main();

src/lib/ChatCompletionStream.ts

@@ -157,20 +157,26 @@ export class ChatCompletionStream
     for (const { delta, finish_reason, index, logprobs = null, ...other } of chunk.choices) {
       let choice = snapshot.choices[index];
       if (!choice) {
-        choice = snapshot.choices[index] = { finish_reason, index, message: {}, logprobs, ...other };
+        choice = snapshot.choices[index] = {
+          finish_reason,
+          index,
+          message: {},
+          logprobs: { token_ids: [], token_logprobs: [], tokens: [] },
+          ...other,
+        };
       }
 
       if (logprobs) {
+        console.log({ logprobs });
         if (!choice.logprobs) {
-          choice.logprobs = Object.assign({}, logprobs);
-        } else {
-          const { content, ...rest } = logprobs;
-          Object.assign(choice.logprobs, rest);
-          if (content) {
-            choice.logprobs.content ??= [];
-            choice.logprobs.content.push(...content);
-          }
+          choice.logprobs = { token_ids: [], token_logprobs: [], tokens: [] };
         }
+        choice.logprobs.token_ids ??= [];
+        choice.logprobs.token_ids.push(delta.token_id ?? null);
+        choice.logprobs.token_logprobs ??= [];
+        choice.logprobs.token_logprobs.push(logprobs ?? null);
+        choice.logprobs.tokens ??= [];
+        choice.logprobs.tokens.push(delta.content ?? null);
       }
 
       if (finish_reason) choice.finish_reason = finish_reason;

src/resources/completions.ts

@@ -66,17 +66,17 @@ export interface LogProbs {
   /**
    * List of token IDs corresponding to the logprobs
    */
-  token_ids?: Array<number>;
+  token_ids?: Array<number | null>;
 
   /**
    * List of token log probabilities
    */
-  token_logprobs?: Array<number>;
+  token_logprobs?: Array<number | null>;
 
   /**
    * List of token strings
    */
-  tokens?: Array<string>;
+  tokens?: Array<string | null>;
 }
 
 export interface ToolChoice {