@2kw/ai-mcp-server 4.0.0-dev.2

package/dist/tools/experiments.js

@@ -0,0 +1,231 @@
+import { z } from "zod";
+import { formatErrorForMcp } from "../errors.js";
+export function register(server, client) {
+    // ── list_experiments ────────────────────────────────────────────
+    server.tool("2kw_list_experiments", "List experiments in the organization with optional search, status filter, and pagination.", {
+        search: z.string().optional().describe("Filter experiments by name"),
+        status: z
+            .enum(["DRAFT", "RUNNING", "COMPLETED", "FAILED"])
+            .optional()
+            .describe("Filter by experiment status"),
+        page: z.number().optional().default(0).describe("Page number (0-based)"),
+        size: z.number().optional().default(20).describe("Page size"),
+        sort: z.string().optional().describe("Sort field and direction"),
+    }, async ({ search, status, page, size, sort }) => {
+        try {
+            const { data } = await client.GET("/v1/experiments", {
+                params: {
+                    query: {
+                        search,
+                        status,
+                        pageable: {
+                            page,
+                            size,
+                            sort: sort ? [sort] : undefined,
+                        },
+                    },
+                },
+            });
+            const result = data;
+            const lines = (result.content ?? []).map((e) => `- ${e.name} (id: ${e.id}, status: ${e.status ?? "DRAFT"}, type: ${e.type ?? "—"})${e.description ? ` — ${e.description}` : ""}`);
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Experiments (${result.totalElements} total, page ${(result.number ?? 0) + 1}/${result.totalPages}):\n${lines.join("\n") || "(none)"}`,
+                    },
+                ],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── get_experiment ──────────────────────────────────────────────
+    server.tool("2kw_get_experiment", "Get an experiment by its ID, including status and dataset version reference.", {
+        experimentId: z.string().describe("The experiment ID"),
+    }, async ({ experimentId }) => {
+        try {
+            const { data } = await client.GET("/v1/experiments/{id}", {
+                params: { path: { id: experimentId } },
+            });
+            return {
+                content: [{ type: "text", text: JSON.stringify(data, null, 2) }],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── create_experiment ───────────────────────────────────────────
+    server.tool("2kw_create_experiment", "Create a new experiment, optionally linked to a dataset version.", {
+        name: z.string().min(1).describe("Experiment name"),
+        description: z.string().optional().describe("Experiment description"),
+        type: z.string().optional().describe("Experiment type"),
+        datasetVersionId: z
+            .string()
+            .optional()
+            .describe("Dataset version ID to evaluate against"),
+        metadata: z.unknown().optional().describe("Arbitrary metadata (JSON object)"),
+    }, async ({ name, description, type, datasetVersionId, metadata }) => {
+        try {
+            const body = { name };
+            if (description !== undefined)
+                body.description = description;
+            if (type !== undefined)
+                body.type = type;
+            if (datasetVersionId !== undefined)
+                body.datasetVersionId = datasetVersionId;
+            if (metadata !== undefined)
+                body.metadata = metadata;
+            const { data } = await client.POST("/v1/experiments", {
+                body: body,
+            });
+            return {
+                content: [{ type: "text", text: JSON.stringify(data, null, 2) }],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── add_variant ─────────────────────────────────────────────────
+    server.tool("2kw_add_variant", "Add a variant to an experiment with a task type and configuration.", {
+        experimentId: z.string().describe("The experiment ID"),
+        name: z.string().min(1).describe("Variant name"),
+        taskType: z.string().min(1).describe("Task type (e.g., EXTRACTION, CLASSIFICATION)"),
+        configuration: z.unknown().describe("Variant configuration (JSON object)"),
+        description: z.string().optional().describe("Variant description"),
+        sortOrder: z.number().optional().describe("Sort order for display"),
+    }, async ({ experimentId, name, taskType, configuration, description, sortOrder }) => {
+        try {
+            const body = { name, taskType, configuration };
+            if (description !== undefined)
+                body.description = description;
+            if (sortOrder !== undefined)
+                body.sortOrder = sortOrder;
+            const { data } = await client.POST("/v1/experiments/{id}/variants", {
+                params: { path: { id: experimentId } },
+                body: body,
+            });
+            return {
+                content: [{ type: "text", text: JSON.stringify(data, null, 2) }],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── run_experiment ──────────────────────────────────────────────
+    server.tool("2kw_run_experiment", "Start an experiment run. Returns immediately with the run ID — use 2kw_get_experiment_runs to poll for completion.", {
+        experimentId: z.string().describe("The experiment ID"),
+    }, async ({ experimentId }) => {
+        try {
+            const { data } = await client.POST("/v1/experiments/{id}/runs", {
+                params: { path: { id: experimentId } },
+            });
+            const run = data;
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Run started (id: ${run?.id}, status: ${run?.status ?? "PENDING"}).\n\nUse 2kw_get_experiment_runs to check progress, then 2kw_get_run_results to retrieve results once completed.\n\n${JSON.stringify(run, null, 2)}`,
+                    },
+                ],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── get_experiment_runs ─────────────────────────────────────────
+    server.tool("2kw_get_experiment_runs", "List runs for an experiment with status, progress, and timing information.", {
+        experimentId: z.string().describe("The experiment ID"),
+        page: z.number().optional().default(0).describe("Page number (0-based)"),
+        size: z.number().optional().default(20).describe("Page size"),
+    }, async ({ experimentId, page, size }) => {
+        try {
+            const { data } = await client.GET("/v1/experiments/{id}/runs", {
+                params: {
+                    path: { id: experimentId },
+                    query: { pageable: { page, size } },
+                },
+            });
+            const result = data;
+            const lines = (result.content ?? []).map((r) => {
+                const progress = r.itemsTotal
+                    ? `${r.itemsCompleted ?? 0}/${r.itemsTotal}${r.itemsFailed ? ` (${r.itemsFailed} failed)` : ""}`
+                    : "—";
+                return `- Run ${r.id} [${r.status}] variant=${r.variantName ?? r.variantId} progress=${progress}${r.completedAt ? ` completed=${r.completedAt}` : ""}`;
+            });
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Runs (${result.totalElements} total, page ${(result.number ?? 0) + 1}/${result.totalPages}):\n${lines.join("\n") || "(none)"}`,
+                    },
+                ],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── get_run_results ─────────────────────────────────────────────
+    server.tool("2kw_get_run_results", "Get detailed results for an experiment run, including output, token usage, cost, and errors per dataset item.", {
+        experimentId: z.string().describe("The experiment ID"),
+        runId: z.string().describe("The run ID"),
+        page: z.number().optional().default(0).describe("Page number (0-based)"),
+        size: z.number().optional().default(20).describe("Page size"),
+    }, async ({ experimentId, runId, page, size }) => {
+        try {
+            const { data } = await client.GET("/v1/experiments/{id}/runs/{runId}/results", {
+                params: {
+                    path: { id: experimentId, runId },
+                    query: { pageable: { page, size } },
+                },
+            });
+            const result = data;
+            const lines = (result.content ?? []).map((r) => {
+                const cost = r.estimatedCost != null ? `$${r.estimatedCost}` : "—";
+                const tokens = `${r.inputTokens ?? 0}in/${r.outputTokens ?? 0}out`;
+                const duration = r.durationMs != null ? `${r.durationMs}ms` : "—";
+                const error = r.error ? ` ERROR: ${JSON.stringify(r.error)}` : "";
+                return `- Item ${r.datasetItemId}: ${duration}, ${tokens}, ${cost}${error}`;
+            });
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Results (${result.totalElements} total, page ${(result.number ?? 0) + 1}/${result.totalPages}):\n${lines.join("\n") || "(none)"}`,
+                    },
+                ],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+}
+//# sourceMappingURL=experiments.js.map

package/dist/tools/extraction.d.ts

@@ -0,0 +1,4 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import type { ApiClient } from "../client.js";
+export declare function register(server: McpServer, client: ApiClient): void;
+//# sourceMappingURL=extraction.d.ts.map

package/dist/tools/extraction.js

@@ -0,0 +1,245 @@
+import { z } from "zod";
+import { readFile } from "node:fs/promises";
+import { basename, extname } from "node:path";
+import { formatErrorForMcp } from "../errors.js";
+const IMAGE_MIME_TYPES = {
+    ".png": "image/png",
+    ".jpg": "image/jpeg",
+    ".jpeg": "image/jpeg",
+    ".gif": "image/gif",
+    ".webp": "image/webp",
+};
+export function register(server, client) {
+    // ── create_extraction ───────────────────────────────────────────────────
+    server.tool("2kw_create_extraction", "Extract structured data from text and/or images using a schema and AI model. Supports text-only, vision (images), or combined extraction. Use a vision-capable model (e.g. openai/gpt-4o) when providing images. IMPORTANT: inputText accepts only plain text/markdown strings. For binary files (PDF, DOCX, XLSX, etc.), first convert them using 2kw_convert_document to get markdown, then pass the mdContent here as inputText. Use async=true for large documents to avoid gateway timeouts.", {
+        schemaId: z.string().describe("The schema ID to extract with"),
+        schemaVersionId: z
+            .string()
+            .optional()
+            .describe("Specific schema version ID (uses latest active version if omitted)"),
+        inputText: z
+            .string()
+            .optional()
+            .describe("The text to extract structured data from (required unless inputImages provided)"),
+        inputImages: z
+            .array(z.object({
+            type: z.enum(["base64", "file"]).describe("Source type: 'base64' for pre-encoded data, 'file' for local file path"),
+            data: z.string().optional().describe("Base64-encoded image data (when type='base64')"),
+            path: z.string().optional().describe("Local file path to image (when type='file')"),
+            mimeType: z.string().optional().describe("MIME type (auto-detected for files, required for base64)"),
+        }))
+            .optional()
+            .describe("Images for vision extraction (max 10). Requires a vision-capable model."),
+        model: z.string().describe("AI model in 'provider/model' format"),
+        async: z
+            .boolean()
+            .optional()
+            .default(false)
+            .describe("If true, submit as async extraction and return immediately (not supported with images)"),
+    }, async ({ schemaId, schemaVersionId, inputText, inputImages, model, async: isAsync }) => {
+        try {
+            // Resolve images: read local files and convert to base64
+            const resolvedImages = [];
+            if (inputImages?.length) {
+                for (const img of inputImages) {
+                    if (img.type === "file") {
+                        if (!img.path)
+                            throw new Error("path is required for type='file'");
+                        const fileBuffer = await readFile(img.path);
+                        const ext = extname(img.path).toLowerCase();
+                        const mimeType = img.mimeType ?? IMAGE_MIME_TYPES[ext];
+                        if (!mimeType)
+                            throw new Error(`Cannot determine MIME type for ${basename(img.path)}. Supported: ${Object.keys(IMAGE_MIME_TYPES).join(", ")}`);
+                        resolvedImages.push({ data: fileBuffer.toString("base64"), mimeType });
+                    }
+                    else {
+                        if (!img.data)
+                            throw new Error("data is required for type='base64'");
+                        if (!img.mimeType)
+                            throw new Error("mimeType is required for type='base64'");
+                        resolvedImages.push({ data: img.data, mimeType: img.mimeType });
+                    }
+                }
+            }
+            const body = {
+                schemaId,
+                model,
+                ...(inputText && { inputText }),
+                ...(resolvedImages.length > 0 && { inputImages: resolvedImages }),
+                ...(schemaVersionId && { schemaVersionId }),
+            };
+            if (isAsync) {
+                const { data } = await client.POST("/v1/extractions/async", { body });
+                const result = data;
+                return {
+                    content: [
+                        {
+                            type: "text",
+                            text: `Extraction submitted (id: ${result.id}, status: ${result.status}). Poll with 2kw_get_extraction to check progress.`,
+                        },
+                    ],
+                };
+            }
+            const { data } = await client.POST("/v1/extractions", { body });
+            const result = data;
+            if (result.status === "PENDING" || result.status === "PROCESSING") {
+                return {
+                    content: [
+                        {
+                            type: "text",
+                            text: `Extraction submitted (id: ${result.id}, status: ${result.status}). Poll with 2kw_get_extraction to check progress.`,
+                        },
+                    ],
+                };
+            }
+            return { content: formatExtractionResult(result) };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── get_extraction ──────────────────────────────────────────────────────
+    server.tool("2kw_get_extraction", "Get an extraction by its ID. Use this to check the status of async extractions or retrieve results.", {
+        extractionId: z.string().describe("The extraction ID"),
+    }, async ({ extractionId }) => {
+        try {
+            const { data } = await client.GET("/v1/extractions/{id}", {
+                params: { path: { id: extractionId } },
+            });
+            return { content: formatExtractionResult(data) };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── list_extractions ────────────────────────────────────────────────────
+    server.tool("2kw_list_extractions", "List extractions in the organization with optional filtering by status, schema version, or search term.", {
+        status: z
+            .enum(["PENDING", "PROCESSING", "COMPLETED", "FAILED"])
+            .optional()
+            .describe("Filter by extraction status"),
+        schemaVersionId: z.string().optional().describe("Filter by schema version ID"),
+        search: z.string().optional().describe("Search by model name"),
+        page: z.number().optional().default(0).describe("Page number (0-based)"),
+        size: z.number().optional().default(20).describe("Page size"),
+        sort: z.string().optional().describe("Sort field and direction (default: createdAt,desc)"),
+    }, async ({ status, schemaVersionId, search, page, size, sort }) => {
+        try {
+            const { data } = await client.GET("/v1/extractions", {
+                params: {
+                    query: {
+                        status,
+                        schemaVersionId,
+                        search,
+                        pageable: {
+                            page,
+                            size,
+                            sort: sort ? [sort] : undefined,
+                        },
+                    },
+                },
+            });
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            const result = data;
+            const lines = (result?.content ?? []).map(
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            (e) => `- ${e.id} [${e.status}] model=${e.model} schema=${e.schemaId}${e.errorMessage ? ` error: ${e.errorMessage}` : ""}`);
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Extractions (${result?.totalElements} total, page ${(result?.number ?? 0) + 1}/${result?.totalPages}):\n${lines.join("\n") || "(none)"}`,
+                    },
+                ],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── estimate_tokens ─────────────────────────────────────────────────────
+    server.tool("2kw_estimate_tokens", "Estimate token usage for an extraction without executing it. Useful for cost estimation.", {
+        schemaId: z.string().describe("The schema ID"),
+        schemaVersionId: z.string().optional().describe("Specific schema version ID"),
+        inputText: z.string().describe("The text to estimate tokens for"),
+    }, async ({ schemaId, schemaVersionId, inputText }) => {
+        try {
+            const { data } = await client.POST("/v1/extractions/estimate", {
+                body: {
+                    schemaId,
+                    inputText,
+                    ...(schemaVersionId && { schemaVersionId }),
+                },
+            });
+            // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            const result = data;
+            return {
+                content: [
+                    {
+                        type: "text",
+                        text: `Token estimate:\n  Input tokens: ${result?.inputTokens}\n  Estimated output tokens: ${result?.estimatedOutputTokens}\n  Strategy: ${result?.strategy}`,
+                    },
+                ],
+            };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+    // ── rerun_extraction ────────────────────────────────────────────────────
+    server.tool("2kw_rerun_extraction", "Re-run an existing extraction with the same configuration. Creates a new extraction.", {
+        extractionId: z.string().describe("The extraction ID to re-run"),
+    }, async ({ extractionId }) => {
+        try {
+            const { data } = await client.POST("/v1/extractions/{id}/rerun", {
+                params: { path: { id: extractionId } },
+            });
+            return { content: formatExtractionResult(data) };
+        }
+        catch (error) {
+            return {
+                content: [{ type: "text", text: formatErrorForMcp(error) }],
+                isError: true,
+            };
+        }
+    });
+}
+function formatExtractionResult(extraction) {
+    const parts = [];
+    parts.push({
+        type: "text",
+        text: `Extraction ${extraction.id} [${extraction.status}]\nModel: ${extraction.model}${extraction.strategy ? `\nStrategy: ${extraction.strategy}` : ""}`,
+    });
+    if (extraction.status === "COMPLETED" && extraction.result) {
+        parts.push({
+            type: "text",
+            text: `Result:\n${JSON.stringify(extraction.result, null, 2)}`,
+        });
+    }
+    if (extraction.status === "FAILED" && extraction.errorMessage) {
+        parts.push({
+            type: "text",
+            text: `Error: ${extraction.errorMessage}`,
+        });
+    }
+    if (extraction.estimatedInputTokens != null) {
+        parts.push({
+            type: "text",
+            text: `Tokens: ~${extraction.estimatedInputTokens} input, ~${extraction.estimatedOutputTokens} output`,
+        });
+    }
+    return parts;
+}
+//# sourceMappingURL=extraction.js.map

package/dist/tools/prompts.d.ts

@@ -0,0 +1,4 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+import type { ApiClient } from "../client.js";
+export declare function register(server: McpServer, client: ApiClient): void;
+//# sourceMappingURL=prompts.d.ts.map