Skip to content

Commit 2b2d851

Browse files
os-zhuangclaude
andcommitted
fix(service-ai): estimate non-zero token usage in MemoryLLMAdapter
The in-memory adapter stands in for a real LLM in dev/CI/local-E2E, but returned a flat `usage: {0,0,0}` on every path. That made every token-metering feature — quota guardrails, lifetime caps, usage dashboards, cost stops — impossible to exercise without a paid provider key, i.e. without spending real money. Estimate usage crudely (~4 chars/token, the standard ballpark) from the prompt messages and echoed output across chat/complete/streamChat/ generateObject. Intentionally NOT provider-accurate — it exists so usage-driven behaviour is testable money-free, not to bill anyone. Adds a unit test asserting non-zero, additive (total = prompt+completion), monotonic-in-input-size usage. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent cb352a0 commit 2b2d851

2 files changed

Lines changed: 48 additions & 11 deletions

File tree

packages/services/service-ai/src/__tests__/ai-service.test.ts

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,21 @@ describe('MemoryLLMAdapter', () => {
4545
expect(result.usage).toBeDefined();
4646
});
4747

48+
it('estimates non-zero, monotonic token usage so metering is testable money-free', async () => {
49+
const short = await adapter.chat([{ role: 'user', content: 'hi' }]);
50+
const long = await adapter.chat([
51+
{ role: 'system', content: 'You are a helpful assistant with a long preamble. '.repeat(20) },
52+
{ role: 'user', content: 'Tell me everything about my data. '.repeat(20) },
53+
]);
54+
// Non-zero (the whole point — a flat 0 made quota/guardrail features untestable).
55+
expect(short.usage!.totalTokens).toBeGreaterThan(0);
56+
// total = prompt + completion, and grows with input size (monotonic).
57+
expect(short.usage!.totalTokens).toBe(
58+
short.usage!.promptTokens + short.usage!.completionTokens,
59+
);
60+
expect(long.usage!.promptTokens).toBeGreaterThan(short.usage!.promptTokens);
61+
});
62+
4863
it('should handle no user message in chat()', async () => {
4964
const messages: ModelMessage[] = [{ role: 'system', content: 'System only' }];
5065
const result = await adapter.chat(messages);

packages/services/service-ai/src/adapters/memory-adapter.ts

Lines changed: 33 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,27 @@ import type { LLMAdapter } from '@objectstack/spec/contracts';
1818
* Always echoes back the last user message prefixed with "[memory] ".
1919
* Useful for unit tests, CI pipelines, and local dev without an LLM key.
2020
*/
21+
/**
22+
* Rough token estimate for the in-memory adapter. The memory adapter stands in
23+
* for a real LLM in dev/CI/local-E2E; returning a flat `0` usage made every
24+
* token-metering feature (quota guardrails, usage dashboards, cost caps)
25+
* impossible to exercise without a paid provider key. This is intentionally
26+
* crude — ~4 chars/token, the standard ballpark — NOT provider-accurate. It
27+
* exists so usage-driven behaviour can be tested money-free, not to bill anyone.
28+
*/
29+
function estimateTokens(text: string): number {
30+
return text ? Math.max(1, Math.ceil(text.length / 4)) : 0;
31+
}
32+
33+
function estimateUsage(messages: ModelMessage[], output = ''): AIResult['usage'] {
34+
const promptText = messages
35+
.map(m => (typeof m.content === 'string' ? m.content : JSON.stringify(m.content)))
36+
.join('\n');
37+
const promptTokens = estimateTokens(promptText);
38+
const completionTokens = estimateTokens(output);
39+
return { promptTokens, completionTokens, totalTokens: promptTokens + completionTokens };
40+
}
41+
2142
export class MemoryLLMAdapter implements LLMAdapter {
2243
readonly name = 'memory';
2344

@@ -66,7 +87,7 @@ export class MemoryLLMAdapter implements LLMAdapter {
6687
return {
6788
content: '',
6889
model: options?.model ?? 'memory',
69-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
90+
usage: estimateUsage(messages),
7091
toolCalls: [
7192
{
7293
type: 'tool-call',
@@ -89,7 +110,7 @@ export class MemoryLLMAdapter implements LLMAdapter {
89110
return {
90111
content: '',
91112
model: options?.model ?? 'memory',
92-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
113+
usage: estimateUsage(messages),
93114
toolCalls: [
94115
{
95116
type: 'tool-call',
@@ -126,13 +147,13 @@ export class MemoryLLMAdapter implements LLMAdapter {
126147
return {
127148
content: `[memory] action ${payload.action ?? ''} failed: ${payload.error}`,
128149
model: options?.model ?? 'memory',
129-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
150+
usage: estimateUsage(messages),
130151
};
131152
}
132153
return {
133154
content: `[memory] ${payload.message ?? 'Action executed.'} (${payload.action ?? 'action'})`,
134155
model: options?.model ?? 'memory',
135-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
156+
usage: estimateUsage(messages),
136157
};
137158
}
138159

@@ -163,15 +184,15 @@ export class MemoryLLMAdapter implements LLMAdapter {
163184
return {
164185
content: `[memory] query_data failed: ${payload.error}`,
165186
model: options?.model ?? 'memory',
166-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
187+
usage: estimateUsage(messages),
167188
};
168189
}
169190
const records = payload.records ?? [];
170191
const count = payload.count ?? records.length;
171192
return {
172193
content: `[memory] Found ${count} record${count === 1 ? '' : 's'} for "${userText}".`,
173194
model: options?.model ?? 'memory',
174-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
195+
usage: estimateUsage(messages),
175196
};
176197
}
177198

@@ -182,15 +203,16 @@ export class MemoryLLMAdapter implements LLMAdapter {
182203
return {
183204
content,
184205
model: options?.model ?? 'memory',
185-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
206+
usage: estimateUsage(messages),
186207
};
187208
}
188209

189210
async complete(prompt: string, options?: AIRequestOptions): Promise<AIResult> {
211+
const content = `[memory] ${prompt}`;
190212
return {
191-
content: `[memory] ${prompt}`,
213+
content,
192214
model: options?.model ?? 'memory',
193-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
215+
usage: estimateUsage([{ role: 'user', content: prompt }], content),
194216
};
195217
}
196218

@@ -208,7 +230,7 @@ export class MemoryLLMAdapter implements LLMAdapter {
208230
yield {
209231
type: 'finish',
210232
finishReason: 'stop' as const,
211-
totalUsage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
233+
totalUsage: estimateUsage(messages, result.content),
212234
rawFinishReason: 'stop',
213235
} as unknown as TextStreamPart<ToolSet>;
214236
}
@@ -310,7 +332,7 @@ export class MemoryLLMAdapter implements LLMAdapter {
310332
return {
311333
object: result.data,
312334
model: options?.model ?? 'memory',
313-
usage: { promptTokens: 0, completionTokens: 0, totalTokens: 0 },
335+
usage: estimateUsage(messages),
314336
};
315337
}
316338
}

0 commit comments

Comments
 (0)