mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
Use openrouter stream_options include_usage (#1905)
This commit is contained in:
parent
92dee31634
commit
2597347e91
2 changed files with 25 additions and 60 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue