Use openrouter stream_options include_usage (#1905)

This commit is contained in:
Matt Rubens 2025-03-23 01:52:43 -04:00 committed by GitHub
parent 92dee31634
commit 2597347e91
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 25 additions and 60 deletions

View file

@ -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

View file

@ -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<OpenAI.Chat.Completions.ChatCompletionChunk>) {
// 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,
}
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 (chunk.usage) {
lastUsage = chunk.usage
}
}
if (lastError) {
console.error(
`Failed to fetch OpenRouter generation details after attempt #${attempt} (${Date.now() - startTime}ms) [${genId}]`,
lastError,
)
if (lastUsage) {
yield this.processUsageMetrics(lastUsage)
}
}
processUsageMetrics(usage: any): ApiStreamUsageChunk {
return {
type: "usage",
inputTokens: usage?.prompt_tokens || 0,
outputTokens: usage?.completion_tokens || 0,
totalCost: usage?.cost || 0,
}
}