mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
feat: migrate OpenAiHandler to AI SDK (#11351)
Co-authored-by: Roo Code <roomote@roocode.com>
This commit is contained in:
Daniel
GitHub
parent
7a38b99232
commit
a4914c438c
5 changed files with 796 additions and 1294 deletions
|
|
@ -4,15 +4,27 @@ import OpenAI from "openai"
|
|||
|
||||
import { OpenAiHandler } from "../openai"
|
||||
|
||||
vi.mock("@ai-sdk/openai-compatible", () => ({
|
||||
createOpenAICompatible: vi.fn(() => vi.fn((modelId: string) => ({ modelId, provider: "openai-compatible" }))),
|
||||
}))
|
||||
|
||||
vi.mock("@ai-sdk/azure", () => ({
|
||||
createAzure: vi.fn(() => ({
|
||||
chat: vi.fn((modelId: string) => ({ modelId, provider: "azure.chat" })),
|
||||
})),
|
||||
}))
|
||||
|
||||
describe("OpenAiHandler native tools", () => {
|
||||
it("includes tools in request when tools are provided via metadata (regression test)", async () => {
|
||||
const mockCreate = vi.fn().mockImplementationOnce(() => ({
|
||||
[Symbol.asyncIterator]: async function* () {
|
||||
yield {
|
||||
choices: [{ delta: { content: "Test response" } }],
|
||||
}
|
||||
},
|
||||
}))
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test response" }
|
||||
}
|
||||
|
||||
mockStreamText.mockReturnValueOnce({
|
||||
fullStream: mockFullStream(),
|
||||
usage: Promise.resolve({ inputTokens: 10, outputTokens: 5 }),
|
||||
providerMetadata: Promise.resolve(undefined),
|
||||
})
|
||||
|
||||
// Set openAiCustomModelInfo without any tool capability flags; tools should
|
||||
// still be passed whenever metadata.tools is present.
|
||||
|
|
@ -26,16 +38,6 @@ describe("OpenAiHandler native tools", () => {
|
|||
},
|
||||
} as unknown as import("../../../shared/api").ApiHandlerOptions)
|
||||
|
||||
// Patch the OpenAI client call
|
||||
const mockClient = {
|
||||
chat: {
|
||||
completions: {
|
||||
create: mockCreate,
|
||||
},
|
||||
},
|
||||
} as unknown as OpenAI
|
||||
;(handler as unknown as { client: OpenAI }).client = mockClient
|
||||
|
||||
const tools: OpenAI.Chat.ChatCompletionTool[] = [
|
||||
{
|
||||
type: "function",
|
||||
|
|
@ -53,17 +55,12 @@ describe("OpenAiHandler native tools", () => {
|
|||
})
|
||||
await stream.next()
|
||||
|
||||
expect(mockCreate).toHaveBeenCalledWith(
|
||||
expect(mockStreamText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
tools: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
type: "function",
|
||||
function: expect.objectContaining({ name: "test_tool" }),
|
||||
}),
|
||||
]),
|
||||
parallel_tool_calls: true,
|
||||
tools: expect.objectContaining({
|
||||
test_tool: expect.anything(),
|
||||
}),
|
||||
}),
|
||||
expect.anything(),
|
||||
)
|
||||
})
|
||||
})
|
||||
|
|
@ -92,6 +89,10 @@ vi.mock("@ai-sdk/openai", () => ({
|
|||
modelId: "gpt-4o",
|
||||
provider: "openai.responses",
|
||||
}))
|
||||
;(provider as any).chat = vi.fn((modelId: string) => ({
|
||||
modelId,
|
||||
provider: "openai.chat",
|
||||
}))
|
||||
return provider
|
||||
}),
|
||||
}))
|
||||
|
|
|
|||
|
|
@ -3,51 +3,34 @@
|
|||
import { OpenAiHandler } from "../openai"
|
||||
import { ApiHandlerOptions } from "../../../shared/api"
|
||||
|
||||
// Mock the timeout config utility
|
||||
vitest.mock("../utils/timeout-config", () => ({
|
||||
getApiRequestTimeout: vitest.fn(),
|
||||
const mockCreateOpenAI = vi.hoisted(() => vi.fn())
|
||||
const mockCreateOpenAICompatible = vi.hoisted(() => vi.fn())
|
||||
const mockCreateAzure = vi.hoisted(() => vi.fn())
|
||||
|
||||
vi.mock("@ai-sdk/openai", () => ({
|
||||
createOpenAI: mockCreateOpenAI.mockImplementation(() => ({
|
||||
chat: vi.fn(() => ({ modelId: "test", provider: "openai.chat" })),
|
||||
})),
|
||||
}))
|
||||
|
||||
import { getApiRequestTimeout } from "../utils/timeout-config"
|
||||
vi.mock("@ai-sdk/openai-compatible", () => ({
|
||||
createOpenAICompatible: mockCreateOpenAICompatible.mockImplementation(() =>
|
||||
vi.fn((modelId: string) => ({ modelId, provider: "openai-compatible" })),
|
||||
),
|
||||
}))
|
||||
|
||||
// Mock OpenAI and AzureOpenAI
|
||||
const mockOpenAIConstructor = vitest.fn()
|
||||
const mockAzureOpenAIConstructor = vitest.fn()
|
||||
vi.mock("@ai-sdk/azure", () => ({
|
||||
createAzure: mockCreateAzure.mockImplementation(() => ({
|
||||
chat: vi.fn((modelId: string) => ({ modelId, provider: "azure.chat" })),
|
||||
})),
|
||||
}))
|
||||
|
||||
vitest.mock("openai", () => {
|
||||
return {
|
||||
__esModule: true,
|
||||
default: vitest.fn().mockImplementation((config) => {
|
||||
mockOpenAIConstructor(config)
|
||||
return {
|
||||
chat: {
|
||||
completions: {
|
||||
create: vitest.fn(),
|
||||
},
|
||||
},
|
||||
}
|
||||
}),
|
||||
AzureOpenAI: vitest.fn().mockImplementation((config) => {
|
||||
mockAzureOpenAIConstructor(config)
|
||||
return {
|
||||
chat: {
|
||||
completions: {
|
||||
create: vitest.fn(),
|
||||
},
|
||||
},
|
||||
}
|
||||
}),
|
||||
}
|
||||
})
|
||||
|
||||
describe("OpenAiHandler timeout configuration", () => {
|
||||
describe("OpenAiHandler provider configuration", () => {
|
||||
beforeEach(() => {
|
||||
vitest.clearAllMocks()
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
it("should use default timeout for standard OpenAI", () => {
|
||||
;(getApiRequestTimeout as any).mockReturnValue(600000)
|
||||
|
||||
it("should use createOpenAI for standard OpenAI endpoints", () => {
|
||||
const options: ApiHandlerOptions = {
|
||||
apiModelId: "gpt-4",
|
||||
openAiModelId: "gpt-4",
|
||||
|
|
@ -56,19 +39,15 @@ describe("OpenAiHandler timeout configuration", () => {
|
|||
|
||||
new OpenAiHandler(options)
|
||||
|
||||
expect(getApiRequestTimeout).toHaveBeenCalled()
|
||||
expect(mockOpenAIConstructor).toHaveBeenCalledWith(
|
||||
expect(mockCreateOpenAI).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: "https://api.openai.com/v1",
|
||||
apiKey: "test-key",
|
||||
timeout: 600000, // 600 seconds in milliseconds
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should use custom timeout for OpenAI-compatible providers", () => {
|
||||
;(getApiRequestTimeout as any).mockReturnValue(1800000) // 30 minutes
|
||||
|
||||
it("should use createOpenAI for custom OpenAI-compatible providers", () => {
|
||||
const options: ApiHandlerOptions = {
|
||||
apiModelId: "custom-model",
|
||||
openAiModelId: "custom-model",
|
||||
|
|
@ -78,17 +57,14 @@ describe("OpenAiHandler timeout configuration", () => {
|
|||
|
||||
new OpenAiHandler(options)
|
||||
|
||||
expect(mockOpenAIConstructor).toHaveBeenCalledWith(
|
||||
expect(mockCreateOpenAI).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
baseURL: "http://localhost:8080/v1",
|
||||
timeout: 1800000, // 1800 seconds in milliseconds
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should use timeout for Azure OpenAI", () => {
|
||||
;(getApiRequestTimeout as any).mockReturnValue(900000) // 15 minutes
|
||||
|
||||
it("should use createAzure for Azure OpenAI", () => {
|
||||
const options: ApiHandlerOptions = {
|
||||
apiModelId: "gpt-4",
|
||||
openAiModelId: "gpt-4",
|
||||
|
|
@ -99,16 +75,16 @@ describe("OpenAiHandler timeout configuration", () => {
|
|||
|
||||
new OpenAiHandler(options)
|
||||
|
||||
expect(mockAzureOpenAIConstructor).toHaveBeenCalledWith(
|
||||
expect(mockCreateAzure).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
timeout: 900000, // 900 seconds in milliseconds
|
||||
baseURL: "https://myinstance.openai.azure.com/openai",
|
||||
apiKey: "test-key",
|
||||
useDeploymentBasedUrls: true,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should use timeout for Azure AI Inference", () => {
|
||||
;(getApiRequestTimeout as any).mockReturnValue(1200000) // 20 minutes
|
||||
|
||||
it("should use createOpenAICompatible for Azure AI Inference", () => {
|
||||
const options: ApiHandlerOptions = {
|
||||
apiModelId: "deepseek",
|
||||
openAiModelId: "deepseek",
|
||||
|
|
@ -118,26 +94,32 @@ describe("OpenAiHandler timeout configuration", () => {
|
|||
|
||||
new OpenAiHandler(options)
|
||||
|
||||
expect(mockOpenAIConstructor).toHaveBeenCalledWith(
|
||||
expect(mockCreateOpenAICompatible).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
timeout: 1200000, // 1200 seconds in milliseconds
|
||||
baseURL: "https://myinstance.services.ai.azure.com/models",
|
||||
apiKey: "test-key",
|
||||
queryParams: expect.objectContaining({
|
||||
"api-version": expect.any(String),
|
||||
}),
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should handle zero timeout (no timeout)", () => {
|
||||
;(getApiRequestTimeout as any).mockReturnValue(0)
|
||||
|
||||
it("should include custom headers in provider configuration", () => {
|
||||
const options: ApiHandlerOptions = {
|
||||
apiModelId: "gpt-4",
|
||||
openAiModelId: "gpt-4",
|
||||
openAiApiKey: "test-key",
|
||||
openAiHeaders: { "X-Custom": "value" },
|
||||
}
|
||||
|
||||
new OpenAiHandler(options)
|
||||
|
||||
expect(mockOpenAIConstructor).toHaveBeenCalledWith(
|
||||
expect(mockCreateOpenAI).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
timeout: 0, // No timeout
|
||||
headers: expect.objectContaining({
|
||||
"X-Custom": "value",
|
||||
}),
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
|
|
|||
|
|
@ -5,89 +5,38 @@ import { Anthropic } from "@anthropic-ai/sdk"
|
|||
import { ApiHandlerOptions } from "../../../shared/api"
|
||||
import { OpenAiHandler } from "../openai"
|
||||
|
||||
const mockCreate = vitest.fn()
|
||||
const { mockStreamText } = vi.hoisted(() => ({
|
||||
mockStreamText: vi.fn(),
|
||||
}))
|
||||
|
||||
vitest.mock("openai", () => {
|
||||
vi.mock("ai", async (importOriginal) => {
|
||||
const actual = await importOriginal<typeof import("ai")>()
|
||||
return {
|
||||
__esModule: true,
|
||||
default: vitest.fn().mockImplementation(() => ({
|
||||
chat: {
|
||||
completions: {
|
||||
create: mockCreate.mockImplementation(async (options) => {
|
||||
if (!options.stream) {
|
||||
return {
|
||||
id: "test-completion",
|
||||
choices: [
|
||||
{
|
||||
message: { role: "assistant", content: "Test response", refusal: null },
|
||||
finish_reason: "stop",
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usage: {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 5,
|
||||
total_tokens: 15,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// Return a stream with multiple chunks that include usage metrics
|
||||
return {
|
||||
[Symbol.asyncIterator]: async function* () {
|
||||
// First chunk with partial usage
|
||||
yield {
|
||||
choices: [
|
||||
{
|
||||
delta: { content: "Test " },
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usage: {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 2,
|
||||
total_tokens: 12,
|
||||
},
|
||||
}
|
||||
|
||||
// Second chunk with updated usage
|
||||
yield {
|
||||
choices: [
|
||||
{
|
||||
delta: { content: "response" },
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usage: {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 4,
|
||||
total_tokens: 14,
|
||||
},
|
||||
}
|
||||
|
||||
// Final chunk with complete usage
|
||||
yield {
|
||||
choices: [
|
||||
{
|
||||
delta: {},
|
||||
index: 0,
|
||||
},
|
||||
],
|
||||
usage: {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 5,
|
||||
total_tokens: 15,
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
}),
|
||||
},
|
||||
},
|
||||
})),
|
||||
...actual,
|
||||
streamText: mockStreamText,
|
||||
generateText: vi.fn(),
|
||||
}
|
||||
})
|
||||
|
||||
vi.mock("@ai-sdk/openai", () => ({
|
||||
createOpenAI: vi.fn(() => ({
|
||||
chat: vi.fn(() => ({
|
||||
modelId: "gpt-4",
|
||||
provider: "openai.chat",
|
||||
})),
|
||||
})),
|
||||
}))
|
||||
|
||||
vi.mock("@ai-sdk/openai-compatible", () => ({
|
||||
createOpenAICompatible: vi.fn(() => vi.fn((modelId: string) => ({ modelId, provider: "openai-compatible" }))),
|
||||
}))
|
||||
|
||||
vi.mock("@ai-sdk/azure", () => ({
|
||||
createAzure: vi.fn(() => ({
|
||||
chat: vi.fn((modelId: string) => ({ modelId, provider: "azure.chat" })),
|
||||
})),
|
||||
}))
|
||||
|
||||
describe("OpenAiHandler with usage tracking fix", () => {
|
||||
let handler: OpenAiHandler
|
||||
let mockOptions: ApiHandlerOptions
|
||||
|
|
@ -99,7 +48,7 @@ describe("OpenAiHandler with usage tracking fix", () => {
|
|||
openAiBaseUrl: "https://api.openai.com/v1",
|
||||
}
|
||||
handler = new OpenAiHandler(mockOptions)
|
||||
mockCreate.mockClear()
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
describe("usage metrics with streaming", () => {
|
||||
|
|
@ -117,19 +66,31 @@ describe("OpenAiHandler with usage tracking fix", () => {
|
|||
]
|
||||
|
||||
it("should only yield usage metrics once at the end of the stream", async () => {
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test " }
|
||||
yield { type: "text-delta", text: "response" }
|
||||
}
|
||||
|
||||
mockStreamText.mockReturnValueOnce({
|
||||
fullStream: mockFullStream(),
|
||||
usage: Promise.resolve({
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
}),
|
||||
providerMetadata: Promise.resolve(undefined),
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages)
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Check we have text chunks
|
||||
const textChunks = chunks.filter((chunk) => chunk.type === "text")
|
||||
expect(textChunks).toHaveLength(2)
|
||||
expect(textChunks[0].text).toBe("Test ")
|
||||
expect(textChunks[1].text).toBe("response")
|
||||
|
||||
// Check we only have one usage chunk and it's the last one
|
||||
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
|
||||
expect(usageChunks).toHaveLength(1)
|
||||
expect(usageChunks[0]).toEqual({
|
||||
|
|
@ -138,49 +99,25 @@ describe("OpenAiHandler with usage tracking fix", () => {
|
|||
outputTokens: 5,
|
||||
})
|
||||
|
||||
// Check the usage chunk is the last one reported from the API
|
||||
const lastChunk = chunks[chunks.length - 1]
|
||||
expect(lastChunk.type).toBe("usage")
|
||||
expect(lastChunk.inputTokens).toBe(10)
|
||||
expect(lastChunk.outputTokens).toBe(5)
|
||||
})
|
||||
|
||||
it("should handle case where usage is only in the final chunk", async () => {
|
||||
// Override the mock for this specific test
|
||||
mockCreate.mockImplementationOnce(async (options) => {
|
||||
if (!options.stream) {
|
||||
return {
|
||||
id: "test-completion",
|
||||
choices: [{ message: { role: "assistant", content: "Test response" } }],
|
||||
usage: { prompt_tokens: 10, completion_tokens: 5, total_tokens: 15 },
|
||||
}
|
||||
}
|
||||
it("should handle case where usage is provided after stream completes", async () => {
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test " }
|
||||
yield { type: "text-delta", text: "response" }
|
||||
}
|
||||
|
||||
return {
|
||||
[Symbol.asyncIterator]: async function* () {
|
||||
// First chunk with no usage
|
||||
yield {
|
||||
choices: [{ delta: { content: "Test " }, index: 0 }],
|
||||
usage: null,
|
||||
}
|
||||
|
||||
// Second chunk with no usage
|
||||
yield {
|
||||
choices: [{ delta: { content: "response" }, index: 0 }],
|
||||
usage: null,
|
||||
}
|
||||
|
||||
// Final chunk with usage data
|
||||
yield {
|
||||
choices: [{ delta: {}, index: 0 }],
|
||||
usage: {
|
||||
prompt_tokens: 10,
|
||||
completion_tokens: 5,
|
||||
total_tokens: 15,
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
mockStreamText.mockReturnValueOnce({
|
||||
fullStream: mockFullStream(),
|
||||
usage: Promise.resolve({
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
}),
|
||||
providerMetadata: Promise.resolve(undefined),
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages)
|
||||
|
|
@ -189,7 +126,6 @@ describe("OpenAiHandler with usage tracking fix", () => {
|
|||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Check usage metrics
|
||||
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
|
||||
expect(usageChunks).toHaveLength(1)
|
||||
expect(usageChunks[0]).toEqual({
|
||||
|
|
@ -200,28 +136,14 @@ describe("OpenAiHandler with usage tracking fix", () => {
|
|||
})
|
||||
|
||||
it("should handle case where no usage is provided", async () => {
|
||||
// Override the mock for this specific test
|
||||
mockCreate.mockImplementationOnce(async (options) => {
|
||||
if (!options.stream) {
|
||||
return {
|
||||
id: "test-completion",
|
||||
choices: [{ message: { role: "assistant", content: "Test response" } }],
|
||||
usage: null,
|
||||
}
|
||||
}
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test response" }
|
||||
}
|
||||
|
||||
return {
|
||||
[Symbol.asyncIterator]: async function* () {
|
||||
yield {
|
||||
choices: [{ delta: { content: "Test response" }, index: 0 }],
|
||||
usage: null,
|
||||
}
|
||||
yield {
|
||||
choices: [{ delta: {}, index: 0 }],
|
||||
usage: null,
|
||||
}
|
||||
},
|
||||
}
|
||||
mockStreamText.mockReturnValueOnce({
|
||||
fullStream: mockFullStream(),
|
||||
usage: Promise.resolve(undefined),
|
||||
providerMetadata: Promise.resolve(undefined),
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages)
|
||||
|
|
@ -230,9 +152,81 @@ describe("OpenAiHandler with usage tracking fix", () => {
|
|||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Check we don't have any usage chunks
|
||||
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
|
||||
expect(usageChunks).toHaveLength(0)
|
||||
})
|
||||
|
||||
it("should include reasoningTokens from usage.details", async () => {
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test response" }
|
||||
}
|
||||
|
||||
mockStreamText.mockReturnValueOnce({
|
||||
fullStream: mockFullStream(),
|
||||
usage: Promise.resolve({
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
details: {
|
||||
reasoningTokens: 3,
|
||||
},
|
||||
}),
|
||||
providerMetadata: Promise.resolve(undefined),
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages)
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
|
||||
expect(usageChunks).toHaveLength(1)
|
||||
expect(usageChunks[0]).toEqual(
|
||||
expect.objectContaining({
|
||||
type: "usage",
|
||||
inputTokens: 10,
|
||||
outputTokens: 5,
|
||||
reasoningTokens: 3,
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
it("should extract cache and reasoning tokens from providerMetadata", async () => {
|
||||
async function* mockFullStream() {
|
||||
yield { type: "text-delta", text: "Test response" }
|
||||
}
|
||||
|
||||
mockStreamText.mockReturnValueOnce({
|
||||
fullStream: mockFullStream(),
|
||||
usage: Promise.resolve({
|
||||
inputTokens: 100,
|
||||
outputTokens: 50,
|
||||
}),
|
||||
providerMetadata: Promise.resolve({
|
||||
openai: {
|
||||
cachedPromptTokens: 80,
|
||||
reasoningTokens: 20,
|
||||
},
|
||||
}),
|
||||
})
|
||||
|
||||
const stream = handler.createMessage(systemPrompt, messages)
|
||||
const chunks: any[] = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
const usageChunks = chunks.filter((chunk) => chunk.type === "usage")
|
||||
expect(usageChunks).toHaveLength(1)
|
||||
expect(usageChunks[0]).toEqual(
|
||||
expect.objectContaining({
|
||||
type: "usage",
|
||||
inputTokens: 100,
|
||||
outputTokens: 50,
|
||||
cacheReadTokens: 80,
|
||||
reasoningTokens: 20,
|
||||
}),
|
||||
)
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1,5 +1,8 @@
|
|||
import { Anthropic } from "@anthropic-ai/sdk"
|
||||
import OpenAI, { AzureOpenAI } from "openai"
|
||||
import { createOpenAI } from "@ai-sdk/openai"
|
||||
import { createOpenAICompatible } from "@ai-sdk/openai-compatible"
|
||||
import { createAzure } from "@ai-sdk/azure"
|
||||
import { streamText, generateText, ToolSet, LanguageModel } from "ai"
|
||||
import axios from "axios"
|
||||
|
||||
import {
|
||||
|
|
@ -7,31 +10,35 @@ import {
|
|||
azureOpenAiDefaultApiVersion,
|
||||
openAiModelInfoSaneDefaults,
|
||||
DEEP_SEEK_DEFAULT_TEMPERATURE,
|
||||
OPENAI_AZURE_AI_INFERENCE_PATH,
|
||||
} from "@roo-code/types"
|
||||
|
||||
import type { ApiHandlerOptions } from "../../shared/api"
|
||||
|
||||
import { TagMatcher } from "../../utils/tag-matcher"
|
||||
|
||||
import { convertToOpenAiMessages } from "../transform/openai-format"
|
||||
import { convertToR1Format } from "../transform/r1-format"
|
||||
import {
|
||||
convertToAiSdkMessages,
|
||||
convertToolsForAiSdk,
|
||||
processAiSdkStreamPart,
|
||||
mapToolChoice,
|
||||
handleAiSdkError,
|
||||
} from "../transform/ai-sdk"
|
||||
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
|
||||
import { getModelParams } from "../transform/model-params"
|
||||
|
||||
import { DEFAULT_HEADERS } from "./constants"
|
||||
import { BaseProvider } from "./base-provider"
|
||||
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
|
||||
import { getApiRequestTimeout } from "./utils/timeout-config"
|
||||
import { handleOpenAIError } from "./utils/openai-error-handler"
|
||||
|
||||
// TODO: Rename this to OpenAICompatibleHandler. Also, I think the
|
||||
// `OpenAINativeHandler` can subclass from this, since it's obviously
|
||||
// compatible with the OpenAI API. We can also rename it to `OpenAIHandler`.
|
||||
export class OpenAiHandler extends BaseProvider implements SingleCompletionHandler {
|
||||
protected options: ApiHandlerOptions
|
||||
protected client: OpenAI
|
||||
private readonly providerName = "OpenAI"
|
||||
private readonly isAzureAiInference: boolean
|
||||
private readonly isAzureOpenAi: boolean
|
||||
private readonly languageModelFactory: (modelId: string) => LanguageModel
|
||||
|
||||
constructor(options: ApiHandlerOptions) {
|
||||
super()
|
||||
|
|
@ -39,243 +46,258 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
|
|||
|
||||
const baseURL = this.options.openAiBaseUrl || "https://api.openai.com/v1"
|
||||
const apiKey = this.options.openAiApiKey ?? "not-provided"
|
||||
const isAzureAiInference = this._isAzureAiInference(this.options.openAiBaseUrl)
|
||||
const urlHost = this._getUrlHost(this.options.openAiBaseUrl)
|
||||
const isAzureOpenAi = urlHost === "azure.com" || urlHost.endsWith(".azure.com") || options.openAiUseAzure
|
||||
this.isAzureAiInference = this._isAzureAiInference(baseURL)
|
||||
const urlHost = this._getUrlHost(baseURL)
|
||||
this.isAzureOpenAi =
|
||||
!this.isAzureAiInference &&
|
||||
(urlHost === "azure.com" || urlHost.endsWith(".azure.com") || !!options.openAiUseAzure)
|
||||
|
||||
const headers = {
|
||||
...DEFAULT_HEADERS,
|
||||
...(this.options.openAiHeaders || {}),
|
||||
}
|
||||
|
||||
const timeout = getApiRequestTimeout()
|
||||
|
||||
if (isAzureAiInference) {
|
||||
// Azure AI Inference Service (e.g., for DeepSeek) uses a different path structure
|
||||
this.client = new OpenAI({
|
||||
baseURL,
|
||||
if (this.isAzureAiInference) {
|
||||
const provider = createOpenAICompatible({
|
||||
name: "OpenAI",
|
||||
baseURL: `${baseURL}/models`,
|
||||
apiKey,
|
||||
defaultHeaders: headers,
|
||||
defaultQuery: { "api-version": this.options.azureApiVersion || "2024-05-01-preview" },
|
||||
timeout,
|
||||
headers,
|
||||
queryParams: { "api-version": this.options.azureApiVersion || "2024-05-01-preview" },
|
||||
})
|
||||
} else if (isAzureOpenAi) {
|
||||
// Azure API shape slightly differs from the core API shape:
|
||||
// https://github.com/openai/openai-node?tab=readme-ov-file#microsoft-azure-openai
|
||||
this.client = new AzureOpenAI({
|
||||
baseURL,
|
||||
this.languageModelFactory = (modelId: string) => provider(modelId)
|
||||
} else if (this.isAzureOpenAi) {
|
||||
const azureBaseURL = baseURL.endsWith("/openai") ? baseURL : `${baseURL}/openai`
|
||||
const provider = createAzure({
|
||||
baseURL: azureBaseURL,
|
||||
apiKey,
|
||||
apiVersion: this.options.azureApiVersion || azureOpenAiDefaultApiVersion,
|
||||
defaultHeaders: headers,
|
||||
timeout,
|
||||
headers,
|
||||
useDeploymentBasedUrls: true,
|
||||
})
|
||||
this.languageModelFactory = (modelId: string) => provider.chat(modelId)
|
||||
} else {
|
||||
this.client = new OpenAI({
|
||||
const provider = createOpenAI({
|
||||
baseURL,
|
||||
apiKey,
|
||||
defaultHeaders: headers,
|
||||
timeout,
|
||||
headers,
|
||||
})
|
||||
this.languageModelFactory = (modelId: string) => provider.chat(modelId)
|
||||
}
|
||||
}
|
||||
|
||||
protected getLanguageModel(): LanguageModel {
|
||||
const { id } = this.getModel()
|
||||
return this.languageModelFactory(id)
|
||||
}
|
||||
|
||||
override async *createMessage(
|
||||
systemPrompt: string,
|
||||
messages: Anthropic.Messages.MessageParam[],
|
||||
metadata?: ApiHandlerCreateMessageMetadata,
|
||||
): ApiStream {
|
||||
const { info: modelInfo, reasoning } = this.getModel()
|
||||
const modelUrl = this.options.openAiBaseUrl ?? ""
|
||||
const { info: modelInfo, temperature, reasoning } = this.getModel()
|
||||
const modelId = this.options.openAiModelId ?? ""
|
||||
const enabledR1Format = this.options.openAiR1FormatEnabled ?? false
|
||||
const isAzureAiInference = this._isAzureAiInference(modelUrl)
|
||||
const deepseekReasoner = modelId.includes("deepseek-reasoner") || enabledR1Format
|
||||
const isO3Family = modelId.includes("o1") || modelId.includes("o3") || modelId.includes("o4")
|
||||
|
||||
if (modelId.includes("o1") || modelId.includes("o3") || modelId.includes("o4")) {
|
||||
yield* this.handleO3FamilyMessage(modelId, systemPrompt, messages, metadata)
|
||||
return
|
||||
const languageModel = this.getLanguageModel()
|
||||
|
||||
const aiSdkMessages = convertToAiSdkMessages(messages)
|
||||
|
||||
const openAiTools = this.convertToolsForOpenAI(metadata?.tools)
|
||||
const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined
|
||||
|
||||
let effectiveSystemPrompt: string | undefined = systemPrompt
|
||||
let effectiveTemperature: number | undefined =
|
||||
this.options.modelTemperature ?? (deepseekReasoner ? DEEP_SEEK_DEFAULT_TEMPERATURE : (temperature ?? 0))
|
||||
|
||||
const providerOptions: Record<string, any> = {}
|
||||
|
||||
if (isO3Family) {
|
||||
effectiveSystemPrompt = `Formatting re-enabled\n${systemPrompt}`
|
||||
effectiveTemperature = undefined
|
||||
|
||||
const openaiOpts: Record<string, unknown> = {
|
||||
systemMessageMode: "developer",
|
||||
parallelToolCalls: metadata?.parallelToolCalls ?? true,
|
||||
}
|
||||
|
||||
const effort = modelInfo.reasoningEffort as string | undefined
|
||||
if (effort) {
|
||||
openaiOpts.reasoningEffort = effort
|
||||
}
|
||||
|
||||
providerOptions.openai = openaiOpts
|
||||
} else if (reasoning?.reasoning_effort) {
|
||||
providerOptions.openai = {
|
||||
reasoningEffort: reasoning.reasoning_effort,
|
||||
parallelToolCalls: metadata?.parallelToolCalls ?? true,
|
||||
}
|
||||
}
|
||||
|
||||
let systemMessage: OpenAI.Chat.ChatCompletionSystemMessageParam = {
|
||||
role: "system",
|
||||
content: systemPrompt,
|
||||
if (deepseekReasoner) {
|
||||
effectiveSystemPrompt = undefined
|
||||
aiSdkMessages.unshift({ role: "user", content: systemPrompt })
|
||||
}
|
||||
|
||||
if (this.options.openAiStreamingEnabled ?? true) {
|
||||
let convertedMessages
|
||||
|
||||
if (deepseekReasoner) {
|
||||
convertedMessages = convertToR1Format([{ role: "user", content: systemPrompt }, ...messages])
|
||||
} else {
|
||||
if (modelInfo.supportsPromptCache) {
|
||||
systemMessage = {
|
||||
role: "system",
|
||||
content: [
|
||||
{
|
||||
type: "text",
|
||||
text: systemPrompt,
|
||||
// @ts-ignore-next-line
|
||||
cache_control: { type: "ephemeral" },
|
||||
},
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
convertedMessages = [systemMessage, ...convertToOpenAiMessages(messages)]
|
||||
|
||||
if (modelInfo.supportsPromptCache) {
|
||||
// Note: the following logic is copied from openrouter:
|
||||
// Add cache_control to the last two user messages
|
||||
// (note: this works because we only ever add one user message at a time, but if we added multiple we'd need to mark the user message before the last assistant message)
|
||||
const lastTwoUserMessages = convertedMessages.filter((msg) => msg.role === "user").slice(-2)
|
||||
|
||||
lastTwoUserMessages.forEach((msg) => {
|
||||
if (typeof msg.content === "string") {
|
||||
msg.content = [{ type: "text", text: msg.content }]
|
||||
}
|
||||
|
||||
if (Array.isArray(msg.content)) {
|
||||
// NOTE: this is fine since env details will always be added at the end. but if it weren't there, and the user added a image_url type message, it would pop a text part before it and then move it after to the end.
|
||||
let lastTextPart = msg.content.filter((part) => part.type === "text").pop()
|
||||
|
||||
if (!lastTextPart) {
|
||||
lastTextPart = { type: "text", text: "..." }
|
||||
msg.content.push(lastTextPart)
|
||||
}
|
||||
|
||||
// @ts-ignore-next-line
|
||||
lastTextPart["cache_control"] = { type: "ephemeral" }
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
const isGrokXAI = this._isGrokXAI(this.options.openAiBaseUrl)
|
||||
|
||||
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = {
|
||||
model: modelId,
|
||||
temperature: this.options.modelTemperature ?? (deepseekReasoner ? DEEP_SEEK_DEFAULT_TEMPERATURE : 0),
|
||||
messages: convertedMessages,
|
||||
stream: true as const,
|
||||
...(isGrokXAI ? {} : { stream_options: { include_usage: true } }),
|
||||
...(reasoning && reasoning),
|
||||
tools: this.convertToolsForOpenAI(metadata?.tools),
|
||||
tool_choice: metadata?.tool_choice,
|
||||
parallel_tool_calls: metadata?.parallelToolCalls ?? true,
|
||||
}
|
||||
|
||||
// Add max_tokens if needed
|
||||
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
|
||||
let stream
|
||||
try {
|
||||
stream = await this.client.chat.completions.create(
|
||||
requestOptions,
|
||||
isAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {},
|
||||
)
|
||||
} catch (error) {
|
||||
throw handleOpenAIError(error, this.providerName)
|
||||
}
|
||||
|
||||
const matcher = new TagMatcher(
|
||||
"think",
|
||||
(chunk) =>
|
||||
({
|
||||
type: chunk.matched ? "reasoning" : "text",
|
||||
text: chunk.data,
|
||||
}) as const,
|
||||
yield* this.handleStreaming(
|
||||
languageModel,
|
||||
effectiveSystemPrompt,
|
||||
aiSdkMessages,
|
||||
effectiveTemperature,
|
||||
aiSdkTools,
|
||||
metadata,
|
||||
providerOptions,
|
||||
modelInfo,
|
||||
)
|
||||
} else {
|
||||
yield* this.handleNonStreaming(
|
||||
languageModel,
|
||||
effectiveSystemPrompt,
|
||||
aiSdkMessages,
|
||||
effectiveTemperature,
|
||||
aiSdkTools,
|
||||
metadata,
|
||||
providerOptions,
|
||||
modelInfo,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
let lastUsage
|
||||
const activeToolCallIds = new Set<string>()
|
||||
private async *handleStreaming(
|
||||
languageModel: LanguageModel,
|
||||
systemPrompt: string | undefined,
|
||||
messages: ReturnType<typeof convertToAiSdkMessages>,
|
||||
temperature: number | undefined,
|
||||
tools: ToolSet | undefined,
|
||||
metadata: ApiHandlerCreateMessageMetadata | undefined,
|
||||
providerOptions: Record<string, any>,
|
||||
modelInfo: ModelInfo,
|
||||
): ApiStream {
|
||||
const result = streamText({
|
||||
model: languageModel,
|
||||
system: systemPrompt,
|
||||
messages,
|
||||
temperature,
|
||||
maxOutputTokens: this.getMaxOutputTokens(),
|
||||
tools,
|
||||
toolChoice: mapToolChoice(metadata?.tool_choice),
|
||||
providerOptions: Object.keys(providerOptions).length > 0 ? providerOptions : undefined,
|
||||
})
|
||||
|
||||
for await (const chunk of stream) {
|
||||
const delta = chunk.choices?.[0]?.delta ?? {}
|
||||
const finishReason = chunk.choices?.[0]?.finish_reason
|
||||
const matcher = new TagMatcher(
|
||||
"think",
|
||||
(chunk) =>
|
||||
({
|
||||
type: chunk.matched ? "reasoning" : "text",
|
||||
text: chunk.data,
|
||||
}) as const,
|
||||
)
|
||||
|
||||
if (delta.content) {
|
||||
for (const chunk of matcher.update(delta.content)) {
|
||||
try {
|
||||
for await (const part of result.fullStream) {
|
||||
for (const chunk of processAiSdkStreamPart(part)) {
|
||||
if (chunk.type === "text") {
|
||||
for (const matchedChunk of matcher.update(chunk.text)) {
|
||||
yield matchedChunk
|
||||
}
|
||||
} else {
|
||||
yield chunk
|
||||
}
|
||||
}
|
||||
|
||||
if ("reasoning_content" in delta && delta.reasoning_content) {
|
||||
yield {
|
||||
type: "reasoning",
|
||||
text: (delta.reasoning_content as string | undefined) || "",
|
||||
}
|
||||
}
|
||||
|
||||
yield* this.processToolCalls(delta, finishReason, activeToolCallIds)
|
||||
|
||||
if (chunk.usage) {
|
||||
lastUsage = chunk.usage
|
||||
}
|
||||
}
|
||||
|
||||
for (const chunk of matcher.final()) {
|
||||
yield chunk
|
||||
}
|
||||
|
||||
if (lastUsage) {
|
||||
yield this.processUsageMetrics(lastUsage, modelInfo)
|
||||
}
|
||||
} else {
|
||||
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = {
|
||||
model: modelId,
|
||||
messages: deepseekReasoner
|
||||
? convertToR1Format([{ role: "user", content: systemPrompt }, ...messages])
|
||||
: [systemMessage, ...convertToOpenAiMessages(messages)],
|
||||
// Tools are always present (minimum ALWAYS_AVAILABLE_TOOLS)
|
||||
tools: this.convertToolsForOpenAI(metadata?.tools),
|
||||
tool_choice: metadata?.tool_choice,
|
||||
parallel_tool_calls: metadata?.parallelToolCalls ?? true,
|
||||
const usage = await result.usage
|
||||
const providerMetadata = await result.providerMetadata
|
||||
if (usage) {
|
||||
yield this.processUsageMetrics(usage, modelInfo, providerMetadata as any)
|
||||
}
|
||||
} catch (error) {
|
||||
throw handleAiSdkError(error, this.providerName)
|
||||
}
|
||||
}
|
||||
|
||||
// Add max_tokens if needed
|
||||
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
private async *handleNonStreaming(
|
||||
languageModel: LanguageModel,
|
||||
systemPrompt: string | undefined,
|
||||
messages: ReturnType<typeof convertToAiSdkMessages>,
|
||||
temperature: number | undefined,
|
||||
tools: ToolSet | undefined,
|
||||
metadata: ApiHandlerCreateMessageMetadata | undefined,
|
||||
providerOptions: Record<string, any>,
|
||||
modelInfo: ModelInfo,
|
||||
): ApiStream {
|
||||
try {
|
||||
const { text, toolCalls, usage, providerMetadata } = await generateText({
|
||||
model: languageModel,
|
||||
system: systemPrompt,
|
||||
messages,
|
||||
temperature,
|
||||
maxOutputTokens: this.getMaxOutputTokens(),
|
||||
tools,
|
||||
toolChoice: mapToolChoice(metadata?.tool_choice),
|
||||
providerOptions: Object.keys(providerOptions).length > 0 ? providerOptions : undefined,
|
||||
})
|
||||
|
||||
let response
|
||||
try {
|
||||
response = await this.client.chat.completions.create(
|
||||
requestOptions,
|
||||
this._isAzureAiInference(modelUrl) ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {},
|
||||
)
|
||||
} catch (error) {
|
||||
throw handleOpenAIError(error, this.providerName)
|
||||
}
|
||||
|
||||
const message = response.choices?.[0]?.message
|
||||
|
||||
if (message?.tool_calls) {
|
||||
for (const toolCall of message.tool_calls) {
|
||||
if (toolCall.type === "function") {
|
||||
yield {
|
||||
type: "tool_call",
|
||||
id: toolCall.id,
|
||||
name: toolCall.function.name,
|
||||
arguments: toolCall.function.arguments,
|
||||
}
|
||||
if (toolCalls && toolCalls.length > 0) {
|
||||
for (const toolCall of toolCalls) {
|
||||
yield {
|
||||
type: "tool_call",
|
||||
id: toolCall.toolCallId,
|
||||
name: toolCall.toolName,
|
||||
arguments: JSON.stringify((toolCall as any).args),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
yield {
|
||||
type: "text",
|
||||
text: message?.content || "",
|
||||
text: text || "",
|
||||
}
|
||||
|
||||
yield this.processUsageMetrics(response.usage, modelInfo)
|
||||
if (usage) {
|
||||
yield this.processUsageMetrics(usage, modelInfo, providerMetadata as any)
|
||||
}
|
||||
} catch (error) {
|
||||
throw handleAiSdkError(error, this.providerName)
|
||||
}
|
||||
}
|
||||
|
||||
protected processUsageMetrics(usage: any, _modelInfo?: ModelInfo): ApiStreamUsageChunk {
|
||||
protected processUsageMetrics(
|
||||
usage: {
|
||||
inputTokens?: number
|
||||
outputTokens?: number
|
||||
details?: {
|
||||
cachedInputTokens?: number
|
||||
reasoningTokens?: number
|
||||
}
|
||||
},
|
||||
_modelInfo?: ModelInfo,
|
||||
providerMetadata?: {
|
||||
openai?: {
|
||||
cachedPromptTokens?: number
|
||||
reasoningTokens?: number
|
||||
}
|
||||
},
|
||||
): ApiStreamUsageChunk {
|
||||
// Extract cache and reasoning metrics from OpenAI's providerMetadata when available,
|
||||
// falling back to usage.details for standard AI SDK fields.
|
||||
const cacheReadTokens = providerMetadata?.openai?.cachedPromptTokens ?? usage.details?.cachedInputTokens
|
||||
const reasoningTokens = providerMetadata?.openai?.reasoningTokens ?? usage.details?.reasoningTokens
|
||||
|
||||
return {
|
||||
type: "usage",
|
||||
inputTokens: usage?.prompt_tokens || 0,
|
||||
outputTokens: usage?.completion_tokens || 0,
|
||||
cacheWriteTokens: usage?.cache_creation_input_tokens || undefined,
|
||||
cacheReadTokens: usage?.cache_read_input_tokens || undefined,
|
||||
inputTokens: usage.inputTokens || 0,
|
||||
outputTokens: usage.outputTokens || 0,
|
||||
cacheReadTokens,
|
||||
reasoningTokens,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -292,208 +314,37 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
|
|||
return { id, info, ...params }
|
||||
}
|
||||
|
||||
protected getMaxOutputTokens(): number | undefined {
|
||||
if (this.options.includeMaxTokens !== true) {
|
||||
return undefined
|
||||
}
|
||||
const { info } = this.getModel()
|
||||
return this.options.modelMaxTokens || info.maxTokens || undefined
|
||||
}
|
||||
|
||||
async completePrompt(prompt: string): Promise<string> {
|
||||
try {
|
||||
const isAzureAiInference = this._isAzureAiInference(this.options.openAiBaseUrl)
|
||||
const model = this.getModel()
|
||||
const modelInfo = model.info
|
||||
const { temperature } = this.getModel()
|
||||
const languageModel = this.getLanguageModel()
|
||||
|
||||
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = {
|
||||
model: model.id,
|
||||
messages: [{ role: "user", content: prompt }],
|
||||
}
|
||||
const { text } = await generateText({
|
||||
model: languageModel,
|
||||
prompt,
|
||||
maxOutputTokens: this.getMaxOutputTokens(),
|
||||
temperature: this.options.modelTemperature ?? temperature ?? 0,
|
||||
})
|
||||
|
||||
// Add max_tokens if needed
|
||||
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
|
||||
let response
|
||||
try {
|
||||
response = await this.client.chat.completions.create(
|
||||
requestOptions,
|
||||
isAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {},
|
||||
)
|
||||
} catch (error) {
|
||||
throw handleOpenAIError(error, this.providerName)
|
||||
}
|
||||
|
||||
return response.choices?.[0]?.message.content || ""
|
||||
return text
|
||||
} catch (error) {
|
||||
if (error instanceof Error) {
|
||||
throw new Error(`${this.providerName} completion error: ${error.message}`)
|
||||
}
|
||||
|
||||
throw error
|
||||
}
|
||||
}
|
||||
|
||||
private async *handleO3FamilyMessage(
|
||||
modelId: string,
|
||||
systemPrompt: string,
|
||||
messages: Anthropic.Messages.MessageParam[],
|
||||
metadata?: ApiHandlerCreateMessageMetadata,
|
||||
): ApiStream {
|
||||
const modelInfo = this.getModel().info
|
||||
const methodIsAzureAiInference = this._isAzureAiInference(this.options.openAiBaseUrl)
|
||||
|
||||
if (this.options.openAiStreamingEnabled ?? true) {
|
||||
const isGrokXAI = this._isGrokXAI(this.options.openAiBaseUrl)
|
||||
|
||||
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = {
|
||||
model: modelId,
|
||||
messages: [
|
||||
{
|
||||
role: "developer",
|
||||
content: `Formatting re-enabled\n${systemPrompt}`,
|
||||
},
|
||||
...convertToOpenAiMessages(messages),
|
||||
],
|
||||
stream: true,
|
||||
...(isGrokXAI ? {} : { stream_options: { include_usage: true } }),
|
||||
reasoning_effort: modelInfo.reasoningEffort as "low" | "medium" | "high" | undefined,
|
||||
temperature: undefined,
|
||||
// Tools are always present (minimum ALWAYS_AVAILABLE_TOOLS)
|
||||
tools: this.convertToolsForOpenAI(metadata?.tools),
|
||||
tool_choice: metadata?.tool_choice,
|
||||
parallel_tool_calls: metadata?.parallelToolCalls ?? true,
|
||||
}
|
||||
|
||||
// O3 family models do not support the deprecated max_tokens parameter
|
||||
// but they do support max_completion_tokens (the modern OpenAI parameter)
|
||||
// This allows O3 models to limit response length when includeMaxTokens is enabled
|
||||
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
|
||||
let stream
|
||||
try {
|
||||
stream = await this.client.chat.completions.create(
|
||||
requestOptions,
|
||||
methodIsAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {},
|
||||
)
|
||||
} catch (error) {
|
||||
throw handleOpenAIError(error, this.providerName)
|
||||
}
|
||||
|
||||
yield* this.handleStreamResponse(stream)
|
||||
} else {
|
||||
const requestOptions: OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming = {
|
||||
model: modelId,
|
||||
messages: [
|
||||
{
|
||||
role: "developer",
|
||||
content: `Formatting re-enabled\n${systemPrompt}`,
|
||||
},
|
||||
...convertToOpenAiMessages(messages),
|
||||
],
|
||||
reasoning_effort: modelInfo.reasoningEffort as "low" | "medium" | "high" | undefined,
|
||||
temperature: undefined,
|
||||
// Tools are always present (minimum ALWAYS_AVAILABLE_TOOLS)
|
||||
tools: this.convertToolsForOpenAI(metadata?.tools),
|
||||
tool_choice: metadata?.tool_choice,
|
||||
parallel_tool_calls: metadata?.parallelToolCalls ?? true,
|
||||
}
|
||||
|
||||
// O3 family models do not support the deprecated max_tokens parameter
|
||||
// but they do support max_completion_tokens (the modern OpenAI parameter)
|
||||
// This allows O3 models to limit response length when includeMaxTokens is enabled
|
||||
this.addMaxTokensIfNeeded(requestOptions, modelInfo)
|
||||
|
||||
let response
|
||||
try {
|
||||
response = await this.client.chat.completions.create(
|
||||
requestOptions,
|
||||
methodIsAzureAiInference ? { path: OPENAI_AZURE_AI_INFERENCE_PATH } : {},
|
||||
)
|
||||
} catch (error) {
|
||||
throw handleOpenAIError(error, this.providerName)
|
||||
}
|
||||
|
||||
const message = response.choices?.[0]?.message
|
||||
if (message?.tool_calls) {
|
||||
for (const toolCall of message.tool_calls) {
|
||||
if (toolCall.type === "function") {
|
||||
yield {
|
||||
type: "tool_call",
|
||||
id: toolCall.id,
|
||||
name: toolCall.function.name,
|
||||
arguments: toolCall.function.arguments,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
yield {
|
||||
type: "text",
|
||||
text: message?.content || "",
|
||||
}
|
||||
yield this.processUsageMetrics(response.usage)
|
||||
}
|
||||
}
|
||||
|
||||
private async *handleStreamResponse(stream: AsyncIterable<OpenAI.Chat.Completions.ChatCompletionChunk>): ApiStream {
|
||||
const activeToolCallIds = new Set<string>()
|
||||
|
||||
for await (const chunk of stream) {
|
||||
const delta = chunk.choices?.[0]?.delta
|
||||
const finishReason = chunk.choices?.[0]?.finish_reason
|
||||
|
||||
if (delta) {
|
||||
if (delta.content) {
|
||||
yield {
|
||||
type: "text",
|
||||
text: delta.content,
|
||||
}
|
||||
}
|
||||
|
||||
yield* this.processToolCalls(delta, finishReason, activeToolCallIds)
|
||||
}
|
||||
|
||||
if (chunk.usage) {
|
||||
yield {
|
||||
type: "usage",
|
||||
inputTokens: chunk.usage.prompt_tokens || 0,
|
||||
outputTokens: chunk.usage.completion_tokens || 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Helper generator to process tool calls from a stream chunk.
|
||||
* Tracks active tool call IDs and yields tool_call_partial and tool_call_end events.
|
||||
* @param delta - The delta object from the stream chunk
|
||||
* @param finishReason - The finish_reason from the stream chunk
|
||||
* @param activeToolCallIds - Set to track active tool call IDs (mutated in place)
|
||||
*/
|
||||
private *processToolCalls(
|
||||
delta: OpenAI.Chat.Completions.ChatCompletionChunk.Choice.Delta | undefined,
|
||||
finishReason: string | null | undefined,
|
||||
activeToolCallIds: Set<string>,
|
||||
): Generator<
|
||||
| { type: "tool_call_partial"; index: number; id?: string; name?: string; arguments?: string }
|
||||
| { type: "tool_call_end"; id: string }
|
||||
> {
|
||||
if (delta?.tool_calls) {
|
||||
for (const toolCall of delta.tool_calls) {
|
||||
if (toolCall.id) {
|
||||
activeToolCallIds.add(toolCall.id)
|
||||
}
|
||||
yield {
|
||||
type: "tool_call_partial",
|
||||
index: toolCall.index,
|
||||
id: toolCall.id,
|
||||
name: toolCall.function?.name,
|
||||
arguments: toolCall.function?.arguments,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Emit tool_call_end events when finish_reason is "tool_calls"
|
||||
// This ensures tool calls are finalized even if the stream doesn't properly close
|
||||
if (finishReason === "tool_calls" && activeToolCallIds.size > 0) {
|
||||
for (const id of activeToolCallIds) {
|
||||
yield { type: "tool_call_end", id }
|
||||
}
|
||||
activeToolCallIds.clear()
|
||||
}
|
||||
override isAiSdkProvider(): boolean {
|
||||
return true
|
||||
}
|
||||
|
||||
protected _getUrlHost(baseUrl?: string): string {
|
||||
|
|
@ -504,34 +355,10 @@ export class OpenAiHandler extends BaseProvider implements SingleCompletionHandl
|
|||
}
|
||||
}
|
||||
|
||||
private _isGrokXAI(baseUrl?: string): boolean {
|
||||
const urlHost = this._getUrlHost(baseUrl)
|
||||
return urlHost.includes("x.ai")
|
||||
}
|
||||
|
||||
protected _isAzureAiInference(baseUrl?: string): boolean {
|
||||
const urlHost = this._getUrlHost(baseUrl)
|
||||
return urlHost.endsWith(".services.ai.azure.com")
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds max_completion_tokens to the request body if needed based on provider configuration
|
||||
* Note: max_tokens is deprecated in favor of max_completion_tokens as per OpenAI documentation
|
||||
* O3 family models handle max_tokens separately in handleO3FamilyMessage
|
||||
*/
|
||||
protected addMaxTokensIfNeeded(
|
||||
requestOptions:
|
||||
| OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming
|
||||
| OpenAI.Chat.Completions.ChatCompletionCreateParamsNonStreaming,
|
||||
modelInfo: ModelInfo,
|
||||
): void {
|
||||
// Only add max_completion_tokens if includeMaxTokens is true
|
||||
if (this.options.includeMaxTokens === true) {
|
||||
// Use user-configured modelMaxTokens if available, otherwise fall back to model's default maxTokens
|
||||
// Using max_completion_tokens as max_tokens is deprecated
|
||||
requestOptions.max_completion_tokens = this.options.modelMaxTokens || modelInfo.maxTokens
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
export async function getOpenAiModels(baseUrl?: string, apiKey?: string, openAiHeaders?: Record<string, string>) {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue