update store setup for new plugin, not sure it works

Author: Anemy

packages/compass-telemetry/src/telemetry-events.ts

@@ -2632,6 +2632,7 @@ type ScreenEvent = ConnectionScopedEvent<{
       | 'my_queries'
       | 'performance'
       | 'schema'
+      | 'vector_visualizer'
       | 'validation'
       | 'confirm_new_pipeline_modal'
       | 'create_collection_modal'

packages/compass-vector-embedding-visualizer/src/components/vector-visualizer.tsx

@@ -1,18 +1,21 @@
 import React, { useEffect, useState } from 'react';
+import { connect } from 'react-redux';
 import Plotly from 'plotly.js';
 const PCA = require('ml-pca');
-import { Binary } from 'mongodb';
+import type { Binary } from 'mongodb';
+import type { Document } from 'bson';
+
+import type { VectorEmbeddingVisualizerState } from '../stores/reducer';
+import { loadDocuments } from '../stores/visualization';
+import { ErrorSummary } from '@mongodb-js/compass-components';
 
 type HoverInfo = { x: number; y: number; text: string } | null;
 
 export interface VectorVisualizerProps {
-  dataService: {
-    find: (
-      ns: string,
-      filter: Record<string, unknown>,
-      options?: { limit?: number }
-    ) => Promise<any[]>;
-  };
+  onFetchDocs: (namespace: string) => void;
+  docs: Document[];
+  loadingDocumentsState: 'initial' | 'loading' | 'loaded' | 'error';
+  loadingDocumentsError: Error | null;
   collection: { namespace: string };
 }
 
@@ -23,24 +26,33 @@ function normalizeTo2D(vectors: Binary[]): { x: number; y: number }[] {
   return reduced.map(([x, y]: [number, number]) => ({ x, y }));
 }
 
-export const VectorVisualizer: React.FC<VectorVisualizerProps> = ({
-  dataService,
+const VectorVisualizer: React.FC<VectorVisualizerProps> = ({
+  onFetchDocs,
+  docs,
+  loadingDocumentsState,
+  loadingDocumentsError,
   collection,
 }) => {
   const [hoverInfo, setHoverInfo] = useState<HoverInfo>(null);
 
+  useEffect(() => {
+    if (loadingDocumentsState === 'initial') {
+      // Fetch the documents when the component mounts when they aren't already loaded.
+      onFetchDocs(collection.namespace);
+    }
+  }, [loadingDocumentsState, onFetchDocs, collection.namespace]);
+
   useEffect(() => {
     const container = document.getElementById('vector-plot');
     if (!container) return;
 
-    let isMounted = true;
+    const abortController = new AbortController();
 
     const plot = async () => {
       try {
         const ns = collection?.namespace;
-        if (!ns || !dataService) return;
+        if (!ns) return;
 
-        const docs = await dataService.find(ns, {}, { limit: 1000 });
         const vectors = docs.map((doc) => doc.review_vec).filter(Boolean);
 
         if (!vectors.length) return;
@@ -120,13 +132,16 @@ export const VectorVisualizer: React.FC<VectorVisualizerProps> = ({
     void plot();
 
     return () => {
-      isMounted = false;
+      abortController.abort();
     };
-  }, [collection?.namespace, dataService]);
+  }, [docs, collection.namespace]);
 
   return (
     <div style={{ position: 'relative', width: '100%', height: '100%' }}>
       <div id="vector-plot" style={{ width: '100%', height: '100%' }} />
+      {loadingDocumentsError && (
+        <ErrorSummary errors={loadingDocumentsError.message} />
+      )}
       {hoverInfo && (
         <div
           style={{
@@ -148,3 +163,14 @@ export const VectorVisualizer: React.FC<VectorVisualizerProps> = ({
     </div>
   );
 };
+
+export default connect(
+  (state: VectorEmbeddingVisualizerState) => ({
+    docs: state.visualization.docs,
+    loadingDocumentsState: state.visualization.loadingDocumentsState,
+    loadingDocumentsError: state.visualization.loadingDocumentsError,
+  }),
+  {
+    onFetchDocs: loadDocuments,
+  }
+)(VectorVisualizer);

packages/compass-vector-embedding-visualizer/src/index.ts

@@ -2,15 +2,17 @@ import React from 'react';
 import { registerHadronPlugin } from 'hadron-app-registry';
 import { createLoggerLocator } from '@mongodb-js/compass-logging/provider';
 import { collectionModelLocator } from '@mongodb-js/compass-app-stores/provider';
-import { dataServiceLocator } from '@mongodb-js/compass-connections/provider';
-import { createStore } from 'redux';
+import {
+  dataServiceLocator,
+  type DataServiceLocator,
+} from '@mongodb-js/compass-connections/provider';
 
-import { VectorVisualizer } from './components/vector-visualizer';
-import { activateVectorPlugin } from './stores/store';
-
-function reducer(state = {}, _action: any) {
-  return state;
-}
+import VectorVisualizer from './components/vector-visualizer';
+import { VectorsTabTitle } from './plugin-tab-title';
+import {
+  activateVectorPlugin,
+  type VectorDataServiceProps,
+} from './stores/store';
 
 export const CompassVectorPluginProvider = registerHadronPlugin(
   {
@@ -21,25 +23,25 @@ export const CompassVectorPluginProvider = registerHadronPlugin(
     activate: activateVectorPlugin,
   },
   {
-    dataService: dataServiceLocator,
+    dataService:
+      dataServiceLocator as DataServiceLocator<VectorDataServiceProps>,
     collection: collectionModelLocator,
     logger: createLoggerLocator('COMPASS-VECTOR-VISUALIZER'),
   }
 );
 
-const VectorVisualizerWrapper = (props: {
-  dataService: any;
-  collection: any;
-}) => {
-  return React.createElement(VectorVisualizer, props);
-};
+// const VectorVisualizerWrapper = (props: {
+//   dataService: any;
+//   collection: any;
+// }) => {
+//   return React.createElement(VectorVisualizer, props);
+// };
 
 export const CompassVectorPlugin = {
   name: 'Vector Visualizer' as const,
-  type: 'CollectionTab' as const,
   provider: CompassVectorPluginProvider,
-  content: VectorVisualizerWrapper,
-  header: () => React.createElement('div', null, 'Vector Embeddings'),
+  content: VectorVisualizer, // VectorVisualizerWrapper,
+  header: VectorsTabTitle,
 };
 
 export default CompassVectorPluginProvider;

packages/compass-vector-embedding-visualizer/src/plugin-tab-title.tsx

@@ -0,0 +1,5 @@
+import React from 'react';
+
+export function VectorsTabTitle() {
+  return <div>Vectors</div>;
+}

packages/compass-vector-embedding-visualizer/src/stores/reducer.ts

@@ -0,0 +1,35 @@
+import type { AnyAction } from 'redux';
+import { combineReducers } from 'redux';
+import type {
+  VisualizationActions,
+  VisualizationActionTypes,
+} from './visualization';
+import type { ThunkAction } from 'redux-thunk';
+import { visualizationReducer } from './visualization';
+import type { VectorPluginServices } from './store';
+
+const reducer = combineReducers({
+  visualization: visualizationReducer,
+});
+
+export type VectorEmbeddingVisualizerActions = VisualizationActions;
+
+export type VectorEmbeddingVisualizerActionTypes = VisualizationActionTypes;
+
+export type VectorEmbeddingVisualizerState = ReturnType<typeof reducer>;
+
+export type VectorEmbeddingVisualizerExtraArgs = VectorPluginServices & {
+  cancelControllerRef: { current: AbortController | null };
+};
+
+export type VectorEmbeddingVisualizerThunkAction<
+  R,
+  A extends AnyAction
+> = ThunkAction<
+  R,
+  VectorEmbeddingVisualizerState,
+  VectorEmbeddingVisualizerExtraArgs,
+  A
+>;
+
+export default reducer;

packages/compass-vector-embedding-visualizer/src/stores/store.ts

@@ -1,17 +1,19 @@
-import type { Store } from 'redux';
-import { createStore } from 'redux';
+import { applyMiddleware, createStore } from 'redux';
 import type { DataService } from 'mongodb-data-service';
-import type { ConnectionInfoRef } from '@mongodb-js/compass-connections/provider';
+// import type { ConnectionInfoRef } from '@mongodb-js/compass-connections/provider';
 import type {
   Collection,
-  MongoDBInstance,
+  // MongoDBInstance,
 } from '@mongodb-js/compass-app-stores/provider';
-import type AppRegistry from 'hadron-app-registry';
+// import type AppRegistry from 'hadron-app-registry';
 import type { Logger } from '@mongodb-js/compass-logging';
-import type { TrackFunction } from '@mongodb-js/compass-telemetry';
-import type { AtlasService } from '@mongodb-js/atlas-service/provider';
-import type { PreferencesAccess } from 'compass-preferences-model';
+// import type { TrackFunction } from '@mongodb-js/compass-telemetry';
+// import type { AtlasService } from '@mongodb-js/atlas-service/provider';
+// import type { PreferencesAccess } from 'compass-preferences-model';
 import type { ActivateHelpers } from 'hadron-app-registry';
+import thunk from 'redux-thunk';
+
+import reducer from './reducer';
 
 export type VectorDataServiceProps =
   | 'find'
@@ -23,37 +25,39 @@ export type VectorDataService = Pick<DataService, VectorDataServiceProps>;
 
 export type VectorPluginServices = {
   dataService: VectorDataService;
-  connectionInfoRef: ConnectionInfoRef;
-  instance: MongoDBInstance;
-  localAppRegistry: Pick<AppRegistry, 'on' | 'emit' | 'removeListener'>;
-  globalAppRegistry: Pick<AppRegistry, 'on' | 'emit' | 'removeListener'>;
   logger: Logger;
   collection: Collection;
-  track: TrackFunction;
-  atlasService: AtlasService;
-  preferences: PreferencesAccess;
-};
 
-export type VectorPluginOptions = {
-  namespace: string;
-  serverVersion: string;
-  isReadonly: boolean;
-};
+  // Note(Rhys): If we want more of these services, we can uncomment,
+  // and then in ../index.ts, we add them as well.
 
-export type VectorPluginStore = Store;
+  // connectionInfoRef: ConnectionInfoRef;
+  // instance: MongoDBInstance;
+  // localAppRegistry: Pick<AppRegistry, 'on' | 'emit' | 'removeListener'>;
+  // globalAppRegistry: Pick<AppRegistry, 'on' | 'emit' | 'removeListener'>;
+  // track: TrackFunction;
+  // atlasService: AtlasService;
+  // preferences: PreferencesAccess;
+};
 
-function reducer(state = {}, _action: any) {
-  return state;
-}
+// export type VectorPluginOptions = {
+//   namespace: string;
+//   serverVersion: string;
+//   isReadonly: boolean;
+// };
+export type VectorPluginOptions = Record<string, unknown>;
 
 export function activateVectorPlugin(
   _options: VectorPluginOptions,
-  _services: VectorPluginServices,
+  services: VectorPluginServices,
   { cleanup }: ActivateHelpers
 ) {
-  const store = createStore((state = {}) => state);
-  return {
-    store,
-    deactivate: () => cleanup(),
-  };
+  const cancelControllerRef = { current: null };
+  const store = createStore(
+    reducer,
+    applyMiddleware(
+      thunk.withExtraArgument({ ...services, cancelControllerRef })
+    )
+  );
+  return { store, deactivate: cleanup };
 }

packages/compass-vector-embedding-visualizer/src/stores/util.ts

@@ -0,0 +1,12 @@
+import type { AnyAction } from 'redux';
+import type {
+  VectorEmbeddingVisualizerActions,
+  VectorEmbeddingVisualizerActionTypes,
+} from './reducer';
+
+export function isAction<T extends VectorEmbeddingVisualizerActionTypes>(
+  action: AnyAction,
+  type: T
+): action is Extract<VectorEmbeddingVisualizerActions, { type: T }> {
+  return action.type === type;
+}