@startsimpli/llm 0.4.22 → 0.4.24
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/package.json +1 -1
- package/src/__tests__/anthropic-provider.test.ts +11 -0
- package/src/__tests__/openai-provider.test.ts +53 -0
- package/src/__tests__/usage-tally.test.ts +22 -0
- package/src/index.ts +3 -0
- package/src/providers/anthropic.ts +17 -2
- package/src/providers/openai.ts +51 -6
- package/src/types.ts +22 -0
- package/src/utils/usage-tally.ts +56 -0
package/package.json
CHANGED
|
@@ -169,3 +169,14 @@ describe('AnthropicProvider — recovering from a parameter the model rejects',
|
|
|
169
169
|
expect(sent).toHaveLength(1);
|
|
170
170
|
});
|
|
171
171
|
});
|
|
172
|
+
|
|
173
|
+
describe('images', () => {
|
|
174
|
+
it('sends images as base64 image blocks before the text', async () => {
|
|
175
|
+
await provider('claude-opus-5-5').generate('what is this?', 's', { images: [{ data: 'AAAA', mediaType: 'image/png' }] });
|
|
176
|
+
const body = sent[0] as unknown as { messages: Array<{ content: Array<Record<string, unknown>> }> };
|
|
177
|
+
const content = body.messages[0]!.content;
|
|
178
|
+
expect(content[0]).toEqual({ type: 'image', source: { type: 'base64', media_type: 'image/png', data: 'AAAA' } });
|
|
179
|
+
expect(content[1]!.type).toBe('text');
|
|
180
|
+
expect(String(content[1]!.text)).toContain('what is this?');
|
|
181
|
+
});
|
|
182
|
+
});
|
|
@@ -33,6 +33,8 @@ interface Sent {
|
|
|
33
33
|
max_tokens?: number;
|
|
34
34
|
max_completion_tokens?: number;
|
|
35
35
|
reasoning_effort?: string;
|
|
36
|
+
prompt_cache_key?: string;
|
|
37
|
+
prompt_cache_retention?: string;
|
|
36
38
|
}
|
|
37
39
|
|
|
38
40
|
const OK_BODY = {
|
|
@@ -265,3 +267,54 @@ describe('generateStream', () => {
|
|
|
265
267
|
await expect(provider('gpt-6-astra').generateStream!('u', 's', () => {})).rejects.toThrow(/Empty response/);
|
|
266
268
|
});
|
|
267
269
|
});
|
|
270
|
+
|
|
271
|
+
describe('images', () => {
|
|
272
|
+
it('sends images as image_url parts after the text, and plain text without them', async () => {
|
|
273
|
+
const p = provider('gpt-6-astra');
|
|
274
|
+
await p.generate('what is this?', 's', { images: [{ data: 'AAAA', mediaType: 'image/png' }] });
|
|
275
|
+
await p.generate('plain', 's');
|
|
276
|
+
const withImage = sent[0] as unknown as { messages: Array<{ role: string; content: unknown }> };
|
|
277
|
+
expect(withImage.messages[1]!.content).toEqual([
|
|
278
|
+
{ type: 'text', text: 'what is this?' },
|
|
279
|
+
{ type: 'image_url', image_url: { url: 'data:image/png;base64,AAAA' } },
|
|
280
|
+
]);
|
|
281
|
+
const plain = sent[1] as unknown as { messages: Array<{ role: string; content: unknown }> };
|
|
282
|
+
expect(plain.messages[1]!.content).toBe('plain');
|
|
283
|
+
});
|
|
284
|
+
});
|
|
285
|
+
|
|
286
|
+
describe('prompt caching', () => {
|
|
287
|
+
// Measured against api.openai.com 2026-10-15 (gpt-6-astra, gpt-6-luna): both
|
|
288
|
+
// parameters accepted; a repeat of a 2,020-token prompt came back with
|
|
289
|
+
// prompt_tokens_details.cached_tokens 2,017.
|
|
290
|
+
it('sends the cache key and retention when asked, and nothing when not', async () => {
|
|
291
|
+
const p = provider('gpt-6-astra');
|
|
292
|
+
await p.generate('u', 's', { cacheKey: 'crochet-decompose', cacheRetention: '24h' });
|
|
293
|
+
await p.generate('u', 's');
|
|
294
|
+
expect(sent[0]!.prompt_cache_key).toBe('crochet-decompose');
|
|
295
|
+
expect(sent[0]!.prompt_cache_retention).toBe('24h');
|
|
296
|
+
expect(sent[1]!.prompt_cache_key).toBeUndefined();
|
|
297
|
+
expect(sent[1]!.prompt_cache_retention).toBeUndefined();
|
|
298
|
+
});
|
|
299
|
+
|
|
300
|
+
it('reports cached prompt tokens', async () => {
|
|
301
|
+
fetchMock.mockImplementationOnce(async (_url: unknown, init: RequestInit) => {
|
|
302
|
+
sent.push(JSON.parse(String(init.body)) as Sent);
|
|
303
|
+
return { ok: true, status: 200, json: async () => ({ ...OK_BODY, usage: { prompt_tokens: 2020, completion_tokens: 4, total_tokens: 2024, prompt_tokens_details: { cached_tokens: 2017 } } }) } as unknown as Response;
|
|
304
|
+
});
|
|
305
|
+
const res = await provider('gpt-6-astra').generate('u', 's', { cacheKey: 'k' });
|
|
306
|
+
expect(res.usage.cachedTokens).toBe(2017);
|
|
307
|
+
});
|
|
308
|
+
|
|
309
|
+
it('drops the cache parameters for a model that rejects them, and remembers', async () => {
|
|
310
|
+
const p = provider('some-old-model');
|
|
311
|
+
reject({ code: 'unknown_parameter', param: 'prompt_cache_retention', message: "Unrecognized request argument supplied: prompt_cache_retention" });
|
|
312
|
+
const res = await p.generate('u', 's', { cacheKey: 'k', cacheRetention: '24h' });
|
|
313
|
+
expect(res.content).toBe('{"ok":true}');
|
|
314
|
+
expect(sent[1]!.prompt_cache_retention).toBeUndefined();
|
|
315
|
+
expect(sent[1]!.prompt_cache_key).toBeUndefined();
|
|
316
|
+
await p.generate('u', 's', { cacheKey: 'k', cacheRetention: '24h' });
|
|
317
|
+
expect(sent).toHaveLength(3);
|
|
318
|
+
expect(sent[2]!.prompt_cache_retention).toBeUndefined();
|
|
319
|
+
});
|
|
320
|
+
});
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
import { describe, it, expect } from 'vitest';
|
|
2
|
+
import { newUsageTally, tallyingService } from '../utils/usage-tally';
|
|
3
|
+
|
|
4
|
+
describe('tallyingService', () => {
|
|
5
|
+
it('adds up every call per model, cached tokens included, and passes other members through', async () => {
|
|
6
|
+
const svc = {
|
|
7
|
+
name: 'svc',
|
|
8
|
+
async generateChatResponse(_u: string, _s: string, opts?: { model?: string }) {
|
|
9
|
+
return { content: '{}', model: opts?.model ?? 'a', usage: { promptTokens: 100, completionTokens: 40, totalTokens: 140, cachedTokens: 60 } };
|
|
10
|
+
},
|
|
11
|
+
};
|
|
12
|
+
const tally = newUsageTally();
|
|
13
|
+
const counted = tallyingService(svc, tally);
|
|
14
|
+
await counted.generateChatResponse('u', 's', { model: 'astra' });
|
|
15
|
+
await counted.generateChatResponse('u', 's', { model: 'astra' });
|
|
16
|
+
await counted.generateChatResponse('u', 's', { model: 'luna' });
|
|
17
|
+
expect(counted.name).toBe('svc');
|
|
18
|
+
expect(tally).toMatchObject({ calls: 3, promptTokens: 300, completionTokens: 120, cachedTokens: 180 });
|
|
19
|
+
expect(tally.byModel.astra).toEqual({ calls: 2, promptTokens: 200, completionTokens: 80, cachedTokens: 120 });
|
|
20
|
+
expect(tally.byModel.luna.calls).toBe(1);
|
|
21
|
+
});
|
|
22
|
+
});
|
package/src/index.ts
CHANGED
|
@@ -8,6 +8,7 @@ export type {
|
|
|
8
8
|
LLMResponse,
|
|
9
9
|
GenerateOptions,
|
|
10
10
|
ReasoningEffort,
|
|
11
|
+
LLMImage,
|
|
11
12
|
ProviderConfig,
|
|
12
13
|
JsonExtractionResult,
|
|
13
14
|
ValidationResult,
|
|
@@ -35,6 +36,8 @@ export type { MockBehavior } from './providers/mock';
|
|
|
35
36
|
// Utilities
|
|
36
37
|
export { extractJson } from './utils/json-extraction'
|
|
37
38
|
export { llmErrorToHttpStatus } from './utils/http-status';
|
|
39
|
+
export { newUsageTally, addUsage, tallyingService } from './utils/usage-tally';
|
|
40
|
+
export type { UsageTally, UsageCounts } from './utils/usage-tally';
|
|
38
41
|
export {
|
|
39
42
|
buildSchemaErrorCorrectionPrompt,
|
|
40
43
|
buildParseErrorCorrectionPrompt,
|
|
@@ -11,9 +11,13 @@ import { LLM_DEFAULTS } from '../types';
|
|
|
11
11
|
import { LLMProviderError } from '../errors';
|
|
12
12
|
import type { LLMErrorCode } from '../errors';
|
|
13
13
|
|
|
14
|
+
type AnthropicContentBlock =
|
|
15
|
+
| { type: 'text'; text: string }
|
|
16
|
+
| { type: 'image'; source: { type: 'base64'; media_type: string; data: string } };
|
|
17
|
+
|
|
14
18
|
interface AnthropicMessage {
|
|
15
19
|
role: 'user' | 'assistant';
|
|
16
|
-
content: string;
|
|
20
|
+
content: string | AnthropicContentBlock[];
|
|
17
21
|
}
|
|
18
22
|
|
|
19
23
|
interface AnthropicRequest {
|
|
@@ -199,7 +203,18 @@ Please respond with a valid JSON object only. Do not include any text before or
|
|
|
199
203
|
const requestBody: AnthropicRequest = {
|
|
200
204
|
model,
|
|
201
205
|
max_tokens: maxTokens,
|
|
202
|
-
messages: [{
|
|
206
|
+
messages: [{
|
|
207
|
+
role: 'user',
|
|
208
|
+
content: options?.images?.length
|
|
209
|
+
? [
|
|
210
|
+
...options.images.map((img) => ({
|
|
211
|
+
type: 'image' as const,
|
|
212
|
+
source: { type: 'base64' as const, media_type: img.mediaType, data: img.data },
|
|
213
|
+
})),
|
|
214
|
+
{ type: 'text' as const, text: jsonWrappedPrompt },
|
|
215
|
+
]
|
|
216
|
+
: jsonWrappedPrompt,
|
|
217
|
+
}],
|
|
203
218
|
system: systemPrompt,
|
|
204
219
|
};
|
|
205
220
|
// Only for models known to take it. The Claude 5 family answers a request
|
package/src/providers/openai.ts
CHANGED
|
@@ -18,14 +18,29 @@ import type {
|
|
|
18
18
|
NativeChatMessage,
|
|
19
19
|
NativeToolResponse,
|
|
20
20
|
ReasoningEffort,
|
|
21
|
+
LLMImage,
|
|
22
|
+
TokenUsage,
|
|
21
23
|
} from '../types';
|
|
22
24
|
import { LLM_DEFAULTS } from '../types';
|
|
23
25
|
import { LLMProviderError } from '../errors';
|
|
24
26
|
import type { LLMErrorCode } from '../errors';
|
|
25
27
|
|
|
28
|
+
type OpenAIContentPart =
|
|
29
|
+
| { type: 'text'; text: string }
|
|
30
|
+
| { type: 'image_url'; image_url: { url: string } };
|
|
31
|
+
|
|
26
32
|
interface OpenAIMessage {
|
|
27
33
|
role: 'system' | 'user' | 'assistant';
|
|
28
|
-
content: string;
|
|
34
|
+
content: string | OpenAIContentPart[];
|
|
35
|
+
}
|
|
36
|
+
|
|
37
|
+
/** The user message: plain text, or text plus images for vision models. */
|
|
38
|
+
function userContent(prompt: string, images?: LLMImage[]): string | OpenAIContentPart[] {
|
|
39
|
+
if (!images?.length) return prompt;
|
|
40
|
+
return [
|
|
41
|
+
{ type: 'text', text: prompt },
|
|
42
|
+
...images.map((img) => ({ type: 'image_url' as const, image_url: { url: `data:${img.mediaType};base64,${img.data}` } })),
|
|
43
|
+
];
|
|
29
44
|
}
|
|
30
45
|
|
|
31
46
|
interface OpenAIRequest {
|
|
@@ -37,6 +52,8 @@ interface OpenAIRequest {
|
|
|
37
52
|
/** Reasoning models only; dropped automatically for models that reject it. */
|
|
38
53
|
reasoning_effort?: ReasoningEffort;
|
|
39
54
|
response_format?: { type: 'json_object' | 'text' };
|
|
55
|
+
prompt_cache_key?: string;
|
|
56
|
+
prompt_cache_retention?: 'in_memory' | '24h';
|
|
40
57
|
}
|
|
41
58
|
|
|
42
59
|
/**
|
|
@@ -71,6 +88,23 @@ interface ModelQuirks {
|
|
|
71
88
|
efforts?: ReasoningEffort[];
|
|
72
89
|
/** Largest completion-token cap this model accepts. */
|
|
73
90
|
maxOutput?: number;
|
|
91
|
+
/** The model rejected prompt-cache parameters; send none. */
|
|
92
|
+
noPromptCache?: boolean;
|
|
93
|
+
}
|
|
94
|
+
|
|
95
|
+
/** Prompt-cache parameters for a request, unless this model has rejected them. */
|
|
96
|
+
function cacheParams(q: ModelQuirks, options?: GenerateOptions): Pick<OpenAIRequest, 'prompt_cache_key' | 'prompt_cache_retention'> {
|
|
97
|
+
if (q.noPromptCache) return {};
|
|
98
|
+
return {
|
|
99
|
+
...(options?.cacheKey ? { prompt_cache_key: options.cacheKey } : {}),
|
|
100
|
+
...(options?.cacheRetention ? { prompt_cache_retention: options.cacheRetention } : {}),
|
|
101
|
+
};
|
|
102
|
+
}
|
|
103
|
+
|
|
104
|
+
/** Cached prompt tokens from an OpenAI usage block, when it reports them. */
|
|
105
|
+
function cachedTokens(usage: { prompt_tokens_details?: { cached_tokens?: number } } | undefined): { cachedTokens?: number } {
|
|
106
|
+
const n = usage?.prompt_tokens_details?.cached_tokens;
|
|
107
|
+
return typeof n === 'number' ? { cachedTokens: n } : {};
|
|
74
108
|
}
|
|
75
109
|
|
|
76
110
|
const EFFORT_ORDER: ReasoningEffort[] = ['none', 'minimal', 'low', 'medium', 'high', 'xhigh'];
|
|
@@ -89,6 +123,11 @@ function nearestEffort(want: ReasoningEffort, supported: ReasoningEffort[]): Rea
|
|
|
89
123
|
function learnFromError(q: ModelQuirks, err: OpenAIErrorResponse['error'] | undefined, sent: { effort?: ReasoningEffort; maxTokens: number }): boolean {
|
|
90
124
|
if (!err?.message) return false;
|
|
91
125
|
const msg = err.message;
|
|
126
|
+
if (err.param === 'prompt_cache_key' || err.param === 'prompt_cache_retention' || /prompt_cache_(key|retention)/.test(msg)) {
|
|
127
|
+
if (q.noPromptCache) return false;
|
|
128
|
+
q.noPromptCache = true;
|
|
129
|
+
return true;
|
|
130
|
+
}
|
|
92
131
|
if (err.param === 'temperature' || /'temperature'/.test(msg)) {
|
|
93
132
|
if (q.noTemperature) return false;
|
|
94
133
|
q.noTemperature = true;
|
|
@@ -169,6 +208,7 @@ interface OpenAIResponse {
|
|
|
169
208
|
prompt_tokens: number;
|
|
170
209
|
completion_tokens: number;
|
|
171
210
|
total_tokens: number;
|
|
211
|
+
prompt_tokens_details?: { cached_tokens?: number };
|
|
172
212
|
};
|
|
173
213
|
}
|
|
174
214
|
|
|
@@ -202,7 +242,7 @@ interface OpenAIToolsResponse {
|
|
|
202
242
|
message: { role: string; content: string | null; tool_calls?: OpenAIToolCall[] };
|
|
203
243
|
finish_reason: string;
|
|
204
244
|
}>;
|
|
205
|
-
usage: { prompt_tokens: number; completion_tokens: number; total_tokens: number };
|
|
245
|
+
usage: { prompt_tokens: number; completion_tokens: number; total_tokens: number; prompt_tokens_details?: { cached_tokens?: number } };
|
|
206
246
|
}
|
|
207
247
|
|
|
208
248
|
const OPENAI_MODELS = [
|
|
@@ -324,7 +364,7 @@ export class OpenAIProvider implements ILLMProvider {
|
|
|
324
364
|
model,
|
|
325
365
|
messages: [
|
|
326
366
|
{ role: 'system' as const, content: systemPrompt },
|
|
327
|
-
{ role: 'user' as const, content: userPrompt },
|
|
367
|
+
{ role: 'user' as const, content: userContent(userPrompt, options?.images) },
|
|
328
368
|
],
|
|
329
369
|
response_format: { type: 'json_object' as const },
|
|
330
370
|
stream: true,
|
|
@@ -337,13 +377,14 @@ export class OpenAIProvider implements ILLMProvider {
|
|
|
337
377
|
const response = await this.postAdaptivelyRaw(model, (q) => ({
|
|
338
378
|
...base,
|
|
339
379
|
...samplingParams(q, temperature, maxTokens, effort),
|
|
380
|
+
...cacheParams(q, options),
|
|
340
381
|
}), { effort, maxTokens }, controller.signal);
|
|
341
382
|
|
|
342
383
|
let content = '';
|
|
343
384
|
let reportedModel = model;
|
|
344
385
|
let finishReason: string | undefined;
|
|
345
386
|
let responseId: string | undefined;
|
|
346
|
-
const usage = { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
|
|
387
|
+
const usage: TokenUsage = { promptTokens: 0, completionTokens: 0, totalTokens: 0 };
|
|
347
388
|
const reader = response.body!.getReader();
|
|
348
389
|
const decoder = new TextDecoder();
|
|
349
390
|
let buffer = '';
|
|
@@ -362,7 +403,7 @@ export class OpenAIProvider implements ILLMProvider {
|
|
|
362
403
|
id?: string;
|
|
363
404
|
model?: string;
|
|
364
405
|
choices?: Array<{ delta?: { content?: string | null }; finish_reason?: string | null }>;
|
|
365
|
-
usage?: { prompt_tokens: number; completion_tokens: number; total_tokens: number };
|
|
406
|
+
usage?: { prompt_tokens: number; completion_tokens: number; total_tokens: number; prompt_tokens_details?: { cached_tokens?: number } };
|
|
366
407
|
};
|
|
367
408
|
if (chunk.id) responseId = chunk.id;
|
|
368
409
|
if (chunk.model) reportedModel = chunk.model;
|
|
@@ -377,6 +418,7 @@ export class OpenAIProvider implements ILLMProvider {
|
|
|
377
418
|
usage.promptTokens = chunk.usage.prompt_tokens;
|
|
378
419
|
usage.completionTokens = chunk.usage.completion_tokens;
|
|
379
420
|
usage.totalTokens = chunk.usage.total_tokens;
|
|
421
|
+
Object.assign(usage, cachedTokens(chunk.usage));
|
|
380
422
|
}
|
|
381
423
|
}
|
|
382
424
|
}
|
|
@@ -424,7 +466,7 @@ export class OpenAIProvider implements ILLMProvider {
|
|
|
424
466
|
model,
|
|
425
467
|
messages: [
|
|
426
468
|
{ role: 'system' as const, content: systemPrompt },
|
|
427
|
-
{ role: 'user' as const, content: userPrompt },
|
|
469
|
+
{ role: 'user' as const, content: userContent(userPrompt, options?.images) },
|
|
428
470
|
],
|
|
429
471
|
response_format: { type: 'json_object' as const },
|
|
430
472
|
};
|
|
@@ -436,6 +478,7 @@ export class OpenAIProvider implements ILLMProvider {
|
|
|
436
478
|
const data = await this.postAdaptively<OpenAIResponse>(model, (q) => ({
|
|
437
479
|
...base,
|
|
438
480
|
...samplingParams(q, temperature, maxTokens, effort),
|
|
481
|
+
...cacheParams(q, options),
|
|
439
482
|
}), { effort, maxTokens }, controller.signal);
|
|
440
483
|
|
|
441
484
|
clearTimeout(timeoutId);
|
|
@@ -456,6 +499,7 @@ export class OpenAIProvider implements ILLMProvider {
|
|
|
456
499
|
promptTokens: data.usage.prompt_tokens,
|
|
457
500
|
completionTokens: data.usage.completion_tokens,
|
|
458
501
|
totalTokens: data.usage.total_tokens,
|
|
502
|
+
...cachedTokens(data.usage),
|
|
459
503
|
},
|
|
460
504
|
model: data.model,
|
|
461
505
|
provider: this.name,
|
|
@@ -554,6 +598,7 @@ export class OpenAIProvider implements ILLMProvider {
|
|
|
554
598
|
const data = await this.postAdaptively<OpenAIToolsResponse>(model, (q) => ({
|
|
555
599
|
...base,
|
|
556
600
|
...samplingParams(q, temperature, maxTokens, effort),
|
|
601
|
+
...cacheParams(q, options),
|
|
557
602
|
}), { effort, maxTokens }, controller.signal);
|
|
558
603
|
clearTimeout(timeoutId);
|
|
559
604
|
const choice = data.choices?.[0];
|
package/src/types.ts
CHANGED
|
@@ -7,6 +7,8 @@ export interface TokenUsage {
|
|
|
7
7
|
promptTokens: number;
|
|
8
8
|
completionTokens: number;
|
|
9
9
|
totalTokens: number;
|
|
10
|
+
/** Of promptTokens, how many were served from the provider's prompt cache (billed at the cached rate). */
|
|
11
|
+
cachedTokens?: number;
|
|
10
12
|
}
|
|
11
13
|
|
|
12
14
|
/**
|
|
@@ -30,6 +32,12 @@ export interface LLMResponse {
|
|
|
30
32
|
/**
|
|
31
33
|
* Options passed through to the provider for a single generation call.
|
|
32
34
|
*/
|
|
35
|
+
/** An image attached to a prompt: base64 data (no data: prefix) and its media type. */
|
|
36
|
+
export interface LLMImage {
|
|
37
|
+
data: string;
|
|
38
|
+
mediaType: 'image/png' | 'image/jpeg' | 'image/gif' | 'image/webp';
|
|
39
|
+
}
|
|
40
|
+
|
|
33
41
|
export interface GenerateOptions {
|
|
34
42
|
/** Model to use (provider-specific) */
|
|
35
43
|
model?: string;
|
|
@@ -45,6 +53,20 @@ export interface GenerateOptions {
|
|
|
45
53
|
* provider adds headroom for the reasoning on top.
|
|
46
54
|
*/
|
|
47
55
|
reasoningEffort?: ReasoningEffort;
|
|
56
|
+
/**
|
|
57
|
+
* Images sent with the user prompt, for vision-capable models (e.g. to have
|
|
58
|
+
* a model look at a rendering of what it designed).
|
|
59
|
+
*/
|
|
60
|
+
images?: LLMImage[];
|
|
61
|
+
/**
|
|
62
|
+
* Prompt caching (OpenAI). Requests that share a long identical opening
|
|
63
|
+
* (instructions first, per-request data last) are billed at the cached rate
|
|
64
|
+
* when they reach the same cache: `cacheKey` routes them there, and
|
|
65
|
+
* `cacheRetention: '24h'` keeps the cache for a day instead of minutes, which
|
|
66
|
+
* is what makes it pay off for low-traffic features.
|
|
67
|
+
*/
|
|
68
|
+
cacheKey?: string;
|
|
69
|
+
cacheRetention?: 'in_memory' | '24h';
|
|
48
70
|
/** Arbitrary extra data the caller wants to thread through (e.g. domain preferences) */
|
|
49
71
|
extra?: Record<string, unknown>;
|
|
50
72
|
}
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
/**
|
|
2
|
+
* Token usage added up across every call a feature makes, per model, so what a
|
|
3
|
+
* request cost can be recorded and priced (a crochet design makes 3-7 calls on
|
|
4
|
+
* two models). Wrap any service whose methods resolve to `{ model, usage }`.
|
|
5
|
+
*/
|
|
6
|
+
import type { TokenUsage } from '../types';
|
|
7
|
+
|
|
8
|
+
export interface UsageCounts {
|
|
9
|
+
calls: number;
|
|
10
|
+
promptTokens: number;
|
|
11
|
+
completionTokens: number;
|
|
12
|
+
/** Of promptTokens, how many were served from the prompt cache. */
|
|
13
|
+
cachedTokens: number;
|
|
14
|
+
}
|
|
15
|
+
|
|
16
|
+
export interface UsageTally extends UsageCounts {
|
|
17
|
+
byModel: Record<string, UsageCounts>;
|
|
18
|
+
}
|
|
19
|
+
|
|
20
|
+
const zero = (): UsageCounts => ({ calls: 0, promptTokens: 0, completionTokens: 0, cachedTokens: 0 });
|
|
21
|
+
|
|
22
|
+
export function newUsageTally(): UsageTally {
|
|
23
|
+
return { ...zero(), byModel: {} };
|
|
24
|
+
}
|
|
25
|
+
|
|
26
|
+
/** Add one response's usage to the tally. Returns the response unchanged. */
|
|
27
|
+
export function addUsage<R extends { model?: string; usage?: TokenUsage }>(tally: UsageTally, response: R): R {
|
|
28
|
+
const u = response.usage;
|
|
29
|
+
const m = (tally.byModel[response.model ?? 'unknown'] ??= zero());
|
|
30
|
+
for (const c of [tally, m]) {
|
|
31
|
+
c.calls++;
|
|
32
|
+
c.promptTokens += u?.promptTokens ?? 0;
|
|
33
|
+
c.completionTokens += u?.completionTokens ?? 0;
|
|
34
|
+
c.cachedTokens += u?.cachedTokens ?? 0;
|
|
35
|
+
}
|
|
36
|
+
return response;
|
|
37
|
+
}
|
|
38
|
+
|
|
39
|
+
/**
|
|
40
|
+
* The same service, with every call's usage added to `tally`. Methods in
|
|
41
|
+
* `methods` that return a promise of a response are counted; everything else
|
|
42
|
+
* passes through untouched.
|
|
43
|
+
*/
|
|
44
|
+
export function tallyingService<S extends object>(
|
|
45
|
+
service: S,
|
|
46
|
+
tally: UsageTally,
|
|
47
|
+
methods: readonly string[] = ['generateChatResponse', 'generateChatResponseStream', 'generate'],
|
|
48
|
+
): S {
|
|
49
|
+
return new Proxy(service, {
|
|
50
|
+
get(target, prop, receiver) {
|
|
51
|
+
const value = Reflect.get(target, prop, receiver);
|
|
52
|
+
if (typeof prop !== 'string' || !methods.includes(prop) || typeof value !== 'function') return value;
|
|
53
|
+
return (...args: unknown[]) => (value as (...a: unknown[]) => Promise<{ model?: string; usage?: TokenUsage }>).apply(target, args).then((r) => addUsage(tally, r));
|
|
54
|
+
},
|
|
55
|
+
});
|
|
56
|
+
}
|