[AiAssistance] Add UI elements for supporting screenshots as input

This CL covers the screenshot mechanism, how the screenshot will be
displayed in input area and how user can remove it.

Bug: 393036589
Change-Id: Ifc147f8ed2c960fb0fa5695b8a0f3c163a9a0670
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6225119
Commit-Queue: Samiya Caur <[email protected]>
Auto-Submit: Samiya Caur <[email protected]>
Reviewed-by: Alex Rudenko <[email protected]>

Author: Samiya Caur

front_end/panels/ai_assistance/AiAssistancePanel.test.ts

@@ -11,7 +11,8 @@ import * as Workspace from '../../models/workspace/workspace.js';
 import {createAiAssistancePanel, detachPanels, mockAidaClient} from '../../testing/AiAssistanceHelpers.js';
 import {findMenuItemWithLabel, getMenu} from '../../testing/ContextMenuHelpers.js';
 import {dispatchClickEvent} from '../../testing/DOMHelpers.js';
-import {describeWithEnvironment, registerNoopActions} from '../../testing/EnvironmentHelpers.js';
+import {createTarget, registerNoopActions} from '../../testing/EnvironmentHelpers.js';
+import {describeWithMockConnection} from '../../testing/MockConnection.js';
 import * as UI from '../../ui/legacy/legacy.js';
 import * as ElementsPanel from '../elements/elements.js';
 import * as NetworkPanel from '../network/network.js';
@@ -27,7 +28,7 @@ async function drainMicroTasks() {
   await new Promise(resolve => setTimeout(resolve, 0));
 }
 
-describeWithEnvironment('AI Assistance Panel', () => {
+describeWithMockConnection('AI Assistance Panel', () => {
   beforeEach(() => {
     registerNoopActions(['elements.toggle-element-search']);
     UI.Context.Context.instance().setFlavor(ElementsPanel.ElementsPanel.ElementsPanel, null);
@@ -1092,4 +1093,63 @@ describeWithEnvironment('AI Assistance Panel', () => {
       },
     ]);
   });
+
+  describe('multimodal input', () => {
+    function mockScreenshotModel() {
+      const target = createTarget();
+      const screenCaptureModel = target.model(SDK.ScreenCaptureModel.ScreenCaptureModel);
+      assert.exists(screenCaptureModel);
+      return {
+        captureScreenshotStub:
+            sinon.stub(screenCaptureModel, 'captureScreenshot').returns(Promise.resolve('imageInput')),
+      };
+    }
+
+    it('multimodal related functions unavailable when multimodal is disabled', async () => {
+      Object.assign(Root.Runtime.hostConfig, {
+        devToolsFreestyler: {
+          enabled: true,
+          multimodal: false,
+        },
+      });
+      UI.Context.Context.instance().setFlavor(
+          ElementsPanel.ElementsPanel.ElementsPanel,
+          sinon.createStubInstance(ElementsPanel.ElementsPanel.ElementsPanel));
+      const {
+        view,
+      } = createAiAssistancePanel();
+
+      assert.isFalse(view.lastCall.args[0].multimodalInputEnabled);
+      assert.notExists(view.lastCall.args[0].onTakeScreenshot);
+      assert.notExists(view.lastCall.args[0].onRemoveImageInput);
+      assert.isEmpty(view.lastCall.args[0].imageInput);
+    });
+
+    it('adds an image input and then removes it', async () => {
+      const {captureScreenshotStub} = mockScreenshotModel();
+      Object.assign(Root.Runtime.hostConfig, {
+        devToolsFreestyler: {
+          enabled: true,
+          multimodal: true,
+        },
+      });
+      UI.Context.Context.instance().setFlavor(
+          ElementsPanel.ElementsPanel.ElementsPanel,
+          sinon.createStubInstance(ElementsPanel.ElementsPanel.ElementsPanel));
+      const {
+        view,
+      } = createAiAssistancePanel();
+
+      assert.isTrue(view.lastCall.args[0].multimodalInputEnabled);
+
+      await view.lastCall.args[0].onTakeScreenshot?.();
+
+      expect(captureScreenshotStub.calledOnce);
+      assert.deepEqual(view.lastCall.args[0].imageInput, 'imageInput');
+
+      await view.lastCall.args[0].onRemoveImageInput?.();
+
+      assert.isEmpty(view.lastCall.args[0].imageInput);
+    });
+  });
 });

front_end/panels/ai_assistance/AiAssistancePanel.ts

@@ -10,6 +10,7 @@ import * as i18n from '../../core/i18n/i18n.js';
 import type * as Platform from '../../core/platform/platform.js';
 import * as Root from '../../core/root/root.js';
 import * as SDK from '../../core/sdk/sdk.js';
+import * as Protocol from '../../generated/protocol.js';
 import * as Workspace from '../../models/workspace/workspace.js';
 import * as UI from '../../ui/legacy/legacy.js';
 import * as Lit from '../../ui/lit/lit.js';
@@ -58,6 +59,7 @@ const {html} = Lit;
 
 const AI_ASSISTANCE_SEND_FEEDBACK = 'https://crbug.com/364805393' as Platform.DevToolsPath.UrlString;
 const AI_ASSISTANCE_HELP = 'https://goo.gle/devtools-ai-assistance' as Platform.DevToolsPath.UrlString;
+const SCREENSHOT_QUALITY = 100;
 
 const UIStrings = {
   /**
@@ -254,6 +256,7 @@ export class AiAssistancePanel extends UI.Panel.Panel {
   #project?: Workspace.Workspace.Project;
   #patchSuggestion?: string;
   #patchSuggestionLoading?: boolean;
+  #imageInput: string = '';
 
   constructor(private view: View = defaultView, {aidaClient, aidaAvailability, syncInfo}: {
     aidaClient: Host.AidaClient.AidaClient,
@@ -639,7 +642,11 @@ export class AiAssistancePanel extends UI.Panel.Panel {
           inspectElementToggled: this.#toggleSearchElementAction.toggled(),
           userInfo: this.#userInfo,
           canShowFeedbackForm: this.#serverSideLoggingEnabled,
+          multimodalInputEnabled:
+              isAiAssistanceMultimodalInputEnabled() && this.#currentAgent?.type === AgentType.STYLING,
+          imageInput: this.#imageInput,
           onTextSubmit: (text: string) => {
+            this.#imageInput = '';
             void this.#startConversation(text);
             Host.userMetrics.actionTaken(Host.UserMetrics.Action.AiAssistanceQuerySubmitted);
           },
@@ -651,6 +658,9 @@ export class AiAssistancePanel extends UI.Panel.Panel {
           onCancelCrossOriginChat: this.#blockedByCrossOrigin && this.#previousSameOriginContext ?
               this.#handleCrossOriginChatCancellation.bind(this) :
               undefined,
+          onTakeScreenshot: isAiAssistanceMultimodalInputEnabled() ? this.#handleTakeScreenshot.bind(this) : undefined,
+          onRemoveImageInput: isAiAssistanceMultimodalInputEnabled() ? this.#handleRemoveImageInput.bind(this) :
+                                                                       undefined,
           onApplyToWorkspace: this.#onApplyToWorkspace.bind(this)
         },
         this.#viewOutput, this.#contentContainer);
@@ -834,6 +844,31 @@ export class AiAssistancePanel extends UI.Panel.Panel {
     }
   }
 
+  async #handleTakeScreenshot(): Promise<void> {
+    const mainTarget = SDK.TargetManager.TargetManager.instance().primaryPageTarget();
+    if (!mainTarget) {
+      throw new Error('Could not find main target');
+    }
+    const model = mainTarget.model(SDK.ScreenCaptureModel.ScreenCaptureModel);
+    if (!model) {
+      throw new Error('Could not find model');
+    }
+    const bytes = await model.captureScreenshot(
+        Protocol.Page.CaptureScreenshotRequestFormat.Jpeg,
+        SCREENSHOT_QUALITY,
+        SDK.ScreenCaptureModel.ScreenshotMode.FROM_VIEWPORT,
+    );
+    if (bytes) {
+      this.#imageInput = bytes;
+      void this.doUpdate();
+    }
+  }
+
+  #handleRemoveImageInput(): void {
+    this.#imageInput = '';
+    void this.doUpdate();
+  }
+
   #runAbortController = new AbortController();
   #cancel(): void {
     this.#runAbortController.abort();
@@ -1122,6 +1157,11 @@ export class ActionDelegate implements UI.ActionRegistration.ActionDelegate {
   }
 }
 
+function isAiAssistanceMultimodalInputEnabled(): boolean {
+  const {hostConfig} = Root.Runtime;
+  return Boolean(hostConfig.devToolsFreestyler?.multimodal);
+}
+
 function isAiAssistancePatchingEnabled(): boolean {
   const {hostConfig} = Root.Runtime;
   return Boolean(hostConfig.devToolsFreestyler?.patching);

front_end/panels/ai_assistance/components/ChatView.ts

@@ -268,6 +268,14 @@ const UIStringsNotTranslate = {
    *@description Button text for staging changes to workspace.
    */
   applyToWorkspace: 'Apply to workspace',
+  /**
+   *@description Title for the take screenshot button.
+   */
+  takeScreenshotButtonTitle: 'Take screenshot',
+  /**
+   *@description Title for the remove image input button.
+   */
+  removeImageInputButtonTitle: 'Remove image input',
 };
 
 const str_ = i18n.i18n.registerUIStrings('panels/ai_assistance/components/ChatView.ts', UIStrings);
@@ -324,6 +332,8 @@ export interface Props {
   onContextClick: () => void | Promise<void>;
   onNewConversation: () => void;
   onCancelCrossOriginChat?: () => void;
+  onTakeScreenshot?: () => Promise<void>;
+  onRemoveImageInput?: () => void;
   inspectElementToggled: boolean;
   state: State;
   aidaAvailability: Host.AidaClient.AidaAccessPreconditions;
@@ -340,6 +350,8 @@ export interface Props {
   patchSuggestion?: string;
   patchSuggestionLoading?: boolean;
   projectName?: string;
+  multimodalInputEnabled?: boolean;
+  imageInput?: string;
   onApplyToWorkspace?: () => void;
 }
 
@@ -738,14 +750,18 @@ export class ChatView extends HTMLElement {
                 state: this.#props.state,
                 selectedContext: this.#props.selectedContext,
                 inspectElementToggled: this.#props.inspectElementToggled,
+                multimodalInputEnabled: this.#props.multimodalInputEnabled,
                 agentType: this.#props.agentType,
+                imageInput: this.#props.imageInput,
                 onContextClick: this.#props.onContextClick,
                 onInspectElementClick: this.#props.onInspectElementClick,
                 onSubmit: this.#handleSubmit,
                 onTextAreaKeyDown: this.#handleTextAreaKeyDown,
                 onCancel: this.#handleCancel,
                 onNewConversation: this.#props.onNewConversation,
                 onCancelCrossOriginChat: this.#props.onCancelCrossOriginChat,
+                onTakeScreenshot: this.#props.onTakeScreenshot,
+                onRemoveImageInput: this.#props.onRemoveImageInput,
               })
           }
         </main>
@@ -1420,6 +1436,65 @@ function renderChatInputButtons(
   ></devtools-button>`;
 }
 
+function renderTakeScreenshotButton({
+  multimodalInputEnabled,
+  isTextInputDisabled,
+  onTakeScreenshot,
+}: {
+  isTextInputDisabled: boolean,
+  multimodalInputEnabled?: boolean,
+  onTakeScreenshot?: () => Promise<void>,
+}): Lit.LitTemplate {
+    if (!multimodalInputEnabled) {
+      return Lit.nothing;
+    }
+    return html`<devtools-button
+      class="chat-input-button"
+      aria-label=${lockedString(UIStringsNotTranslate.takeScreenshotButtonTitle)}
+      @click=${onTakeScreenshot}
+      .data=${
+        {
+          variant: Buttons.Button.Variant.ICON,
+          size: Buttons.Button.Size.REGULAR,
+          disabled: isTextInputDisabled,
+          iconName: 'photo-camera',
+          title: lockedString(UIStringsNotTranslate.takeScreenshotButtonTitle),
+          jslogContext: 'take-screenshot',
+        } as Buttons.Button.ButtonData
+      }
+    ></devtools-button>`;
+  }
+
+function renderImageInput({
+  multimodalInputEnabled,
+  imageInput,
+  onRemoveImageInput,
+}: {
+  multimodalInputEnabled?: boolean,
+  imageInput?: string,
+  onRemoveImageInput?: () => void,
+}): Lit.LitTemplate {
+    if (!multimodalInputEnabled || !imageInput || imageInput==='') {
+      return Lit.nothing;
+    }
+    return  html`
+    <div class="image-input-container">
+      <devtools-button
+      aria-label=${lockedString(UIStringsNotTranslate.removeImageInputButtonTitle)}
+      @click=${onRemoveImageInput}
+      .data=${
+        {
+          variant: Buttons.Button.Variant.ICON,
+          size: Buttons.Button.Size.MICRO,
+          iconName: 'cross',
+          title: lockedString(UIStringsNotTranslate.removeImageInputButtonTitle),
+        } as Buttons.Button.ButtonData
+      }
+    ></devtools-button>
+      <img src="data:image/jpeg;base64, ${imageInput}" alt="Screenshot input" />
+    </div>`;
+  }
+
 function renderChatInput({
   isLoading,
   blockedByCrossOrigin,
@@ -1428,14 +1503,18 @@ function renderChatInput({
   state,
   selectedContext,
   inspectElementToggled,
+  multimodalInputEnabled,
   agentType,
+  imageInput,
   onContextClick,
   onInspectElementClick,
   onSubmit,
   onTextAreaKeyDown,
   onCancel,
   onNewConversation,
   onCancelCrossOriginChat,
+  onTakeScreenshot,
+  onRemoveImageInput,
 }: {
   isLoading: boolean,
   blockedByCrossOrigin: boolean,
@@ -1444,24 +1523,33 @@ function renderChatInput({
   state: State,
   selectedContext: ConversationContext<unknown> | null,
   inspectElementToggled: boolean,
+  multimodalInputEnabled?: boolean,
   agentType?: AgentType,
+  imageInput?: string,
   onContextClick: () => void | Promise<void>,
   onInspectElementClick: () => void,
   onSubmit: (ev: SubmitEvent) => void,
   onTextAreaKeyDown: (ev: KeyboardEvent) => void,
   onCancel: (ev: SubmitEvent) => void,
   onNewConversation: () => void,
   onCancelCrossOriginChat?: () => void,
+  onTakeScreenshot?: () => Promise<void>,
+  onRemoveImageInput?: () => void,
 }): Lit.LitTemplate {
   if (!agentType) {
     return Lit.nothing;
   }
 
-  const cls = Lit.Directives.classMap({
+  const chatInputCls = Lit.Directives.classMap({
     'chat-input': true,
     'two-big-buttons': blockedByCrossOrigin,
   });
 
+    const chatInputContainerCls = Lit.Directives.classMap({
+      'chat-input-container': true,
+      disabled: isTextInputDisabled,
+    });
+
   // clang-format off
   return html`
   <form class="input-form" @submit=${onSubmit}>
@@ -1481,8 +1569,11 @@ function renderChatInput({
         </div>
       </div>
     ` : Lit.nothing}
-    <div class="chat-input-container">
-      <textarea class=${cls}
+    <div class=${chatInputContainerCls}>
+      ${renderImageInput(
+        {multimodalInputEnabled, imageInput, onRemoveImageInput}
+      )}
+      <textarea class=${chatInputCls}
         .disabled=${isTextInputDisabled}
         wrap="hard"
         maxlength="10000"
@@ -1491,6 +1582,9 @@ function renderChatInput({
         jslog=${VisualLogging.textField('query').track({ keydown: 'Enter' })}
       ></textarea>
       <div class="chat-input-buttons">
+        ${renderTakeScreenshotButton({
+          multimodalInputEnabled, isTextInputDisabled, onTakeScreenshot
+        })}
         ${renderChatInputButtons({ isLoading, blockedByCrossOrigin, isTextInputDisabled, onCancel, onNewConversation, onCancelCrossOriginChat })}
       </div>
     </div>