Merge pull request #3374 from github/robertbrignull/sentToLLMMethods

Add sentToLLM to ModelingStore

Author: robertbrignull

extensions/ql-vscode/src/common/interface-types.ts

@@ -543,6 +543,11 @@ interface SetInProgressMethodsMessage {
   methods: string[];
 }
 
+interface SetProcessedByAutoModelMethodsMessage {
+  t: "setProcessedByAutoModelMethods";
+  methods: string[];
+}
+
 interface SwitchModeMessage {
   t: "switchMode";
   mode: Mode;
@@ -626,6 +631,7 @@ export type ToModelEditorMessage =
   | SetModeledMethodsMessage
   | SetModifiedMethodsMessage
   | SetInProgressMethodsMessage
+  | SetProcessedByAutoModelMethodsMessage
   | RevealMethodMessage
   | SetAccessPathSuggestionsMessage;
 

extensions/ql-vscode/src/model-editor/auto-modeler.ts

@@ -159,6 +159,12 @@ export class AutoModeler {
             this.databaseItem,
             candidateSignatures,
           );
+
+          // Let the UI know which methods have been sent to the LLM
+          this.modelingStore.addProcessedByAutoModelMethods(
+            this.databaseItem,
+            candidateSignatures,
+          );
         }
       } finally {
         // Clear out in progress methods in case anything went wrong

extensions/ql-vscode/src/model-editor/model-editor-view.ts

@@ -835,6 +835,19 @@ export class ModelEditorView extends AbstractWebview<
       }),
     );
 
+    this.push(
+      this.modelingEvents.onProcessedByAutoModelMethodsChanged(
+        async (event) => {
+          if (event.dbUri === this.databaseItem.databaseUri.toString()) {
+            await this.postMessage({
+              t: "setProcessedByAutoModelMethods",
+              methods: Array.from(event.methods),
+            });
+          }
+        },
+      ),
+    );
+
     this.push(
       this.modelingEvents.onRevealInModelEditor(async (event) => {
         if (event.dbUri === this.databaseItem.databaseUri.toString()) {

extensions/ql-vscode/src/model-editor/modeling-events.ts

@@ -42,13 +42,19 @@ interface SelectedMethodChangedEvent {
   readonly modeledMethods: readonly ModeledMethod[];
   readonly isModified: boolean;
   readonly isInProgress: boolean;
+  readonly processedByAutoModel: boolean;
 }
 
 interface InProgressMethodsChangedEvent {
   readonly dbUri: string;
   readonly methods: ReadonlySet<string>;
 }
 
+interface ProcessedByAutoModelMethodsChangedEvent {
+  readonly dbUri: string;
+  readonly methods: ReadonlySet<string>;
+}
+
 interface RevealInModelEditorEvent {
   dbUri: string;
   method: Method;
@@ -69,6 +75,7 @@ export class ModelingEvents extends DisposableObject {
   public readonly onModifiedMethodsChanged: AppEvent<ModifiedMethodsChangedEvent>;
   public readonly onSelectedMethodChanged: AppEvent<SelectedMethodChangedEvent>;
   public readonly onInProgressMethodsChanged: AppEvent<InProgressMethodsChangedEvent>;
+  public readonly onProcessedByAutoModelMethodsChanged: AppEvent<ProcessedByAutoModelMethodsChangedEvent>;
   public readonly onRevealInModelEditor: AppEvent<RevealInModelEditorEvent>;
   public readonly onFocusModelEditor: AppEvent<FocusModelEditorEvent>;
 
@@ -82,6 +89,7 @@ export class ModelingEvents extends DisposableObject {
   private readonly onModifiedMethodsChangedEventEmitter: AppEventEmitter<ModifiedMethodsChangedEvent>;
   private readonly onSelectedMethodChangedEventEmitter: AppEventEmitter<SelectedMethodChangedEvent>;
   private readonly onInProgressMethodsChangedEventEmitter: AppEventEmitter<InProgressMethodsChangedEvent>;
+  private readonly onProcessedByAutoModelMethodsChangedEventEmitter: AppEventEmitter<ProcessedByAutoModelMethodsChangedEvent>;
   private readonly onRevealInModelEditorEventEmitter: AppEventEmitter<RevealInModelEditorEvent>;
   private readonly onFocusModelEditorEventEmitter: AppEventEmitter<FocusModelEditorEvent>;
 
@@ -141,6 +149,12 @@ export class ModelingEvents extends DisposableObject {
     this.onInProgressMethodsChanged =
       this.onInProgressMethodsChangedEventEmitter.event;
 
+    this.onProcessedByAutoModelMethodsChangedEventEmitter = this.push(
+      app.createEventEmitter<ProcessedByAutoModelMethodsChangedEvent>(),
+    );
+    this.onProcessedByAutoModelMethodsChanged =
+      this.onProcessedByAutoModelMethodsChangedEventEmitter.event;
+
     this.onRevealInModelEditorEventEmitter = this.push(
       app.createEventEmitter<RevealInModelEditorEvent>(),
     );
@@ -226,6 +240,7 @@ export class ModelingEvents extends DisposableObject {
     modeledMethods: ModeledMethod[],
     isModified: boolean,
     isInProgress: boolean,
+    processedByAutoModel: boolean,
   ) {
     this.onSelectedMethodChangedEventEmitter.fire({
       databaseItem,
@@ -234,6 +249,7 @@ export class ModelingEvents extends DisposableObject {
       modeledMethods,
       isModified,
       isInProgress,
+      processedByAutoModel,
     });
   }
 
@@ -247,6 +263,16 @@ export class ModelingEvents extends DisposableObject {
     });
   }
 
+  public fireProcessedByAutoModelMethodsChangedEvent(
+    dbUri: string,
+    methods: ReadonlySet<string>,
+  ) {
+    this.onProcessedByAutoModelMethodsChangedEventEmitter.fire({
+      dbUri,
+      methods,
+    });
+  }
+
   public fireRevealInModelEditorEvent(dbUri: string, method: Method) {
     this.onRevealInModelEditorEventEmitter.fire({
       dbUri,

extensions/ql-vscode/src/model-editor/modeling-store.ts

@@ -14,6 +14,7 @@ interface InternalDbModelingState {
   modeledMethods: Record<string, ModeledMethod[]>;
   modifiedMethodSignatures: Set<string>;
   inProgressMethods: Set<string>;
+  processedByAutoModelMethods: Set<string>;
   selectedMethod: Method | undefined;
   selectedUsage: Usage | undefined;
 }
@@ -26,6 +27,7 @@ interface DbModelingState {
   readonly modeledMethods: Readonly<Record<string, readonly ModeledMethod[]>>;
   readonly modifiedMethodSignatures: ReadonlySet<string>;
   readonly inProgressMethods: ReadonlySet<string>;
+  readonly processedByAutoModelMethods: ReadonlySet<string>;
   readonly selectedMethod: Method | undefined;
   readonly selectedUsage: Usage | undefined;
 }
@@ -37,6 +39,7 @@ interface SelectedMethodDetails {
   readonly modeledMethods: readonly ModeledMethod[];
   readonly isModified: boolean;
   readonly isInProgress: boolean;
+  readonly processedByAutoModel: boolean;
 }
 
 export class ModelingStore extends DisposableObject {
@@ -59,6 +62,7 @@ export class ModelingStore extends DisposableObject {
       mode,
       modeledMethods: {},
       modifiedMethodSignatures: new Set(),
+      processedByAutoModelMethods: new Set(),
       selectedMethod: undefined,
       selectedUsage: undefined,
       inProgressMethods: new Set(),
@@ -301,13 +305,17 @@ export class ModelingStore extends DisposableObject {
     const modeledMethods = dbState.modeledMethods[method.signature] ?? [];
     const isModified = dbState.modifiedMethodSignatures.has(method.signature);
     const isInProgress = dbState.inProgressMethods.has(method.signature);
+    const processedByAutoModel = dbState.processedByAutoModelMethods.has(
+      method.signature,
+    );
     this.modelingEvents.fireSelectedMethodChangedEvent(
       dbItem,
       method,
       usage,
       modeledMethods,
       isModified,
       isInProgress,
+      processedByAutoModel,
     );
   }
 
@@ -336,6 +344,18 @@ export class ModelingStore extends DisposableObject {
     });
   }
 
+  public addProcessedByAutoModelMethods(
+    dbItem: DatabaseItem,
+    processedByAutoModelMethods: string[],
+  ) {
+    this.changeProcessedByAutoModelMethods(dbItem, (state) => {
+      state.processedByAutoModelMethods = new Set([
+        ...state.processedByAutoModelMethods,
+        ...processedByAutoModelMethods,
+      ]);
+    });
+  }
+
   public getSelectedMethodDetails(): SelectedMethodDetails | undefined {
     const dbState = this.getInternalStateForActiveDb();
     if (!dbState) {
@@ -356,6 +376,9 @@ export class ModelingStore extends DisposableObject {
         selectedMethod.signature,
       ),
       isInProgress: dbState.inProgressMethods.has(selectedMethod.signature),
+      processedByAutoModel: dbState.processedByAutoModelMethods.has(
+        selectedMethod.signature,
+      ),
     };
   }
 
@@ -412,4 +435,18 @@ export class ModelingStore extends DisposableObject {
       state.inProgressMethods,
     );
   }
+
+  private changeProcessedByAutoModelMethods(
+    dbItem: DatabaseItem,
+    updateState: (state: InternalDbModelingState) => void,
+  ) {
+    const state = this.getState(dbItem);
+
+    updateState(state);
+
+    this.modelingEvents.fireProcessedByAutoModelMethodsChangedEvent(
+      dbItem.databaseUri.toString(),
+      state.processedByAutoModelMethods,
+    );
+  }
 }

extensions/ql-vscode/src/view/model-editor/ModelEditor.tsx

@@ -148,6 +148,10 @@ export function ModelEditor({
             setInProgressMethods(new Set(msg.methods));
             break;
           }
+          case "setProcessedByAutoModelMethods": {
+            // TODO: set state
+            break;
+          }
           case "revealMethod":
             setRevealedMethodSignature(msg.methodSignature);
             break;

extensions/ql-vscode/test/__mocks__/model-editor/modelingEventsMock.ts

@@ -10,6 +10,7 @@ export function createMockModelingEvents({
   onModeledMethodsChanged = jest.fn(),
   onModifiedMethodsChanged = jest.fn(),
   onInProgressMethodsChanged = jest.fn(),
+  onProcessedByAutoModelMethodsChanged = jest.fn(),
   onRevealInModelEditor = jest.fn(),
   onFocusModelEditor = jest.fn(),
 }: {
@@ -21,6 +22,7 @@ export function createMockModelingEvents({
   onModeledMethodsChanged?: ModelingEvents["onModeledMethodsChanged"];
   onModifiedMethodsChanged?: ModelingEvents["onModifiedMethodsChanged"];
   onInProgressMethodsChanged?: ModelingEvents["onInProgressMethodsChanged"];
+  onProcessedByAutoModelMethodsChanged?: ModelingEvents["onProcessedByAutoModelMethodsChanged"];
   onRevealInModelEditor?: ModelingEvents["onRevealInModelEditor"];
   onFocusModelEditor?: ModelingEvents["onFocusModelEditor"];
 } = {}): ModelingEvents {
@@ -33,6 +35,7 @@ export function createMockModelingEvents({
     onModeledMethodsChanged,
     onModifiedMethodsChanged,
     onInProgressMethodsChanged,
+    onProcessedByAutoModelMethodsChanged,
     onRevealInModelEditor,
     onFocusModelEditor,
   });