tool-calling.ts

// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

/**
 * Tool Calling Example — Copilot SDK + Foundry Local
 *
 * Demonstrates multiple custom tools that the model can invoke:
 *   - calculate: Evaluate math expressions
 *   - get_system_info: Return local system details
 *   - lookup_definition: Look up programming term definitions
 *
 * Run:  npm run tools
 */

import { CopilotClient, defineTool, approveAll } from "@github/copilot-sdk";
import { FoundryLocalManager } from "foundry-local-sdk";
import { z } from "zod";
import * as os from "os";

const alias = "phi-4-mini";
const endpointUrl = "http://localhost:6543";

// Timeout for each model turn (ms).  Override with FOUNDRY_TIMEOUT_MS env var.
// Local models on CPU can be slow — increase this on less powerful hardware.
const TIMEOUT_MS = Number(process.env.FOUNDRY_TIMEOUT_MS) || 120_000;

// ---------------------------------------------------------------------------
// Helper: send a message and wait for the assistant's full reply.
// Uses sendAndWait with a fallback: if the session emits an error (e.g.
// missing finish_reason from the local model), we catch it and continue.
// ---------------------------------------------------------------------------
async function sendMessage(
    session: Awaited<ReturnType<CopilotClient["createSession"]>>,
    prompt: string,
    timeoutMs = TIMEOUT_MS,
) {
    try {
        await session.sendAndWait({ prompt }, timeoutMs);
    } catch (err: any) {
        // Foundry Local streaming may omit finish_reason, causing a
        // session.error that rejects sendAndWait. Treat as non-fatal.
        console.error(`\n[sendMessage error: ${err?.message ?? err}]`);
    }
}

type Model = Awaited<ReturnType<FoundryLocalManager["catalog"]["getModel"]>>;

// ---------------------------------------------------------------------------
// Tool definitions
// ---------------------------------------------------------------------------

function defineCalculateTool() {
    return defineTool("calculate", {
        description:
            "Evaluate a math expression and return the numeric result. " +
            "Supports +, -, *, /, parentheses, and Math.* functions like Math.sqrt, Math.pow.",
        parameters: z.object({
            expression: z.string().describe('Math expression to evaluate, e.g. "2 + 2" or "Math.sqrt(144)"'),
        }),
        handler: async (args) => {
            try {
                // Only allow safe math characters and Math.* calls
                const sanitized = args.expression.replace(/[^0-9+\-*/().,%\s]|Math\.\w+/g, (m) =>
                    m.startsWith("Math.") ? m : "",
                );
                const result = new Function(`"use strict"; return (${sanitized})`)();
                console.log(`\n    → calculate("${args.expression}") = ${result}`);
                return { expression: args.expression, result: Number(result) };
            } catch {
                return { expression: args.expression, error: "Could not evaluate expression" };
            }
        },
    });
}

function defineLookupTool() {
    const glossary: Record<string, string> = {
        "byok": "Bring Your Own Key — a pattern where you supply your own API credentials to route requests to a custom endpoint instead of the default provider.",
        "onnx": "Open Neural Network Exchange — an open format for representing machine learning models, enabling interoperability between frameworks.",
        "rag": "Retrieval-Augmented Generation — a technique that combines a retrieval system with a generative model so responses are grounded in external documents.",
        "json-rpc": "JSON Remote Procedure Call — a lightweight protocol for calling methods on a remote server using JSON-encoded messages.",
        "streaming": "A technique where the server sends response tokens incrementally as they are generated, rather than waiting for the full response.",
    };

    return defineTool("lookup_definition", {
        description:
            "Look up the definition of a programming or AI term. " +
            "Available terms: " + Object.keys(glossary).join(", "),
        parameters: z.object({
            term: z.string().describe("The term to look up (case-insensitive)"),
        }),
        handler: async (args) => {
            const key = args.term.toLowerCase().trim();
            const definition = glossary[key];
            console.log(`\n    → lookup_definition("${args.term}") → ${definition ? "found" : "not found"}`);
            if (definition) {
                return { term: args.term, definition };
            }
            return { term: args.term, error: `Term not found. Available: ${Object.keys(glossary).join(", ")}` };
        },
    });
}

function defineSystemInfoTool(modelId: string, endpoint: string) {
    return defineTool("get_system_info", {
        description: "Get information about the local system: OS, architecture, memory, CPU count, and the running model.",
        parameters: z.object({}),
        handler: async () => {
            const info = {
                platform: os.platform(),
                arch: os.arch(),
                cpus: os.cpus().length,
                totalMemory: `${Math.round(os.totalmem() / 1024 ** 3)} GB`,
                freeMemory: `${Math.round(os.freemem() / 1024 ** 3)} GB`,
                nodeVersion: process.version,
                model: modelId,
                endpoint,
            };
            console.log(`\n    → get_system_info() → ${JSON.stringify(info)}`);
            return info;
        },
    });
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

async function main() {
    let manager: FoundryLocalManager | undefined;
    let model: Model | undefined;
    let client: CopilotClient | undefined;
    let session: Awaited<ReturnType<CopilotClient["createSession"]>> | undefined;

    try {
        console.log("Initializing Foundry Local...");
        manager = FoundryLocalManager.create({
            appName: "foundry_local_samples",
            webServiceUrls: endpointUrl,
        });

        // Download and register all execution providers.
        let currentEp = '';
        await manager.downloadAndRegisterEps((epName, percent) => {
            if (epName !== currentEp) {
                if (currentEp !== '') process.stdout.write('\n');
                currentEp = epName;
            }
            process.stdout.write(`\r  ${epName.padEnd(30)}  ${percent.toFixed(1).padStart(5)}%`);
        });
        if (currentEp !== '') process.stdout.write('\n');

        model = await manager.catalog.getModel(alias);
        await model.download();
        await model.load();
        console.log(`Model: ${model.id}`);

        manager.startWebService();
        const endpoint = endpointUrl + "/v1";
        console.log(`Endpoint: ${endpoint}\n`);

        const calculate = defineCalculateTool();
        const lookupDefinition = defineLookupTool();
        const getSystemInfo = defineSystemInfoTool(model.id, endpoint);

        client = new CopilotClient();

        session = await client.createSession({
            onPermissionRequest: approveAll,
            model: model.id,
            provider: {
                type: "openai",
                baseUrl: endpoint,
                apiKey: "local",
                wireApi: "completions",
            },
            streaming: true,
            tools: [calculate, lookupDefinition, getSystemInfo],
            systemMessage: {
                content:
                    "You are a helpful AI assistant running locally via Foundry Local. " +
                    "You have access to tools. ALWAYS use the appropriate tool when the user asks you to " +
                    "calculate something, look up a term, or get system information. " +
                    "Do not guess — call the tool and report its result. " +
                    "Keep responses concise.",
            },
        });

        // Stream assistant text to stdout
        session.on("assistant.message_delta", (event) => {
            process.stdout.write(event.data.deltaContent);
        });
        session.on("tool.execution_start", (event) => {
            console.log(`\n  [Tool called: ${(event as any).data?.toolName ?? "unknown"}]`);
        });

        // --- Turn 1: Calculator tool ---
        console.log("=== Turn 1: Calculator ===\n");
        process.stdout.write("User: What is the square root of 144 plus 8 times 3?\n\nAssistant: ");
        await sendMessage(
            session,
            "Use the calculate tool to compute: Math.sqrt(144) + 8 * 3",
        );
        console.log("\n");

        // --- Turn 2: Glossary lookup tool ---
        console.log("=== Turn 2: Glossary Lookup ===\n");
        process.stdout.write("User: What does BYOK mean? And what about RAG?\n\nAssistant: ");
        await sendMessage(
            session,
            "Use the lookup_definition tool to look up 'byok' and 'rag', then explain both.",
        );
        console.log("\n");

        // --- Turn 3: System info tool ---
        console.log("=== Turn 3: System Info ===\n");
        process.stdout.write("User: What system am I running on?\n\nAssistant: ");
        await sendMessage(
            session,
            "Use the get_system_info tool to check what system this is running on, then summarize.",
        );
        console.log("\n");

        console.log("Done!");
    } finally {
        // Clean up resources in reverse order of creation
        if (session) {
            await session.destroy().catch(() => {});
        }
        if (client) {
            await client.stop().catch(() => {});
        }
        if (model) {
            console.log("Unloading model...");
            await model.unload().catch((e) => {
                console.warn("Warning: failed to unload model:", e);
            });
        }
        if (manager) {
            console.log("Stopping web service...");
            try {
                manager.stopWebService();
            } catch (e) {
                console.warn("Warning: failed to stop web service:", e);
            }
        }
    }
}

main().catch(console.error);