Blame for backend/src/lib/assistantRouting/classifier.ts

import type { RunTree } from "langsmith";
import { executeCompleteText } from "../inference/execute";
import type { UserApiKeys } from "../llm/types";
import type { AssistantRoutingContextV2, AssistantRouterOutputV2 } from "./types";
import { parseAssistantRouterOutput } from "./schema";
import {
    AssistantRoutingValidationError,
} from "./types";
import { buildAssistantRouterMessages } from "../../prompts/chat/routing/assistant-turn-router";
import { getRouterModelCandidates, type AssistantModelCandidate } from "./modelCandidates";

/** Minimum remaining deadline required to attempt a fallback router call. */
export const MIN_ROUTER_FALLBACK_REMAINING_MS = 1_500;

export interface ClassifyAssistantTurnResult {
    output: AssistantRouterOutputV2;
    usedFallback: boolean;
    model: string | null;
    failureReasons: string[];
}

/**
 * Injected inference executor. Defaults to the real {@link executeCompleteText}
 * but may be replaced in tests with a deterministic mock.
 */
export type RouterInferenceExecutor = (params: {
    model: string;
    featureId: string;
    systemPrompt: string;
    user: string;
    maxTokens: number;
    temperature: number;
    apiKeys?: UserApiKeys;
    trace?: RunTree | null;
    signal?: AbortSignal;
}) => Promise<string>;

const defaultExecutor: RouterInferenceExecutor = async (p) => {
    return executeCompleteText({
        model: p.model,
        featureId: p.featureId,
        systemPrompt: p.systemPrompt,
        user: p.user,
        maxTokens: p.maxTokens,
        temperature: p.temperature,
        lightweight: true,
        apiKeys: p.apiKeys,
        trace: p.trace,
    });
};

/**
 * Build the deterministic safe fallback router output.
 *
 * Safe fallback requirements (§15):
 * - no web research;
 * - no comment/edit/generation/workflow capability (only a general answer may
 *   be produced when safe);
 * - otherwise require clarification;
 * - if the deterministic mutation hint is true, the fallback signals an
 *   ambiguous/unsafe state rather than ignoring possible mutation (the policy
 *   compiler turns this into clarification/confirmation).
 */
export function buildSafeFallbackOutput(
    context: AssistantRoutingContextV2,
): AssistantRouterOutputV2 {
    const ambiguities: AssistantRouterOutputV2["ambiguities"] = [];
    const reasonCodes: AssistantRouterOutputV2["reasonCodes"] = ["ambiguous_request"];
    if (!context.latestUserMessage.trim()) {
        ambiguities.push("empty_request");
    }
    return Object.freeze({
        schemaVersion: 2,
        operations: ["answer"],
        webNeed: "not_needed",
        targetScopes: ["conversation"],
        ambiguities,
        reasonCodes,
    }) as AssistantRouterOutputV2;
}

function remainingFromNow(deadlineAt: number): number {
    return Math.max(0, deadlineAt - Date.now());
}

/**
 * Run the single semantic routing stage with at most one independent fallback.
 *
 * Rules (§15):
 * 1. Call the primary routing model once (temperature 0, bounded output).
 * 2. Parse through the strict validator.
 * 3. On transport failure / timeout / empty / invalid output, check remaining
 *    deadline and configured fallback candidates; call at most one fallback.
 * 4. The fallback receives the original routing context, not the first output.
 * 5. The fallback does not critique or refine the first result.
 * 6. If all candidates fail, return a deterministic safe fallback.
 *
 * Never chains a second semantic decision from a first router output.
 */
export async function classifyAssistantTurn(params: {
    context: AssistantRoutingContextV2;
    apiKeys?: UserApiKeys;
    signal?: AbortSignal;
    deadlineAt: number;
    trace?: RunTree | null;
    executor?: RouterInferenceExecutor;
    candidates?: AssistantModelCandidate[];
}): Promise<ClassifyAssistantTurnResult> {
    const {
        context,
        apiKeys,
        signal,
        deadlineAt,
        trace,
        executor = defaultExecutor,
        candidates = getRouterModelCandidates(),
    } = params;

    const failureReasons: string[] = [];
    const messages = buildAssistantRouterMessages(context);
    const systemPrompt = messages[0].content;
    const userPayload = messages[1].content;

    if (candidates.length === 0) {
        return {
            output: buildSafeFallbackOutput(context),
            usedFallback: true,
            model: null,
            failureReasons: ["no_router_candidates"],
        };
    }

    async function attempt(
        candidate: AssistantModelCandidate,
    ): Promise<AssistantRouterOutputV2 | null> {
        if (signal?.aborted) throw new DOMException("Aborted", "AbortError");
        try {
            const raw = await executor({
                model: candidate.model,
                featureId: candidate.featureId,
                systemPrompt,
                user: userPayload,
                maxTokens: 512,
                temperature: 0,
                apiKeys,
                trace,
                signal,
            });
            if (!raw || !raw.trim()) {
                failureReasons.push(`${candidate.model}:empty_output`);
                return null;
            }
            return parseAssistantRouterOutput(raw);
        } catch (err) {
            if (err instanceof DOMException && err.name === "AbortError") throw err;
            if (err instanceof AssistantRoutingValidationError) {
                failureReasons.push(`${candidate.model}:invalid_output`);
            } else {
                failureReasons.push(`${candidate.model}:transport_error`);
            }
            return null;
        }
    }

    // Primary.
    try {
        const primary = await attempt(candidates[0]);
        if (primary) {
            return {
                output: primary,
                usedFallback: false,
                model: candidates[0].model,
                failureReasons,
            };
        }
    } catch (err) {
        if (err instanceof DOMException && err.name === "AbortError") {
            return {
                output: buildSafeFallbackOutput(context),
                usedFallback: true,
                model: null,
                failureReasons: [...failureReasons, "cancelled"],
            };
        }
        throw err;
    }

    // At most one fallback, if remaining deadline permits and a distinct
    // candidate exists.
    const fallbackCandidate = candidates.find(
        (c) => c.model !== candidates[0].model,
    );
    if (
        fallbackCandidate &&
        remainingFromNow(deadlineAt) >= MIN_ROUTER_FALLBACK_REMAINING_MS &&
        !signal?.aborted
    ) {
        try {
            const fallback = await attempt(fallbackCandidate);
            if (fallback) {
                return {
                    output: fallback,
                    usedFallback: true,
                    model: fallbackCandidate.model,
                    failureReasons,
                };
            }
        } catch (err) {
            if (err instanceof DOMException && err.name === "AbortError") {
                return {
                    output: buildSafeFallbackOutput(context),
                    usedFallback: true,
                    model: null,
                    failureReasons: [...failureReasons, "cancelled"],
                };
            }
        }
    }

    // All candidates failed: deterministic safe fallback.
    return {
        output: buildSafeFallbackOutput(context),
        usedFallback: true,
        model: null,
        failureReasons,
    };
}