| 1 | import { BaseCallbackHandler } from '@langchain/core/callbacks/base' |
| 2 | import type { BaseMessage } from '@langchain/core/messages' |
| 3 | import type { LLMResult } from '@langchain/core/outputs' |
| 4 | import log from 'electron-log/main.js' |
| 5 | |
| 6 | export interface ExtractedModelUsage { |
| 7 | inputTokens: number |
| 8 | outputTokens: number |
| 9 | totalTokens: number |
| 10 | source: 'provider' | 'estimated' |
| 11 | } |
| 12 | |
| 13 | export interface ModelUsageEntry extends ExtractedModelUsage { |
| 14 | provider: string |
| 15 | model: string |
| 16 | modelConfigId?: string |
| 17 | } |
| 18 | |
| 19 | /** Application-provided persistence capability. Runtime never imports a database implementation. */ |
| 20 | export interface ModelUsageRecorder { |
| 21 | record(entry: ModelUsageEntry): Promise<void> |
| 22 | } |
| 23 | |
| 24 | export interface ModelRuntimeConfig { |
| 25 | recorder: ModelUsageRecorder | null |
| 26 | } |
| 27 | |
| 28 | type UnknownRecord = Record<string, unknown> |
| 29 | |
| 30 | const asRecord = (value: unknown): UnknownRecord | null => |
| 31 | value && typeof value === 'object' && !Array.isArray(value) |
| 32 | ? (value as UnknownRecord) |
| 33 | : null |
| 34 | |
| 35 | const readNumber = (record: UnknownRecord | null, keys: string[]): number | null => { |
| 36 | if (!record) return null |
| 37 | for (const key of keys) { |
| 38 | const value = record[key] |
| 39 | if (typeof value === 'number' && Number.isFinite(value) && value >= 0) { |
| 40 | return Math.floor(value) |
| 41 | } |
| 42 | } |
| 43 | return null |
| 44 | } |
| 45 | |
| 46 | const readUsageRecord = (value: unknown): ExtractedModelUsage | null => { |
| 47 | const record = asRecord(value) |
| 48 | if (!record) return null |
| 49 | const inputTokens = readNumber(record, ['input_tokens', 'inputTokens', 'promptTokens']) |
| 50 | const outputTokens = readNumber(record, [ |
| 51 | 'output_tokens', |
| 52 | 'outputTokens', |
| 53 | 'completionTokens' |
| 54 | ]) |
| 55 | const totalTokens = readNumber(record, ['total_tokens', 'totalTokens']) |
| 56 | if (inputTokens === null && outputTokens === null && totalTokens === null) return null |
| 57 | |
| 58 | const input = inputTokens ?? Math.max(0, (totalTokens ?? 0) - (outputTokens ?? 0)) |
| 59 | const output = outputTokens ?? Math.max(0, (totalTokens ?? 0) - input) |
| 60 | return { |
| 61 | inputTokens: input, |
| 62 | outputTokens: output, |
| 63 | totalTokens: totalTokens ?? input + output, |
| 64 | source: 'provider' |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | const countEstimatedTokens = (value: string): number => { |
| 69 | if (!value) return 0 |
| 70 | return Math.max(1, Math.ceil(value.length / 4)) |
| 71 | } |
| 72 | |
| 73 | const serializeMessages = (messages: BaseMessage[][]): string => |
| 74 | JSON.stringify( |
| 75 | messages.map((batch) => |
| 76 | batch.map((message) => ({ |
| 77 | role: message.getType(), |
| 78 | content: message.content, |
| 79 | additionalKwargs: message.additional_kwargs |
| 80 | })) |
| 81 | ) |
| 82 | ) |
| 83 | |
| 84 | export const extractModelUsage = (output: LLMResult): ExtractedModelUsage => { |
| 85 | const generationUsages = output.generations |
| 86 | .flat() |
| 87 | .map((generation) => { |
| 88 | const message = asRecord((generation as unknown as UnknownRecord).message) |
| 89 | return ( |
| 90 | readUsageRecord(message?.usage_metadata) || |
| 91 | readUsageRecord(asRecord(message?.response_metadata)?.tokenUsage) |
| 92 | ) |
| 93 | }) |
| 94 | .filter((usage): usage is ExtractedModelUsage => usage !== null) |
| 95 | |
| 96 | if (generationUsages.length > 0) { |
| 97 | return generationUsages.reduce<ExtractedModelUsage>( |
| 98 | (total, usage) => ({ |
| 99 | inputTokens: total.inputTokens + usage.inputTokens, |
| 100 | outputTokens: total.outputTokens + usage.outputTokens, |
| 101 | totalTokens: total.totalTokens + usage.totalTokens, |
| 102 | source: 'provider' |
| 103 | }), |
| 104 | { inputTokens: 0, outputTokens: 0, totalTokens: 0, source: 'provider' } |
| 105 | ) |
| 106 | } |
| 107 | |
| 108 | const llmOutput = asRecord(output.llmOutput) |
| 109 | return ( |
| 110 | readUsageRecord(llmOutput?.tokenUsage) || |
| 111 | readUsageRecord(llmOutput?.usage) || |
| 112 | { inputTokens: 0, outputTokens: 0, totalTokens: 0, source: 'estimated' } |
| 113 | ) |
| 114 | } |
| 115 | |
| 116 | export class ModelUsageCallbackHandler extends BaseCallbackHandler { |
| 117 | name = 'ohmyppt-model-usage' |
| 118 | private readonly estimatedInputByRun = new Map<string, number>() |
| 119 | |
| 120 | constructor( |
| 121 | private readonly context: { |
| 122 | provider: string |
| 123 | model: string |
| 124 | modelConfigId?: string |
| 125 | }, |
| 126 | private readonly recorder: ModelUsageRecorder | null |
| 127 | ) { |
| 128 | super({ _awaitHandler: true }) |
| 129 | } |
| 130 | |
| 131 | copy(): this { |
| 132 | return this |
| 133 | } |
| 134 | |
| 135 | handleLLMStart(_llm: unknown, prompts: string[], runId: string): void { |
| 136 | this.estimatedInputByRun.set(runId, countEstimatedTokens(prompts.join('\n'))) |
| 137 | } |
| 138 | |
| 139 | handleChatModelStart(_llm: unknown, messages: BaseMessage[][], runId: string): void { |
| 140 | this.estimatedInputByRun.set(runId, countEstimatedTokens(serializeMessages(messages))) |
| 141 | } |
| 142 | |
| 143 | handleLLMError(_error: unknown, runId: string): void { |
| 144 | this.estimatedInputByRun.delete(runId) |
| 145 | } |
| 146 | |
| 147 | async handleLLMEnd(output: LLMResult, runId: string): Promise<void> { |
| 148 | if (!this.recorder) { |
| 149 | this.estimatedInputByRun.delete(runId) |
| 150 | return |
| 151 | } |
| 152 | const providerUsage = extractModelUsage(output) |
| 153 | const usage = |
| 154 | providerUsage.source === 'provider' |
| 155 | ? providerUsage |
| 156 | : (() => { |
| 157 | const inputTokens = this.estimatedInputByRun.get(runId) ?? 0 |
| 158 | const outputTokens = countEstimatedTokens( |
| 159 | output.generations.flat().map((generation) => generation.text).join('\n') |
| 160 | ) |
| 161 | return { |
| 162 | inputTokens, |
| 163 | outputTokens, |
| 164 | totalTokens: inputTokens + outputTokens, |
| 165 | source: 'estimated' as const |
| 166 | } |
| 167 | })() |
| 168 | this.estimatedInputByRun.delete(runId) |
| 169 | try { |
| 170 | await this.recorder.record({ |
| 171 | ...this.context, |
| 172 | ...usage |
| 173 | }) |
| 174 | } catch (error) { |
| 175 | log.warn('[model-usage] failed to persist usage', { |
| 176 | provider: this.context.provider, |
| 177 | model: this.context.model, |
| 178 | message: error instanceof Error ? error.message : String(error) |
| 179 | }) |
| 180 | } |
| 181 | } |
| 182 | } |
| 183 |