rag

import { text, tool, type Tool, type ToolCallContext, type ToolsProvider } from "@lmstudio/sdk";
import { z } from "zod";
import { readdir, stat } from "fs/promises";
import { join, basename } from "path";
import { pluginConfigSchematics } from "./config";
import { getEmbedFn } from "./embedder";
import { chunkText } from "./chunker";
import { parseFile, resolvePath } from "./parser";
import {
  getIndex, insertChunks, queryChunks,
  deleteBySource, listSources, indexStats,
} from "./vectorStore";

function json(obj: unknown): string {
  return JSON.stringify(obj, null, 2);
}

function safe_impl<T extends Record<string, unknown>>(
  name: string,
  fn: (params: T, ctx: ToolCallContext) => Promise<string>,
): (params: T, ctx: ToolCallContext) => Promise<string> {
  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: unknown) {
      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

export const toolsProvider: ToolsProvider = async (ctl) => {
  const cfg = ctl.getPluginConfig(pluginConfigSchematics);

  const dataPath = () => {
    const p = cfg.get("dataPath").trim();
    return p || 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: Tool[] = [

    tool({
      name: "rag_ingest",
      description: 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: z.string().describe("File or directory path"),
        collection: z.string().default(DEFAULT_COLLECTION).describe("Collection name (default: \"default\")"),
        recursive: z.boolean().default(false).describe("Recurse into subdirectories"),
        chunk_size: z.coerce.number().int().min(100).max(8000).optional().describe("Override chunk size"),
        chunk_overlap: 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 getEmbedFn(ctl.client, embeddingId());

        // Collect files to ingest
        const fullPath = await resolvePath(path, ws, 9999); // size check per file below
        const info = await stat(fullPath);
        const files: string[] = [];

        if (info.isDirectory()) {
          const collect = async (dir: string) => {
            const entries = await readdir(dir, { withFileTypes: true });
            for (const e of entries) {
              const fp = 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 = getIndex(dp, collection);
        const results: Array<{ file: string; chunks: number; status: string }> = [];

        for (const filePath of files) {
          ctx.status(`Ingesting ${basename(filePath)}…`);
          try {
            const fileStat = await stat(filePath);
            if (fileStat.size / (1024 * 1024) > maxMb) {
              results.push({ file: filePath, chunks: 0, status: `skipped — exceeds ${maxMb} MB` });
              continue;
            }

            const { text: docText } = await 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 deleteBySource(idx, filePath);

            const chunks = 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: number[][] = [];

            // Embed in batches
            for (let i = 0; i < chunks.length; i += BATCH_SIZE) {
              const batch = chunks.slice(i, i + BATCH_SIZE);
              ctx.status(`Embedding ${basename(filePath)} (${i + batch.length}/${chunks.length})…`);
              const vecs = await embed(batch.map(c => c.text));
              allVectors.push(...vecs);
            }

            await insertChunks(idx, chunks.map((c, i) => ({
              vector: allVectors[i],
              metadata: {
                text: c.text,
                source: filePath,
                fileName: 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: unknown) {
            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 });
      }),
    }),

    tool({
      name: "rag_query",
      description: 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: z.string().describe("Search query"),
        collection: z.string().default(DEFAULT_COLLECTION).describe("Collection to search"),
        top_k: 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 getEmbedFn(ctl.client, embeddingId());
        const [queryVec] = await embed([query]);

        ctx.status("Searching index…");
        const idx = getIndex(dataPath(), collection);
        const results = await 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,
          })),
        });
      }),
    }),

    tool({
      name: "rag_list",
      description: text`
        List all documents ingested into a collection with chunk counts and ingest timestamps.
        collection: which collection to inspect (default: "default")
      `,
      parameters: {
        collection: z.string().default(DEFAULT_COLLECTION).describe("Collection name"),
      },
      implementation: safe_impl("rag_list", async ({ collection }) => {
        const idx = getIndex(dataPath(), collection);
        const sources = await listSources(idx);
        return json({ collection, documentCount: sources.length, documents: sources });
      }),
    }),

    tool({
      name: "rag_delete",
      description: 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: z.string().describe("Exact file path of the document to remove"),
        collection: z.string().default(DEFAULT_COLLECTION).describe("Collection name"),
      },
      implementation: safe_impl("rag_delete", async ({ source_path, collection }, ctx) => {
        ctx.status(`Deleting chunks for ${basename(source_path)}…`);
        const idx = getIndex(dataPath(), collection);
        const deleted = await deleteBySource(idx, source_path);
        return json({ deleted: deleted > 0, chunksRemoved: deleted, source: source_path, collection });
      }),
    }),

    tool({
      name: "rag_stats",
      description: 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: string[] = [];
        try {
          const entries = await 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: Array<{ collection: string; documents: number; chunks: number }> = [];
        let totalDocs = 0;
        let totalChunks = 0;

        for (const col of collections) {
          const idx = getIndex(dp, col);
          const s = await 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;
};