@theia/ai-core 1.60.1 → 1.61.0-next.8

package/src/common/token-usage-service.ts

@@ -0,0 +1,56 @@
+// *****************************************************************************
+// Copyright (C) 2025 EclipseSource GmbH.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0.
+//
+// This Source Code may also be made available under the following Secondary
+// Licenses when the conditions for such availability set forth in the Eclipse
+// Public License v. 2.0 are satisfied: GNU General Public License, version 2
+// with the GNU Classpath Exception which is available at
+// https://www.gnu.org/software/classpath/license.html.
+//
+// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
+// *****************************************************************************
+
+import { TokenUsageServiceClient } from './protocol';
+
+export const TokenUsageService = Symbol('TokenUsageService');
+
+export interface TokenUsage {
+    /** The input token count */
+    inputTokens: number;
+    /** The output token count */
+    outputTokens: number;
+    /** The model identifier */
+    model: string;
+    /** The timestamp of when the tokens were used */
+    timestamp: Date;
+    /** Request identifier */
+    requestId: string;
+}
+
+export interface TokenUsageParams {
+    /** The input token count */
+    inputTokens: number;
+    /** The output token count */
+    outputTokens: number;
+    /** Request identifier */
+    requestId: string;
+}
+
+export interface TokenUsageService {
+    /**
+     * Records token usage for a model interaction.
+     *
+     * @param model The identifier of the model that was used
+     * @param params Object containing token usage information
+     * @returns A promise that resolves when the token usage has been recorded
+     */
+    recordTokenUsage(model: string, params: TokenUsageParams): Promise<void>;
+
+    getTokenUsages(): Promise<TokenUsage[]>;
+
+    setClient(tokenUsageClient: TokenUsageServiceClient): void;
+}

package/src/node/ai-core-backend-module.ts

@@ -36,9 +36,13 @@ import {
     LanguageModelRegistryFrontendDelegate,
     languageModelDelegatePath,
     languageModelRegistryDelegatePath,
-    LanguageModelRegistryClient
+    LanguageModelRegistryClient,
+    TokenUsageService,
+    TokenUsageServiceClient,
+    TOKEN_USAGE_SERVICE_PATH
 } from '../common';
 import { BackendLanguageModelRegistry } from './backend-language-model-registry';
+import { TokenUsageServiceImpl } from './token-usage-service-impl';
 
 // We use a connection module to handle AI services separately for each frontend.
 const aiCoreConnectionModule = ConnectionContainerModule.create(({ bind, bindBackendService, bindFrontendService }) => {
@@ -46,6 +50,22 @@ const aiCoreConnectionModule = ConnectionContainerModule.create(({ bind, bindBac
     bind(BackendLanguageModelRegistry).toSelf().inSingletonScope();
     bind(LanguageModelRegistry).toService(BackendLanguageModelRegistry);
 
+    bind(TokenUsageService).to(TokenUsageServiceImpl).inSingletonScope();
+
+    bind(ConnectionHandler)
+        .toDynamicValue(
+            ({ container }) =>
+                new RpcConnectionHandler<TokenUsageServiceClient>(
+                    TOKEN_USAGE_SERVICE_PATH,
+                    client => {
+                        const service = container.get<TokenUsageService>(TokenUsageService);
+                        service.setClient(client);
+                        return service;
+                    }
+                )
+        )
+        .inSingletonScope();
+
     bind(LanguageModelRegistryFrontendDelegate).to(LanguageModelRegistryFrontendDelegateImpl).inSingletonScope();
     bind(ConnectionHandler)
         .toDynamicValue(

package/src/node/language-model-frontend-delegate.ts

@@ -19,7 +19,6 @@ import { CancellationToken, CancellationTokenSource, ILogger, generateUuid } fro
 import {
     LanguageModelMetaData,
     LanguageModelRegistry,
-    LanguageModelRequest,
     isLanguageModelStreamResponse,
     isLanguageModelTextResponse,
     LanguageModelStreamResponsePart,
@@ -29,6 +28,7 @@ import {
     LanguageModelResponseDelegate,
     LanguageModelRegistryClient,
     isLanguageModelParsedResponse,
+    UserRequest,
 } from '../common';
 import { BackendLanguageModelRegistry } from './backend-language-model-registry';
 
@@ -69,7 +69,7 @@ export class LanguageModelFrontendDelegateImpl implements LanguageModelFrontendD
 
     async request(
         modelId: string,
-        request: LanguageModelRequest,
+        request: UserRequest,
         requestId: string,
         cancellationToken?: CancellationToken
     ): Promise<LanguageModelResponseDelegate> {

package/src/node/token-usage-service-impl.ts

@@ -0,0 +1,65 @@
+// *****************************************************************************
+// Copyright (C) 2025 EclipseSource GmbH.
+//
+// This program and the accompanying materials are made available under the
+// terms of the Eclipse Public License v. 2.0 which is available at
+// http://www.eclipse.org/legal/epl-2.0.
+//
+// This Source Code may also be made available under the following Secondary
+// Licenses when the conditions for such availability set forth in the Eclipse
+// Public License v. 2.0 are satisfied: GNU General Public License, version 2
+// with the GNU Classpath Exception which is available at
+// https://www.gnu.org/software/classpath/license.html.
+//
+// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
+// *****************************************************************************
+
+import { injectable } from '@theia/core/shared/inversify';
+import { TokenUsage, TokenUsageParams, TokenUsageService } from '../common/token-usage-service';
+import { TokenUsageServiceClient } from '../common/protocol';
+
+@injectable()
+export class TokenUsageServiceImpl implements TokenUsageService {
+    private client: TokenUsageServiceClient | undefined;
+
+    /**
+     * Sets the client to notify about token usage changes
+     */
+    setClient(client: TokenUsageServiceClient | undefined): void {
+        this.client = client;
+    }
+
+    private readonly tokenUsages: TokenUsage[] = [];
+
+    /**
+     * Records token usage for a model interaction.
+     *
+     * @param model The model identifier
+     * @param params Token usage parameters
+     * @returns A promise that resolves when the token usage has been recorded
+     */
+    async recordTokenUsage(model: string, params: TokenUsageParams): Promise<void> {
+        const usage: TokenUsage = {
+            inputTokens: params.inputTokens,
+            outputTokens: params.outputTokens,
+            model,
+            timestamp: new Date(),
+            requestId: params.requestId
+        };
+
+        this.tokenUsages.push(usage);
+        this.client?.notifyTokenUsage(usage);
+
+        console.log(`Input Tokens: ${params.inputTokens}; Output Tokens: ${params.outputTokens}; Model: ${model}${params.requestId ? `; RequestId: ${params.requestId}` : ''}`);
+        // For now we just store in memory
+        // In the future, this could be persisted to disk, a database, or sent to a service
+        return Promise.resolve();
+    }
+
+    /**
+     * Gets all token usage records stored in this service.
+     */
+    async getTokenUsages(): Promise<TokenUsage[]> {
+        return [...this.tokenUsages];
+    }
+}