Press n or j to go to the next uncovered block, b, p or k for the previous block.
| 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 4x 1x 1x 1x 4x 4x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 2x 1x 1x 1x 1x 2x 2x 2x 2x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x | /**
* LLM client usage recording middleware and wrappers.
*
* - createUsageRecordingMiddleware: StreamMiddleware that auto-records on stream completion
* - trackedCall: wraps llm.call() with automatic usage recording
* - trackedEmbed: wraps llm.embed() with automatic usage recording
*
* IMPORTANT: This module is transitively imported by client components via
* llm.ts → worksheet-parsing → PhotoViewerEditor → SummaryClient.
* All imports of ./helpers (which pulls in @/db via ./record) MUST be
* dynamic import() to avoid bundling node:http into the client.
*/
import type {
StreamMiddleware,
StreamEvent,
LLMClient,
LLMRequest,
LLMResponse,
EmbeddingRequest,
EmbeddingResponse,
} from '@soroban/llm-client'
import type { z } from 'zod'
import type { AiFeatureValue } from './features'
export interface UsageRecordingContext {
userId: string
feature: AiFeatureValue
backgroundTaskId?: string
}
/**
* StreamMiddleware that records usage from the 'complete' event.
*
* Intercepts the stream, watches for the `complete` event, and calls
* recordAiUsage() with token counts. All events are yielded downstream
* unchanged.
*/
export function createUsageRecordingMiddleware(
context: UsageRecordingContext,
/** Provider name — passed at construction since stream events don't carry it */
provider?: string,
/** Model name — passed at construction since stream events don't carry it */
model?: string
): StreamMiddleware {
return {
async *wrap<T>(
stream: AsyncGenerator<StreamEvent<T>, void, unknown>
): AsyncGenerator<StreamEvent<T>, void, unknown> {
for await (const event of stream) {
if (event.type === 'complete') {
const { recordLlmClientStreamUsage } = await import('./helpers')
recordLlmClientStreamUsage(event.usage, provider ?? 'openai', model ?? 'unknown', context)
}
yield event
}
},
}
}
/**
* Call llm.call() and record usage automatically.
*/
export async function trackedCall<T extends z.ZodType>(
llm: LLMClient,
request: LLMRequest<T>,
context: UsageRecordingContext
): Promise<LLMResponse<z.infer<T>>> {
const response = await llm.call(request)
const { recordLlmClientUsage } = await import('./helpers')
recordLlmClientUsage(response, context)
return response
}
/**
* Call llm.embed() and record usage automatically.
*/
export async function trackedEmbed(
llm: LLMClient,
request: EmbeddingRequest,
context: UsageRecordingContext
): Promise<EmbeddingResponse> {
const response = await llm.embed(request)
const { recordEmbeddingUsage } = await import('./helpers')
recordEmbeddingUsage(response, context)
return response
}
|