Skip to content

Commit 8d34507

Browse files
eneufeldeneufeld
committed
Add initial support for tracking token usage
1 parent 53a070b commit 8d34507

26 files changed

+801
-47
lines changed

‎package-lock.json

+11
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/ai-anthropic/src/node/anthropic-language-model.ts

+35-6
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,10 @@ import {
2121
LanguageModelResponse,
2222
LanguageModelStreamResponse,
2323
LanguageModelStreamResponsePart,
24-
LanguageModelTextResponse
24+
LanguageModelTextResponse,
25+
TokenUsageService,
26+
TokenUsageParams,
27+
UserRequest
2528
} from '@theia/ai-core';
2629
import { CancellationToken, isArray } from '@theia/core';
2730
import { Anthropic } from '@anthropic-ai/sdk';
@@ -100,14 +103,15 @@ export class AnthropicModel implements LanguageModel {
100103
public model: string,
101104
public enableStreaming: boolean,
102105
public apiKey: () => string | undefined,
103-
public maxTokens: number = DEFAULT_MAX_TOKENS
106+
public maxTokens: number = DEFAULT_MAX_TOKENS,
107+
protected readonly tokenUsageService?: TokenUsageService
104108
) { }
105109

106110
protected getSettings(request: LanguageModelRequest): Readonly<Record<string, unknown>> {
107111
return request.settings ?? {};
108112
}
109113

110-
async request(request: LanguageModelRequest, cancellationToken?: CancellationToken): Promise<LanguageModelResponse> {
114+
async request(request: UserRequest, cancellationToken?: CancellationToken): Promise<LanguageModelResponse> {
111115
if (!request.messages?.length) {
112116
throw new Error('Request must contain at least one message');
113117
}
@@ -144,7 +148,7 @@ export class AnthropicModel implements LanguageModel {
144148

145149
protected async handleStreamingRequest(
146150
anthropic: Anthropic,
147-
request: LanguageModelRequest,
151+
request: UserRequest,
148152
cancellationToken?: CancellationToken,
149153
toolMessages?: readonly Anthropic.Messages.MessageParam[]
150154
): Promise<LanguageModelStreamResponse> {
@@ -173,6 +177,7 @@ export class AnthropicModel implements LanguageModel {
173177
const toolCalls: ToolCallback[] = [];
174178
let toolCall: ToolCallback | undefined;
175179
const currentMessages: Message[] = [];
180+
let currentMessage: Message | undefined = undefined;
176181

177182
for await (const event of stream) {
178183
if (event.type === 'content_block_start') {
@@ -217,6 +222,21 @@ export class AnthropicModel implements LanguageModel {
217222
}
218223
} else if (event.type === 'message_start') {
219224
currentMessages.push(event.message);
225+
currentMessage = event.message;
226+
} else if (event.type === 'message_stop') {
227+
if (currentMessage) {
228+
yield { input_tokens: currentMessage.usage.input_tokens, output_tokens: currentMessage.usage.output_tokens };
229+
// Record token usage if token usage service is available
230+
if (that.tokenUsageService && currentMessage.usage) {
231+
const tokenUsageParams: TokenUsageParams = {
232+
inputTokens: currentMessage.usage.input_tokens,
233+
outputTokens: currentMessage.usage.output_tokens,
234+
requestId: request.requestId
235+
};
236+
await that.tokenUsageService.recordTokenUsage(that.id, tokenUsageParams);
237+
}
238+
}
239+
220240
}
221241
}
222242
if (toolCalls.length > 0) {
@@ -278,10 +298,9 @@ export class AnthropicModel implements LanguageModel {
278298

279299
protected async handleNonStreamingRequest(
280300
anthropic: Anthropic,
281-
request: LanguageModelRequest
301+
request: UserRequest
282302
): Promise<LanguageModelTextResponse> {
283303
const settings = this.getSettings(request);
284-
285304
const { messages, systemMessage } = transformToAnthropicParams(request.messages);
286305

287306
const params: Anthropic.MessageCreateParams = {
@@ -296,6 +315,16 @@ export class AnthropicModel implements LanguageModel {
296315
const response = await anthropic.messages.create(params);
297316
const textContent = response.content[0];
298317

318+
// Record token usage if token usage service is available
319+
if (this.tokenUsageService && response.usage) {
320+
const tokenUsageParams: TokenUsageParams = {
321+
inputTokens: response.usage.input_tokens,
322+
outputTokens: response.usage.output_tokens,
323+
requestId: request.requestId
324+
};
325+
await this.tokenUsageService.recordTokenUsage(this.id, tokenUsageParams);
326+
}
327+
299328
if (textContent?.type === 'text') {
300329
return { text: textContent.text };
301330
}

‎packages/ai-anthropic/src/node/anthropic-language-models-manager-impl.ts

+6-2
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
1515
// *****************************************************************************
1616

17-
import { LanguageModelRegistry } from '@theia/ai-core';
17+
import { LanguageModelRegistry, TokenUsageService } from '@theia/ai-core';
1818
import { inject, injectable } from '@theia/core/shared/inversify';
1919
import { AnthropicModel, DEFAULT_MAX_TOKENS } from './anthropic-language-model';
2020
import { AnthropicLanguageModelsManager, AnthropicModelDescription } from '../common';
@@ -27,6 +27,9 @@ export class AnthropicLanguageModelsManagerImpl implements AnthropicLanguageMode
2727
@inject(LanguageModelRegistry)
2828
protected readonly languageModelRegistry: LanguageModelRegistry;
2929

30+
@inject(TokenUsageService)
31+
protected readonly tokenUsageService: TokenUsageService;
32+
3033
get apiKey(): string | undefined {
3134
return this._apiKey ?? process.env.ANTHROPIC_API_KEY;
3235
}
@@ -64,7 +67,8 @@ export class AnthropicLanguageModelsManagerImpl implements AnthropicLanguageMode
6467
modelDescription.model,
6568
modelDescription.enableStreaming,
6669
apiKeyProvider,
67-
modelDescription.maxTokens
70+
modelDescription.maxTokens,
71+
this.tokenUsageService
6872
)
6973
]);
7074
}

‎packages/ai-chat/src/common/chat-agents.ts

+4
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ import {
2727
isTextResponsePart,
2828
isThinkingResponsePart,
2929
isToolCallResponsePart,
30+
isUsageResponsePart,
3031
LanguageModel,
3132
LanguageModelMessage,
3233
LanguageModelRequirement,
@@ -442,6 +443,9 @@ export abstract class AbstractStreamParsingChatAgent extends AbstractChatAgent {
442443
if (isThinkingResponsePart(token)) {
443444
return new ThinkingChatResponseContentImpl(token.thought, token.signature);
444445
}
446+
if (isUsageResponsePart(token)) {
447+
return [];
448+
}
445449
return this.defaultContentFactory.create('', request);
446450
}
447451

‎packages/ai-chat/src/common/chat-session-naming-service.ts

+10-13
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,10 @@ import {
2020
CommunicationRecordingService,
2121
getTextOfResponse,
2222
LanguageModelRegistry,
23-
LanguageModelRequest,
2423
LanguageModelRequirement,
2524
PromptService,
26-
PromptTemplate
25+
PromptTemplate,
26+
UserRequest
2727
} from '@theia/ai-core';
2828
import { inject, injectable } from '@theia/core/shared/inversify';
2929
import { ChatSession } from './chat-service';
@@ -103,22 +103,19 @@ export class ChatSessionNamingAgent implements Agent {
103103
throw new Error('Unable to create prompt message for generating chat session name');
104104
}
105105

106-
const request: LanguageModelRequest = {
106+
const sessionId = generateUuid();
107+
const requestId = generateUuid();
108+
const request: UserRequest = {
107109
messages: [{
108110
actor: 'user',
109111
text: message,
110112
type: 'text'
111-
}]
112-
};
113-
114-
const sessionId = generateUuid();
115-
const requestId = generateUuid();
116-
this.recordingService.recordRequest({
117-
agentId: this.id,
118-
sessionId,
113+
}],
119114
requestId,
120-
...request
121-
});
115+
sessionId,
116+
agentId: this.id
117+
};
118+
this.recordingService.recordRequest(request);
122119

123120
const result = await lm.request(request);
124121
const response = await getTextOfResponse(result);

‎packages/ai-core/src/browser/ai-core-frontend-module.ts

+15-1
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,10 @@ import {
3535
PromptCustomizationService,
3636
PromptService,
3737
PromptServiceImpl,
38-
ToolProvider
38+
ToolProvider,
39+
TokenUsageService,
40+
TOKEN_USAGE_SERVICE_PATH,
41+
TokenUsageServiceClient
3942
} from '../common';
4043
import {
4144
FrontendLanguageModelRegistryImpl,
@@ -65,6 +68,8 @@ import { AiCoreCommandContribution } from './ai-core-command-contribution';
6568
import { PromptVariableContribution } from '../common/prompt-variable-contribution';
6669
import { LanguageModelService } from '../common/language-model-service';
6770
import { FrontendLanguageModelServiceImpl } from './frontend-language-model-service';
71+
import { TokenUsageFrontendService } from './token-usage-frontend-service';
72+
import { TokenUsageFrontendServiceImpl, TokenUsageServiceClientImpl } from './token-usage-frontend-service-impl';
6873

6974
export default new ContainerModule(bind => {
7075
bindContributionProvider(bind, LanguageModelProvider);
@@ -144,4 +149,13 @@ export default new ContainerModule(bind => {
144149
bind(CommandContribution).toService(AiCoreCommandContribution);
145150
bind(FrontendLanguageModelServiceImpl).toSelf().inSingletonScope();
146151
bind(LanguageModelService).toService(FrontendLanguageModelServiceImpl);
152+
153+
bind(TokenUsageFrontendService).to(TokenUsageFrontendServiceImpl).inSingletonScope();
154+
bind(TokenUsageServiceClient).to(TokenUsageServiceClientImpl).inSingletonScope();
155+
156+
bind(TokenUsageService).toDynamicValue(ctx => {
157+
const connection = ctx.container.get<ServiceConnectionProvider>(RemoteConnectionProvider);
158+
const client = ctx.container.get<TokenUsageServiceClient>(TokenUsageServiceClient);
159+
return connection.createProxy<TokenUsageService>(TOKEN_USAGE_SERVICE_PATH, client);
160+
}).inSingletonScope();
147161
});

‎packages/ai-core/src/browser/token-usage-frontend-service-impl.ts

+117
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
// *****************************************************************************
2+
// Copyright (C) 2025 EclipseSource GmbH.
3+
//
4+
// This program and the accompanying materials are made available under the
5+
// terms of the Eclipse Public License v. 2.0 which is available at
6+
// http://www.eclipse.org/legal/epl-2.0.
7+
//
8+
// This Source Code may also be made available under the following Secondary
9+
// Licenses when the conditions for such availability set forth in the Eclipse
10+
// Public License v. 2.0 are satisfied: GNU General Public License, version 2
11+
// with the GNU Classpath Exception which is available at
12+
// https://www.gnu.org/software/classpath/license.html.
13+
//
14+
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
15+
// *****************************************************************************
16+
17+
import { injectable, inject, postConstruct } from '@theia/core/shared/inversify';
18+
import { Emitter } from '@theia/core';
19+
import { ModelTokenUsageData, TokenUsageFrontendService } from './token-usage-frontend-service';
20+
import { TokenUsage, TokenUsageService } from '../common/token-usage-service';
21+
import { TokenUsageServiceClient } from '../common/protocol';
22+
23+
@injectable()
24+
export class TokenUsageServiceClientImpl implements TokenUsageServiceClient {
25+
private readonly _onTokenUsageUpdated = new Emitter<TokenUsage>();
26+
readonly onTokenUsageUpdated = this._onTokenUsageUpdated.event;
27+
28+
notifyTokenUsage(usage: TokenUsage): void {
29+
this._onTokenUsageUpdated.fire(usage);
30+
}
31+
32+
}
33+
34+
@injectable()
35+
export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService {
36+
37+
@inject(TokenUsageServiceClient)
38+
protected readonly tokenUsageServiceClient: TokenUsageServiceClient;
39+
40+
@inject(TokenUsageService)
41+
protected readonly tokenUsageService: TokenUsageService;
42+
43+
private readonly _onTokenUsageUpdated = new Emitter<ModelTokenUsageData[]>();
44+
readonly onTokenUsageUpdated = this._onTokenUsageUpdated.event;
45+
46+
private cachedUsageData: ModelTokenUsageData[] = [];
47+
48+
@postConstruct()
49+
protected init(): void {
50+
this.tokenUsageServiceClient.onTokenUsageUpdated(() => {
51+
this.getTokenUsageData().then(data => {
52+
this._onTokenUsageUpdated.fire(data);
53+
});
54+
});
55+
}
56+
57+
/**
58+
* Gets the current token usage data for all models
59+
*/
60+
async getTokenUsageData(): Promise<ModelTokenUsageData[]> {
61+
try {
62+
const usages = await this.tokenUsageService.getTokenUsages();
63+
this.cachedUsageData = this.aggregateTokenUsages(usages);
64+
return this.cachedUsageData;
65+
} catch (error) {
66+
console.error('Failed to get token usage data:', error);
67+
return [];
68+
}
69+
}
70+
71+
/**
72+
* Aggregates token usages by model
73+
*/
74+
private aggregateTokenUsages(usages: TokenUsage[]): ModelTokenUsageData[] {
75+
// Group by model
76+
const modelMap = new Map<string, {
77+
inputTokens: number;
78+
outputTokens: number;
79+
lastUsed?: Date;
80+
}>();
81+
82+
// Process each usage record
83+
for (const usage of usages) {
84+
const existing = modelMap.get(usage.model);
85+
86+
if (existing) {
87+
existing.inputTokens += usage.inputTokens;
88+
existing.outputTokens += usage.outputTokens;
89+
90+
// Update last used if this usage is more recent
91+
if (!existing.lastUsed || (usage.timestamp && usage.timestamp > existing.lastUsed)) {
92+
existing.lastUsed = usage.timestamp;
93+
}
94+
} else {
95+
modelMap.set(usage.model, {
96+
inputTokens: usage.inputTokens,
97+
outputTokens: usage.outputTokens,
98+
lastUsed: usage.timestamp
99+
});
100+
}
101+
}
102+
103+
// Convert map to array of model usage data
104+
const result: ModelTokenUsageData[] = [];
105+
106+
for (const [modelId, data] of modelMap.entries()) {
107+
result.push({
108+
modelId,
109+
inputTokens: data.inputTokens,
110+
outputTokens: data.outputTokens,
111+
lastUsed: data.lastUsed
112+
});
113+
}
114+
115+
return result;
116+
}
117+
}

0 commit comments

Comments
 (0)