diff --git a/src/org/labkey/test/components/domain/CalculatedColumnAssistantDialog.java b/src/org/labkey/test/components/domain/CalculatedColumnAssistantDialog.java new file mode 100644 index 0000000000..4556966f28 --- /dev/null +++ b/src/org/labkey/test/components/domain/CalculatedColumnAssistantDialog.java @@ -0,0 +1,233 @@ +package org.labkey.test.components.domain; + +import org.labkey.test.Locator; +import org.labkey.test.WebDriverWrapper; +import org.labkey.test.components.bootstrap.ModalDialog; +import org.openqa.selenium.WebElement; + +import java.util.List; +import java.util.stream.Collectors; + +/** + * Modal that opens when the user clicks the "AI Assistant" button inside the Calculation field options. + * Provides a chat-style interface where the user enters prompts and the assistant suggests expressions. + */ +public class CalculatedColumnAssistantDialog extends ModalDialog +{ + public static final String TITLE = "Expression AI Assistant"; + + private final DomainFieldRow _row; + + public CalculatedColumnAssistantDialog(DomainFieldRow row, ModalDialogFinder finder) + { + super(finder); + _row = row; + } + + public CalculatedColumnAssistantDialog(DomainFieldRow row) + { + this(row, new ModalDialogFinder(row.getDriver()).withTitle(TITLE)); + } + + /** + * Type the prompt into the textarea. The submit button stays disabled until non-empty text is present. + */ + public CalculatedColumnAssistantDialog setPrompt(String prompt) + { + getWrapper().setFormElement(elementCache().promptInput, prompt); + WebDriverWrapper.waitFor(() -> elementCache().promptSubmitButton.isEnabled(), + "Prompt submit button did not become enabled.", 2_000); + return this; + } + + public String getPrompt() + { + return getWrapper().getFormElement(elementCache().promptInput); + } + + /** + * Click the submit (arrow) button. First waits for the "Thinking..." spinner to disappear (up to 60s) + * and then for a new assistant response to render (up to 10s). + */ + public CalculatedColumnAssistantDialog submitPrompt() + { + int previousCount = getAssistantResponses().size(); + elementCache().promptSubmitButton.click(); + waitForThinkingSpinnerToDisappear(); + WebDriverWrapper.waitFor(() -> getAssistantResponses().size() > previousCount, + "No new assistant response appeared in chat history.", 10_000); + return this; + } + + private void waitForThinkingSpinnerToDisappear() + { + Locator spinner = Locator.tagWithClass("i", "fa-spinner"); + // Spinner may not appear if the response is instantaneous; that's fine. + WebDriverWrapper.waitFor(() -> !spinner.existsIn(this), 60_000); + } + + /** + * Convenience: type the prompt and submit it. + */ + public CalculatedColumnAssistantDialog sendPrompt(String prompt) + { + return setPrompt(prompt).submitPrompt(); + } + + /** + * @return one entry per assistant response bubble (concatenated text of all its {@code .assistant-text} blocks), + * in chat order. Suggested-expression SQL is not included here — see {@link #getSuggestedExpressions()}. + */ + public List getAssistantResponses() + { + return Locator.tagWithClass("div", "chat-item").withClass("assistant-response") + .findElements(this).stream() + .map(WebElement::getText) + .collect(Collectors.toList()); + } + + /** + * @return text of the most recent assistant response, or empty string if there are none. + */ + public String getLastAssistantResponse() + { + List responses = getAssistantResponses(); + return responses.isEmpty() ? "" : responses.get(responses.size() - 1); + } + + /** + * @return every applicable SQL expression suggested in the most recent assistant response, in display + * order. Only counts {@code .assistant-expression} blocks that include an "Apply Expression" button — read-only + * SQL the assistant shows for illustration (e.g. an alternative custom-query example) is excluded, since the user + * can't accept it as the field's calculation. + */ + public List getSuggestedExpressions() + { + WebElement lastResponse = lastAssistantResponseElement(); + if (lastResponse == null) + return List.of(); + return Locator.tagWithClass("div", "assistant-expression") + .withDescendant(Locator.tagWithClass("button", "clickable-text")) + .descendant(Locator.tag("code")) + .findElements(lastResponse).stream() + .map(WebElement::getText) + .collect(Collectors.toList()); + } + + /** + * @return the first SQL expression in the most recent assistant response, or empty string if none. + */ + public String getFirstSuggestedExpression() + { + List expressions = getSuggestedExpressions(); + return expressions.isEmpty() ? "" : expressions.get(0); + } + + /** + * Click "Apply Expression" on the first suggestion in the most recent assistant response. + * Returns the underlying field row (the dialog stays open; call {@link #clickEndChat()} to close it). + */ + public DomainFieldRow applyFirstSuggestedExpression() + { + return applySuggestedExpression(0); + } + + /** + * Click "Apply Expression" on the suggestion at the given index in the most recent assistant response. + */ + public DomainFieldRow applySuggestedExpression(int index) + { + WebElement lastResponse = lastAssistantResponseElement(); + if (lastResponse == null) + throw new IllegalStateException("No assistant response is available to apply."); + List buttons = Locator.tagWithClass("div", "assistant-expression") + .descendant(Locator.tagWithClass("button", "clickable-text")) + .findElements(lastResponse); + if (index < 0 || index >= buttons.size()) + throw new IndexOutOfBoundsException( + "Requested expression index " + index + " but only " + buttons.size() + " expression(s) available."); + buttons.get(index).click(); + return _row; + } + + /** + * @return text of the first assistant response in the chat history, or empty string if there are none. Useful + * for asserting the intro message in NEW / CHANGE / VALIDATE entry modes. + */ + public String getFirstAssistantResponse() + { + List responses = getAssistantResponses(); + return responses.isEmpty() ? "" : responses.get(0); + } + + /** + * @return true while the dialog is waiting for an AI response (the "Thinking..." pending bubble is shown). + */ + public boolean isPending() + { + return Locator.tagWithClass("div", "chat-item").withClass("pending").existsIn(this); + } + + /** + * Click the stop button to abort an in-flight AI request. The submit button toggles to a stop button (fa-stop) + * while the dialog is in the pending state; calling this method when no request is pending will fail. + */ + public void clickStop() + { + Locator.tagWithClass("button", "prompt-button") + .withDescendant(Locator.tagWithClass("i", "fa-stop")) + .findElement(this) + .click(); + } + + /** + * Click submit without waiting for the response. Useful for tests that need to interrupt or otherwise observe + * the pending state before the response arrives. Prefer {@link #submitPrompt()} when the caller wants to wait. + */ + public void clickSubmitWithoutWaiting() + { + elementCache().promptSubmitButton.click(); + } + + private WebElement lastAssistantResponseElement() + { + List responses = Locator.tagWithClass("div", "chat-item").withClass("assistant-response") + .findElements(this); + return responses.isEmpty() ? null : responses.get(responses.size() - 1); + } + + /** + * Click "End Chat" to close the dialog. + */ + public DomainFieldRow clickEndChat() + { + elementCache().endChatButton.click(); + waitForClose(); + return _row; + } + + @Override + protected ElementCache newElementCache() + { + return new ElementCache(); + } + + @Override + protected ElementCache elementCache() + { + return (ElementCache) super.elementCache(); + } + + protected class ElementCache extends ModalDialog.ElementCache + { + final WebElement endChatButton = Locator.tagWithClass("button", "btn") + .withText("End Chat") + .refindWhenNeeded(this); + + final WebElement promptInput = Locator.tagWithClass("textarea", "prompt-input") + .refindWhenNeeded(this); + + final WebElement promptSubmitButton = Locator.tagWithClass("button", "prompt-button") + .refindWhenNeeded(this); + } +} diff --git a/src/org/labkey/test/components/domain/DomainFieldRow.java b/src/org/labkey/test/components/domain/DomainFieldRow.java index bf3487677a..b52fb6ec43 100644 --- a/src/org/labkey/test/components/domain/DomainFieldRow.java +++ b/src/org/labkey/test/components/domain/DomainFieldRow.java @@ -1085,6 +1085,30 @@ public String getValueExpression() return getWrapper().getFormElement(elementCache().expressionInput); } + /** + * Click the "AI Assistant" button in the expanded Calculation field options and return the resulting dialog. + */ + public CalculatedColumnAssistantDialog openAIAssistant() + { + if (!isExpanded()) + expand(); + elementCache().aiAssistantButton.click(); + return new CalculatedColumnAssistantDialog(this); + } + + /** + * @return true if the "AI Assistant" button is present in the expanded Calculation field options. + * The button is only available when the {@code professional} module is enabled. + */ + public boolean hasAIAssistantButton() + { + if (!isExpanded()) + expand(); + return Locator.tagWithClass("button", "btn") + .withText("AI Assistant") + .findElementOrNull(this) != null; + } + // advanced settings public DomainFieldRow showFieldOnDefaultView(boolean checked) @@ -1778,6 +1802,7 @@ protected class ElementCache extends WebDriverComponent.ElementCache public final WebElement expressionStatusError = expressionStatusMsgLoc.descendant(Locator.tagWithClass("span", "error")).refindWhenNeeded(this); public final WebElement expressionStatusMsg = expressionStatusMsgLoc.childTag("div").refindWhenNeeded(this); public final WebElement expressionValidateLink = expressionStatusMsgLoc.child(Locator.tagWithClass("div", "validate-link")).refindWhenNeeded(this); + public final WebElement aiAssistantButton = Locator.tagWithClass("button", "btn").withText("AI Assistant").refindWhenNeeded(this); Locator.XPathLocator aliquotWarningAlert = Locator.tagWithClassContaining("div", "aliquot-alert-warning");