Skip to content

Use openrouter stream_options include_usage #1905

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
mrubens merged 1 commit into from
Mar 23, 2025
Merged

Use openrouter stream_options include_usage #1905

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 10 additions & 12 deletions src/api/providers/__tests__/openrouter.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
},
}
},
}

Expand All @@ -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" }]

Expand All @@ -153,7 +152,6 @@ describe("OpenRouterHandler", () => {
inputTokens: 10,
outputTokens: 20,
totalCost: 0.001,
fullResponseText: "test response",
})

// Verify OpenAI client was called with correct parameters
Expand Down
61 changes: 14 additions & 47 deletions src/api/providers/openrouter.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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 && {
Expand All @@ -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<OpenAI.Chat.Completions.ChatCompletionChunk>) {
// OpenRouter returns an error object instead of the OpenAI SDK throwing an error.
Expand All @@ -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) {
Expand All @@ -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 {
Copy link
Contributor

@ellipsis-dev ellipsis-dev bot Mar 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Consider replacing the any type in processUsageMetrics with a properly defined interface for the usage object. This helps ensure type safety and clarity for the expected structure.

return {
type: "usage",
inputTokens: usage?.prompt_tokens || 0,
outputTokens: usage?.completion_tokens || 0,
totalCost: usage?.cost || 0,
}
}

Expand Down