agent-v1

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

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") {
                    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("\nDocument 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 runDocSubAgent(
    ctl: PluginCapableCtl,
    args: { question: string; file: FileHandle },
): Promise<{ trace: string; answer: string }> {
    const { question, file } = args;

    const tools = buildDocSubAgentTools(ctl, file);
    const model = await ctl.client.llm.model();
    const log: string[] = [];

    const systemPrompt =
        `You are a document query agent working with the file "${file.name}". ` +
        `You have two tools:\n` +
        `- search_document_chunks: for specific facts, quotes, or details.\n` +
        `- get_full_document_text: for summaries, overviews, or "what is this about" questions.\n\n` +
        `CRITICAL: You must call one of these tools to look at the document — never answer from ` +
        `assumptions or prior conversation memory alone. If search_document_chunks returns nothing ` +
        `relevant, try get_full_document_text before giving up. Pick the tool that best matches the ` +
        `shape of the question: specific detail → search_document_chunks; broad/summary → ` +
        `get_full_document_text.`;

    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 a tool.
    if (!sawToolCall) {
        log.push("(Sub-agent skipped tools 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 a tool — you answered without looking at the document. ` +
                    `Call search_document_chunks or get_full_document_text now, then answer using its result.`,
            },
        ];
        ({ finalAnswer, sawToolCall } = await runOnce(model, history, tools, ctl, log));
    }

    if (!sawToolCall) {
        return {
            trace: log.join("\n"),
            answer: `The sub-agent was unable to read "${file.name}" after two attempts. Please try rephrasing the question.`,
        };
    }

    return {
        trace: log.join("\n"),
        answer: finalAnswer || `No result was returned from querying "${file.name}".`,
    };
}