mirror of
https://github.com/FJurmanovic/sugarr.git
synced 2026-11-10 23:31:08 +00:00
Track cached prompt tokens through to Langfuse
OpenAI's chat completions response includes prompt_tokens_details. cached_tokens when prompt caching kicks in - previously dropped entirely. Now extracted in LlmClient and forwarded to Langfuse as the flat cache_read_input_tokens usage key. Used the flat key rather than OpenAI's nested prompt_tokens_details shape: the installed @langfuse/tracing SDK's usageDetails type is Record<string, number> | OpenAiUsage (a 3-field camelCase shape, no _details nesting) even though the wire schema in @langfuse/core supports nesting - a gap in the SDK's public types, not something worth fighting with a cast. Verified: full suite (82/82) passing, tsc/eslint clean.
This commit is contained in:
@@ -16,6 +16,14 @@ import type {
|
||||
* OpenAI's native snake_case schema (prompt_tokens/completion_tokens/...);
|
||||
* our own ChatResult.usage uses camelCase for our internal contract, so it
|
||||
* needs translating at this boundary or Langfuse silently shows zero usage.
|
||||
*
|
||||
* Cached tokens use the flat schema's cache_read_input_tokens key rather
|
||||
* than OpenAI's nested prompt_tokens_details.cached_tokens: the installed
|
||||
* @langfuse/tracing SDK's usageDetails type is `Record<string, number> |
|
||||
* OpenAiUsage` (a 3-field camelCase shape with no _details nesting) even
|
||||
* though the wire schema (@langfuse/core's OpenAiCompletionUsageSchema)
|
||||
* supports the nested form — keeping everything as flat numbers avoids
|
||||
* fighting that gap in the SDK's public types.
|
||||
*/
|
||||
function toUsageDetails(usage: GenerationOutput['usage']): Record<string, number> | undefined {
|
||||
if (!usage) return undefined;
|
||||
@@ -23,6 +31,7 @@ function toUsageDetails(usage: GenerationOutput['usage']): Record<string, number
|
||||
if (usage.promptTokens !== undefined) details.prompt_tokens = usage.promptTokens;
|
||||
if (usage.completionTokens !== undefined) details.completion_tokens = usage.completionTokens;
|
||||
if (usage.totalTokens !== undefined) details.total_tokens = usage.totalTokens;
|
||||
if (usage.cachedTokens !== undefined) details.cache_read_input_tokens = usage.cachedTokens;
|
||||
return Object.keys(details).length > 0 ? details : undefined;
|
||||
}
|
||||
|
||||
|
||||
@@ -13,7 +13,12 @@ export interface ChatOptions {
|
||||
|
||||
export interface ChatResult {
|
||||
content: string;
|
||||
usage?: { promptTokens?: number; completionTokens?: number; totalTokens?: number };
|
||||
usage?: {
|
||||
promptTokens?: number;
|
||||
completionTokens?: number;
|
||||
totalTokens?: number;
|
||||
cachedTokens?: number;
|
||||
};
|
||||
}
|
||||
|
||||
export class LlmClient {
|
||||
@@ -43,6 +48,7 @@ export class LlmClient {
|
||||
promptTokens: response.usage.prompt_tokens,
|
||||
completionTokens: response.usage.completion_tokens,
|
||||
totalTokens: response.usage.total_tokens,
|
||||
cachedTokens: response.usage.prompt_tokens_details?.cached_tokens,
|
||||
}
|
||||
: undefined,
|
||||
};
|
||||
|
||||
@@ -7,7 +7,12 @@ export interface GenerationInput {
|
||||
|
||||
export interface GenerationOutput {
|
||||
output: unknown;
|
||||
usage?: { promptTokens?: number; completionTokens?: number; totalTokens?: number };
|
||||
usage?: {
|
||||
promptTokens?: number;
|
||||
completionTokens?: number;
|
||||
totalTokens?: number;
|
||||
cachedTokens?: number;
|
||||
};
|
||||
metadata?: Record<string, unknown>;
|
||||
}
|
||||
|
||||
|
||||
@@ -78,6 +78,27 @@ describe('LangfuseTracer', () => {
|
||||
expect(generationEnd).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
it('maps cachedTokens to the flat cache_read_input_tokens key', () => {
|
||||
const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
|
||||
const handle = tracer.startTrace('my-pipeline');
|
||||
const gen = handle.startGeneration({ name: 'step-1', model: 'gpt-4o-mini', input: ['hi'] });
|
||||
gen.end({
|
||||
output: 'done',
|
||||
usage: { promptTokens: 100, completionTokens: 20, totalTokens: 120, cachedTokens: 80 },
|
||||
});
|
||||
|
||||
expect(generationUpdate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
usageDetails: {
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 20,
|
||||
total_tokens: 120,
|
||||
cache_read_input_tokens: 80,
|
||||
},
|
||||
})
|
||||
);
|
||||
});
|
||||
|
||||
it('omits usageDetails entirely when no usage is reported', () => {
|
||||
const tracer = new LangfuseTracer({ publicKey: 'pk', secretKey: 'sk', baseUrl: 'http://lf' });
|
||||
const handle = tracer.startTrace('my-pipeline');
|
||||
|
||||
@@ -32,6 +32,28 @@ describe('LlmClient', () => {
|
||||
});
|
||||
});
|
||||
|
||||
it('extracts cachedTokens from prompt_tokens_details when present', async () => {
|
||||
createMock.mockResolvedValue({
|
||||
choices: [{ message: { content: 'hello' } }],
|
||||
usage: {
|
||||
prompt_tokens: 100,
|
||||
completion_tokens: 20,
|
||||
total_tokens: 120,
|
||||
prompt_tokens_details: { cached_tokens: 80 },
|
||||
},
|
||||
});
|
||||
|
||||
const client = new LlmClient({ baseUrl: 'http://localhost:11434/v1' });
|
||||
const result = await client.chat([{ role: 'user', content: 'hi' }], { model: 'gpt-4o-mini' });
|
||||
|
||||
expect(result.usage).toEqual({
|
||||
promptTokens: 100,
|
||||
completionTokens: 20,
|
||||
totalTokens: 120,
|
||||
cachedTokens: 80,
|
||||
});
|
||||
});
|
||||
|
||||
it('throws when the response has no choices', async () => {
|
||||
createMock.mockResolvedValue({ choices: [] });
|
||||
|
||||
|
||||
Reference in New Issue
Block a user