diff --git a/src/api/providers/__tests__/openrouter.test.ts b/src/api/providers/__tests__/openrouter.test.ts index 981c9ad096f..996644b07f7 100644 --- a/src/api/providers/__tests__/openrouter.test.ts +++ b/src/api/providers/__tests__/openrouter.test.ts @@ -112,6 +112,16 @@ describe("OpenRouterHandler", () => { }, ], } + // Add usage information in the stream response + yield { + id: "test-id", + choices: [{ delta: {} }], + usage: { + prompt_tokens: 10, + completion_tokens: 20, + cost: 0.001, + }, + } }, } @@ -121,17 +131,6 @@ describe("OpenRouterHandler", () => { completions: { create: mockCreate }, } as any - // Mock axios.get for generation details - ;(axios.get as jest.Mock).mockResolvedValue({ - data: { - data: { - native_tokens_prompt: 10, - native_tokens_completion: 20, - total_cost: 0.001, - }, - }, - }) - const systemPrompt = "test system prompt" const messages: Anthropic.Messages.MessageParam[] = [{ role: "user" as const, content: "test message" }] @@ -153,7 +152,6 @@ describe("OpenRouterHandler", () => { inputTokens: 10, outputTokens: 20, totalCost: 0.001, - fullResponseText: "test response", }) // Verify OpenAI client was called with correct parameters diff --git a/src/api/providers/openrouter.ts b/src/api/providers/openrouter.ts index 3f215bfc7c3..6d93d0b180b 100644 --- a/src/api/providers/openrouter.ts +++ b/src/api/providers/openrouter.ts @@ -24,11 +24,6 @@ type OpenRouterChatCompletionParams = OpenAI.Chat.ChatCompletionCreateParams & { thinking?: BetaThinkingConfigParam } -// Add custom interface for OpenRouter usage chunk. -interface OpenRouterApiStreamUsageChunk extends ApiStreamUsageChunk { - fullResponseText: string -} - export class OpenRouterHandler extends BaseProvider implements SingleCompletionHandler { protected options: ApiHandlerOptions private client: OpenAI @@ -110,7 +105,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH top_p: topP, messages: openAiMessages, stream: true, - include_reasoning: true, + stream_options: { include_usage: true }, // Only include provider if openRouterSpecificProvider is not "[default]". ...(this.options.openRouterSpecificProvider && this.options.openRouterSpecificProvider !== OPENROUTER_DEFAULT_PROVIDER_NAME && { @@ -122,7 +117,7 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH const stream = await this.client.chat.completions.create(completionParams) - let genId: string | undefined + let lastUsage for await (const chunk of stream as unknown as AsyncIterable) { // OpenRouter returns an error object instead of the OpenAI SDK throwing an error. @@ -132,10 +127,6 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH throw new Error(`OpenRouter API Error ${error?.code}: ${error?.message}`) } - if (!genId && chunk.id) { - genId = chunk.id - } - const delta = chunk.choices[0]?.delta if ("reasoning" in delta && delta.reasoning) { @@ -146,47 +137,23 @@ export class OpenRouterHandler extends BaseProvider implements SingleCompletionH fullResponseText += delta.content yield { type: "text", text: delta.content } as ApiStreamChunk } - } - - const endpoint = `${this.client.baseURL}/generation?id=${genId}` - const config: AxiosRequestConfig = { - headers: { Authorization: `Bearer ${this.options.openRouterApiKey}` }, - timeout: 3_000, + if (chunk.usage) { + lastUsage = chunk.usage + } } - let attempt = 0 - let lastError: Error | undefined - const startTime = Date.now() - - while (attempt++ < 10) { - await delay(attempt * 100) // Give OpenRouter some time to produce the generation metadata. - - try { - const response = await axios.get(endpoint, config) - const generation = response.data?.data - - yield { - type: "usage", - inputTokens: generation?.native_tokens_prompt || 0, - outputTokens: generation?.native_tokens_completion || 0, - totalCost: generation?.total_cost || 0, - fullResponseText, - } as OpenRouterApiStreamUsageChunk - - break - } catch (error: unknown) { - if (error instanceof Error) { - lastError = error - } - } + if (lastUsage) { + yield this.processUsageMetrics(lastUsage) } + } - if (lastError) { - console.error( - `Failed to fetch OpenRouter generation details after attempt #${attempt} (${Date.now() - startTime}ms) [${genId}]`, - lastError, - ) + processUsageMetrics(usage: any): ApiStreamUsageChunk { + return { + type: "usage", + inputTokens: usage?.prompt_tokens || 0, + outputTokens: usage?.completion_tokens || 0, + totalCost: usage?.cost || 0, } }