|
| 1 | +package org.labkey.test.components.domain; |
| 2 | + |
| 3 | +import org.labkey.test.Locator; |
| 4 | +import org.labkey.test.WebDriverWrapper; |
| 5 | +import org.labkey.test.components.bootstrap.ModalDialog; |
| 6 | +import org.openqa.selenium.WebElement; |
| 7 | + |
| 8 | +import java.util.ArrayList; |
| 9 | +import java.util.List; |
| 10 | +import java.util.stream.Collectors; |
| 11 | + |
| 12 | +/** |
| 13 | + * Modal that opens when the user clicks the "AI Assistant" button inside the Calculation field options. |
| 14 | + */ |
| 15 | +public class CalculatedColumnAssistantDialog extends ModalDialog |
| 16 | +{ |
| 17 | + public static final String TITLE = "Expression AI Assistant"; |
| 18 | + |
| 19 | + private final DomainFieldRow _row; |
| 20 | + |
| 21 | + public CalculatedColumnAssistantDialog(DomainFieldRow row, ModalDialogFinder finder) |
| 22 | + { |
| 23 | + super(finder); |
| 24 | + _row = row; |
| 25 | + } |
| 26 | + |
| 27 | + public CalculatedColumnAssistantDialog(DomainFieldRow row) |
| 28 | + { |
| 29 | + this(row, new ModalDialogFinder(row.getDriver()).withTitle(TITLE)); |
| 30 | + } |
| 31 | + |
| 32 | + /** |
| 33 | + * Type the prompt into the textarea. The submit button stays disabled until non-empty text is present. |
| 34 | + */ |
| 35 | + public CalculatedColumnAssistantDialog setPrompt(String prompt) |
| 36 | + { |
| 37 | + getWrapper().setFormElement(elementCache().promptInput, prompt); |
| 38 | + WebDriverWrapper.waitFor(() -> elementCache().promptSubmitButton.isEnabled(), |
| 39 | + "Prompt submit button did not become enabled.", 2_000); |
| 40 | + return this; |
| 41 | + } |
| 42 | + |
| 43 | + public String getPrompt() |
| 44 | + { |
| 45 | + return getWrapper().getFormElement(elementCache().promptInput); |
| 46 | + } |
| 47 | + |
| 48 | + /** |
| 49 | + * Click the submit (arrow) button. First waits for the "Thinking..." spinner to disappear (up to 60s) |
| 50 | + * and then for a new assistant response to render (up to 10s). |
| 51 | + */ |
| 52 | + public CalculatedColumnAssistantDialog submitPrompt() |
| 53 | + { |
| 54 | + int previousCount = getAssistantResponses().size(); |
| 55 | + elementCache().promptSubmitButton.click(); |
| 56 | + waitForThinkingSpinnerToDisappear(); |
| 57 | + WebDriverWrapper.waitFor(() -> getAssistantResponses().size() > previousCount, |
| 58 | + "No new assistant response appeared in chat history.", 10_000); |
| 59 | + return this; |
| 60 | + } |
| 61 | + |
| 62 | + private void waitForThinkingSpinnerToDisappear() |
| 63 | + { |
| 64 | + WebDriverWrapper.waitFor(() -> !Locators.thinkingSpinner.existsIn(this), 60_000); |
| 65 | + } |
| 66 | + |
| 67 | + /** |
| 68 | + * Convenience: type the prompt and submit it. |
| 69 | + */ |
| 70 | + public CalculatedColumnAssistantDialog sendPrompt(String prompt) |
| 71 | + { |
| 72 | + return setPrompt(prompt).submitPrompt(); |
| 73 | + } |
| 74 | + |
| 75 | + /** |
| 76 | + * @return one entry per assistant response bubble (concatenated text of all its {@code .assistant-text} blocks), |
| 77 | + * in chat order. Suggested-expression SQL is not included here — see {@link #getSuggestedExpressions()}. |
| 78 | + */ |
| 79 | + public List<String> getAssistantResponses() |
| 80 | + { |
| 81 | + return Locators.assistantResponse.findElements(this).stream() |
| 82 | + .map(WebElement::getText) |
| 83 | + .collect(Collectors.toList()); |
| 84 | + } |
| 85 | + |
| 86 | + /** |
| 87 | + * @return text of the most recent assistant response, or empty string if there are none. |
| 88 | + */ |
| 89 | + public String getLastAssistantResponse() |
| 90 | + { |
| 91 | + List<String> responses = getAssistantResponses(); |
| 92 | + return responses.isEmpty() ? "" : responses.get(responses.size() - 1); |
| 93 | + } |
| 94 | + |
| 95 | + /** |
| 96 | + * @return every <em>applicable</em> SQL expression suggested in the most recent assistant response, in display |
| 97 | + * order. Only counts {@code .assistant-expression} blocks that include an "Apply Expression" button — read-only |
| 98 | + * SQL the assistant shows for illustration (e.g. an alternative custom-query example) is excluded, since the user |
| 99 | + * can't accept it as the field's calculation. |
| 100 | + */ |
| 101 | + public List<String> getSuggestedExpressions() |
| 102 | + { |
| 103 | + WebElement lastResponse = lastAssistantResponseElement(); |
| 104 | + if (lastResponse == null) |
| 105 | + return List.of(); |
| 106 | + return Locators.applicableSqlCode.findElements(lastResponse).stream() |
| 107 | + .map(WebElement::getText) |
| 108 | + .collect(Collectors.toList()); |
| 109 | + } |
| 110 | + |
| 111 | + /** |
| 112 | + * @return the first SQL expression in the most recent assistant response, or empty string if none. |
| 113 | + */ |
| 114 | + public String getFirstSuggestedExpression() |
| 115 | + { |
| 116 | + List<String> expressions = getSuggestedExpressions(); |
| 117 | + return expressions.isEmpty() ? "" : expressions.get(0); |
| 118 | + } |
| 119 | + |
| 120 | + /** |
| 121 | + * Click "Apply Expression" on the first suggestion in the most recent assistant response. |
| 122 | + * Returns the underlying field row (the dialog stays open; call {@link #clickEndChat()} to close it). |
| 123 | + */ |
| 124 | + public DomainFieldRow applyFirstSuggestedExpression() |
| 125 | + { |
| 126 | + return applySuggestedExpression(0); |
| 127 | + } |
| 128 | + |
| 129 | + /** |
| 130 | + * Click "Apply Expression" on the suggestion at the given index in the most recent assistant response. Waits |
| 131 | + * up to 5 seconds for at least one applicable expression to render — the spinner disappears as soon as the |
| 132 | + * bubble exists, but the inner {@code assistant-expression} block sometimes finishes rendering a moment later. |
| 133 | + */ |
| 134 | + public DomainFieldRow applySuggestedExpression(int index) |
| 135 | + { |
| 136 | + List<WebElement> buttons = new ArrayList<>(); |
| 137 | + WebDriverWrapper.waitFor(() -> { |
| 138 | + buttons.clear(); |
| 139 | + WebElement last = lastAssistantResponseElement(); |
| 140 | + if (last != null) |
| 141 | + buttons.addAll(Locators.applyButton.findElements(last)); |
| 142 | + return !buttons.isEmpty(); |
| 143 | + }, |
| 144 | + "No applicable expression rendered in the assistant response.", |
| 145 | + 5_000); |
| 146 | + |
| 147 | + if (index >= buttons.size()) |
| 148 | + throw new IndexOutOfBoundsException( |
| 149 | + "Requested expression index " + index + " but only " + buttons.size() + " expression(s) available."); |
| 150 | + buttons.get(index).click(); |
| 151 | + return _row; |
| 152 | + } |
| 153 | + |
| 154 | + /** |
| 155 | + * @return text of the first assistant response in the chat history, or empty string if there are none. Useful |
| 156 | + * for asserting the intro message in NEW / CHANGE / VALIDATE entry modes. |
| 157 | + */ |
| 158 | + public String getFirstAssistantResponse() |
| 159 | + { |
| 160 | + List<String> responses = getAssistantResponses(); |
| 161 | + return responses.isEmpty() ? "" : responses.get(0); |
| 162 | + } |
| 163 | + |
| 164 | + /** |
| 165 | + * @return true while the dialog is waiting for an AI response (the "Thinking..." pending bubble is shown). |
| 166 | + */ |
| 167 | + public boolean isPending() |
| 168 | + { |
| 169 | + return Locators.pendingBubble.existsIn(this); |
| 170 | + } |
| 171 | + |
| 172 | + /** |
| 173 | + * Click the stop button to abort an in-flight AI request. The submit button toggles to a stop button (fa-stop) |
| 174 | + * while the dialog is in the pending state; calling this method when no request is pending will fail. |
| 175 | + */ |
| 176 | + public void clickStop() |
| 177 | + { |
| 178 | + Locators.stopButton.findElement(this).click(); |
| 179 | + } |
| 180 | + |
| 181 | + /** |
| 182 | + * Click submit without waiting for the response. Useful for tests that need to interrupt or otherwise observe |
| 183 | + * the pending state before the response arrives. Prefer {@link #submitPrompt()} when the caller wants to wait. |
| 184 | + */ |
| 185 | + public void clickSubmitWithoutWaiting() |
| 186 | + { |
| 187 | + elementCache().promptSubmitButton.click(); |
| 188 | + } |
| 189 | + |
| 190 | + private WebElement lastAssistantResponseElement() |
| 191 | + { |
| 192 | + List<WebElement> responses = Locators.assistantResponse.findElements(this); |
| 193 | + return responses.isEmpty() ? null : responses.get(responses.size() - 1); |
| 194 | + } |
| 195 | + |
| 196 | + /** |
| 197 | + * Click "End Chat" to close the dialog. |
| 198 | + */ |
| 199 | + public DomainFieldRow clickEndChat() |
| 200 | + { |
| 201 | + elementCache().endChatButton.click(); |
| 202 | + waitForClose(); |
| 203 | + return _row; |
| 204 | + } |
| 205 | + |
| 206 | + @Override |
| 207 | + protected ElementCache newElementCache() |
| 208 | + { |
| 209 | + return new ElementCache(); |
| 210 | + } |
| 211 | + |
| 212 | + @Override |
| 213 | + protected ElementCache elementCache() |
| 214 | + { |
| 215 | + return (ElementCache) super.elementCache(); |
| 216 | + } |
| 217 | + |
| 218 | + public static class Locators |
| 219 | + { |
| 220 | + public static final Locator.XPathLocator assistantResponse = Locator.tagWithClass("div", "chat-item").withClass("assistant-response"); |
| 221 | + |
| 222 | + public static final Locator.XPathLocator pendingBubble = Locator.tagWithClass("div", "chat-item").withClass("pending"); |
| 223 | + |
| 224 | + public static final Locator.XPathLocator thinkingSpinner = Locator.tagWithClass("i", "fa-spinner"); |
| 225 | + |
| 226 | + public static final Locator.XPathLocator applyButton = Locator.tagWithClass("div", "assistant-expression") |
| 227 | + .descendant(Locator.tagWithClass("button", "clickable-text")); |
| 228 | + |
| 229 | + public static final Locator.XPathLocator applicableSqlCode = Locator.tagWithClass("div", "assistant-expression") |
| 230 | + .withDescendant(Locator.tagWithClass("button", "clickable-text")) |
| 231 | + .descendant(Locator.tag("code")); |
| 232 | + |
| 233 | + public static final Locator.XPathLocator stopButton = Locator.tagWithClass("button", "prompt-button") |
| 234 | + .withDescendant(Locator.tagWithClass("i", "fa-stop")); |
| 235 | + } |
| 236 | + |
| 237 | + protected class ElementCache extends ModalDialog.ElementCache |
| 238 | + { |
| 239 | + final WebElement endChatButton = Locator.tagWithClass("button", "btn").withText("End Chat").findWhenNeeded(this); |
| 240 | + |
| 241 | + final WebElement promptInput = Locator.tagWithClass("textarea", "prompt-input").findWhenNeeded(this); |
| 242 | + |
| 243 | + final WebElement promptSubmitButton = Locator.tagWithClass("button", "prompt-button").refindWhenNeeded(this); |
| 244 | + } |
| 245 | +} |
0 commit comments