Add openai.finish_reason span attribute for OpenAIChatModel (#507)

This adds an attribute (`openai.finish_reason`) that `OpenAIChatModel` uses to indicate how the stream ended.
fixie-ai · Nov 21, 2023 · 1b14e53 · 1b14e53 · vercel · Nov 21, 2023
1 parent 7325135
commit 1b14e53
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 3 deletions.
diff --git a/packages/ai-jsx/package.json b/packages/ai-jsx/package.json
@@ -4,7 +4,7 @@
   "repository": "fixie-ai/ai-jsx",
   "bugs": "https://github.com/fixie-ai/ai-jsx/issues",
   "homepage": "https://ai-jsx.com",
-  "version": "0.28.0",
+  "version": "0.28.1",
   "volta": {
     "extends": "../../package.json"
   },

diff --git a/packages/ai-jsx/src/lib/openai.tsx b/packages/ai-jsx/src/lib/openai.tsx
@@ -546,6 +546,7 @@ export async function* OpenAIChatModel(
 
     throw ex;
   }
+  let finishReason: string | undefined = undefined;
   const iterator = chatResponse[Symbol.asyncIterator]();
 
   // We have a single response iterator, but we'll wrap tokens _within_ the structure of <AssistantMessage> or <FunctionCall>
@@ -567,6 +568,10 @@ export async function* OpenAIChatModel(
     } while (next.value.choices.length == 0);
 
     logger.trace({ deltaMessage: next.value }, 'Got delta message');
+
+    if (next.value.choices[0].finish_reason) {
+      finishReason = next.value.choices[0].finish_reason;
+    }
     return next.value.choices[0].delta;
   }
 
@@ -667,6 +672,12 @@ export async function* OpenAIChatModel(
     }
   }
 
+  // TS doesn't realize that the advance closure can set `finishReason`.
+  // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
+  if (finishReason) {
+    logger.setAttribute('openai.finish_reason', finishReason);
+  }
+
   // Render the completion conversation to log it.
   await renderToConversation(outputMessages, render, logger, 'completion', tokenCountForConversationMessage);
   return AI.AppendOnlyStream;

diff --git a/packages/docs/docs/changelog.md b/packages/docs/docs/changelog.md
@@ -1,6 +1,10 @@
 # Changelog
 
-## 0.28.0
+## 0.28.1
+
+- Add `openai.finish_reason` span attribute for `OpenAIChatModel`
+
+## [0.28.0](https://github.com/fixie-ai/ai-jsx/commit/73251358a1121c98e9059c57d3c905b6156447c4)
 
 - Improved completion/prompt logging to include explicit message text
 

diff --git a/packages/examples/test/core/completion.tsx b/packages/examples/test/core/completion.tsx
@@ -119,6 +119,7 @@ describe('OpenTelemetry', () => {
           {"hello"}
         </UserMessage>
       </OpenAIChatModel>",
+          "openai.finish_reason": "stop",
         },
         {
           "ai.jsx.result": "opentel response from OpenAI",
@@ -257,6 +258,7 @@ describe('OpenTelemetry', () => {
           {"hello"}
         </UserMessage>
       </OpenAIChatModel>",
+          "openai.finish_reason": "stop",
         },
         {
           "ai.jsx.result": "opentel response from OpenAI",

diff --git a/packages/examples/test/lib/openai.tsx b/packages/examples/test/lib/openai.tsx
@@ -21,7 +21,7 @@ describe('OpenAIChatModel', () => {
                   }),
                 ]);
 
-                yield { choices: [{ delta: { role: 'assistant', content: 'Hi!' } }] };
+                yield { choices: [{ delta: { role: 'assistant', content: 'Hi!' }, finish_reason: 'stop' }] };
               },
             },
           },
-Original file line number
+Diff line change
@@ Expand Up / @@ -21,7 +21,7 @@ describe('OpenAIChatModel', () => { @@
                       }),
                     ]);
-                    yield { choices: [{ delta: { role: 'assistant', content: 'Hi!' } }] };
+                    yield { choices: [{ delta: { role: 'assistant', content: 'Hi!' }, finish_reason: 'stop' }] };
                   },
                 },
               },
@@ Expand Down @@