agent-v1

import { tool } from "@lmstudio/sdk";
import { z } from "zod";
import type { FileHandle } from "@lmstudio/sdk";
import { type PluginCapableCtl } from "../../utils/shared/pluginCtl";
import { parseFile } from "../../utils/docs/parser/parseFile";

const FULL_TEXT_CHAR_CAP = 10000;

/**
 * Builds the two tools the doc sub-agent chooses between:
 * - search_document_chunks: semantic search, best for specific facts/quotes
 * - get_full_document_text: full (possibly truncated) text, best for
 *   summarization / overview questions that don't match any single passage
 */
export function buildDocSubAgentTools(ctl: PluginCapableCtl, file: FileHandle) {
    const searchTool = tool({
        name: "search_document_chunks",
        description:
            `Semantic search for passages similar to a specific query within "${file.name}". ` +
            `Best for specific facts, quotes, numbers, or named details. Returns the most relevant ` +
            `passages, not the whole document. If this returns nothing relevant, try ` +
            `get_full_document_text instead.`,
        parameters: {
            query: z.string().describe("A specific, targeted search query — not a vague request like 'summarize'."),
        },
        implementation: async ({ query }: { query: string }) => {
            const model = await ctl.client.embedding.model("nomic-ai/nomic-embed-text-v1.5-GGUF", {
                signal: ctl.abortSignal,
            });

            const retrievalLimit = ctl.getConfigValue("retrievalLimit");
            const retrievalAffinityThreshold = ctl.getConfigValue("retrievalAffinityThreshold");

            const result = await ctl.client.files.retrieve(query, [file], {
                embeddingModel: model,
                limit: retrievalLimit,
                signal: ctl.abortSignal,
            });

            const entries = result.entries
                .filter((e: any) => e.score > retrievalAffinityThreshold)
                .sort((a: any, b: any) => b.score - a.score)
                .slice(0, retrievalLimit);

            if (entries.length === 0) {
                return (
                    `No passages in "${file.name}" scored above the relevance threshold for this query. ` +
                    `If you need a broad overview instead of a specific fact, call get_full_document_text.`
                );
            }

            return entries
                .map((e: any, i: number) => `[${i + 1}] ${(e.content)}`)
                .join("\n\n");
        },
    });

    const fullTextTool = tool({
        name: "get_full_document_text",
        description:
            `Returns the full text of "${file.name}" (truncated if very long). Use this for ` +
            `summarization, overview, or "what is this about" style questions where no single ` +
            `passage answers the question — search_document_chunks won't find a good match for those.`,
        parameters: {},
        implementation: async () => {
            const parsed = await parseFile(ctl, file);
            const cleaned = parsed.content;
            if (cleaned.length <= FULL_TEXT_CHAR_CAP) {
                return cleaned;
            }
            return cleaned.slice(0, FULL_TEXT_CHAR_CAP) + "\n\n(...document truncated, content continues beyond this point...)";
        },
    });

    return [searchTool, fullTextTool];
}