Aubin/configuration assistant builder

front/components/assistant_builder/actions/MCPAction.tsx

@@ -1,20 +1,29 @@
+import { MIME_TYPES } from "@dust-tt/client";
 import {
+  Button,
   DropdownMenu,
   DropdownMenuContent,
   DropdownMenuItem,
   DropdownMenuTrigger,
 } from "@dust-tt/sparkle";
-import { Button } from "@dust-tt/sparkle";
-import { useState } from "react";
+import { useEffect, useState } from "react";
 
+import AssistantBuilderDataSourceModal from "@app/components/assistant_builder/AssistantBuilderDataSourceModal";
+import DataSourceSelectionSection from "@app/components/assistant_builder/DataSourceSelectionSection";
 import type {
   AssistantBuilderActionConfiguration,
   AssistantBuilderMCPServerConfiguration,
 } from "@app/components/assistant_builder/types";
 import { AVAILABLE_INTERNAL_MCPSERVER_IDS } from "@app/lib/actions/constants";
-import type { LightWorkspaceType, SpaceType } from "@app/types";
+import type { InternalMCPServerIdType } from "@app/lib/actions/mcp";
+import { useInternalMcpServerTools } from "@app/lib/swr/mcp";
+import type {
+  DataSourceViewSelectionConfigurations,
+  LightWorkspaceType,
+  SpaceType,
+} from "@app/types";
 
-type ActionMCPProps = {
+interface ActionMCPProps {
   owner: LightWorkspaceType;
   allowedSpaces: SpaceType[];
   actionConfiguration: AssistantBuilderMCPServerConfiguration;
@@ -24,27 +33,84 @@ type ActionMCPProps = {
     ) => AssistantBuilderMCPServerConfiguration
   ) => void;
   setEdited: (edited: boolean) => void;
-};
+}
 
 export function ActionMCP({
+  owner,
+  allowedSpaces,
   actionConfiguration,
   updateAction,
   setEdited,
 }: ActionMCPProps) {
-  const [selectedInternalMCPServerId, setSelectedInternalMCPServerId] =
-    useState<(typeof AVAILABLE_INTERNAL_MCPSERVER_IDS)[number] | null>(
-      actionConfiguration.internalMCPServerId
-    );
+  const [showDataSourcesModal, setShowDataSourcesModal] = useState(false);
+
+  const { tools } = useInternalMcpServerTools({
+    owner,
+    serverId: actionConfiguration.internalMCPServerId,
+  });
+
+  useEffect(() => {
+    updateAction((previousAction) => ({
+      ...previousAction,
+      resources: {
+        dataSourceConfigurations: tools?.some(
+          (r) => r.inputSchema.mimeType === MIME_TYPES.DATA_SOURCE_VIEW
+        )
+          ? previousAction.resources?.dataSourceConfigurations || {}
+          : undefined,
+      },
+    }));
+  }, [tools, setEdited, updateAction]);
+
+  const handleServerSelection = (serverId: InternalMCPServerIdType) => {
+    setEdited(true);
+    updateAction((previousAction) => ({
+      ...previousAction,
+      serverType: "internal",
+      internalMCPServerId: serverId,
+    }));
+  };
+
+  const handleDataSourceConfigUpdate = (
+    dsConfigs: DataSourceViewSelectionConfigurations
+  ) => {
+    setEdited(true);
+    updateAction((previousAction) => ({
+      ...previousAction,
+      resources: {
+        ...previousAction.resources,
+        dataSourceConfigurations: dsConfigs,
+      },
+    }));
+  };
 
   return (
     <>
+      {actionConfiguration.resources.dataSourceConfigurations && (
+        <AssistantBuilderDataSourceModal
+          isOpen={showDataSourcesModal}
+          setOpen={setShowDataSourcesModal}
+          owner={owner}
+          onSave={handleDataSourceConfigUpdate}
+          initialDataSourceConfigurations={
+            actionConfiguration.resources.dataSourceConfigurations
+          }
+          allowedSpaces={allowedSpaces}
+          viewType="document"
+        />
+      )}
+
       <div>Will expose all the tools available via an MCP Server.</div>
       <div>For testing purposes, pick an internal server</div>
+
       <DropdownMenu>
         <DropdownMenuTrigger asChild>
           <Button
             isSelect
-            label={selectedInternalMCPServerId ?? "Select a internal server"}
+            label={
+              actionConfiguration.internalMCPServerId ??
+              "Select a internal server"
+            }
             className="w-48"
           />
         </DropdownMenuTrigger>
@@ -53,19 +119,23 @@ export function ActionMCP({
             <DropdownMenuItem
               key={id}
               label={id}
-              onClick={() => {
-                setSelectedInternalMCPServerId(id);
-                updateAction((previousAction) => ({
-                  ...previousAction,
-                  serverType: "internal",
-                  internalMCPServerId: id,
-                }));
-                setEdited(true);
-              }}
+              onClick={() => handleServerSelection(id)}
             />
           ))}
         </DropdownMenuContent>
       </DropdownMenu>
+
+      {actionConfiguration.resources.dataSourceConfigurations && (
+        <DataSourceSelectionSection
+          owner={owner}
+          dataSourceConfigurations={
+            actionConfiguration.resources.dataSourceConfigurations
+          }
+          openDataSourceModal={() => setShowDataSourcesModal(true)}
+          onSave={handleDataSourceConfigUpdate}
+          viewType="document"
+        />
+      )}
     </>
   );
 }

front/lib/actions/constants.ts

@@ -4,7 +4,10 @@
 // future based on user feedback.
 export const PROCESS_ACTION_TOP_K = 768;
 
-export const AVAILABLE_INTERNAL_MCPSERVER_IDS = ["helloworld"] as const;
+export const AVAILABLE_INTERNAL_MCPSERVER_IDS = [
+  "helloworld",
+  "data-source-utils",
+] as const;
 
 export const DEFAULT_BROWSE_ACTION_NAME = "browse";
 export const DEFAULT_BROWSE_ACTION_DESCRIPTION =

front/lib/actions/mcp_actions.ts

@@ -64,16 +64,43 @@ const Schema = z.union([
   EmbeddedResourceSchema,
 ]);
 
+export type ToolType = ListToolsResult["tools"][number];
+
+export interface MCPServerConnectionDetails {
+  serverType: MCPServerConfigurationType["serverType"];
+  internalMCPServerId?:
+    | MCPServerConfigurationType["internalMCPServerId"]
+    | null;
+  remoteMCPServerId?: MCPServerConfigurationType["remoteMCPServerId"] | null;
+}
+
 export type MCPToolResultContent = z.infer<typeof Schema>;
 
+// The type provided in the sdk is widely incomplete, below is an attempt to normalize the inputSchema.
+// TODO(mcp): fix InputSchemaType accordingly.
+export function normalizeInputSchemaProperties(
+  inputSchema: ListToolsResult["tools"][number]["inputSchema"]
+) {
+  const result: Record<string, any> = {};
+
+  for (const key in inputSchema.properties) {
+    const field = inputSchema.properties[key];
+    if (typeof field === "object" && field !== null && "const" in field) {
+      result[key] = field?.const;
+    }
+  }
+
+  return result;
+}
+
 function makeMCPConfigurations({
   config,
   listToolsResult,
 }: {
   config: MCPServerConfigurationType;
-  listToolsResult: ListToolsResult;
+  listToolsResult: ToolType[];
 }): MCPToolConfigurationType[] {
-  return listToolsResult.tools.map((tool) => {
+  return listToolsResult.map((tool) => {
     return {
       id: config.id,
       sId: generateRandomModelSId(),
@@ -349,7 +376,11 @@ export async function tryGetMCPTools(
         // Connect to the MCP server.
         const mcpClient = await connectToMCPServer(action);
 
-        const r: ListToolsResult = await mcpClient.listTools();
+        const r: ToolType[] = await listMCPServerTools({
+          serverType: action.serverType,
+          internalMCPServerId: action.internalMCPServerId,
+          remoteMCPServerId: action.remoteMCPServerId,
+        });
 
         // Close immediately after listing tools.
         await mcpClient.close();
@@ -374,3 +405,21 @@ export async function tryGetMCPTools(
 
   return configurations.flat();
 }
+
+export async function listMCPServerTools(
+  connectionDetails: MCPServerConnectionDetails
+): Promise<ToolType[]> {
+  const mcpClient = await connectToMCPServer(connectionDetails);
+
+  let allTools: ToolType[] = [];
+  let nextPageCursor;
+  do {
+    const { tools, nextCursor } = await mcpClient.listTools();
+    nextPageCursor = nextCursor;
+    allTools = [...allTools, ...tools];
+  } while (nextPageCursor);
+
+  await mcpClient.close();
+
+  return allTools;
+}

front/lib/actions/mcp_internal_actions/data_source_utils.ts

@@ -0,0 +1,39 @@
+import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
+
+import { DataSourceViewInputSchema } from "@app/lib/actions/mcp_internal_actions/input_schemas";
+
+export const createServer = (): McpServer => {
+  const server = new McpServer(
+    {
+      name: "data-source-utils",
+      version: "1.0.0",
+      description:
+        "Demo server showing a basic interaction with a data source configuration.",
+    },
+    {
+      capabilities: {
+        resources: {},
+        tools: {},
+      },
+    }
+  );
+
+  server.tool(
+    "pass-through",
+    "Super useful tool that should be used at all time.",
+    DataSourceViewInputSchema.shape.inputSchema.shape,
+    ({ uri, mimeType }) => {
+      return {
+        isError: false,
+        content: [
+          {
+            type: "text",
+            text: `Got the URI: ${uri} and the mimeType ${mimeType}`,
+          },
+        ],
+      };
+    }
+  );
+
+  return server;
+};

front/lib/actions/mcp_internal_actions/index.ts

@@ -4,6 +4,7 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
 import type { MCPServerConfigurationType } from "@app/lib/actions/mcp";
 import { assertNever } from "@app/types";
 
+import { createServer as createDataSourceUtilsServer } from "./data_source_utils";
 import { createServer as createHelloWorldServer } from "./helloworld";
 
 export const connectToInternalMCPServer = async (
@@ -18,6 +19,10 @@ export const connectToInternalMCPServer = async (
     case "helloworld":
       server = createHelloWorldServer();
       break;
+    // TODO(mcp): add a variable for these server IDs.
+    case "data-source-utils":
+      server = createDataSourceUtilsServer();
+      break;
     default:
       assertNever(internalMCPServerId);
   }

front/lib/swr/mcp.ts

@@ -0,0 +1,30 @@
+import { useMemo } from "react";
+import type { Fetcher } from "swr";
+
+import { fetcher, useSWRWithDefaults } from "@app/lib/swr/swr";
+import type { GetMCPServerToolsResponseBody } from "@app/pages/api/w/[wId]/mcp/[serverId]/tools";
+import type { LightWorkspaceType } from "@app/types";
+
+export function useInternalMcpServerTools({
+  owner,
+  serverId,
+}: {
+  owner: LightWorkspaceType;
+  serverId: string | null;
+}) {
+  const configFetcher: Fetcher<GetMCPServerToolsResponseBody> = fetcher;
+
+  const { data, error, mutate } = useSWRWithDefaults(
+    serverId ? `/api/w/${owner.sId}/mcp/${serverId}/tools` : null,
+    configFetcher
+  );
+
+  const tools = useMemo(() => (data ? data.tools : null), [data]);
+
+  return {
+    tools,
+    isLoading: !error && !data,
+    isError: error,
+    mutate,
+  };
+}