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..ca94aa9f17 --- /dev/null +++ b/src/org/labkey/test/components/domain/CalculatedColumnAssistantDialog.java @@ -0,0 +1,180 @@ +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 SQL expression suggested in the most recent assistant response, in display order. + * Usually a single entry, occasionally more. + */ + public List getSuggestedExpressions() + { + WebElement lastResponse = lastAssistantResponseElement(); + if (lastResponse == null) + return List.of(); + return Locator.tagWithClass("div", "assistant-expression") + .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() + { + WebElement lastResponse = lastAssistantResponseElement(); + if (lastResponse == null) + throw new IllegalStateException("No assistant response is available to apply."); + Locator.tagWithClass("div", "assistant-expression") + .descendant(Locator.tagWithClass("button", "clickable-text")) + .findElement(lastResponse) + .click(); + return _row; + } + + 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") + .findWhenNeeded(this); + + final WebElement promptInput = Locator.tagWithClass("textarea", "prompt-input") + .findWhenNeeded(this); + + final WebElement promptSubmitButton = Locator.tagWithClass("button", "prompt-button") + .findWhenNeeded(this); + } +} diff --git a/src/org/labkey/test/components/domain/DomainFieldRow.java b/src/org/labkey/test/components/domain/DomainFieldRow.java index 107df90f36..d05c6fc9aa 100644 --- a/src/org/labkey/test/components/domain/DomainFieldRow.java +++ b/src/org/labkey/test/components/domain/DomainFieldRow.java @@ -1070,6 +1070,17 @@ 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); + } + // advanced settings public DomainFieldRow showFieldOnDefaultView(boolean checked) @@ -1763,6 +1774,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");