agent-v1

import type { FileHandle } from "@lmstudio/sdk";
import { buildExcelTools } from "./buildExcelTools";
import { type PluginCapableCtl } from "../../utils/shared/pluginCtl";



async function runOnce(model: any, history: any[], tools: any[], ctl: PluginCapableCtl, log: string[]) {
    let finalAnswer = "";
    let sawToolCall = false;

    await model.act(history as any, tools, {
        signal: ctl.abortSignal,
        contextOverflowPolicy: "truncateMiddle",
        onMessage: (message: any) => {
            const data = message?.data;
            if (!data) return;
            const blocks: any[] = Array.isArray(data.content) ? data.content : [];
            for (const block of blocks) {
                if (block.type === "toolCallResult") {
                    sawToolCall = true;
                    log.push(`Step: got result — ${String(block.content ?? "").slice(0, 120)}`);
                } else if (block.type === "toolCallRequest" || block.type === "toolCall") {
                    sawToolCall = true;

                    const req = block.toolCallRequest ?? {};
                    const name = req.name ?? "unknown_tool";
                    const argsPreview = JSON.stringify(req.arguments ?? {}).slice(0, 120);
                    log.push(`Step: calling ${name}(${argsPreview})`);
                }
            }
            console.log("\nExcel Sub Agent Processed the following")
            console.log(log)
        },
        onPredictionCompleted: (result: any) => {
            if (result.content.trim().length > 0) {
                finalAnswer = result.content
                    .replace(/<\|channel>thought[\s\S]*?<channel\|>\n?/g, "")
                    .replace(/<think>[\s\S]*?<\/think>\n?/g, "")
                    .trim();
            }
        },
    });

    return { finalAnswer, sawToolCall };
}

export async function runExcelSubAgent(
    ctl: PluginCapableCtl,
    args: { question: string; targetFiles: FileHandle[]; schemaSummary: string },
): Promise<{ trace: string; answer: string }> {
    const { question, targetFiles, schemaSummary } = args;

    const tools = await buildExcelTools(targetFiles);
    const model = await ctl.client.llm.model();
    const log: string[] = [];
    const CLARIFICATION_MARKER = "NEEDS_CLARIFICATION:";

    const systemPrompt =
        `You are a spreadsheet query agent. Schema for the relevant file(s) has already been ` +
        `fetched:\n${schemaSummary}\n\n` +
        `CRITICAL: You must call query_spreadsheet to answer ANY question — do not reason about or ` +
        `guess at data you have not queried.\n\n` +
        `EXCEPTION: If the question is genuinely ambiguous — e.g. it asks for a subjective judgment ` +
        `("best," "recommend," "not recommend") without specifying which metric(s) should decide it, and ` +
        `there is no reasonable default — do not guess a metric on your own. Instead, respond with ONLY:\n` +
        `${CLARIFICATION_MARKER} <a short, specific question about what criteria to use>\n` +
        `Only use this if you truly cannot proceed — if a reasonable default metric exists, just query it.\n\n` +
        `Use query_spreadsheet directly with this schema — you should rarely need to call ` +
        `list_spreadsheet_tables again. Write correct SQL using the exact column names shown above. ` +
        `If a query errors on column/table names, read the corrected names from the error and retry immediately.\n\n` +
        `IMPORTANT — matching text/categorical values: users often refer to a row's identifying value ` +
        `in shortened or reworded form. Never filter text columns with exact equality; use ` +
        `case-insensitive LIKE matching instead. Retry once with a looser pattern before reporting no data found.`;

    let history: any[] = [
        { role: "system", content: systemPrompt },
        { role: "user", content: question },
    ];

    let { finalAnswer, sawToolCall } = await runOnce(model, history, tools, ctl, log);

    // Retry once, more forcefully, if the model never actually called the tool.
    if (!sawToolCall) {
        log.push("(Sub-agent skipped the tool on first attempt — retrying with a stronger instruction.)");
        history = [
            { role: "system", content: systemPrompt },
            { role: "user", content: question },
            { role: "assistant", content: finalAnswer },
            {
                role: "user",
                content:
                    `You did not call query_spreadsheet — you only wrote SQL as text. Call the tool now ` +
                    `with that exact query and use its real result to answer.`,
            },
        ];
        ({ finalAnswer, sawToolCall } = await runOnce(model, history, tools, ctl, log));
    }

    if (!sawToolCall) {
        log.push("(Sub-agent skipped the tool on first attempt — retrying with a stronger instruction.)");
        history = [
            { role: "system", content: systemPrompt },
            { role: "user", content: question },
            { role: "assistant", content: finalAnswer },
            {
                role: "user",
                content:
                    `You answered without calling query_spreadsheet and without asking a clarifying ` +
                    `question. If the question is ambiguous, respond with "${CLARIFICATION_MARKER} <question> based on schema". ` +
                    `Otherwise, call query_spreadsheet now with a real query and use its result to answer.`,
            },
        ];
        ({ finalAnswer, sawToolCall } = await runOnce(model, history, tools, ctl, log));

        if (finalAnswer.trim().startsWith(CLARIFICATION_MARKER)) {
            const clarification = finalAnswer.replace(CLARIFICATION_MARKER, "").trim();
            return {
                trace: log.join("\n"),
                answer: `I need more information to answer this: ${clarification}`,
            };
        }
    }

    return {
        trace: log.join("\n"),
        answer: finalAnswer || "No result was returned from the spreadsheet query.",
    };
}