Tool calls typescript

bindings/typescript/examples/toolCalls.ts

@@ -0,0 +1,82 @@
+import {
+    ChatConfig,
+    ChatMessage,
+    ChatReplyConfig,
+    Engine,
+    EngineConfig,
+    SamplingMethodGreedy,
+    SamplingPolicyCustom,
+    uzuToolFunction,
+} from '@trymirai/uzu';
+import * as z from 'zod';
+
+
+const Coordinate = z.object({
+    latitude: z.number().describe('Latitude in decimal degrees.'),
+    longitude: z.number().describe('Longitude in decimal degrees.'),
+});
+
+type Coordinate = z.infer<typeof Coordinate>;
+
+
+const getCurrentLocation = uzuToolFunction({
+    name: 'get_location',
+    description: 'Return the current location in coordinates',
+    parameters: z.object({}),
+    returns: Coordinate,
+    handler: (): Coordinate => ({
+        latitude: 51.5074,
+        longitude: -0.1278,
+    }),
+});
+
+
+async function calculateCurrentTemperature({latitude, longitude}: Coordinate): Promise<number> {
+    if (!Number.isFinite(Math.hypot(latitude, longitude))) {
+        throw new RangeError('Coordinates must be finite');
+    }
+    return 25;
+}
+
+const getCurrentTemperature = uzuToolFunction({
+    name: 'get_current_temperature',
+    description: 'Return the temperature at the provided coordinates',
+    parameters: Coordinate,
+    returns: z.number(),
+    handler: calculateCurrentTemperature,
+});
+
+
+async function main() {
+    const engine = await Engine.create(EngineConfig.create());
+    const model = await engine.model('mlx-community/Qwen3.5-9B-MLX-8bit');
+    if (!model) {
+        throw new Error('Model not found');
+    }
+
+    for await (const update of await engine.download(model)) {
+        console.log('Download progress:', update.progress);
+    }
+
+    const session = await engine.chat(model, ChatConfig.create());
+    await session.addTool(getCurrentLocation);
+    await session.addTool(getCurrentTemperature);
+
+    const messages = [
+        ChatMessage.system().withText('You are a helpful assistant'),
+        ChatMessage.user().withText('What temperature is it now at my location?'),
+    ];
+    const config = ChatReplyConfig.create().withSamplingPolicy(
+        new SamplingPolicyCustom(new SamplingMethodGreedy()),
+    );
+    const replies = await session.reply(messages, config);
+    const message = replies[replies.length - 1]?.message;
+    if (message) {
+        console.log('Reasoning:', message.reasoning ?? '');
+        console.log('Text:', message.text ?? '');
+    }
+}
+
+main().catch((error: unknown) => {
+    console.error(error);
+});

bindings/typescript/package.json

@@ -54,7 +54,7 @@
         "publint": "^0.2.12",
         "ts-jest": "^29.1.0",
         "ts-node": "^10.5.0",
-        "tsc-multi": "https://github.com/stainless-api/tsc-multi/releases/download/v1.1.9/tsc-multi.tgz",
+        "tsc-multi": "https://github.com/stainless-api/tsc-multi/releases/download/v1.1.11/tsc-multi.tgz",
         "tsconfig-paths": "^4.0.0",
         "tslib": "^2.8.1",
         "typescript": "5.8.3",

bindings/typescript/src/index.ts

@@ -1 +1,2 @@
 export * from './napi/index';
+export * from './tool';

bindings/typescript/src/napi/index.d.mts

@@ -1,6 +1,8 @@
 /* auto-generated by NAPI-RS */
 /* eslint-disable */
 export declare class ChatSession {
+  addTool(tool: NativeTool): Promise<undefined>
+  addTools(tools: Array<NativeTool>): Promise<undefined>
   get state(): Promise<ChatSessionState>
   get messages(): Promise<Array<ChatMessage>>
   reset(): Promise<void>
@@ -35,6 +37,10 @@ export declare class ClassificationSession {
   classify(input: Array<ClassificationMessage>): Promise<ClassificationOutput>
 }
 
+export declare class NativeTool {
+  constructor(definition: ToolFunction, invokeJson: (arg0: string, arg1: string) => Promise<string>, cancel: (arg: string) => void)
+}
+
 export declare class TextToSpeechSession {
   get state(): Promise<TextToSpeechSessionState>
   synthesize(input: string): Promise<TextToSpeechOutput>

bindings/typescript/src/napi/index.d.ts

@@ -1,6 +1,8 @@
 /* auto-generated by NAPI-RS */
 /* eslint-disable */
 export declare class ChatSession {
+  addTool(tool: NativeTool): Promise<undefined>
+  addTools(tools: Array<NativeTool>): Promise<undefined>
   get state(): Promise<ChatSessionState>
   get messages(): Promise<Array<ChatMessage>>
   reset(): Promise<void>
@@ -35,6 +37,10 @@ export declare class ClassificationSession {
   classify(input: Array<ClassificationMessage>): Promise<ClassificationOutput>
 }
 
+export declare class NativeTool {
+  constructor(definition: ToolFunction, invokeJson: (arg0: string, arg1: string) => Promise<string>, cancel: (arg: string) => void)
+}
+
 export declare class TextToSpeechSession {
   get state(): Promise<TextToSpeechSessionState>
   synthesize(input: string): Promise<TextToSpeechOutput>

bindings/typescript/src/napi/index.js

@@ -513,6 +513,7 @@ module.exports.ChatSessionStream = nativeBinding.ChatSessionStream
 module.exports.ChatSessionStreamChunkError = nativeBinding.ChatSessionStreamChunkError
 module.exports.ChatSessionStreamChunkReplies = nativeBinding.ChatSessionStreamChunkReplies
 module.exports.ClassificationSession = nativeBinding.ClassificationSession
+module.exports.NativeTool = nativeBinding.NativeTool
 module.exports.TextToSpeechSession = nativeBinding.TextToSpeechSession
 module.exports.TextToSpeechSessionStream = nativeBinding.TextToSpeechSessionStream
 module.exports.TextToSpeechSessionStreamChunkError = nativeBinding.TextToSpeechSessionStreamChunkError

bindings/typescript/src/napi/index.mjs

@@ -10,6 +10,7 @@ export const ChatSessionStream = cjs.ChatSessionStream;
 export const ChatSessionStreamChunkError = cjs.ChatSessionStreamChunkError;
 export const ChatSessionStreamChunkReplies = cjs.ChatSessionStreamChunkReplies;
 export const ClassificationSession = cjs.ClassificationSession;
+export const NativeTool = cjs.NativeTool;
 export const TextToSpeechSession = cjs.TextToSpeechSession;
 export const TextToSpeechSessionStream = cjs.TextToSpeechSessionStream;
 export const TextToSpeechSessionStreamChunkError = cjs.TextToSpeechSessionStreamChunkError;

bindings/typescript/src/tool.ts

@@ -0,0 +1,144 @@
+import * as z from 'zod';
+
+import { NativeTool, ToolFunction, Value } from './napi/index';
+
+const CANCELLED_INVOCATION_RETENTION_MS = 30_000;
+
+export interface UzuToolContext {
+    readonly signal: AbortSignal;
+}
+
+export interface UzuToolInvokeOptions {
+    readonly signal?: AbortSignal;
+}
+
+export interface UzuToolFunctionOptions<Parameters extends z.ZodObject, ResultSchema extends z.ZodType> {
+    readonly name: string;
+    readonly description?: string;
+    readonly parameters: Parameters;
+    readonly returns: ResultSchema;
+    readonly handler: (
+        parameters: z.output<Parameters>,
+        context: UzuToolContext,
+    ) => z.input<ResultSchema> | Promise<z.input<ResultSchema>>;
+}
+
+export class UzuToolFunction<
+    Parameters extends z.ZodObject,
+    ResultSchema extends z.ZodType,
+> extends NativeTool {
+    readonly name: string;
+    readonly description: string;
+    readonly parameters: Parameters;
+    readonly returns: ResultSchema;
+    readonly parametersSchema: Record<string, unknown>;
+    readonly returnSchema: Record<string, unknown>;
+    readonly handler: UzuToolFunctionOptions<Parameters, ResultSchema>['handler'];
+
+    constructor(options: UzuToolFunctionOptions<Parameters, ResultSchema>) {
+        const name = options.name.trim();
+        if (!name) {
+            throw new TypeError('tool name must not be empty');
+        }
+
+        const description = options.description ?? '';
+        const parametersSchema = z.toJSONSchema(options.parameters);
+        const returnSchema = z.toJSONSchema(options.returns);
+        const activeInvocations = new Map<string, AbortController>();
+        const cancelledInvocations = new Map<string, ReturnType<typeof setTimeout>>();
+
+        const takeCancellation = (invocationId: string): boolean => {
+            const expiration = cancelledInvocations.get(invocationId);
+            if (expiration === undefined) {
+                return false;
+            }
+            clearTimeout(expiration);
+            cancelledInvocations.delete(invocationId);
+            return true;
+        };
+
+        const invokeJson = async (argumentsJson: string, invocationId: string): Promise<string> => {
+            const controller = new AbortController();
+            activeInvocations.set(invocationId, controller);
+            if (takeCancellation(invocationId)) {
+                controller.abort();
+            }
+
+            try {
+                const rawArguments: unknown = JSON.parse(argumentsJson);
+                const parameters = await options.parameters.parseAsync(rawArguments);
+                const rawResult = await options.handler(parameters, {
+                    signal: controller.signal,
+                });
+                const result = await options.returns.parseAsync(rawResult);
+                return serializeResult(result);
+            } finally {
+                activeInvocations.delete(invocationId);
+                takeCancellation(invocationId);
+            }
+        };
+        const cancel = (invocationId: string): void => {
+            const controller = activeInvocations.get(invocationId);
+            if (controller) {
+                controller.abort();
+                return;
+            }
+
+            // Invocation and cancellation use separate nonblocking native callbacks, so
+            // cancellation can arrive before invocation starts. Retain it briefly for
+            // that case, but expire IDs from callbacks that arrive after completion.
+            const expiration = setTimeout(() => {
+                cancelledInvocations.delete(invocationId);
+            }, CANCELLED_INVOCATION_RETENTION_MS);
+            expiration.unref?.();
+            const previousExpiration = cancelledInvocations.get(invocationId);
+            if (previousExpiration !== undefined) {
+                clearTimeout(previousExpiration);
+            }
+            cancelledInvocations.set(invocationId, expiration);
+        };
+
+        super(
+            new ToolFunction(
+                name,
+                description,
+                new Value(JSON.stringify(parametersSchema)),
+                new Value(JSON.stringify(returnSchema)),
+            ),
+            invokeJson,
+            cancel,
+        );
+
+        this.name = name;
+        this.description = description;
+        this.parameters = options.parameters;
+        this.returns = options.returns;
+        this.parametersSchema = parametersSchema;
+        this.returnSchema = returnSchema;
+        this.handler = options.handler;
+    }
+
+    async invoke(
+        input: z.input<Parameters>,
+        options: UzuToolInvokeOptions = {},
+    ): Promise<z.output<ResultSchema>> {
+        const parameters = await this.parameters.parseAsync(input);
+        const signal = options.signal ?? new AbortController().signal;
+        const rawResult = await this.handler(parameters, { signal });
+        return this.returns.parseAsync(rawResult);
+    }
+}
+
+export function uzuToolFunction<Parameters extends z.ZodObject, ResultSchema extends z.ZodType>(
+    options: UzuToolFunctionOptions<Parameters, ResultSchema>,
+): UzuToolFunction<Parameters, ResultSchema> {
+    return new UzuToolFunction(options);
+}
+
+function serializeResult(result: unknown): string {
+    const json = JSON.stringify(result === undefined ? null : result);
+    if (json === undefined) {
+        throw new TypeError('tool result must be JSON serializable');
+    }
+    return json;
+}

bindings/typescript/tests/tool-cancellation.test.ts

@@ -0,0 +1,45 @@
+import * as z from 'zod';
+
+type InvokeJson = (argumentsJson: string, invocationId: string) => Promise<string>;
+type Cancel = (invocationId: string) => void;
+
+jest.mock('../src/napi/index', () => ({
+    NativeTool: class {
+        readonly testInvokeJson: InvokeJson;
+        readonly testCancel: Cancel;
+
+        constructor(_definition: unknown, invokeJson: InvokeJson, cancel: Cancel) {
+            this.testInvokeJson = invokeJson;
+            this.testCancel = cancel;
+        }
+    },
+    ToolFunction: class {},
+    Value: class {},
+}));
+
+import { uzuToolFunction } from '../src/tool';
+
+interface TestNativeTool {
+    readonly testInvokeJson: InvokeJson;
+    readonly testCancel: Cancel;
+}
+
+test('retains pre-start cancellations but expires late cancellations', async () => {
+    jest.useFakeTimers();
+    const tool = uzuToolFunction({
+        name: 'is_cancelled',
+        parameters: z.object({}),
+        returns: z.boolean(),
+        handler: (_parameters, context) => context.signal.aborted,
+    }) as unknown as TestNativeTool;
+
+    tool.testCancel('before-start');
+    await expect(tool.testInvokeJson('{}', 'before-start')).resolves.toBe('true');
+
+    await expect(tool.testInvokeJson('{}', 'already-finished')).resolves.toBe('false');
+    tool.testCancel('already-finished');
+    jest.runOnlyPendingTimers();
+
+    await expect(tool.testInvokeJson('{}', 'already-finished')).resolves.toBe('false');
+    jest.useRealTimers();
+});