Skip to content

Commit f1852f6

Browse files
stainless-botstainless-bot
committed
feat(api): add token logprobs to chat completions (#576)
1 parent 465236d commit f1852f6

File tree

12 files changed

+184
-18
lines changed

12 files changed

+184
-18
lines changed

‎api.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ Types:
3737
- <code><a href="./src/resources/chat/completions.ts">ChatCompletionNamedToolChoice</a></code>
3838
- <code><a href="./src/resources/chat/completions.ts">ChatCompletionRole</a></code>
3939
- <code><a href="./src/resources/chat/completions.ts">ChatCompletionSystemMessageParam</a></code>
40+
- <code><a href="./src/resources/chat/completions.ts">ChatCompletionTokenLogprob</a></code>
4041
- <code><a href="./src/resources/chat/completions.ts">ChatCompletionTool</a></code>
4142
- <code><a href="./src/resources/chat/completions.ts">ChatCompletionToolChoiceOption</a></code>
4243
- <code><a href="./src/resources/chat/completions.ts">ChatCompletionToolMessageParam</a></code>

‎examples/logprobs.ts‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
#!/usr/bin/env -S npm run tsn -T
2+
3+
import OpenAI from 'openai';
4+
5+
// gets API Key from environment variable OPENAI_API_KEY
6+
const openai = new OpenAI();
7+
8+
async function main() {
9+
const stream = await openai.beta.chat.completions
10+
.stream({
11+
model: 'gpt-4',
12+
messages: [{ role: 'user', content: 'Say this is a test' }],
13+
stream: true,
14+
logprobs: true,
15+
})
16+
.on('logprob', (logprob) => {
17+
console.log(logprob);
18+
});
19+
20+
console.dir(await stream.finalChatCompletion(), { depth: null });
21+
}
22+
23+
main();

‎src/index.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -242,6 +242,7 @@ export namespace OpenAI {
242242
export import ChatCompletionNamedToolChoice = API.ChatCompletionNamedToolChoice;
243243
export import ChatCompletionRole = API.ChatCompletionRole;
244244
export import ChatCompletionSystemMessageParam = API.ChatCompletionSystemMessageParam;
245+
export import ChatCompletionTokenLogprob = API.ChatCompletionTokenLogprob;
245246
export import ChatCompletionTool = API.ChatCompletionTool;
246247
export import ChatCompletionToolChoiceOption = API.ChatCompletionToolChoiceOption;
247248
export import ChatCompletionToolMessageParam = API.ChatCompletionToolMessageParam;

‎src/lib/ChatCompletionRunFunctions.test.ts‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,7 @@ function* contentChoiceDeltas(
146146
yield {
147147
index,
148148
finish_reason: i === deltas.length - 1 ? 'stop' : null,
149+
logprobs: null,
149150
delta: {
150151
role,
151152
content: deltas[i] ? `${deltas[i]}${i === deltas.length - 1 ? '' : ' '}` : null,
@@ -593,6 +594,7 @@ describe('resource completions', () => {
593594
{
594595
index: 0,
595596
finish_reason: 'function_call',
597+
logprobs: null,
596598
message: {
597599
role: 'assistant',
598600
content: null,
@@ -645,6 +647,7 @@ describe('resource completions', () => {
645647
{
646648
index: 0,
647649
finish_reason: 'stop',
650+
logprobs: null,
648651
message: {
649652
role: 'assistant',
650653
content: `it's raining`,
@@ -716,6 +719,7 @@ describe('resource completions', () => {
716719
{
717720
index: 0,
718721
finish_reason: 'function_call',
722+
logprobs: null,
719723
message: {
720724
role: 'assistant',
721725
content: null,
@@ -808,6 +812,7 @@ describe('resource completions', () => {
808812
{
809813
index: 0,
810814
finish_reason: 'function_call',
815+
logprobs: null,
811816
message: {
812817
role: 'assistant',
813818
content: null,
@@ -867,6 +872,7 @@ describe('resource completions', () => {
867872
{
868873
index: 0,
869874
finish_reason: 'stop',
875+
logprobs: null,
870876
message: {
871877
role: 'assistant',
872878
content: `there are 3 properties in {"a": 1, "b": 2, "c": 3}`,
@@ -953,6 +959,7 @@ describe('resource completions', () => {
953959
{
954960
index: 0,
955961
finish_reason: 'function_call',
962+
logprobs: null,
956963
message: {
957964
role: 'assistant',
958965
content: null,
@@ -1006,6 +1013,7 @@ describe('resource completions', () => {
10061013
{
10071014
index: 0,
10081015
finish_reason: 'function_call',
1016+
logprobs: null,
10091017
message: {
10101018
role: 'assistant',
10111019
content: null,
@@ -1078,6 +1086,7 @@ describe('resource completions', () => {
10781086
{
10791087
index: 0,
10801088
finish_reason: 'stop',
1089+
logprobs: null,
10811090
message: {
10821091
role: 'assistant',
10831092
content: `there are 3 properties in {"a": 1, "b": 2, "c": 3}`,
@@ -1164,6 +1173,7 @@ describe('resource completions', () => {
11641173
{
11651174
index: 0,
11661175
finish_reason: 'function_call',
1176+
logprobs: null,
11671177
message: {
11681178
role: 'assistant',
11691179
content: null,
@@ -1241,6 +1251,7 @@ describe('resource completions', () => {
12411251
{
12421252
index: 0,
12431253
finish_reason: 'function_call',
1254+
logprobs: null,
12441255
message: {
12451256
role: 'assistant',
12461257
content: null,
@@ -1291,6 +1302,7 @@ describe('resource completions', () => {
12911302
{
12921303
index: 0,
12931304
finish_reason: 'function_call',
1305+
logprobs: null,
12941306
message: {
12951307
role: 'assistant',
12961308
content: null,
@@ -1360,6 +1372,7 @@ describe('resource completions', () => {
13601372
{
13611373
index: 0,
13621374
finish_reason: 'stop',
1375+
logprobs: null,
13631376
message: {
13641377
role: 'assistant',
13651378
content: `it's raining`,
@@ -1436,6 +1449,7 @@ describe('resource completions', () => {
14361449
{
14371450
index: 0,
14381451
finish_reason: 'function_call',
1452+
logprobs: null,
14391453
delta: {
14401454
role: 'assistant',
14411455
content: null,
@@ -2071,6 +2085,7 @@ describe('resource completions', () => {
20712085
{
20722086
index: 0,
20732087
finish_reason: 'function_call',
2088+
logprobs: null,
20742089
delta: {
20752090
role: 'assistant',
20762091
content: null,

‎src/lib/ChatCompletionStream.ts‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -153,13 +153,22 @@ export class ChatCompletionStream
153153
Object.assign(snapshot, rest);
154154
}
155155

156-
for (const { delta, finish_reason, index, ...other } of chunk.choices) {
156+
for (const { delta, finish_reason, index, logprobs = null, ...other } of chunk.choices) {
157157
let choice = snapshot.choices[index];
158158
if (!choice) {
159-
snapshot.choices[index] = { finish_reason, index, message: delta, ...other };
159+
snapshot.choices[index] = { finish_reason, index, message: delta, logprobs, ...other };
160160
continue;
161161
}
162162

163+
if (logprobs) {
164+
if (!choice.logprobs) {
165+
choice.logprobs = logprobs;
166+
} else if (logprobs.content) {
167+
choice.logprobs.content ??= [];
168+
choice.logprobs.content.push(...logprobs.content);
169+
}
170+
}
171+
163172
if (finish_reason) choice.finish_reason = finish_reason;
164173
Object.assign(choice, other);
165174

@@ -242,7 +251,7 @@ function finalizeChatCompletion(snapshot: ChatCompletionSnapshot): ChatCompletio
242251
const { id, choices, created, model } = snapshot;
243252
return {
244253
id,
245-
choices: choices.map(({ message, finish_reason, index }): ChatCompletion.Choice => {
254+
choices: choices.map(({ message, finish_reason, index, logprobs }): ChatCompletion.Choice => {
246255
if (!finish_reason) throw new OpenAIError(`missing finish_reason for choice ${index}`);
247256
const { content = null, function_call, tool_calls } = message;
248257
const role = message.role as 'assistant'; // this is what we expect; in theory it could be different which would make our types a slight lie but would be fine.
@@ -251,12 +260,18 @@ function finalizeChatCompletion(snapshot: ChatCompletionSnapshot): ChatCompletio
251260
const { arguments: args, name } = function_call;
252261
if (args == null) throw new OpenAIError(`missing function_call.arguments for choice ${index}`);
253262
if (!name) throw new OpenAIError(`missing function_call.name for choice ${index}`);
254-
return { message: { content, function_call: { arguments: args, name }, role }, finish_reason, index };
263+
return {
264+
message: { content, function_call: { arguments: args, name }, role },
265+
finish_reason,
266+
index,
267+
logprobs,
268+
};
255269
}
256270
if (tool_calls) {
257271
return {
258272
index,
259273
finish_reason,
274+
logprobs,
260275
message: {
261276
role,
262277
content,
@@ -281,7 +296,7 @@ function finalizeChatCompletion(snapshot: ChatCompletionSnapshot): ChatCompletio
281296
},
282297
};
283298
}
284-
return { message: { content: content, role }, finish_reason, index };
299+
return { message: { content: content, role }, finish_reason, index, logprobs };
285300
}),
286301
created,
287302
model,
@@ -336,6 +351,11 @@ export namespace ChatCompletionSnapshot {
336351
*/
337352
finish_reason: ChatCompletion.Choice['finish_reason'] | null;
338353

354+
/**
355+
* Log probability information for the choice.
356+
*/
357+
logprobs: ChatCompletion.Choice.Logprobs | null;
358+
339359
/**
340360
* The index of the choice in the list of choices.
341361
*/

‎src/resources/beta/threads/runs/steps.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ export interface MessageCreationStepDetails {
180180
message_creation: MessageCreationStepDetails.MessageCreation;
181181

182182
/**
183-
* Always `message_creation``.
183+
* Always `message_creation`.
184184
*/
185185
type: 'message_creation';
186186
}
@@ -269,7 +269,7 @@ export interface RunStep {
269269
metadata: unknown | null;
270270

271271
/**
272-
* The object type, which is always `thread.run.step``.
272+
* The object type, which is always `thread.run.step`.
273273
*/
274274
object: 'thread.run.step';
275275

‎src/resources/chat/chat.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ export namespace Chat {
2323
export import ChatCompletionNamedToolChoice = CompletionsAPI.ChatCompletionNamedToolChoice;
2424
export import ChatCompletionRole = CompletionsAPI.ChatCompletionRole;
2525
export import ChatCompletionSystemMessageParam = CompletionsAPI.ChatCompletionSystemMessageParam;
26+
export import ChatCompletionTokenLogprob = CompletionsAPI.ChatCompletionTokenLogprob;
2627
export import ChatCompletionTool = CompletionsAPI.ChatCompletionTool;
2728
export import ChatCompletionToolChoiceOption = CompletionsAPI.ChatCompletionToolChoiceOption;
2829
export import ChatCompletionToolMessageParam = CompletionsAPI.ChatCompletionToolMessageParam;

0 commit comments

Comments
 (0)