feat(compass-assistant): add basic tool calling & tool cards COMPASS-10144

packages/compass-aggregations/src/components/pipeline-builder-workspace/pipeline-as-text-workspace/pipeline-editor.tsx

@@ -23,6 +23,7 @@ import type { PipelineParserError } from '../../../modules/pipeline-builder/pipe
 import { useAutocompleteFields } from '@mongodb-js/compass-field-store';
 import { useTelemetry } from '@mongodb-js/compass-telemetry/provider';
 import { useConnectionInfoRef } from '@mongodb-js/compass-connections/provider';
+import { useSyncAssistantGlobalState } from '@mongodb-js/compass-assistant';
 
 const containerStyles = css({
   position: 'relative',
@@ -86,6 +87,8 @@ export const PipelineEditor: React.FunctionComponent<PipelineEditorProps> = ({
   const editorInitialValueRef = useRef<string>(pipelineText);
   const editorCurrentValueRef = useCurrentValueRef<string>(pipelineText);
 
+  useSyncAssistantGlobalState('currentAggregation', pipelineText);
+
   const { utmSource, utmMedium } = useRequiredURLSearchParams();
 
   const completer = useMemo(() => {

packages/compass-aggregations/src/components/pipeline-builder-workspace/pipeline-builder-ui-workspace/index.tsx

@@ -17,6 +17,8 @@ import AddStage from '../../add-stage';
 import UseCaseDroppableArea from '../../use-case-droppable-area';
 import type { StageIdAndType } from '../../../modules/pipeline-builder/stage-editor';
 import PipelineBuilderDndWrapper from './dnd-wrapper';
+import { prettify } from '../../../modules/pipeline-builder/pipeline-parser/utils';
+import { useSyncAssistantGlobalState } from '@mongodb-js/compass-assistant';
 
 const pipelineWorkspaceContainerStyles = css({
   position: 'relative',
@@ -36,6 +38,7 @@ const pipelineWorkspaceStyles = css({
 
 export type PipelineBuilderUIWorkspaceProps = {
   stagesIdAndType: StageIdAndType[];
+  pipelineText: string;
   isSidePanelOpen: boolean;
   onStageMoveEnd: (from: number, to: number) => void;
   onStageAddAfterEnd: (after?: number) => void;
@@ -50,11 +53,14 @@ export const PipelineBuilderUIWorkspace: React.FunctionComponent<
   PipelineBuilderUIWorkspaceProps
 > = ({
   stagesIdAndType,
+  pipelineText,
   isSidePanelOpen,
   onStageMoveEnd,
   onStageAddAfterEnd,
   onUseCaseDropped,
 }) => {
+  useSyncAssistantGlobalState('currentAggregation', pipelineText);
+
   return (
     <PipelineBuilderDndWrapper
       stagesIdAndType={stagesIdAndType}
@@ -99,9 +105,34 @@ export const PipelineBuilderUIWorkspace: React.FunctionComponent<
   );
 };
 
+type Stage = {
+  disabled?: boolean;
+  syntaxError: Error | null;
+  stageOperator: string | null;
+  value: string | null;
+};
+
+function stageToString(stage: Stage): string {
+  return `{ ${stage.stageOperator}: ${stage.value} }`;
+}
+
+function getPipelineTextFromStages(stages: Stage[]): string {
+  const code = `[${stages
+    .filter(
+      (stage) =>
+        stage.stageOperator !== null && stage.value !== null && !stage.disabled
+    )
+    .map((stage) => stageToString(stage))
+    .join(',\n')}\n]`;
+  return prettify(code);
+}
+
 const mapState = (state: RootState) => {
   return {
     stagesIdAndType: state.pipelineBuilder.stageEditor.stagesIdAndType,
+    pipelineText: getPipelineTextFromStages(
+      state.pipelineBuilder.stageEditor.stages
+    ),
     isSidePanelOpen: state.sidePanel.isPanelOpen,
   };
 };

packages/compass-assistant/src/@ai-sdk/react/use-chat.ts

@@ -41,6 +41,7 @@ export type UseChatHelpers<UI_MESSAGE extends UIMessage> = {
   | 'stop'
   | 'resumeStream'
   | 'addToolResult'
+  | 'addToolApprovalResponse'
   | 'status'
   | 'messages'
   | 'clearError'
@@ -124,5 +125,6 @@ export function useChat<UI_MESSAGE extends UIMessage = UIMessage>({
     resumeStream: chatRef.current.resumeStream,
     status,
     addToolResult: chatRef.current.addToolResult,
+    addToolApprovalResponse: chatRef.current.addToolApprovalResponse,
   };
 }

packages/compass-assistant/src/assistant-global-state.tsx

@@ -10,8 +10,8 @@ export type GlobalState = {
   activeConnections: ConnectionInfo[];
   activeWorkspace: WorkspaceTab | null;
   activeCollectionMetadata: CollectionMetadata | null;
-  currentQuery: object | null;
-  currentAggregation: object | null;
+  currentQuery: string | null;
+  currentAggregation: string | null;
   activeCollectionSubTab: CollectionSubtab | null;
 };
 

packages/compass-assistant/src/compass-assistant-provider.tsx

@@ -26,6 +26,8 @@ import {
   type ProactiveInsightsContext,
 } from './prompts';
 import {
+  type PreferencesAccess,
+  preferencesLocator,
   useIsAIFeatureEnabled,
   usePreference,
 } from 'compass-preferences-model/provider';
@@ -39,14 +41,24 @@ import {
   type TrackFunction,
   useTelemetry,
 } from '@mongodb-js/compass-telemetry/provider';
-import type { AtlasAiService } from '@mongodb-js/compass-generative-ai/provider';
-import { atlasAiServiceLocator } from '@mongodb-js/compass-generative-ai/provider';
+import type {
+  AtlasAiService,
+  ToolsController,
+} from '@mongodb-js/compass-generative-ai/provider';
+import {
+  atlasAiServiceLocator,
+  toolsControllerLocator,
+} from '@mongodb-js/compass-generative-ai/provider';
 import { buildConversationInstructionsPrompt } from './prompts';
 import { createOpenAI } from '@ai-sdk/openai';
 import {
   AssistantGlobalStateProvider,
   useAssistantGlobalState,
 } from './assistant-global-state';
+import {
+  lastAssistantMessageIsCompleteWithApprovalResponses,
+  type ToolSet,
+} from 'ai';
 
 export const ASSISTANT_DRAWER_ID = 'compass-assistant-drawer';
 
@@ -200,8 +212,10 @@ export const AssistantProvider: React.FunctionComponent<
     appNameForPrompt: string;
     chat: Chat<AssistantMessage>;
     atlasAiService: AtlasAiService;
+    toolsController: ToolsController;
+    preferences: PreferencesAccess;
   }>
-> = ({ chat, atlasAiService, children }) => {
+> = ({ chat, atlasAiService, toolsController, preferences, children }) => {
   const { openDrawer } = useDrawerActions();
   const track = useTelemetry();
 
@@ -272,6 +286,22 @@ export const AssistantProvider: React.FunctionComponent<
         chat.messages = [...chat.messages, contextPrompt];
       }
 
+      const { enableToolCalling } = preferences.getPreferences();
+
+      if (enableToolCalling) {
+        toolsController.setActiveTools(new Set(['compass-ui']));
+        toolsController.setContext({
+          query: assistantGlobalStateRef.current.currentQuery || undefined,
+          aggregation:
+            assistantGlobalStateRef.current.currentAggregation || undefined,
+        });
+      } else {
+        toolsController.setActiveTools(new Set([]));
+        toolsController.setContext({
+          query: undefined,
+          aggregation: undefined,
+        });
+      }
       await chat.sendMessage(message, options);
     };
   });
@@ -340,25 +370,37 @@ export const CompassAssistantProvider = registerCompassPlugin(
       appNameForPrompt,
       chat,
       atlasAiService,
+      toolsController,
+      preferences,
       children,
     }: PropsWithChildren<{
       appNameForPrompt: string;
       originForPrompt: string;
       chat?: Chat<AssistantMessage>;
       atlasAiService?: AtlasAiService;
+      toolsController?: ToolsController;
+      preferences?: PreferencesAccess;
     }>) => {
       if (!chat) {
         throw new Error('Chat was not provided by the state');
       }
       if (!atlasAiService) {
         throw new Error('atlasAiService was not provided by the state');
       }
+      if (!toolsController) {
+        throw new Error('toolsController was not provided by the state');
+      }
+      if (!preferences) {
+        throw new Error('preferences was not provided by the state');
+      }
       return (
         <AssistantGlobalStateProvider>
           <AssistantProvider
             appNameForPrompt={appNameForPrompt}
             chat={chat}
             atlasAiService={atlasAiService}
+            toolsController={toolsController}
+            preferences={preferences}
           >
             {children}
           </AssistantProvider>
@@ -367,7 +409,14 @@ export const CompassAssistantProvider = registerCompassPlugin(
     },
     activate: (
       { chat: initialChat, originForPrompt, appNameForPrompt },
-      { atlasService, atlasAiService, logger, track }
+      {
+        atlasService,
+        atlasAiService,
+        toolsController,
+        preferences,
+        logger,
+        track,
+      }
     ) => {
       const chat =
         initialChat ??
@@ -377,10 +426,13 @@ export const CompassAssistantProvider = registerCompassPlugin(
           atlasService,
           logger,
           track,
+          getTools: () => toolsController.getActiveTools(),
         });
 
       return {
-        store: { state: { chat, atlasAiService } },
+        store: {
+          state: { chat, atlasAiService, toolsController, preferences },
+        },
         deactivate: () => {},
       };
     },
@@ -389,8 +441,10 @@ export const CompassAssistantProvider = registerCompassPlugin(
     atlasService: atlasServiceLocator,
     atlasAiService: atlasAiServiceLocator,
     atlasAuthService: atlasAuthServiceLocator,
+    toolsController: toolsControllerLocator,
     track: telemetryLocator,
     logger: createLoggerLocator('COMPASS-ASSISTANT'),
+    preferences: preferencesLocator,
   }
 );
 
@@ -401,6 +455,7 @@ export function createDefaultChat({
   logger,
   track,
   options,
+  getTools,
 }: {
   originForPrompt: string;
   appNameForPrompt: string;
@@ -410,13 +465,15 @@ export function createDefaultChat({
   options?: {
     transport: Chat<AssistantMessage>['transport'];
   };
+  getTools?: () => ToolSet;
 }): Chat<AssistantMessage> {
   const initialBaseUrl = 'http://PLACEHOLDER_BASE_URL_TO_BE_REPLACED.invalid';
-  return new Chat({
+  return new Chat<AssistantMessage>({
     transport:
       options?.transport ??
       new DocsProviderTransport({
         origin: originForPrompt,
+        getTools,
         instructions: buildConversationInstructionsPrompt({
           target: appNameForPrompt,
         }),
@@ -453,6 +510,7 @@ export function createDefaultChat({
           },
         }).responses('mongodb-chat-latest'),
       }),
+    sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses,
     onError: (err: Error) => {
       logger.log.error(
         logger.mongoLogId(1_001_000_370),