llmorbiter/src/LLMOrbiter.js

'use strict';

import { DirectReturnError } from './Tool.js';

/** Always produce a proper DOMException('...', 'AbortError') from any abort reason. */
function ensureAbortError(reason) {
    if (reason instanceof DOMException && reason.name === 'AbortError') return reason;
    const err = new DOMException(
        typeof reason === 'string' ? reason : (reason?.message || 'LLMOrbiter run aborted'),
        'AbortError',
    );
    if (reason != null && typeof reason !== 'string') err.cause = reason;
    return err;
}

/**
 * LLMOrbiter — main orchestrator.
 *
 * Usage:
 *   const orbiter = new LLMOrbiter(provider, { onStatus, onToolCall });
 *   const result  = await orbiter.run({ prompt, messages, model, tools });
 *   orbiter.abort(); // abort the active request/tool loop
 *
 * `messages` is the conversation transcript (the system prompt is simply its first
 * entry, if any) and is returned on the result for the next run to continue from.
 *
 * The run loop:
 *   1. Send messages + tool specs to the provider
 *   2. If the LLM returns tool_calls → invoke local implementations, feed results back
 *   3. Repeat until the LLM finishes (no more tool_calls) or maxIterations is hit
 *   4. Return { content, toolResults, usage }
 */
class LLMOrbiter {
    /**
     * @param {Object} provider  — must implement chatCompletions({ model, messages, tools, signal, onStatus, onContent, onReasoning })
     * @param {Object} [opts]
     * @param {Function} [opts.onLLMCall]    — called before LLM request ({ model, toolCount, promptSize, messageCount, iteration })
     * @param {Function} [opts.onUsage]      — called with usage info ({ cached, input, output, cost })
     * @param {Function} [opts.onToolCall]   — called with tool call info ({ name, args })
     * @param {Function} [opts.onToolError]  — called with tool error info ({ name, error, args })
     * @param {Function} [opts.onToolResult] — called with tool result info ({ name, result })
     * @param {Function} [opts.onStatus]     — called with status string (e.g. 'connected', 'streaming')
     * @param {Function} [opts.onContent]    — called with streamed content chunks
     * @param {Function} [opts.onReasoning]  — called with streamed reasoning text
     */
    constructor(provider, opts = {}) {
        if (!provider) throw new Error('LLMOrbiter requires a provider');
        this.provider = provider;
        this._activeController = null;
        Object.assign(this, opts);
    }

    /**
     * Abort the currently active run, if any.
     * Note: with concurrent runs, only the most recent run is aborted.
     */
    abort(reason) {
        this._activeController?.abort(ensureAbortError(reason));
    }

    /**
     * Run an LLM conversation with tool-calling loop.
     *
     * @param {Object}   opts
     * @param {string}   opts.prompt        — user message to append (omit/null to continue
     *                                        from `messages` without adding one)
     * @param {Array}    [opts.messages]    — conversation transcript to continue from (e.g. a
     *                                        previous result.messages). Used as-is; a leading
     *                                        { role: 'system', content } entry is the system
     *                                        prompt. Omit to start a fresh conversation.
     * @param {string}   [opts.model]       — model identifier (provider-specific default if omitted)
     * @param {Array}    [opts.tools]       — array of Tool objects from defineTool()
     * @param {number}   [opts.maxIterations=10] — safety limit for tool-call loops
     * @param {Object}   [opts.extraBody]   — extra fields merged into the request body (e.g. reasoning)
     * @param {AbortSignal} [opts.signal]   — optional external abort signal
     * @returns {Promise<RunResult>}
     */
    async run(opts) {
        const {
            prompt,
            messages: seedMessages,
            model,
            tools = [],
            maxIterations = 10,
            extraBody = {},
            signal,
        } = opts;

        if (seedMessages != null && !Array.isArray(seedMessages)) {
            throw new TypeError('run({ messages }): messages must be an array, or omit it for a fresh conversation');
        }

        const { onLLMCall, onUsage, onToolCall, onToolError, onToolResult, onStatus, onContent, onReasoning } = this;

        // Keep an internal controller so callers can cancel a run through orbiter.abort().
        // The external signal remains supported and is forwarded into the same controller.
        const controller = new AbortController();
        const abortFromExternalSignal = () => controller.abort(ensureAbortError(signal.reason));
        if (signal) {
            if (signal.aborted) abortFromExternalSignal();
            else signal.addEventListener('abort', abortFromExternalSignal, { once: true });
        }
        this._activeController = controller;
        const runSignal = controller.signal;

        const call = async (fn, ...args) => {
            if (fn) {
                try { await fn(...args); } catch (_) { }
            }
        };

        // Record a tool response into the transcript. Every tool_call in an assistant
        // message must have a matching tool message, so this is the single place that
        // writes them (including the direct-return/skip cases).
        const recordToolResult = (callItem, result) => {
            const content = typeof result === 'string' ? result : JSON.stringify(result);
            messages.push({ role: 'tool', tool_call_id: callItem.id, content });
        };

        // Build tool lookup: name -> invoke
        const toolMap = new Map();
        const toolSpecs = [];
        for (const t of tools) {
            toolMap.set(t.name, t.invoke);
            toolSpecs.push(t.spec);
        }

        // Seed messages: continue from a prior transcript, or start fresh.
        // A leading { role: 'system' } entry in the seed is the system prompt.
        const messages = seedMessages?.length ? [...seedMessages] : [];
        if (prompt != null) {
            messages.push({ role: 'user', content: prompt });
        }

        const allToolResults = {};   // toolName -> [result, ...]
        const usage = {
            input: 0,
            output: 0,
            cached: 0,
            tokenIn: 0,
            tokenOut: 0,
            tokenCached: 0,
            costTotal: 0,
        };
        const rawUsages = [];
        let finalContent = null;
        let iteration = 0;
        let directReturn = null;
        let hasDirectReturn = false;

        const promptStr = typeof prompt === 'string' ? prompt : JSON.stringify(prompt || '');
        const promptSize = promptStr.length;
        const toolCount = toolSpecs.length;

        try {
            while (iteration < maxIterations) {
                if (runSignal.aborted) throw ensureAbortError(runSignal.reason);
                iteration++;

                const messageFlow = messages.flatMap((m) => {
                    if (m.role === 'assistant' && m.tool_calls?.length) {
                        return ['assistant', 'tool'];
                    }
                    if (m.role === 'tool') {
                        return ['tool_response'];
                    }
                    return [m.role];
                }).join(',');

                await call(onLLMCall, {
                    model,
                    toolCount,
                    tools: toolCount,
                    promptSize,
                    messageCount: messages.length,
                    messages: messageFlow,
                    iteration,
                });

                const completion = await this.provider.chatCompletions({
                    model,
                    messages,
                    tools: toolSpecs.length ? toolSpecs : undefined,
                    signal: runSignal,
                    extraBody,
                    onStatus: (s) => call(onStatus, s),
                    onContent: (...args) => call(onContent, ...args),
                    onReasoning: (r) => call(onReasoning, r),
                });

                // Accumulate usage & emit event per LLM call completed
                const callUsage = this._addUsage(usage, completion.usage);
                if (completion.usage) {
                    rawUsages.push(completion.usage);
                }
                await call(onUsage, callUsage);

                const toolCalls = completion.tool_calls;
                const content = completion.content;
                const hasToolCalls = Array.isArray(toolCalls) && toolCalls.length > 0;
                const hasContent = typeof content === 'string' && content.trim().length > 0;

                // When there is no response (neither tool calls nor content), the loop shall continue.
                if (!hasToolCalls && !hasContent) {
                    continue;
                }

                // Push assistant message
                messages.push({
                    role: 'assistant',
                    content: content || null,
                    ...(hasToolCalls ? { tool_calls: toolCalls } : {}),
                    ...(completion.reasoning_details?.length ? { reasoning_details: completion.reasoning_details } : {}),
                    ...(completion.reasoning ? { reasoning: completion.reasoning } : {}),
                });

                // No tool calls → done
                if (!hasToolCalls) {
                    finalContent = content;
                    break;
                }

                const pendingAppendedMessages = [];
                for (let callIdx = 0; callIdx < toolCalls.length; callIdx++) {
                    const callItem = toolCalls[callIdx];
                    const fnName = callItem.function.name;
                    const invokeFn = toolMap.get(fnName);
                    let result;
                    let isError = false;
                    let errorMsg = null;

                    if (!invokeFn) {
                        isError = true;
                        errorMsg = `Unknown tool: ${fnName}`;
                        result = { error: errorMsg };
                        await call(onToolError, { name: fnName, error: errorMsg });
                    } else {
                        let args;
                        try {
                            args = JSON.parse(callItem.function.arguments || '{}');
                        } catch (e) {
                            isError = true;
                            errorMsg = `Invalid JSON arguments: ${e.message}`;
                            result = { error: errorMsg };
                            await call(onToolError, { name: fnName, error: errorMsg });
                        }

                        if (!isError) {
                            await call(onToolCall, { name: fnName, args, id: callItem.id, callItem, tool_calls: toolCalls, reasoning_details: completion.reasoning_details, reasoning: completion.reasoning, content: content || null });

                            try {
                                result = await invokeFn(args, {
                                    signal: runSignal,
                                    messages,
                                    appendMessage: (msg) => pendingAppendedMessages.push(msg),
                                });
                            } catch (err) {
                                // An abort must propagate, not be recorded as a tool error.
                                if (runSignal.aborted) throw ensureAbortError(runSignal.reason);
                                if (err instanceof DirectReturnError) {
                                    // Stop the loop and return the data to the caller.
                                    directReturn = err.data;
                                    hasDirectReturn = true;

                                    // Record this call's response and close out any remaining
                                    // calls in the batch with synthetic tool messages, so the
                                    // returned transcript stays valid to send back to the API
                                    // (every tool_call must have a matching tool response).
                                    if (err.recordResult !== false) {
                                        recordToolResult(callItem, err.data);
                                        for (let j = callIdx + 1; j < toolCalls.length; j++) {
                                            const pending = toolCalls[j];
                                            recordToolResult(pending, 'Skipped: run terminated by direct return');
                                        }
                                    }
                                    break;
                                }
                                isError = true;
                                if (err instanceof Error) {
                                    errorMsg = err.message
                                        ? `${err.name !== 'Error' && err.name !== 'ToolError' ? `${err.name}: ` : ''}${err.message}`
                                        : `${err.name || 'Error'} (unspecified crash)`;
                                } else if (err != null && String(err).trim()) {
                                    errorMsg = `Tool crashed: ${String(err)}`;
                                } else {
                                    errorMsg = 'Tool crashed due to an unspecified internal error';
                                }
                                result = { error: errorMsg };
                                await call(onToolError, { name: fnName, error: errorMsg, args, rawError: err, id: callItem.id, callItem });
                            }

                            // Also handle tools that return { error: ... } instead of throwing
                            if (!isError && result && typeof result === 'object' && 'error' in result && result.error) {
                                isError = true;
                                errorMsg = typeof result.error === 'string' ? result.error : JSON.stringify(result.error);
                                await call(onToolError, { name: fnName, error: errorMsg, args });
                            }

                            if (result && typeof result === 'object' && Array.isArray(result.appendMessages)) {
                                for (const msg of result.appendMessages) {
                                    pendingAppendedMessages.push(msg);
                                }
                            }
                        }
                    }

                    if (hasDirectReturn) break;

                    // Record result
                    if (!allToolResults[fnName]) allToolResults[fnName] = [];
                    allToolResults[fnName].push(result);

                    if (!isError) {
                        await call(onToolResult, { name: fnName, result, id: callItem.id, callItem });
                    }

                    // Feed result back to conversation
                    recordToolResult(callItem, result);
                }

                if (hasDirectReturn) break;

                if (pendingAppendedMessages.length > 0) {
                    for (const msg of pendingAppendedMessages) {
                        messages.push(msg);
                    }
                }
            }


            usage.raw = rawUsages.length === 1 ? rawUsages[0] : rawUsages;

            return {
                content: finalContent,
                toolResults: allToolResults,
                usage,
                rawUsage: rawUsages.length === 1 ? rawUsages[0] : rawUsages,
                rawUsages,
                iterations: iteration,
                messages,
                ...(directReturn !== null ? { directReturn: true, data: directReturn } : {}),
            };
        } finally {
            if (signal) signal.removeEventListener('abort', abortFromExternalSignal);
            if (this._activeController === controller) this._activeController = null;
        }
    }

    /** @private */
    _addUsage(totals, u) {
        const cached = u?.prompt_tokens_details?.cached_tokens || 0;
        const input = (u?.prompt_tokens || 0) - cached;
        const output = u?.completion_tokens || 0;

        totals.tokenCached += cached;
        totals.tokenIn += input;
        totals.tokenOut += output;
        totals.cached = totals.tokenCached;
        totals.input = totals.tokenIn;
        totals.output = totals.tokenOut;

        let cost = 0;
        if (u?.cost != null) {
            cost = u.cost;
        } else if (u?.cost_details) {
            const d = u.cost_details;
            cost = (d.upstream_inference_prompt_cost || 0) + (d.upstream_inference_completions_cost || 0);
        }
        totals.costTotal += cost;
        totals.cost = totals.costTotal;

        const callUsage = { cached, input, output };
        if (cost) callUsage.cost = cost;
        if (u) {
            Object.defineProperty(callUsage, 'raw', {
                value: u,
                enumerable: false,
                writable: true,
                configurable: true,
            });
        }
        return callUsage;
    }
}

export { LLMOrbiter };