Merge pull request #1003 from RooVetGit/o1_streaming

Enable streaming for o1
This commit is contained in:
Matt Rubens 2025-02-14 22:33:00 -05:00 committed by GitHub
commit 36bbfd917d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 20 additions and 10 deletions

View file

@ -130,11 +130,20 @@ describe("OpenAiNativeHandler", () => {
})
mockCreate.mockResolvedValueOnce({
choices: [{ message: { content: null } }],
usage: {
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
[Symbol.asyncIterator]: async function* () {
yield {
choices: [
{
delta: { content: null },
index: 0,
},
],
usage: {
prompt_tokens: 0,
completion_tokens: 0,
total_tokens: 0,
},
}
},
})
@ -144,10 +153,7 @@ describe("OpenAiNativeHandler", () => {
results.push(result)
}
expect(results).toEqual([
{ type: "text", text: "" },
{ type: "usage", inputTokens: 0, outputTokens: 0 },
])
expect(results).toEqual([{ type: "usage", inputTokens: 0, outputTokens: 0 }])
// Verify developer role is used for system prompt with o1 model
expect(mockCreate).toHaveBeenCalledWith({
@ -156,6 +162,8 @@ describe("OpenAiNativeHandler", () => {
{ role: "developer", content: "Formatting re-enabled\n" + systemPrompt },
{ role: "user", content: "Hello!" },
],
stream: true,
stream_options: { include_usage: true },
})
})

View file

@ -56,9 +56,11 @@ export class OpenAiNativeHandler implements ApiHandler, SingleCompletionHandler
},
...convertToOpenAiMessages(messages),
],
stream: true,
stream_options: { include_usage: true },
})
yield* this.yieldResponseData(response)
yield* this.handleStreamResponse(response)
}
private async *handleO3FamilyMessage(