@celljs/ai-core 3.4.2 → 3.4.4

package/src/common/chat/metadata/metadata-protocol.ts

@@ -10,25 +10,29 @@ export interface Usage {
     /**
      * The number of tokens used in the prompt of the AI request.
      */
-    readonly promptTokens: number;
+    readonly promptTokens?: number;
 
     /**
      * The number of tokens returned in the generation (aka completion)
      * of the AI's response.
      */
-    readonly generationTokens: number;
+    readonly generationTokens?: number;
 
     /**
      * The total number of tokens from both the prompt of an AI request
      * and generation of the AI's response.
      */
-    readonly totalTokens: number;
+    readonly totalTokens?: number;
 
 }
 
 export namespace Usage {
-    export function from(promptTokens: number, generationTokens: number): Usage {
-        return { promptTokens, generationTokens, totalTokens: promptTokens + generationTokens };
+    export function createEmpty(): Usage {
+        return { promptTokens: 0, generationTokens: 0, totalTokens: 0 };
+    }
+
+    export function from(promptTokens = 0, generationTokens = 0, totalTokens = 0): Usage {
+        return { promptTokens, generationTokens, totalTokens: totalTokens ? totalTokens : promptTokens + generationTokens };
     }
 }
 
@@ -42,37 +46,37 @@ export interface RateLimit {
      * The maximum number of requests that are permitted before exhausting the
      * rate limit.
      */
-    readonly requestsLimit: number;
+    readonly requestsLimit?: number;
 
     /**
      * The remaining number of requests that are permitted before exhausting the
      * requestsLimit rate limit.
      */
-    readonly requestsRemaining: number;
+    readonly requestsRemaining?: number;
 
     /**
      * The duration time until the rate limit (based on requests) resets
      * to its requestsLimit initial state.
      */
-    readonly requestsReset: number;
+    readonly requestsReset?: number;
 
     /**
      * The maximum number of tokens that are permitted before exhausting the rate
      * limit.
      */
-    readonly tokensLimit: number;
+    readonly tokensLimit?: number;
 
     /**
      * The remaining number of tokens that are permitted before exhausting the
      * tokensLimit rate limit.
      */
-    readonly tokensRemaining: number;
+    readonly tokensRemaining?: number;
 
     /**
      * The duration time until the rate limit (based on tokens) resets to
      * its tokensLimit initial state.
      */
-    readonly tokensReset: number;
+    readonly tokensReset?: number;
 
 }
 
@@ -123,6 +127,12 @@ export namespace PromptMetadata {
     }
 }
 
+export namespace RateLimit {
+    export function createEmpty(): RateLimit {
+        return { requestsLimit: 0, requestsRemaining: 0, requestsReset: 0, tokensLimit: 0, tokensRemaining: 0, tokensReset: 0 };
+    }
+}
+
 export interface ChatResponseMetadata extends ResponseMetadata {
     id?: string;
     model?: string;
@@ -144,25 +154,7 @@ export interface ChatResponseMetadata extends ResponseMetadata {
 
 export class ChatResponseMetadataBuilder {
 
-    private readonly chatResponseMetadata: ChatResponseMetadata = {
-        id: '',
-        model: '',
-        rateLimit: {
-            requestsLimit: 0,
-            requestsRemaining: 0,
-            requestsReset: 0,
-            tokensLimit: 0,
-            tokensRemaining: 0,
-            tokensReset: 0
-        },
-        usage: {
-            promptTokens: 0,
-            generationTokens: 0,
-            totalTokens: 0
-        },
-        promptMetadata: [] as PromptMetadata,
-        extra: {}
-    };
+    private readonly chatResponseMetadata: ChatResponseMetadata = ChatResponseMetadata.createEmpty();
 
     id(id: string): ChatResponseMetadataBuilder {
         this.chatResponseMetadata.id = id;
@@ -207,25 +199,17 @@ export class ChatResponseMetadataBuilder {
 }
 
 export namespace ChatResponseMetadata {
-    export const EMPTY: ChatResponseMetadata = {
-        id: '',
-        model: '',
-        rateLimit: {
-            requestsLimit: 0,
-            requestsRemaining: 0,
-            requestsReset: 0,
-            tokensLimit: 0,
-            tokensRemaining: 0,
-            tokensReset: 0
-        },
-        usage: {
-            promptTokens: 0,
-            generationTokens: 0,
-            totalTokens: 0
-        },
-        promptMetadata: [] as PromptMetadata,
-        extra: {}
-    };
+
+    export function createEmpty(): ChatResponseMetadata {
+        return {
+            id: '',
+            model: '',
+            rateLimit: RateLimit.createEmpty(),
+            usage: Usage.createEmpty(),
+            promptMetadata: [] as PromptMetadata,
+            extra: {}
+        };
+    }
 
     export function builder(): ChatResponseMetadataBuilder {
         return new ChatResponseMetadataBuilder();

package/src/common/chat/model/chat-protocol.ts

@@ -10,23 +10,41 @@ export const ChatModel = Symbol('ChatModel');
  * Represents a response returned by the AI.
  */
 export interface Generation extends ModelResult<AssistantMessage> {
-
+    metadata: ChatGenerationMetadata;
 }
 
 export interface ChatResponse extends ModelResponse<Generation> {
-
+    metadata: ChatResponseMetadata;
 }
 
 export namespace Generation {
-    export function from(output: AssistantMessage, metadata = ChatGenerationMetadata.EMPTY): Generation {
+    export function from(output: AssistantMessage, metadata = ChatGenerationMetadata.from()): Generation {
         return { output, metadata };
     }
 }
 
 export namespace ChatResponse {
-    export function from(results: Generation[], metadata = ChatResponseMetadata.EMPTY): ChatResponse {
+    export function from(results: Generation[], metadata = ChatResponseMetadata.createEmpty()): ChatResponse {
         return { result: results[0], results, metadata };
     }
+
+    export function isToolCall(chatResponse: ChatResponse, toolCallFinishReasons: Set<string>): boolean {
+        if (!chatResponse) {
+            return false;
+        }
+
+        const generations = chatResponse.results;
+        if (!generations || generations.length === 0) {
+            return false;
+        }
+
+        return generations.some(g => doIsToolCall(g, toolCallFinishReasons));
+    }
+
+    function doIsToolCall(generation: Generation, toolCallFinishReasons: Set<string>): boolean {
+        const finishReason = generation.metadata.finishReason ?? '';
+        return generation.output.toolCalls?.length > 0 && toolCallFinishReasons.has(finishReason.toLowerCase());
+    }
 }
 
 export interface StreamingChatModel extends StreamingModel<Prompt, ChatResponse> {

package/src/common/chat/prompt/prompt-protocol.ts

@@ -30,5 +30,5 @@ export interface PromptTemplateContext {
 
 export interface PromptTemplate {
     render(template: string, ctx?: PromptTemplateContext): Promise<string>;
-    create(template: string, ctx?: PromptTemplateContext): Promise<Prompt>;
+    create(template: string | Message[], ctx?: PromptTemplateContext): Promise<Prompt>;
 }

package/src/common/chat/prompt/prompt-template.ts

@@ -1,7 +1,7 @@
 import { Autowired, Component, ExpressionHandler, IllegalArgumentError } from '@celljs/core';
 import { Prompt, PromptTemplate, PromptTemplateContext } from './prompt-protocol';
 import { PromptImpl } from './prompt';
-import { MessageType, SystemMessage } from '../message';
+import { Message, MessageType, SystemMessage } from '../message';
 
 @Component(PromptTemplate)
 export class PromptTemplateImpl implements PromptTemplate {
@@ -15,11 +15,17 @@ export class PromptTemplateImpl implements PromptTemplate {
         }
         return template;
     }
-    async create(template: string, ctx?: PromptTemplateContext): Promise<Prompt> {
-        if (!ctx?.messageType || ctx?.messageType === MessageType.USER) {
-            return new PromptImpl(await this.render(template, ctx), ctx?.chatOptions);
-        } else if (ctx?.messageType === MessageType.SYSTEM) {
-            return new PromptImpl(new SystemMessage(await this.render(template, ctx)), ctx?.chatOptions);
+    async create(template: string | Message[], ctx?: PromptTemplateContext): Promise<Prompt> {
+        if (Array.isArray(template)
+            && template.some(message => Message.isMessage(message))
+        ) {
+            return new PromptImpl(template, ctx?.chatOptions);
+        } else if (typeof template === 'string') {
+            if (!ctx?.messageType || ctx?.messageType === MessageType.USER) {
+                return new PromptImpl(await this.render(template, ctx), ctx?.chatOptions);
+            } else if (ctx?.messageType === MessageType.SYSTEM) {
+                return new PromptImpl(new SystemMessage(await this.render(template, ctx)), ctx?.chatOptions);
+            }
         }
         throw new IllegalArgumentError(`Invalid message type: ${ctx?.messageType}`);
     }

package/src/common/chat/tool/index.ts

@@ -0,0 +1,2 @@
+export * from './tool-handler';
+export * from './tool-protocol';

package/src/common/chat/tool/tool-handler.ts

@@ -0,0 +1,51 @@
+import { Autowired, Component, IllegalStateError } from '@celljs/core';
+import { FunctionCallbackRegister, FunctionCallingOptions, ToolContext } from '../../model';
+import { ChatResponse } from '../model';
+import { Prompt } from '../prompt';
+import { ToolHandler } from './tool-protocol';
+import { AssistantMessage, Message, ToolResponse, ToolResponseMessage } from '../message';
+
+@Component(ToolHandler)
+export class ToolHanderImpl implements ToolHandler {
+
+    @Autowired(FunctionCallbackRegister)
+    protected readonly functionCallbackRegister: FunctionCallbackRegister;
+
+    protected async executeFunctions(assistantMessage: AssistantMessage, toolContext: ToolContext): Promise<ToolResponseMessage> {
+        const toolResponses: ToolResponse[] = [];
+        for (const toolCall of assistantMessage.toolCalls) {
+            const functionResponse = await this.functionCallbackRegister.call(toolCall.name, toolCall.arguments, toolContext);
+            toolResponses.push({
+                id: toolCall.id,
+                name: toolCall.name,
+                responseData: functionResponse
+            });
+        }
+        return new ToolResponseMessage(toolResponses);
+    }
+
+    protected buildToolCallConversation(previousMessages: Message[], assistantMessage: AssistantMessage, toolResponseMessage: ToolResponseMessage): Message[] {
+        const messages = [...previousMessages];
+        messages.push(assistantMessage);
+        messages.push(toolResponseMessage);
+        return messages;
+    }
+
+    async handle(prompt: Prompt, chatResponse: ChatResponse): Promise<Message[]> {
+        const toolCallGeneration = chatResponse.results.find(result => result.output.toolCalls?.length > 0);
+        if (!toolCallGeneration) {
+            throw new IllegalStateError('No tool call generation found in the response!');
+        }
+        const assistantMessage = toolCallGeneration.output;
+        let toolContext: ToolContext = {};
+        const options = prompt.options;
+        if (FunctionCallingOptions.is(options) && options.toolContext) {
+            toolContext = { ...options.toolContext };
+            const toolCallHistory = prompt.copy().instructions;
+            toolContext[ToolContext.TOOL_CALL_HISTORY] = toolCallHistory;
+        }
+        const toolResponseMessage = await this.executeFunctions(assistantMessage, toolContext);
+        const toolConversationHistory = this.buildToolCallConversation(prompt.instructions, assistantMessage, toolResponseMessage);
+        return toolConversationHistory;
+    }
+}

package/src/common/chat/tool/tool-protocol.ts

@@ -0,0 +1,9 @@
+import { Message } from '../message';
+import { ChatResponse } from '../model';
+import { Prompt } from '../prompt';
+
+export const ToolHandler = Symbol('ToolHandler');
+
+export interface ToolHandler {
+    handle(prompt: Prompt, chatResponse: ChatResponse): Promise<Message[]>;
+}

package/src/common/error/error.ts

@@ -5,3 +5,9 @@ export class AIError extends CustomError {
         super(message);
     }
 }
+
+export class NotFoundError extends AIError {
+    constructor(message: string) {
+        super(message);
+    }
+}

package/src/common/model/function/function-callback-register.ts

@@ -0,0 +1,38 @@
+import { Component } from '@celljs/core';
+import { FunctionCallback, FunctionCallbackRegister, ToolContext } from './function-protocol';
+import { NotFoundError } from '../../error';
+
+@Component(FunctionCallbackRegister)
+export class FunctionCallbackRegisterImpl implements FunctionCallbackRegister {
+
+    private functionCallbacks: Map<string, FunctionCallback> = new Map<string, FunctionCallback>();
+
+    register(functionCallback: FunctionCallback): void {
+        this.functionCallbacks.set(functionCallback.name, functionCallback);
+    }
+
+    unregister(functionName: string): void {
+        this.functionCallbacks.delete(functionName);
+    }
+
+    call(functionName: string, functionArguments: string, toolContext: ToolContext): Promise<string> {
+        const functionCallback = this.functionCallbacks.get(functionName);
+        if (functionCallback) {
+            return functionCallback.call(functionArguments, toolContext);
+        }
+        throw new NotFoundError(`Function ${functionName} not found`);
+    }
+
+    resolve(functionNames: string[]): FunctionCallback[] {
+        const functionCallbacks: FunctionCallback[] = [];
+        new Set(functionNames).forEach(functionName => {
+            const functionCallback = this.functionCallbacks.get(functionName);
+            if (functionCallback) {
+                functionCallbacks.push(functionCallback);
+            } else {
+                throw new NotFoundError(`No function callback found for name: ${functionName}`);
+            }
+        });
+        return functionCallbacks;
+    }
+}

package/src/common/model/function/function-protocol.ts

@@ -1,5 +1,13 @@
 import { ChatOptions } from '../../chat/prompt';
 
+export const FunctionCallbackRegister = Symbol('FunctionCallbackRegister');
+
+export interface ToolContext extends Record<string, any> {}
+
+export namespace ToolContext {
+    export const TOOL_CALL_HISTORY = 'TOOL_CALL_HISTORY';
+}
+
 /**
  * Represents a model function call handler. Implementations are registered with the
  * Models and called on prompts that trigger the function call.
@@ -28,9 +36,10 @@ export interface FunctionCallback {
      * @param functionArguments JSON string with the function arguments to be passed to the
      * function. The arguments are defined as JSON schema usually registered with the the
      * model.
+     * @param toolContext Optional tool context to be passed to the function call.
      * @return String containing the function call response.
      */
-    call(functionArguments: string): Promise<string>;
+    call(functionArguments: string, toolContext?: ToolContext): Promise<string>;
 
 }
 
@@ -49,13 +58,61 @@ export interface FunctionCallingOptions {
      * List of function names from the ChatClient registry to be used in the next
      * chat completion requests.
      */
-    functions: Set<String>;
+    functions: Set<string>;
+
+    proxyToolCalls?: boolean;
+
+    toolContext?: ToolContext;
 }
 
 export namespace FunctionCallingOptions {
     export function is(obj: any): obj is FunctionCallingOptions {
         return obj && obj.functionCallbacks !== undefined && obj.functions !== undefined;
     }
+
+    export function getEnabledFunctionsToCall(runtimeFunctionOptions: FunctionCallingOptions): Set<string> {
+        const enabledFunctionsToCall = new Set<string>();
+        if (runtimeFunctionOptions) {
+            // Add the explicitly enabled functions.
+            if (runtimeFunctionOptions.functions) {
+                runtimeFunctionOptions.functions.forEach(fn => {
+                    enabledFunctionsToCall.add(fn);
+                });
+            }
+
+            if (runtimeFunctionOptions.functionCallbacks) {
+                runtimeFunctionOptions.functionCallbacks.forEach(functionCallback => {
+                    // Automatically enable the function, usually from prompt callback.
+                    enabledFunctionsToCall.add(functionCallback.name);
+                });
+            }
+        }
+        return enabledFunctionsToCall;
+    }
+
+    /**
+     * Check if the proxyToolCalls is enabled for the given prompt or the default tool
+     * call options. The prompt options take precedence over the default options. When the
+     * proxyToolCalls is enabled the ChatModel implementation will pass the proxyToolCalls to the
+     * function call.
+     * @param chatOptions the chat options to check.
+     * @param defaultOptions the default tool call options to check.
+     * @return true if the proxyToolCalls is enabled, false otherwise.
+     */
+    export function isProxyToolCalls(chatOptions: ChatOptions, defaultOptions: FunctionCallingOptions): boolean {
+        if (FunctionCallingOptions.is(chatOptions) && chatOptions.proxyToolCalls !== undefined) {
+            return chatOptions.proxyToolCalls;
+        }
+        return defaultOptions.proxyToolCalls ?? false;
+    }
+
+}
+
+export interface FunctionCallbackRegister {
+    resolve(functionNames: string[]): FunctionCallback[];
+    register(functionCallback: FunctionCallback): void;
+    unregister(functionName: string): void;
+    call(functionName: string, functionArguments: string, toolContext: ToolContext): Promise<string>;
 }
 
 export interface PortableFunctionCallingOptions extends FunctionCallingOptions, ChatOptions {

package/src/common/model/function/index.ts

@@ -1 +1,2 @@
 export * from './function-protocol';
+export * from './function-callback-register';

package/src/common/utils/index.ts

@@ -1,2 +1,3 @@
 export * from './prompt-util';
 export * from './sse-util';
+export * from './usage-util';

package/src/common/utils/sse-util.ts

@@ -5,7 +5,7 @@ import { Bytes } from '@celljs/core';
 /**
  * SSE utility class.
  */
-export class SSEUtil {
+export abstract class SSEUtil {
 
     static toObservable<Data>(readableStream: ReadableStream, controller?: AbortController): Observable<StreamEvent<Data>> {
         return new Observable<StreamEvent<Data>>(subscriber => {

package/src/common/utils/usage-util.ts

@@ -0,0 +1,60 @@
+import { ChatResponse } from '../chat';
+import { Usage } from '../chat/metadata';
+
+/**
+ * An utility class to provide support methods handling {@link Usage}.
+ */
+export abstract class UsageUtil {
+
+    /**
+     * Accumulate usage tokens from the previous chat response to the current usage
+     * tokens.
+     * @param currentUsage the current usage.
+     * @param previousChatResponse the previous chat response.
+     * @return accumulated usage.
+     */
+    public static getCumulativeUsage(currentUsage: Usage, previousChatResponse?: ChatResponse): Usage {
+        let usageFromPreviousChatResponse: Usage | undefined = undefined;
+        if (previousChatResponse?.metadata?.usage) {
+            usageFromPreviousChatResponse = previousChatResponse.metadata.usage;
+        } else {
+            return currentUsage;
+        }
+        // For a valid usage from previous chat response, accumulate it to the current
+        // usage.
+        if (!this.isEmpty(currentUsage)) {
+            let promptTokens = currentUsage.promptTokens ?? 0;
+            let generationTokens = currentUsage.generationTokens ?? 0;
+            let totalTokens = currentUsage.totalTokens ?? 0;
+            // Make sure to accumulate the usage from the previous chat response.
+            promptTokens += usageFromPreviousChatResponse.promptTokens ?? 0;
+            generationTokens += usageFromPreviousChatResponse.generationTokens ?? 0;
+            totalTokens += usageFromPreviousChatResponse.totalTokens ?? 0;
+            return Usage.from(promptTokens, generationTokens, totalTokens);
+        }
+        // When current usage is empty, return the usage from the previous chat response.
+        return usageFromPreviousChatResponse;
+    }
+
+    // /**
+    //  * Check if the {@link Usage} is empty. Returns true when the {@link Usage} is null.
+    //  * Returns true when the {@link Usage} has zero tokens.
+    //  * @param usage the usage to check against.
+    //  * @return the boolean value to represent if it is empty.
+    //  */
+    // static isEmpty(usage: Usage | null): boolean {
+    //     if (!usage) {
+    //         return true;
+    //     } else if (usage != null && usage.totalTokens === 0) {
+    //         return true;
+    //     }
+    //     return false;
+    // }
+    static isEmpty(usage: Usage | null): boolean {
+        if (!usage) {
+            return true;
+        }
+
+        return false;
+    }
+}