rag

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.toolsProvider = void 0;
const sdk_1 = require("@lmstudio/sdk");
const zod_1 = require("zod");
const promises_1 = require("fs/promises");
const path_1 = require("path");
const config_1 = require("./config");
const embedder_1 = require("./embedder");
const chunker_1 = require("./chunker");
const parser_1 = require("./parser");
const vectorStore_1 = require("./vectorStore");
function json(obj) {
    return JSON.stringify(obj, null, 2);
}
function safe_impl(name, fn) {
    return async (params, ctx) => {
        if (ctx.signal.aborted)
            return json({ tool_error: true, tool: name, error: "cancelled" });
        try {
            return await fn(params, ctx);
        }
        catch (err) {
            const msg = err instanceof Error ? err.message : String(err);
            return json({ tool_error: true, tool: name, error: msg });
        }
    };
}
const DEFAULT_COLLECTION = "default";
const BATCH_SIZE = 32; // embed this many chunks at a time
const toolsProvider = async (ctl) => {
    const cfg = ctl.getPluginConfig(config_1.pluginConfigSchematics);
    const dataPath = () => {
        const p = cfg.get("dataPath").trim();
        return p || (0, path_1.join)(process.env.HOME ?? "~", "rag-data");
    };
    const workspacePath = () => cfg.get("workspacePath").trim();
    const maxFileSizeMb = () => cfg.get("maxFileSizeMb");
    const chunkSize = () => cfg.get("chunkSize");
    const chunkOverlap = () => cfg.get("chunkOverlap");
    const topKDefault = () => cfg.get("topK");
    const embeddingId = () => cfg.get("embeddingModelIdentifier").trim();
    const tools = [
        (0, sdk_1.tool)({
            name: "rag_ingest",
            description: (0, sdk_1.text) `
        Parse, chunk, embed, and index a document (or all documents in a directory) into the local RAG index.

        path: file or directory path (relative to workspace or absolute)
        collection: named sub-index to store documents in (default: "default")
        recursive: if path is a directory, ingest subdirectories too (default: false)
        chunk_size: override config chunk size for this ingest
        chunk_overlap: override config chunk overlap for this ingest

        Supported formats: PDF, DOCX, XLSX/XLS/ODS/CSV, PPTX/PPT, EPUB, HTML/HTM, JSON/JSONL, TXT/MD and most plain text files.
        Requires an embedding model to be loaded in LM Studio.

        Re-ingesting the same file replaces all its previous chunks.
      `,
            parameters: {
                path: zod_1.z.string().describe("File or directory path"),
                collection: zod_1.z.string().default(DEFAULT_COLLECTION).describe("Collection name (default: \"default\")"),
                recursive: zod_1.z.boolean().default(false).describe("Recurse into subdirectories"),
                chunk_size: zod_1.z.coerce.number().int().min(100).max(8000).optional().describe("Override chunk size"),
                chunk_overlap: zod_1.z.coerce.number().int().min(0).max(500).optional().describe("Override chunk overlap"),
            },
            implementation: safe_impl("rag_ingest", async ({ path, collection, recursive, chunk_size, chunk_overlap }, ctx) => {
                const dp = dataPath();
                const ws = workspacePath();
                const cs = chunk_size ?? chunkSize();
                const co = chunk_overlap ?? chunkOverlap();
                const maxMb = maxFileSizeMb();
                ctx.status("Getting embedding model…");
                const embed = await (0, embedder_1.getEmbedFn)(ctl.client, embeddingId());
                // Collect files to ingest
                const fullPath = await (0, parser_1.resolvePath)(path, ws, 9999); // size check per file below
                const info = await (0, promises_1.stat)(fullPath);
                const files = [];
                if (info.isDirectory()) {
                    const collect = async (dir) => {
                        const entries = await (0, promises_1.readdir)(dir, { withFileTypes: true });
                        for (const e of entries) {
                            const fp = (0, path_1.join)(dir, e.name);
                            if (e.isFile())
                                files.push(fp);
                            else if (e.isDirectory() && recursive)
                                await collect(fp);
                        }
                    };
                    await collect(fullPath);
                }
                else {
                    files.push(fullPath);
                }
                const idx = (0, vectorStore_1.getIndex)(dp, collection);
                const results = [];
                for (const filePath of files) {
                    ctx.status(`Ingesting ${(0, path_1.basename)(filePath)}…`);
                    try {
                        const fileStat = await (0, promises_1.stat)(filePath);
                        if (fileStat.size / (1024 * 1024) > maxMb) {
                            results.push({ file: filePath, chunks: 0, status: `skipped — exceeds ${maxMb} MB` });
                            continue;
                        }
                        const { text: docText } = await (0, parser_1.parseFile)(filePath, 10_000_000);
                        if (!docText.trim()) {
                            results.push({ file: filePath, chunks: 0, status: "skipped — no text extracted" });
                            continue;
                        }
                        // Remove old chunks for this file
                        await (0, vectorStore_1.deleteBySource)(idx, filePath);
                        const chunks = (0, chunker_1.chunkText)(docText, cs, co);
                        if (chunks.length === 0) {
                            results.push({ file: filePath, chunks: 0, status: "skipped — no chunks produced" });
                            continue;
                        }
                        const now = new Date().toISOString();
                        const allVectors = [];
                        // Embed in batches
                        for (let i = 0; i < chunks.length; i += BATCH_SIZE) {
                            const batch = chunks.slice(i, i + BATCH_SIZE);
                            ctx.status(`Embedding ${(0, path_1.basename)(filePath)} (${i + batch.length}/${chunks.length})…`);
                            const vecs = await embed(batch.map(c => c.text));
                            allVectors.push(...vecs);
                        }
                        await (0, vectorStore_1.insertChunks)(idx, chunks.map((c, i) => ({
                            vector: allVectors[i],
                            metadata: {
                                text: c.text,
                                source: filePath,
                                fileName: (0, path_1.basename)(filePath),
                                chunkIndex: c.index,
                                charStart: c.charStart,
                                charEnd: c.charEnd,
                                collection,
                                ingestedAt: now,
                            },
                        })));
                        results.push({ file: filePath, chunks: chunks.length, status: "ok" });
                    }
                    catch (e) {
                        results.push({ file: filePath, chunks: 0, status: `error: ${e instanceof Error ? e.message : String(e)}` });
                    }
                }
                const totalChunks = results.reduce((s, r) => s + r.chunks, 0);
                return json({ collection, filesProcessed: results.length, totalChunks, results });
            }),
        }),
        (0, sdk_1.tool)({
            name: "rag_query",
            description: (0, sdk_1.text) `
        Semantically search the RAG index. Returns the most relevant text chunks for a query.

        query: natural language question or keyword phrase
        collection: which collection to search (default: "default")
        top_k: number of chunks to return (default from config)

        Each result includes the chunk text, source file, similarity score (0–1), and chunk position.
        Use the returned chunks as context when answering questions about ingested documents.
      `,
            parameters: {
                query: zod_1.z.string().describe("Search query"),
                collection: zod_1.z.string().default(DEFAULT_COLLECTION).describe("Collection to search"),
                top_k: zod_1.z.coerce.number().int().min(1).max(20).optional().describe("Number of results"),
            },
            implementation: safe_impl("rag_query", async ({ query, collection, top_k }, ctx) => {
                ctx.status("Embedding query…");
                const embed = await (0, embedder_1.getEmbedFn)(ctl.client, embeddingId());
                const [queryVec] = await embed([query]);
                ctx.status("Searching index…");
                const idx = (0, vectorStore_1.getIndex)(dataPath(), collection);
                const results = await (0, vectorStore_1.queryChunks)(idx, queryVec, top_k ?? topKDefault());
                if (results.length === 0) {
                    return json({ query, collection, results: [], hint: "No documents ingested yet. Use rag_ingest first." });
                }
                return json({
                    query,
                    collection,
                    results: results.map(r => ({
                        score: parseFloat(r.score.toFixed(4)),
                        source: r.metadata.source,
                        fileName: r.metadata.fileName,
                        chunkIndex: r.metadata.chunkIndex,
                        text: r.metadata.text,
                    })),
                });
            }),
        }),
        (0, sdk_1.tool)({
            name: "rag_list",
            description: (0, sdk_1.text) `
        List all documents ingested into a collection with chunk counts and ingest timestamps.
        collection: which collection to inspect (default: "default")
      `,
            parameters: {
                collection: zod_1.z.string().default(DEFAULT_COLLECTION).describe("Collection name"),
            },
            implementation: safe_impl("rag_list", async ({ collection }) => {
                const idx = (0, vectorStore_1.getIndex)(dataPath(), collection);
                const sources = await (0, vectorStore_1.listSources)(idx);
                return json({ collection, documentCount: sources.length, documents: sources });
            }),
        }),
        (0, sdk_1.tool)({
            name: "rag_delete",
            description: (0, sdk_1.text) `
        Remove all chunks for a specific document from the index.
        source_path: the exact file path used when the document was ingested
        collection: which collection to delete from (default: "default")
      `,
            parameters: {
                source_path: zod_1.z.string().describe("Exact file path of the document to remove"),
                collection: zod_1.z.string().default(DEFAULT_COLLECTION).describe("Collection name"),
            },
            implementation: safe_impl("rag_delete", async ({ source_path, collection }, ctx) => {
                ctx.status(`Deleting chunks for ${(0, path_1.basename)(source_path)}…`);
                const idx = (0, vectorStore_1.getIndex)(dataPath(), collection);
                const deleted = await (0, vectorStore_1.deleteBySource)(idx, source_path);
                return json({ deleted: deleted > 0, chunksRemoved: deleted, source: source_path, collection });
            }),
        }),
        (0, sdk_1.tool)({
            name: "rag_stats",
            description: (0, sdk_1.text) `
        Show index statistics: total documents, total chunks, and list of collections.
        Use to understand what's in the index before querying.
      `,
            parameters: {},
            implementation: safe_impl("rag_stats", async (_, ctx) => {
                ctx.status("Reading index stats…");
                const dp = dataPath();
                let collections = [];
                try {
                    const entries = await (0, promises_1.readdir)(dp, { withFileTypes: true });
                    collections = entries.filter(e => e.isDirectory()).map(e => e.name);
                }
                catch {
                    return json({ dataPath: dp, collections: [], totalDocuments: 0, totalChunks: 0, hint: "No documents ingested yet." });
                }
                const stats = [];
                let totalDocs = 0;
                let totalChunks = 0;
                for (const col of collections) {
                    const idx = (0, vectorStore_1.getIndex)(dp, col);
                    const s = await (0, vectorStore_1.indexStats)(idx);
                    stats.push({ collection: col, ...s });
                    totalDocs += s.documents;
                    totalChunks += s.chunks;
                }
                return json({
                    dataPath: dp,
                    embeddingModel: embeddingId() || "(any loaded)",
                    totalDocuments: totalDocs,
                    totalChunks,
                    collections: stats,
                });
            }),
        }),
    ];
    return tools;
};
exports.toolsProvider = toolsProvider;