@lightdash-tools/mcp 0.2.5 → 0.3.1

package/src/tools/ai-agents.ts

@@ -0,0 +1,759 @@
+/**
+ * MCP tools: AI agents (admin + project-scoped), threads, and evaluations.
+ */
+
+import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
+import type { LightdashClient } from '@lightdash-tools/client';
+import { z } from 'zod';
+import {
+  wrapTool,
+  registerToolSafe,
+  READ_ONLY_DEFAULT,
+  WRITE_IDEMPOTENT,
+  WRITE_DESTRUCTIVE,
+} from './shared.js';
+
+export function registerAiAgentTools(server: McpServer, client: LightdashClient): void {
+  // ─── Admin: agents ───────────────────────────────────────────────────────────
+
+  registerToolSafe(
+    server,
+    'list_admin_agents',
+    {
+      title: 'List AI agents (admin)',
+      description: 'List all AI agents across the organization (admin view)',
+      inputSchema: {},
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(client, (c) => async () => {
+      const result = await c.v1.aiAgents.listAdminAgents();
+      return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+    }),
+  );
+
+  // ─── Admin: threads ──────────────────────────────────────────────────────────
+
+  registerToolSafe(
+    server,
+    'list_admin_agent_threads',
+    {
+      title: 'List AI agent threads (admin)',
+      description:
+        'List AI agent conversation threads across the organization with optional filters',
+      inputSchema: {
+        page: z.number().optional().describe('Page number (1-based)'),
+        pageSize: z.number().optional().describe('Number of results per page'),
+        agentUuids: z.array(z.string()).optional().describe('Filter by agent UUIDs'),
+        projectUuids: z.array(z.string()).optional().describe('Filter by project UUIDs'),
+        humanScore: z
+          .number()
+          .optional()
+          .describe('Filter by human score: -1 (negative), 0 (neutral), 1 (positive)'),
+        dateFrom: z.string().optional().describe('Start date filter (YYYY-MM-DD)'),
+        dateTo: z.string().optional().describe('End date filter (YYYY-MM-DD)'),
+        sortField: z.enum(['createdAt', 'title']).optional().describe('Sort field'),
+        sortDirection: z.enum(['asc', 'desc']).optional().describe('Sort direction'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async (params: {
+          page?: number;
+          pageSize?: number;
+          agentUuids?: string[];
+          projectUuids?: string[];
+          humanScore?: number;
+          dateFrom?: string;
+          dateTo?: string;
+          sortField?: 'createdAt' | 'title';
+          sortDirection?: 'asc' | 'desc';
+        }) => {
+          const result = await c.v1.aiAgents.getAdminThreads(params);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  // ─── Admin: settings ─────────────────────────────────────────────────────────
+
+  registerToolSafe(
+    server,
+    'get_ai_organization_settings',
+    {
+      title: 'Get AI organization settings',
+      description: 'Get the AI settings for the current organization (admin)',
+      inputSchema: {},
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(client, (c) => async () => {
+      const result = await c.v1.aiAgents.getAiOrganizationSettings();
+      return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+    }),
+  );
+
+  registerToolSafe(
+    server,
+    'update_ai_organization_settings',
+    {
+      title: 'Update AI organization settings',
+      description: 'Update the AI settings for the current organization (admin)',
+      inputSchema: {
+        aiAgentsVisible: z
+          .boolean()
+          .optional()
+          .describe('Whether AI agents feature is visible to users'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(client, (c) => async (params: { aiAgentsVisible?: boolean }) => {
+      const result = await c.v1.aiAgents.updateAiOrganizationSettings(
+        params as Parameters<typeof c.v1.aiAgents.updateAiOrganizationSettings>[0],
+      );
+      return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+    }),
+  );
+
+  // ─── Project-scoped: agent CRUD ──────────────────────────────────────────────
+
+  registerToolSafe(
+    server,
+    'list_project_agents',
+    {
+      title: 'List agents in a project',
+      description: 'List all AI agents configured for a specific project',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(client, (c) => async ({ projectUuid }: { projectUuid: string }) => {
+      const result = await c.v1.aiAgents.listAgents(projectUuid);
+      return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+    }),
+  );
+
+  registerToolSafe(
+    server,
+    'get_project_agent',
+    {
+      title: 'Get agent',
+      description: 'Get details of a specific AI agent in a project',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({ projectUuid, agentUuid }: { projectUuid: string; agentUuid: string }) => {
+          const result = await c.v1.aiAgents.getAgent(projectUuid, agentUuid);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'create_project_agent',
+    {
+      title: 'Create agent',
+      description: 'Create a new AI agent in a project',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        name: z.string().describe('Agent name'),
+        description: z.string().optional().describe('Agent description'),
+        instruction: z.string().optional().describe('System instruction for the agent'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          name,
+          description,
+          instruction,
+        }: {
+          projectUuid: string;
+          name: string;
+          description?: string;
+          instruction?: string;
+        }) => {
+          const body = {
+            name,
+            projectUuid,
+            ...(description != null ? { description } : {}),
+            ...(instruction != null ? { instruction } : {}),
+          } as Parameters<typeof c.v1.aiAgents.createAgent>[1];
+          const result = await c.v1.aiAgents.createAgent(projectUuid, body);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'update_project_agent',
+    {
+      title: 'Update agent',
+      description: 'Update an existing AI agent',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        name: z.string().optional().describe('New name'),
+        description: z.string().optional().describe('New description'),
+        instruction: z.string().optional().describe('New system instruction'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          ...body
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          name?: string;
+          description?: string;
+          instruction?: string;
+        }) => {
+          const result = await c.v1.aiAgents.updateAgent(
+            projectUuid,
+            agentUuid,
+            body as Parameters<typeof c.v1.aiAgents.updateAgent>[2],
+          );
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'delete_project_agent',
+    {
+      title: 'Delete agent',
+      description: 'Delete an AI agent from a project',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+      },
+      annotations: WRITE_DESTRUCTIVE,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({ projectUuid, agentUuid }: { projectUuid: string; agentUuid: string }) => {
+          await c.v1.aiAgents.deleteAgent(projectUuid, agentUuid);
+          return {
+            content: [{ type: 'text', text: `Agent ${agentUuid} deleted successfully` }],
+          };
+        },
+    ),
+  );
+
+  // ─── Project-scoped: threads ─────────────────────────────────────────────────
+
+  registerToolSafe(
+    server,
+    'list_agent_threads',
+    {
+      title: 'List agent threads',
+      description: 'List all conversation threads for an agent',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({ projectUuid, agentUuid }: { projectUuid: string; agentUuid: string }) => {
+          const result = await c.v1.aiAgents.listAgentThreads(projectUuid, agentUuid);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'get_agent_thread',
+    {
+      title: 'Get agent thread',
+      description: 'Get a conversation thread with all its messages',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        threadUuid: z.string().describe('Thread UUID'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          threadUuid,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          threadUuid: string;
+        }) => {
+          const result = await c.v1.aiAgents.getAgentThread(projectUuid, agentUuid, threadUuid);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'generate_agent_message',
+    {
+      title: 'Generate agent message',
+      description:
+        'Start a new conversation thread and generate the first agent response for a given prompt',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        prompt: z.string().describe('User prompt to send to the agent'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          prompt,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          prompt: string;
+        }) => {
+          const thread = await c.v1.aiAgents.createAgentThread(projectUuid, agentUuid);
+          const result = await c.v1.aiAgents.generateAgentThreadResponse(
+            projectUuid,
+            agentUuid,
+            thread.uuid,
+            { prompt },
+          );
+          return {
+            content: [
+              {
+                type: 'text',
+                text: JSON.stringify({ threadUuid: thread.uuid, ...result }, null, 2),
+              },
+            ],
+          };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'continue_agent_thread',
+    {
+      title: 'Continue agent thread',
+      description: 'Continue an existing conversation thread with a new prompt',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        threadUuid: z.string().describe('Thread UUID to continue'),
+        prompt: z.string().describe('Follow-up prompt'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          threadUuid,
+          prompt,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          threadUuid: string;
+          prompt: string;
+        }) => {
+          const result = await c.v1.aiAgents.generateAgentThreadResponse(
+            projectUuid,
+            agentUuid,
+            threadUuid,
+            { prompt },
+          );
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  // ─── Project-scoped: evaluations ─────────────────────────────────────────────
+
+  registerToolSafe(
+    server,
+    'list_agent_evaluations',
+    {
+      title: 'List agent evaluations',
+      description: 'List all evaluations for an agent',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({ projectUuid, agentUuid }: { projectUuid: string; agentUuid: string }) => {
+          const result = await c.v1.aiAgents.listEvaluations(projectUuid, agentUuid);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'get_agent_evaluation',
+    {
+      title: 'Get agent evaluation',
+      description: 'Get a full evaluation including its test prompts',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        evalUuid: z.string().describe('Evaluation UUID'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          evalUuid,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          evalUuid: string;
+        }) => {
+          const result = await c.v1.aiAgents.getEvaluation(projectUuid, agentUuid, evalUuid);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'create_agent_evaluation',
+    {
+      title: 'Create agent evaluation',
+      description:
+        'Create a new evaluation test suite for an agent with a title and optional prompts',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        title: z.string().describe('Evaluation title'),
+        description: z.string().optional().describe('Evaluation description'),
+        prompts: z
+          .array(
+            z.union([
+              z.object({
+                prompt: z.string().describe('Test prompt text'),
+                expectedResponse: z.string().nullable().describe('Expected response (optional)'),
+              }),
+              z.object({
+                threadUuid: z.string().describe('Existing thread UUID'),
+                promptUuid: z.string().describe('Existing prompt UUID within the thread'),
+                expectedResponse: z.string().nullable().describe('Expected response (optional)'),
+              }),
+            ]),
+          )
+          .optional()
+          .describe('Test prompts for the evaluation'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          title,
+          description,
+          prompts,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          title: string;
+          description?: string;
+          prompts?: Parameters<typeof c.v1.aiAgents.createEvaluation>[2]['prompts'];
+        }) => {
+          const body: Parameters<typeof c.v1.aiAgents.createEvaluation>[2] = {
+            title,
+            prompts: prompts ?? [],
+            ...(description != null ? { description } : {}),
+          };
+          const result = await c.v1.aiAgents.createEvaluation(projectUuid, agentUuid, body);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'update_agent_evaluation',
+    {
+      title: 'Update agent evaluation',
+      description: 'Update an evaluation title, description, or replace its prompts',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        evalUuid: z.string().describe('Evaluation UUID'),
+        title: z.string().optional().describe('New title'),
+        description: z.string().optional().describe('New description'),
+        prompts: z
+          .array(
+            z.union([
+              z.object({
+                prompt: z.string().describe('Test prompt text'),
+                expectedResponse: z.string().nullable().describe('Expected response (optional)'),
+              }),
+              z.object({
+                threadUuid: z.string().describe('Existing thread UUID'),
+                promptUuid: z.string().describe('Existing prompt UUID within the thread'),
+                expectedResponse: z.string().nullable().describe('Expected response (optional)'),
+              }),
+            ]),
+          )
+          .optional()
+          .describe('Replacement prompt list (omit to leave unchanged)'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          evalUuid,
+          ...body
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          evalUuid: string;
+          title?: string;
+          description?: string;
+          prompts?: Parameters<typeof c.v1.aiAgents.updateEvaluation>[3]['prompts'];
+        }) => {
+          const result = await c.v1.aiAgents.updateEvaluation(
+            projectUuid,
+            agentUuid,
+            evalUuid,
+            body as Parameters<typeof c.v1.aiAgents.updateEvaluation>[3],
+          );
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'append_agent_evaluation_prompts',
+    {
+      title: 'Append evaluation prompts',
+      description:
+        'Append additional prompts to an existing evaluation without replacing existing ones',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        evalUuid: z.string().describe('Evaluation UUID'),
+        prompts: z
+          .array(
+            z.union([
+              z.object({
+                prompt: z.string().describe('Test prompt text'),
+                expectedResponse: z.string().nullable().describe('Expected response (optional)'),
+              }),
+              z.object({
+                threadUuid: z.string().describe('Existing thread UUID'),
+                promptUuid: z.string().describe('Existing prompt UUID within the thread'),
+                expectedResponse: z.string().nullable().describe('Expected response (optional)'),
+              }),
+            ]),
+          )
+          .describe('Prompts to append'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          evalUuid,
+          prompts,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          evalUuid: string;
+          prompts: Parameters<typeof c.v1.aiAgents.appendToEvaluation>[3]['prompts'];
+        }) => {
+          await c.v1.aiAgents.appendToEvaluation(projectUuid, agentUuid, evalUuid, { prompts });
+          return {
+            content: [
+              { type: 'text', text: `Prompts appended to evaluation ${evalUuid} successfully` },
+            ],
+          };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'run_agent_evaluation',
+    {
+      title: 'Run agent evaluation',
+      description: 'Trigger a new evaluation run for an agent. Returns the run UUID and status.',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        evalUuid: z.string().describe('Evaluation UUID to run'),
+      },
+      annotations: WRITE_IDEMPOTENT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          evalUuid,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          evalUuid: string;
+        }) => {
+          const result = await c.v1.aiAgents.runEvaluation(projectUuid, agentUuid, evalUuid);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'list_agent_evaluation_runs',
+    {
+      title: 'List evaluation runs',
+      description: 'List all runs for an evaluation with their status and pass/fail counts',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        evalUuid: z.string().describe('Evaluation UUID'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          evalUuid,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          evalUuid: string;
+        }) => {
+          const result = await c.v1.aiAgents.listEvaluationRuns(projectUuid, agentUuid, evalUuid);
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'get_agent_evaluation_run_results',
+    {
+      title: 'Get evaluation run results',
+      description:
+        'Get detailed per-prompt results for a specific evaluation run, including pass/fail and assessments',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        evalUuid: z.string().describe('Evaluation UUID'),
+        runUuid: z.string().describe('Run UUID'),
+      },
+      annotations: READ_ONLY_DEFAULT,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          evalUuid,
+          runUuid,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          evalUuid: string;
+          runUuid: string;
+        }) => {
+          const result = await c.v1.aiAgents.getEvaluationRunResults(
+            projectUuid,
+            agentUuid,
+            evalUuid,
+            runUuid,
+          );
+          return { content: [{ type: 'text', text: JSON.stringify(result, null, 2) }] };
+        },
+    ),
+  );
+
+  registerToolSafe(
+    server,
+    'delete_agent_evaluation',
+    {
+      title: 'Delete agent evaluation',
+      description: 'Delete an evaluation and all its runs',
+      inputSchema: {
+        projectUuid: z.string().describe('Project UUID'),
+        agentUuid: z.string().describe('Agent UUID'),
+        evalUuid: z.string().describe('Evaluation UUID'),
+      },
+      annotations: WRITE_DESTRUCTIVE,
+    },
+    wrapTool(
+      client,
+      (c) =>
+        async ({
+          projectUuid,
+          agentUuid,
+          evalUuid,
+        }: {
+          projectUuid: string;
+          agentUuid: string;
+          evalUuid: string;
+        }) => {
+          await c.v1.aiAgents.deleteEvaluation(projectUuid, agentUuid, evalUuid);
+          return {
+            content: [{ type: 'text', text: `Evaluation ${evalUuid} deleted successfully` }],
+          };
+        },
+    ),
+  );
+}