@theia/ai-core 1.63.0-next.0 → 1.63.0-next.52

package/src/browser/token-usage-frontend-service-impl.ts

@@ -76,6 +76,8 @@ export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService
         const modelMap = new Map<string, {
             inputTokens: number;
             outputTokens: number;
+            cachedInputTokens: number;
+            readCachedInputTokens: number;
             lastUsed?: Date;
         }>();
 
@@ -87,6 +89,16 @@ export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService
                 existing.inputTokens += usage.inputTokens;
                 existing.outputTokens += usage.outputTokens;
 
+                // Add cached tokens if they exist
+                if (usage.cachedInputTokens !== undefined) {
+                    existing.cachedInputTokens += usage.cachedInputTokens;
+                }
+
+                // Add read cached tokens if they exist
+                if (usage.readCachedInputTokens !== undefined) {
+                    existing.readCachedInputTokens += usage.readCachedInputTokens;
+                }
+
                 // Update last used if this usage is more recent
                 if (!existing.lastUsed || (usage.timestamp && usage.timestamp > existing.lastUsed)) {
                     existing.lastUsed = usage.timestamp;
@@ -95,6 +107,8 @@ export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService
                 modelMap.set(usage.model, {
                     inputTokens: usage.inputTokens,
                     outputTokens: usage.outputTokens,
+                    cachedInputTokens: usage.cachedInputTokens || 0,
+                    readCachedInputTokens: usage.readCachedInputTokens || 0,
                     lastUsed: usage.timestamp
                 });
             }
@@ -104,12 +118,23 @@ export class TokenUsageFrontendServiceImpl implements TokenUsageFrontendService
         const result: ModelTokenUsageData[] = [];
 
         for (const [modelId, data] of modelMap.entries()) {
-            result.push({
+            const modelData: ModelTokenUsageData = {
                 modelId,
                 inputTokens: data.inputTokens,
                 outputTokens: data.outputTokens,
                 lastUsed: data.lastUsed
-            });
+            };
+
+            // Only include cache-related fields if they have non-zero values
+            if (data.cachedInputTokens > 0) {
+                modelData.cachedInputTokens = data.cachedInputTokens;
+            }
+
+            if (data.readCachedInputTokens > 0) {
+                modelData.readCachedInputTokens = data.readCachedInputTokens;
+            }
+
+            result.push(modelData);
         }
 
         return result;

package/src/browser/token-usage-frontend-service.ts

@@ -26,6 +26,10 @@ export interface ModelTokenUsageData {
     inputTokens: number;
     /** Number of output tokens used */
     outputTokens: number;
+    /** Number of input tokens written to cache */
+    cachedInputTokens?: number;
+    /** Number of input tokens read from cache */
+    readCachedInputTokens?: number;
     /** Date when the model was last used */
     lastUsed?: Date;
 }

package/src/common/language-model.ts

@@ -19,7 +19,7 @@ import { inject, injectable, named, postConstruct } from '@theia/core/shared/inv
 
 export type MessageActor = 'user' | 'ai' | 'system';
 
-export type LanguageModelMessage = TextMessage | ThinkingMessage | ToolUseMessage | ToolResultMessage;
+export type LanguageModelMessage = TextMessage | ThinkingMessage | ToolUseMessage | ToolResultMessage | ImageMessage;
 export namespace LanguageModelMessage {
 
     export function isTextMessage(obj: LanguageModelMessage): obj is TextMessage {
@@ -34,6 +34,9 @@ export namespace LanguageModelMessage {
     export function isToolResultMessage(obj: LanguageModelMessage): obj is ToolResultMessage {
         return obj.type === 'tool_result';
     }
+    export function isImageMessage(obj: LanguageModelMessage): obj is ImageMessage {
+        return obj.type === 'image';
+    }
 }
 export interface TextMessage {
     actor: MessageActor;
@@ -63,6 +66,22 @@ export interface ToolUseMessage {
     input: unknown;
     name: string;
 }
+export type ImageMimeType = 'image/jpeg' | 'image/png' | 'image/gif' | 'image/webp' | 'image/bmp' | 'image/svg+xml' | string & {};
+export interface UrlImageContent { url: string };
+export interface Base64ImageContent {
+    base64data: string;
+    mimeType: ImageMimeType;
+};
+export type ImageContent = UrlImageContent | Base64ImageContent;
+export namespace ImageContent {
+    export const isUrl = (obj: ImageContent): obj is UrlImageContent => 'url' in obj;
+    export const isBase64 = (obj: ImageContent): obj is Base64ImageContent => 'base64data' in obj && 'mimeType' in obj;
+}
+export interface ImageMessage {
+    actor: 'ai' | 'user';
+    type: 'image';
+    image: ImageContent;
+}
 
 export const isLanguageModelRequestMessage = (obj: unknown): obj is LanguageModelMessage =>
     !!(obj && typeof obj === 'object' &&

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

@@ -23,6 +23,10 @@ export interface TokenUsage {
     inputTokens: number;
     /** The output token count */
     outputTokens: number;
+    /** Input tokens written to cache */
+    cachedInputTokens?: number;
+    /** Input tokens read from cache */
+    readCachedInputTokens?: number;
     /** The model identifier */
     model: string;
     /** The timestamp of when the tokens were used */
@@ -36,6 +40,10 @@ export interface TokenUsageParams {
     inputTokens: number;
     /** The output token count */
     outputTokens: number;
+    /** Input tokens placed in cache */
+    cachedInputTokens?: number;
+    /** Input tokens read from cache */
+    readCachedInputTokens?: number;
     /** Request identifier */
     requestId: string;
 }

package/src/common/tool-invocation-registry.ts

@@ -16,7 +16,7 @@
 
 import { inject, injectable, named, postConstruct, interfaces } from '@theia/core/shared/inversify';
 import { ToolRequest } from './language-model';
-import { ContributionProvider } from '@theia/core';
+import { ContributionProvider, Emitter, Event } from '@theia/core';
 
 export const ToolInvocationRegistry = Symbol('ToolInvocationRegistry');
 
@@ -62,6 +62,11 @@ export interface ToolInvocationRegistry {
      * @param providerName - The name of the tool provider whose tools should be removed (as specificed in the `ToolRequest`).
      */
     unregisterAllTools(providerName: string): void;
+
+    /**
+     * Event that is fired whenever the registry changes (tool registered or unregistered).
+     */
+    onDidChange: Event<void>;
 }
 
 export const ToolProvider = Symbol('ToolProvider');
@@ -80,6 +85,9 @@ export class ToolInvocationRegistryImpl implements ToolInvocationRegistry {
 
     private tools: Map<string, ToolRequest> = new Map<string, ToolRequest>();
 
+    private readonly onDidChangeEmitter = new Emitter<void>();
+    readonly onDidChange: Event<void> = this.onDidChangeEmitter.event;
+
     @inject(ContributionProvider)
     @named(ToolProvider)
     private providers: ContributionProvider<ToolProvider>;
@@ -98,7 +106,15 @@ export class ToolInvocationRegistryImpl implements ToolInvocationRegistry {
                 toolsToRemove.push(id);
             }
         }
-        toolsToRemove.forEach(id => this.tools.delete(id));
+        let changed = false;
+        toolsToRemove.forEach(id => {
+            if (this.tools.delete(id)) {
+                changed = true;
+            }
+        });
+        if (changed) {
+            this.onDidChangeEmitter.fire();
+        }
     }
     getAllFunctions(): ToolRequest[] {
         return Array.from(this.tools.values());
@@ -109,6 +125,7 @@ export class ToolInvocationRegistryImpl implements ToolInvocationRegistry {
             console.warn(`Function with id ${tool.id} is already registered.`);
         } else {
             this.tools.set(tool.id, tool);
+            this.onDidChangeEmitter.fire();
         }
     }
 
@@ -128,3 +145,4 @@ export class ToolInvocationRegistryImpl implements ToolInvocationRegistry {
         return tools;
     }
 }
+

package/src/common/variable-service.ts

@@ -172,6 +172,7 @@ export interface AIVariableService {
     getVariable(name: string): Readonly<AIVariable> | undefined;
     getVariables(): Readonly<AIVariable>[];
     getContextVariables(): Readonly<AIContextVariable>[];
+    registerVariable(variable: AIVariable): Disposable;
     unregisterVariable(name: string): void;
     readonly onDidChangeVariables: Event<void>;
 
@@ -287,12 +288,19 @@ export class DefaultAIVariableService implements AIVariableService {
         return this.getVariables().filter(AIContextVariable.is);
     }
 
-    registerResolver(variable: AIVariable, resolver: AIVariableResolver): Disposable {
+    registerVariable(variable: AIVariable): Disposable {
         const key = this.getKey(variable.name);
         if (!this.variables.get(key)) {
             this.variables.set(key, variable);
             this.onDidChangeVariablesEmitter.fire();
+            return Disposable.create(() => this.unregisterVariable(variable.name));
         }
+        return Disposable.NULL;
+    }
+
+    registerResolver(variable: AIVariable, resolver: AIVariableResolver): Disposable {
+        this.registerVariable(variable);
+        const key = this.getKey(variable.name);
         const resolvers = this.resolvers.get(key) ?? [];
         resolvers.push(resolver);
         this.resolvers.set(key, resolvers);
@@ -315,6 +323,7 @@ export class DefaultAIVariableService implements AIVariableService {
     }
 
     registerArgumentPicker(variable: AIVariable, argPicker: AIVariableArgPicker): Disposable {
+        this.registerVariable(variable);
         const key = this.getKey(variable.name);
         this.argPickers.set(key, argPicker);
         return Disposable.create(() => this.unregisterArgumentPicker(variable, argPicker));
@@ -333,6 +342,7 @@ export class DefaultAIVariableService implements AIVariableService {
     }
 
     registerArgumentCompletionProvider(variable: AIVariable, completionProvider: AIVariableArgCompletionProvider): Disposable {
+        this.registerVariable(variable);
         const key = this.getKey(variable.name);
         this.argCompletionProviders.set(key, completionProvider);
         return Disposable.create(() => this.unregisterArgumentCompletionProvider(variable, completionProvider));

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

@@ -41,6 +41,8 @@ export class TokenUsageServiceImpl implements TokenUsageService {
     async recordTokenUsage(model: string, params: TokenUsageParams): Promise<void> {
         const usage: TokenUsage = {
             inputTokens: params.inputTokens,
+            cachedInputTokens: params.cachedInputTokens,
+            readCachedInputTokens: params.readCachedInputTokens,
             outputTokens: params.outputTokens,
             model,
             timestamp: new Date(),
@@ -50,7 +52,23 @@ export class TokenUsageServiceImpl implements TokenUsageService {
         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}` : ''}`);
+        let logMessage = `Input Tokens: ${params.inputTokens};`;
+
+        if (params.cachedInputTokens) {
+            logMessage += ` Input Tokens written to cache: ${params.cachedInputTokens};`;
+        }
+
+        if (params.readCachedInputTokens) {
+            logMessage += ` Input Tokens read from cache: ${params.readCachedInputTokens};`;
+        }
+
+        logMessage += ` Output Tokens: ${params.outputTokens}; Model: ${model};`;
+
+        if (params.requestId) {
+            logMessage += `; RequestId: ${params.requestId}`;
+        }
+
+        console.debug(logMessage);
         // 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();