Skip to content

Commit ea8b886

Browse files
charlesliencharleslien
committed
add cost to sdk response
1 parent f558f58 commit ea8b886

File tree

2 files changed

+131
-102
lines changed

2 files changed

+131
-102
lines changed

‎scripts/fat-sdk-openrouter-example.ts‎

Lines changed: 49 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,53 @@
1-
import { createOpenAICompatible } from '@ai-sdk/openai-compatible'
1+
import path from 'path'
2+
3+
import {
4+
OpenAICompatibleChatLanguageModel,
5+
VERSION,
6+
} from '@ai-sdk/openai-compatible'
27
import { websiteUrl } from '@codebuff/npm-app/config'
3-
import { generateObject } from 'ai'
4-
import z from 'zod/v4'
8+
import { streamText } from 'ai'
59

6-
const codebuffBackendProvider = createOpenAICompatible({
7-
name: 'codebuff',
8-
apiKey: '12345',
9-
baseURL: websiteUrl + '/api/v1',
10-
// apiKey: process.env.OPEN_ROUTER_API_KEY,
11-
// baseURL: 'https://openrouter.ai/api/v1',
12-
supportsStructuredOutputs: true,
13-
})
10+
const apiKey = '12345'
11+
12+
const codebuffBackendModel = new OpenAICompatibleChatLanguageModel(
13+
'openai/gpt-5',
14+
{
15+
provider: 'codebuff.chat',
16+
url: ({ path: endpoint }) =>
17+
new URL(path.join('/api/v1', endpoint), websiteUrl).toString(),
18+
headers: () => ({
19+
Authorization: `Bearer ${apiKey}`,
20+
'user-agent': `ai-sdk/openai-compatible/${VERSION}`,
21+
}),
22+
metadataExtractor: {
23+
extractMetadata: async (...inputs) => {
24+
console.log(inputs, 'extractMetadata')
25+
return undefined
26+
},
27+
createStreamExtractor: () => ({
28+
processChunk: (...inputs) => {
29+
console.log(
30+
JSON.stringify(inputs, null, 2),
31+
'createStreamExtractor.processChunk',
32+
)
33+
},
34+
buildMetadata: (...inputs) => {
35+
console.log(inputs, 'createStreamExtractor.buildMetadata')
36+
return undefined
37+
},
38+
}),
39+
},
40+
fetch: undefined,
41+
includeUsage: undefined,
42+
supportsStructuredOutputs: true,
43+
},
44+
)
1445

15-
const response = await generateObject({
16-
schema: z.object({ greeting: z.string() }),
46+
const response = streamText({
1747
// const response = await streamText({
1848
// const response = await generateText({
19-
model: codebuffBackendProvider('openai/gpt-5'),
49+
// model: codebuffBackendProvider('openai/gpt-5'),
50+
model: codebuffBackendModel,
2051
messages: [
2152
{
2253
role: 'user',
@@ -44,15 +75,13 @@ const response = await generateObject({
4475
// all these get directly added to the body at the top level
4576
reasoningEffort: 'low',
4677
codebuff_metadata: {
47-
agent_run_id: '19b636d9-bfbf-40ff-b3e9-92dc86f4a8d0',
78+
run_id: '19b636d9-bfbf-40ff-b3e9-92dc86f4a8d0',
4879
client_id: 'test-client-id-123',
49-
client_request_id: 'test-client-session-id-456',
5080
},
5181
},
5282
},
5383
})
5484

55-
console.dir({ response }, { depth: null })
56-
// for await (const chunk of response.fullStream) {
57-
// console.dir({ chunk }, { depth: null })
58-
// }
85+
for await (const chunk of response.fullStream) {
86+
console.dir({ chunk }, { depth: null })
87+
}

‎sdk/src/impl/llm.ts‎

Lines changed: 82 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1-
import { createOpenAICompatible } from '@ai-sdk/openai-compatible'
1+
import path from 'path'
2+
3+
import { OpenAICompatibleChatLanguageModel } from '@ai-sdk/openai-compatible'
24
import { streamText, APICallError, generateText, generateObject } from 'ai'
35

46
import { PROFIT_MARGIN } from '../../../common/src/old-constants'
@@ -20,24 +22,74 @@ import type {
2022
} from '../../../common/src/types/contracts/llm'
2123
import type { ParamsOf } from '../../../common/src/types/function-params'
2224
import type { LanguageModelV2 } from '@ai-sdk/provider'
23-
import type {
24-
OpenRouterProviderOptions,
25-
OpenRouterUsageAccounting,
26-
} from '@openrouter/ai-sdk-provider'
25+
import type { Logger } from '@codebuff/common/types/contracts/logger'
26+
import type { OpenRouterProviderOptions } from '@openrouter/ai-sdk-provider'
2727
import type z from 'zod/v4'
2828

29+
// Forked from https://github.com/OpenRouterTeam/ai-sdk-provider/
30+
type OpenRouterUsageAccounting = {
31+
cost: number | null
32+
costDetails: {
33+
upstreamInferenceCost: number | null
34+
}
35+
}
36+
37+
function calculateUsedCredits(params: { costDollars: number }): number {
38+
const { costDollars } = params
39+
40+
return Math.round(costDollars * (1 + PROFIT_MARGIN) * 100)
41+
}
42+
2943
function getAiSdkModel(params: {
3044
apiKey: string
3145
model: string
46+
logger: Logger
3247
}): LanguageModelV2 {
33-
const { apiKey, model } = params
48+
const { apiKey, model, logger } = params
49+
50+
const openrouterUsage: OpenRouterUsageAccounting = {
51+
cost: null,
52+
costDetails: {
53+
upstreamInferenceCost: null,
54+
},
55+
}
3456

35-
return createOpenAICompatible({
36-
name: 'codebuff',
37-
apiKey,
38-
baseURL: WEBSITE_URL + '/api/v1',
57+
const codebuffBackendModel = new OpenAICompatibleChatLanguageModel(model, {
58+
provider: 'codebuff.chat',
59+
url: ({ path: endpoint }) =>
60+
new URL(path.join('/api/v1', endpoint), WEBSITE_URL).toString(),
61+
headers: () => ({
62+
Authorization: `Bearer ${apiKey}`,
63+
'user-agent': `ai-sdk/codebuff/${process.env.NEXT_PUBLIC_NPM_APP_VERSION || 'unknown-version'}`,
64+
}),
65+
metadataExtractor: {
66+
extractMetadata: async (...inputs) => {
67+
console.log(inputs, 'extractMetadata')
68+
return undefined
69+
},
70+
createStreamExtractor: () => ({
71+
processChunk: (parsedChunk: any) => {
72+
if (typeof parsedChunk?.usage?.cost === 'number') {
73+
openrouterUsage.cost = parsedChunk.usage.cost
74+
}
75+
if (
76+
typeof parsedChunk?.usage?.cost_details?.upstream_inference_cost ===
77+
'number'
78+
) {
79+
openrouterUsage.costDetails.upstreamInferenceCost =
80+
parsedChunk.usage.cost_details.upstream_inference_cost
81+
}
82+
},
83+
buildMetadata: () => {
84+
return { codebuff: { usage: openrouterUsage } }
85+
},
86+
}),
87+
},
88+
fetch: undefined,
89+
includeUsage: undefined,
3990
supportsStructuredOutputs: true,
40-
})(model)
91+
})
92+
return codebuffBackendModel
4193
}
4294

4395
export async function* promptAiSdkStream(
@@ -172,28 +224,12 @@ export async function* promptAiSdkStream(
172224
}
173225

174226
const providerMetadata = (await response.providerMetadata) ?? {}
175-
const usage = await response.usage
176-
let inputTokens = usage.inputTokens || 0
177-
let cacheReadInputTokens: number = 0
178-
let cacheCreationInputTokens: number = 0
227+
179228
let costOverrideDollars: number | undefined
180-
if (providerMetadata.anthropic) {
181-
cacheReadInputTokens =
182-
typeof providerMetadata.anthropic.cacheReadInputTokens === 'number'
183-
? providerMetadata.anthropic.cacheReadInputTokens
184-
: 0
185-
cacheCreationInputTokens =
186-
typeof providerMetadata.anthropic.cacheCreationInputTokens === 'number'
187-
? providerMetadata.anthropic.cacheCreationInputTokens
188-
: 0
189-
}
190-
if (providerMetadata.openrouter) {
191-
if (providerMetadata.openrouter.usage) {
192-
const openrouterUsage = providerMetadata.openrouter
229+
if (providerMetadata.codebuff) {
230+
if (providerMetadata.codebuff.usage) {
231+
const openrouterUsage = providerMetadata.codebuff
193232
.usage as OpenRouterUsageAccounting
194-
cacheReadInputTokens =
195-
openrouterUsage.promptTokensDetails?.cachedTokens ?? 0
196-
inputTokens = openrouterUsage.promptTokens - cacheReadInputTokens
197233

198234
costOverrideDollars =
199235
(openrouterUsage.cost ?? 0) +
@@ -205,8 +241,9 @@ export async function* promptAiSdkStream(
205241

206242
// Call the cost callback if provided
207243
if (params.onCostCalculated && costOverrideDollars) {
208-
const creditsUsed = costOverrideDollars * (1 + PROFIT_MARGIN)
209-
await params.onCostCalculated(creditsUsed)
244+
await params.onCostCalculated(
245+
calculateUsedCredits({ costDollars: costOverrideDollars }),
246+
)
210247
}
211248

212249
return messageId
@@ -229,7 +266,6 @@ export async function promptAiSdk(
229266
return ''
230267
}
231268

232-
const startTime = Date.now()
233269
let aiSDKModel = getAiSdkModel(params)
234270

235271
const response = await generateText({
@@ -248,31 +284,12 @@ export async function promptAiSdk(
248284
})
249285
const content = response.text
250286

251-
const messageId = response.response.id
252287
const providerMetadata = response.providerMetadata ?? {}
253-
const usage = response.usage
254-
let inputTokens = usage.inputTokens || 0
255-
const outputTokens = usage.outputTokens || 0
256-
let cacheReadInputTokens: number = 0
257-
let cacheCreationInputTokens: number = 0
258288
let costOverrideDollars: number | undefined
259-
if (providerMetadata.anthropic) {
260-
cacheReadInputTokens =
261-
typeof providerMetadata.anthropic.cacheReadInputTokens === 'number'
262-
? providerMetadata.anthropic.cacheReadInputTokens
263-
: 0
264-
cacheCreationInputTokens =
265-
typeof providerMetadata.anthropic.cacheCreationInputTokens === 'number'
266-
? providerMetadata.anthropic.cacheCreationInputTokens
267-
: 0
268-
}
269-
if (providerMetadata.openrouter) {
270-
if (providerMetadata.openrouter.usage) {
271-
const openrouterUsage = providerMetadata.openrouter
289+
if (providerMetadata.codebuff) {
290+
if (providerMetadata.codebuff.usage) {
291+
const openrouterUsage = providerMetadata.codebuff
272292
.usage as OpenRouterUsageAccounting
273-
cacheReadInputTokens =
274-
openrouterUsage.promptTokensDetails?.cachedTokens ?? 0
275-
inputTokens = openrouterUsage.promptTokens - cacheReadInputTokens
276293

277294
costOverrideDollars =
278295
(openrouterUsage.cost ?? 0) +
@@ -282,8 +299,9 @@ export async function promptAiSdk(
282299

283300
// Call the cost callback if provided
284301
if (params.onCostCalculated && costOverrideDollars) {
285-
const creditsUsed = costOverrideDollars * (1 + PROFIT_MARGIN)
286-
await params.onCostCalculated(creditsUsed)
302+
await params.onCostCalculated(
303+
calculateUsedCredits({ costDollars: costOverrideDollars }),
304+
)
287305
}
288306

289307
return content
@@ -325,31 +343,12 @@ export async function promptAiSdkStructured<T>(
325343

326344
const content = response.object
327345

328-
const messageId = response.response.id
329346
const providerMetadata = response.providerMetadata ?? {}
330-
const usage = response.usage
331-
let inputTokens = usage.inputTokens || 0
332-
const outputTokens = usage.outputTokens || 0
333-
let cacheReadInputTokens: number = 0
334-
let cacheCreationInputTokens: number = 0
335347
let costOverrideDollars: number | undefined
336-
if (providerMetadata.anthropic) {
337-
cacheReadInputTokens =
338-
typeof providerMetadata.anthropic.cacheReadInputTokens === 'number'
339-
? providerMetadata.anthropic.cacheReadInputTokens
340-
: 0
341-
cacheCreationInputTokens =
342-
typeof providerMetadata.anthropic.cacheCreationInputTokens === 'number'
343-
? providerMetadata.anthropic.cacheCreationInputTokens
344-
: 0
345-
}
346-
if (providerMetadata.openrouter) {
347-
if (providerMetadata.openrouter.usage) {
348-
const openrouterUsage = providerMetadata.openrouter
348+
if (providerMetadata.codebuff) {
349+
if (providerMetadata.codebuff.usage) {
350+
const openrouterUsage = providerMetadata.codebuff
349351
.usage as OpenRouterUsageAccounting
350-
cacheReadInputTokens =
351-
openrouterUsage.promptTokensDetails?.cachedTokens ?? 0
352-
inputTokens = openrouterUsage.promptTokens - cacheReadInputTokens
353352

354353
costOverrideDollars =
355354
(openrouterUsage.cost ?? 0) +
@@ -359,8 +358,9 @@ export async function promptAiSdkStructured<T>(
359358

360359
// Call the cost callback if provided
361360
if (params.onCostCalculated && costOverrideDollars) {
362-
const creditsUsed = costOverrideDollars * (1 + PROFIT_MARGIN)
363-
await params.onCostCalculated(creditsUsed)
361+
await params.onCostCalculated(
362+
calculateUsedCredits({ costDollars: costOverrideDollars }),
363+
)
364364
}
365365

366366
return content

0 commit comments

Comments
 (0)