Iq terms n conditions window

.run/Run_Plugin__force_iq_terms_window_.xml

@@ -0,0 +1,31 @@
+<component name="ProjectRunConfigurationManager">
+  <configuration default="false" name="Run Plugin (force iq terms window)" type="GradleRunConfiguration" factoryName="Gradle">
+    <log_file alias="idea.log" path="$PROJECT_DIR$/build/idea-sandbox/system/log/idea.log" />
+    <ExternalSystemSettings>
+      <option name="env">
+        <map>
+          <entry key="CB_IQ_FORCE_ORG_CHOOSER" value="1" />
+          <entry key="CB_IQ_FORCE_TERMS_DIALOG" value="1" />
+        </map>
+      </option>
+      <option name="executionName" />
+      <option name="externalProjectPath" value="$PROJECT_DIR$" />
+      <option name="externalSystemIdString" value="GRADLE" />
+      <option name="scriptParameters" value="--stacktrace" />
+      <option name="taskDescriptions">
+        <list />
+      </option>
+      <option name="taskNames">
+        <list>
+          <option value="runIde" />
+        </list>
+      </option>
+      <option name="vmOptions" value="-Didea.log.debug.categories=org.intellij.sdk.language.completion" />
+    </ExternalSystemSettings>
+    <ExternalSystemDebugServerProcess>true</ExternalSystemDebugServerProcess>
+    <ExternalSystemReattachDebugProcess>true</ExternalSystemReattachDebugProcess>
+    <DebugAllEnabled>false</DebugAllEnabled>
+    <RunAsTest>false</RunAsTest>
+    <method v="2" />
+  </configuration>
+</component>

src/main/java/com/couchbase/intellij/tree/iq/CapellaApiMethods.java

@@ -2,10 +2,14 @@
 
 import com.couchbase.client.core.deps.com.fasterxml.jackson.databind.DeserializationFeature;
 import com.couchbase.client.core.deps.com.fasterxml.jackson.databind.ObjectMapper;
+import com.couchbase.client.java.json.JsonObject;
 import com.couchbase.intellij.tree.iq.core.CapellaAuth;
+import com.couchbase.intellij.workbench.Log;
 import org.apache.commons.io.IOUtils;
 
 import java.io.IOException;
+import java.io.OutputStream;
+import java.io.StringWriter;
 import java.net.HttpURLConnection;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
@@ -84,4 +88,63 @@ public static void setCapellaDomain(String domain) {
         AUTH_URL = domain + "/sessions";
         ORGANIZATIONS_URL = domain + "/v2/organizations";
     }
+
+    public static JsonObject getOrganization(CapellaAuth auth, String orgId) throws Exception {
+        URL url = new URL(String.format("%s/%s", ORGANIZATIONS_URL, orgId));
+        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
+        connection.setRequestMethod("GET");
+        connection.setRequestProperty("Authorization", "Bearer " + auth.getJwt());
+        connection.setRequestProperty("Content-Type", "application/json");
+
+
+        if (connection.getResponseCode() != 200) {
+            throw new RuntimeException("Failed to get organization list: HTTP error code: " + connection.getResponseCode());
+        }
+
+
+        String result = IOUtils.toString(connection.getInputStream(), StandardCharsets.UTF_8);
+        return JsonObject.fromJson(result).getArray("data").getObject(0);
+    }
+
+    public static void acceptIqTerms(CapellaAuth auth, CapellaOrganization organization) throws Exception {
+        JsonObject original = getOrganization(auth, organization.getId());
+        URL url = new URL(String.format("%s/%s", ORGANIZATIONS_URL, organization.getId()));
+        HttpURLConnection connection = (HttpURLConnection) url.openConnection();
+        connection.setDoInput(true);
+        connection.setDoOutput(true);
+        connection.setRequestProperty("Authorization", "Bearer " + auth.getJwt());
+        connection.setRequestProperty("Content-Type", "application/json");
+        connection.setRequestMethod("PUT");
+        OutputStream os = connection.getOutputStream();
+        ObjectMapper objectMapper = new ObjectMapper()
+                .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
+        StringWriter writer = new StringWriter();
+
+        original.put("iq", JsonObject.create());
+
+        JsonObject origIq = original.getObject("iq");
+        origIq.put("enabled", true);
+        origIq.put("other", JsonObject.create());
+        JsonObject other = origIq.getObject("other");
+        other.put("isTermsAcceptedForOrg", true);
+
+        original.removeKey("modifiedAt");
+        original.removeKey("upsertedAt");
+        original.removeKey("modifiedBy");
+        original.removeKey("modifiedByUserID");
+
+        objectMapper.writeValue(writer, organization);
+        objectMapper.writeValue(os, organization);
+
+        Log.debug(String.format("Submitting the organization: %s", writer.toString()));
+
+        if (connection.getResponseCode() != 200) {
+            throw new RuntimeException("Failed to put organization: HTTP error code: " + connection.getResponseCode());
+        }
+
+
+        String result = IOUtils.toString(connection.getInputStream(), StandardCharsets.UTF_8);
+        Log.debug(String.format("Accepted Capella IQ terms. Updated org: %s", result));
+        return;
+    }
 }

src/main/java/com/couchbase/intellij/tree/iq/CapellaOrganization.java

@@ -1,13 +1,18 @@
 package com.couchbase.intellij.tree.iq;
 
+import com.fasterxml.jackson.annotation.JsonProperty;
 import lombok.Data;
 
 import java.util.List;
+import java.util.Map;
 
 @Data
 public class CapellaOrganization {
     private String id;
     private String name;
+    private String description;
+    private String website;
+    private Map<String, Object> preferences;
 
     private IQ iq;
 
@@ -28,7 +33,7 @@ public void setIsTermsAcceptedForOrg(boolean value) {
             isTermsAcceptedForOrg = value;
         }
 
-        public boolean isTermsAcceptedForOrg() {
+        public boolean getIsTermsAcceptedForOrg() {
             return isTermsAcceptedForOrg;
         }
     }

src/main/java/com/couchbase/intellij/tree/iq/CapellaOrganizationList.java

@@ -26,7 +26,7 @@ public CapellaOrganizationList getOnlyIqEnabledOrgs() {
                         .filter(org -> org.getData().getIq() != null)
                         .filter(org -> org.getData().getIq().isEnabled())
                         .filter(org -> org.getData().getIq().getOther() != null)
-                        .filter(org -> org.getData().getIq().getOther().isTermsAcceptedForOrg())
+                        .filter(org -> org.getData().getIq().getOther().getIsTermsAcceptedForOrg())
                         .collect(Collectors.toList())
         );
         return filteredList;

src/main/java/com/couchbase/intellij/tree/iq/IQWindowContent.java

@@ -5,14 +5,15 @@
 import com.couchbase.intellij.persistence.storage.IQStorage;
 import com.couchbase.intellij.tree.iq.core.IQCredentials;
 import com.couchbase.intellij.tree.iq.settings.OpenAISettingsState;
+import com.couchbase.intellij.tree.iq.ui.CapellaIqTermsDialog;
 import com.couchbase.intellij.tree.iq.ui.ChatPanel;
 import com.couchbase.intellij.tree.iq.ui.LoginPanel;
 import com.couchbase.intellij.tree.iq.ui.action.editor.ActionsUtil;
+import com.couchbase.intellij.tree.iq.ui.view.CapellaOrgSelectorView;
 import com.couchbase.intellij.workbench.Log;
 import com.intellij.notification.Notification;
 import com.intellij.notification.NotificationType;
 import com.intellij.notification.Notifications;
-import com.intellij.openapi.application.ApplicationManager;
 import com.intellij.openapi.progress.ProgressIndicator;
 import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.progress.Task;
@@ -85,28 +86,36 @@ public void onLogin(IQCredentials credentials) {
         this.credentials = credentials;
         this.removeAll();
         try {
-            this.organizationList = CapellaApiMethods.loadOrganizations(credentials.getAuth());
+            this.organizationList = credentials.getOrganizations();
             if (organizationList.getData().isEmpty()) {
                 Notifications.Bus.notify(new Notification(ChatGptBundle.message("group.id"), "No Capella organizations found", "At least one organization is required to use Couchbase IQ. No organizations found.", NotificationType.ERROR));
                 onLogout(null);
                 return;
             }
 
-            if (this.organizationList != null) {
-                this.organizationList = this.organizationList.getOnlyIqEnabledOrgs();
+            CapellaOrganization activeOrg = null;
+            String userSelectedOrg = IQStorage.getInstance().getState().getActiveOrganization();
+            final boolean forceOrgChooser = System.getenv().containsKey("CB_IQ_FORCE_ORG_CHOOSER");
+            if (userSelectedOrg != null && !forceOrgChooser) {
+                // try and load previously selected org
+                activeOrg = organizationList.getData().stream()
+                        .filter(org -> userSelectedOrg.equalsIgnoreCase(org.getData().getId()))
+                        .map(CapellaOrganizationList.Entry::getData)
+                        .findFirst().orElse(null);
             }
 
-            CapellaOrganization activeOrg = organizationList.getData().stream().map(org -> org.getData()).filter(data -> credentials.checkIqIsEnabled(data.getId())).filter(data -> credentials.checkTermsAccepted(data.getId())).findFirst().orElse(null);
-            String orgId = IQStorage.getInstance().getState().getActiveOrganization();
-            if (orgId != null) {
-                activeOrg = organizationList.getData().stream().filter(org -> orgId.equalsIgnoreCase(org.getData().getId())).map(CapellaOrganizationList.Entry::getData).findFirst().orElse(activeOrg);
-            }
-
-            if (activeOrg == null) {
-                Notifications.Bus.notify(new Notification(ChatGptBundle.message("group.id"), "No Capella organizations with iQ enabled found", "At least one organization with enabled iQ feature and accepted terms and conditions is required to use Couchbase IQ. No organizations found.", NotificationType.ERROR));
-                onLogout(null);
-            } else {
+            if (activeOrg != null) {
                 this.onOrgSelected(activeOrg);
+            } else {
+                // try to auto-select an org
+                if (organizationList.getData().size() == 1 && !forceOrgChooser) {
+                    this.onOrgSelected(organizationList.getData().get(0).getData());
+                } else {
+                    this.removeAll();
+                    this.add(new CapellaOrgSelectorView(organizationList, this));
+                    this.updateUI();
+                    return;
+                }
             }
         } catch (Exception e) {
             Log.error("Failed to initialize IQ", e);
@@ -119,26 +128,39 @@ public void onLogin(IQCredentials credentials) {
     public boolean onLogout(@Nullable Throwable reason) {
         this.removeAll();
         this.add(new LoginPanel(credentials, this));
+        this.setEnabled(true);
         this.updateUI();
         return true;
     }
 
     @Override
     public void onOrgSelected(CapellaOrganization organization) {
+        if (organization == null) {
+            this.onLogout(null);
+            return;
+        }
+
         if (!credentials.checkIqIsEnabled(organization.getId())) {
             Notifications.Bus.notify(new Notification(ChatGptBundle.message("group.id"), "Unable to use this organization", "Capella iQ is not enabled for this organization.", NotificationType.ERROR));
             onLogout(null);
             return;
         }
 
-        if (!credentials.checkTermsAccepted(organization.getId())) {
-            Notifications.Bus.notify(new Notification(ChatGptBundle.message("group.id"), "Unable to use this organization", "Capella iQ terms of use have not been accepted for this organization. Please accept terms of use in Capella", NotificationType.ERROR));
-            onLogout(null);
+        final boolean CBIQ_FORCE_TERMS_DIALOG = System.getenv().containsKey("CB_IQ_FORCE_TERMS_DIALOG");
+        if (CBIQ_FORCE_TERMS_DIALOG || !credentials.checkTermsAccepted(organization.getId())) {
+            this.setEnabled(false);
+            CapellaIqTermsDialog termsDialog = getCapellaIqTermsDialog(organization);
+            termsDialog.show();
             return;
         }
+        showIq(organization);
+    }
+
+    private void showIq(CapellaOrganization organization) {
         this.removeAll();
+        this.setEnabled(true);
         this.updateUI();
-        ApplicationManager.getApplication().invokeLater(() -> {
+        SwingUtilities.invokeLater(() -> {
             IQStorage.getInstance().getState().setActiveOrganization(organization.getId());
             final String iqUrl = String.format(IQ_URL.get(), organization.getId());
             iqGptConfig = new OpenAISettingsState.OpenAIConfig();
@@ -149,13 +171,37 @@ public void onOrgSelected(CapellaOrganization organization) {
             iqGptConfig.setModelName("gpt-4");
             iqGptConfig.setApiEndpointUrl(iqUrl);
             iqGptConfig.setEnableCustomApiEndpointUrl(true);
+
             chatPanel = new ChatPanel(project, iqGptConfig.withSystemPrompt(IQWindowContent::systemPrompt), organizationList, organization, this, this);
             ActionsUtil.refreshActions();
             this.add(chatPanel);
             this.updateUI();
         });
     }
 
+    @NotNull
+    private CapellaIqTermsDialog getCapellaIqTermsDialog(CapellaOrganization organization) {
+        CapellaIqTermsDialog termsDialog = new CapellaIqTermsDialog(project);
+        termsDialog.setOnDeactivationAction(() -> {
+            if (termsDialog.isAccepted()) {
+                try {
+                    organization.getIq().getOther().setIsTermsAcceptedForOrg(true);
+                    CapellaApiMethods.acceptIqTerms(credentials.getAuth(), organization);
+                    if (credentials.doLogin()) {
+                        showIq(organization);
+                    } else {
+                        onLogout(null);
+                    }
+                } catch (Exception e) {
+                    throw new RuntimeException(e);
+                }
+            } else {
+                return;
+            }
+        });
+        return termsDialog;
+    }
+
     public static String systemPrompt() {
         try {
             if (cachedPrompt == null) {

src/main/java/com/couchbase/intellij/tree/iq/core/IQCredentials.java

@@ -7,6 +7,7 @@
 import com.intellij.credentialStore.CredentialAttributesKt;
 import com.intellij.credentialStore.Credentials;
 import com.intellij.ide.passwordSafe.PasswordSafe;
+import com.intellij.openapi.progress.ProgressManager;
 import com.intellij.openapi.util.Pass;
 
 import java.io.IOException;
@@ -117,7 +118,7 @@ public boolean checkTermsAccepted(String orgId) {
                 .map(CapellaOrganizationList.Entry::getData)
                 .filter(org -> Objects.equals(orgId, org.getId()))
                 .filter(org -> org.getIq() != null)
-                .anyMatch(org -> org.getIq().getOther().isTermsAcceptedForOrg());
+                .anyMatch(org -> org.getIq().getOther().getIsTermsAcceptedForOrg());
    }
 
     public boolean doLogin() {

src/main/java/com/couchbase/intellij/tree/iq/ui/CapellaIqTermsDialog.java

@@ -0,0 +1,71 @@
+package com.couchbase.intellij.tree.iq.ui;
+
+import com.github.weisj.jsvg.J;
+import com.intellij.openapi.project.Project;
+import com.intellij.openapi.ui.DialogWrapper;
+import com.intellij.util.ui.JBEmptyBorder;
+import com.intellij.util.ui.JBUI;
+import org.jetbrains.annotations.Nullable;
+
+import javax.swing.*;
+import javax.swing.border.Border;
+import javax.swing.border.EmptyBorder;
+import javax.swing.event.HyperlinkEvent;
+import java.awt.*;
+import java.io.IOException;
+import java.net.URI;
+
+public class CapellaIqTermsDialog extends DialogWrapper {
+    private static final String IQ_TERMS_LABEL = "<html>Capella IQ uses a third-party large language model (LLM).<br/>" +
+            "Please do not enter sensitive data into iQ and review its output before using.</html>";
+    private static final String IHAVEREADANDACCEPTED = "I have read and agree to the ";
+    private static final String TERMS_URL = "https://www.couchbase.com/iq-terms/";
+    private static final String SUPPLEMENTALTERMS = "<html><a href='" + TERMS_URL + "'>Capella iQ supplemental terms.</html>";
+
+    private Project project;
+    private JCheckBox acceptCheckBox;
+
+    @Override
+    protected void doOKAction() {
+        super.doOKAction();
+        isOk = true;
+    }
+
+    private boolean isOk;
+
+    public CapellaIqTermsDialog(Project project) {
+        super(project, true);
+        this.project = project;
+        init();
+    }
+
+    @Override
+    protected @Nullable JComponent createCenterPanel() {
+        JPanel mainPanel = new JPanel(new BorderLayout());
+        mainPanel.add(new JLabel(IQ_TERMS_LABEL), BorderLayout.NORTH);
+
+        // todo: check if the user is an admin org
+        acceptCheckBox = new JCheckBox(IHAVEREADANDACCEPTED);
+        mainPanel.add(acceptCheckBox, BorderLayout.WEST);
+        JEditorPane termsLink = new JEditorPane("text/html", SUPPLEMENTALTERMS);
+        termsLink.setEditable(false);
+        termsLink.setBackground(null);
+        termsLink.setCursor(new Cursor(Cursor.HAND_CURSOR));
+        termsLink.setMargin(JBUI.insets(1, 0, 0, 0));
+        termsLink.addHyperlinkListener(e -> {
+            if (e.getEventType() == HyperlinkEvent.EventType.ACTIVATED) {
+                try {
+                    Desktop.getDesktop().browse(URI.create(TERMS_URL));
+                } catch (IOException ex) {
+                    throw new RuntimeException(ex);
+                }
+            }
+        });
+        mainPanel.add(termsLink, BorderLayout.CENTER);
+        return mainPanel;
+    }
+
+    public boolean isAccepted() {
+        return isOk && acceptCheckBox.isSelected();
+    }
+}