[ZEPPELIN-6760] Add an end-to-end test for local assistant tool testing

zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantE2ETest.java

@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.zeppelin.service.assistant;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import org.apache.commons.lang3.StringUtils;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIf;
+import java.util.ArrayList;
+import java.util.Optional;
+import java.util.Set;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+import org.apache.zeppelin.interpreter.InterpreterFactory;
+import org.apache.zeppelin.interpreter.InterpreterResult;
+import org.apache.zeppelin.notebook.AuthorizationService;
+import org.apache.zeppelin.notebook.Note;
+import org.apache.zeppelin.notebook.Notebook;
+import org.apache.zeppelin.notebook.Notebook.NoteProcessor;
+import org.apache.zeppelin.scheduler.Job.Status;
+import org.apache.zeppelin.service.NotebookService;
+import org.apache.zeppelin.user.AuthenticationInfo;
+
+/**
+ * Calls a real model. Skipped unless the API key is passed as a system property:
+ * ./mvnw test -pl zeppelin-server -Dtest=AssistantE2ETest -Dassistant.e2e.apiKey=...
+ */
+@EnabledIf("hasApiKey")
+class AssistantE2ETest {
+  private static final String API_KEY = System.getProperty("assistant.e2e.apiKey");
+  private static final String BASE_URL = "https://api.openai.com/v1";
+  private static final String MODEL = "gpt-6-luna";
+
+  private final ExecutorService executor = Executors.newSingleThreadExecutor();
+
+  static boolean hasApiKey() {
+    return StringUtils.isNotBlank(API_KEY);
+  }
+
+  @AfterEach
+  void tearDown() {
+    executor.shutdownNow();
+  }
+
+  @Test
+  void answersWithMultipleTools() throws Exception {
+    var authInfo = new AuthenticationInfo("holden");
+    var userAndRoles = Set.of("user");
+
+    var note = new Note();
+    note.setInterpreterFactory(mock(InterpreterFactory.class));
+    note.addNewParagraph(authInfo).setText("%md # Intro");
+
+    var paragraph = note.addNewParagraph(authInfo);
+    paragraph.setText("%python\nprint(6 * 7)");
+    paragraph.setStatus(Status.FINISHED);
+    paragraph.setResult(new InterpreterResult(InterpreterResult.Code.SUCCESS, "42"));
+
+    var notebook = mock(Notebook.class);
+    when(notebook.processNote(eq("noteId"), any()))
+        .thenAnswer(invocation -> ((NoteProcessor<?>) invocation.getArgument(1)).process(note));
+    when(notebook.processNote(eq("noteId"), eq(false), any()))
+        .thenAnswer(invocation -> ((NoteProcessor<?>) invocation.getArgument(2)).process(note));
+
+    var authorization = mock(AuthorizationService.class);
+    when(authorization.isReader("noteId", userAndRoles)).thenReturn(true);
+
+    var conversation = Conversation.create("noteId", "e2e", authInfo.getUser());
+    var repository = mock(ConversationRepository.class);
+    when(repository.find("noteId", conversation.getId())).thenReturn(Optional.of(conversation));
+    doAnswer(invocation -> {
+      Conversation target = invocation.getArgument(0);
+      Consumer<Conversation> update = invocation.getArgument(1);
+      update.accept(target);
+      return null;
+    }).when(repository).update(any(), any());
+
+    var userQuery = "What did the second paragraph output?";
+    var toolNames = new ArrayList<String>();
+    var answer = new StringBuilder();
+
+    try (var model = new OpenAiChatModel(BASE_URL, API_KEY, MODEL)) {
+      var sut = new AssistantImpl(
+          notebook,
+          model,
+          new NotebookService(notebook, authorization, null, null),
+          authorization,
+          repository,
+          executor
+      );
+
+      sut.sendMessage("noteId", conversation.getId(),
+          userQuery, authInfo, userAndRoles,
+          (type, payload) -> {
+            if (payload instanceof AssistantEventPayload.ToolCallStarted) {
+              toolNames.add(((AssistantEventPayload.ToolCallStarted) payload).name);
+            } else if (payload instanceof AssistantEventPayload.MessageDone) {
+              answer.append(((AssistantEventPayload.MessageDone) payload).content);
+            } else if (payload instanceof AssistantEventPayload.RunFailed) {
+              answer.append("RUN_FAILED");
+            }
+          }
+      ).get(2, TimeUnit.MINUTES);
+
+      System.out.println(userQuery);
+      System.out.println("Assistant tools: " + toolNames);
+      System.out.println("Assistant response:\n" + answer);
+
+      assertTrue(toolNames.contains("list_paragraphs"), toolNames.toString());
+      assertTrue(toolNames.contains("get_paragraph"), toolNames.toString());
+      assertTrue(answer.toString().contains("42"), answer.toString());
+    }
+  }
+}

zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantTest.java

@@ -36,6 +36,7 @@
 import jakarta.ws.rs.ClientErrorException;
 import jakarta.ws.rs.ForbiddenException;
 import jakarta.ws.rs.ServiceUnavailableException;
+import org.apache.zeppelin.interpreter.InterpreterFactory;
 import org.apache.zeppelin.notebook.AuthorizationService;
 import org.apache.zeppelin.notebook.Note;
 import org.apache.zeppelin.notebook.Notebook;
@@ -361,6 +362,10 @@ void completesToolTurn() throws Exception {
 
     doAnswer(invocation -> {
       Consumer<AssistantEvent> consumer = invocation.getArgument(3);
+      List<ToolSpec> tools = invocation.getArgument(2);
+      var getParagraph = tools.stream().filter(tool -> tool.name.equals("get_paragraph"))
+          .findFirst().orElseThrow();
+      assertEquals(List.of("paragraphId"), getParagraph.parameters.get("required"));
       consumer.accept(new AssistantEvent.ToolCall("call_1", "list_paragraphs", "{}"));
       return null;
     }).doAnswer(invocation -> {
@@ -453,4 +458,70 @@ void guardsRunningConversation() throws Exception {
       finish.countDown();
     }
   }
+
+  @Test
+  void runsToolLoop() throws Exception {
+    var note = new Note();
+    note.setInterpreterFactory(mock(InterpreterFactory.class));
+    var paragraph = note.addNewParagraph(authInfo);
+    paragraph.setText("%md hello");
+    var notebook = mock(Notebook.class);
+    when(notebook.processNote(eq("noteId"), any()))
+        .thenAnswer(invocation -> ((NoteProcessor<?>) invocation.getArgument(1)).process(note));
+    when(notebook.processNote(eq("noteId"), eq(false), any()))
+        .thenAnswer(invocation -> ((NoteProcessor<?>) invocation.getArgument(2)).process(note));
+    var authorization = mock(AuthorizationService.class);
+    when(authorization.isReader("noteId", userAndRoles)).thenReturn(true);
+
+    // Scripted model: list_paragraphs -> get_paragraph -> final answer.
+    var responses = List.<AssistantEvent>of(
+        new AssistantEvent.ToolCall("call_1", "list_paragraphs", "{}"),
+        new AssistantEvent.ToolCall(
+            "call_2", "get_paragraph", "{\"paragraphId\":\"" + paragraph.getId() + "\"}"),
+        new AssistantEvent.TextDelta("done")
+    );
+
+    var capturedHistories = new ArrayList<List<Message>>();
+    ChatModel model = (systemPrompt, history, tools, consumer) -> {
+      capturedHistories.add(List.copyOf(history));
+      consumer.accept(responses.get(capturedHistories.size() - 1));
+    };
+
+    var conversation = Conversation.create("noteId", "test", authInfo.getUser());
+    var repository = mock(ConversationRepository.class);
+
+    when(repository.find("noteId", conversation.getId())).thenReturn(Optional.of(conversation));
+    doAnswer(invocation -> {
+      invocation.<Consumer<Conversation>>getArgument(1).accept(invocation.getArgument(0));
+      return null;
+    }).when(repository).update(any(), any());
+
+    var sut = new AssistantImpl(
+        notebook,
+        model,
+        new NotebookService(notebook, authorization, null, null),
+        authorization,
+        repository,
+        executor
+    );
+
+    var events = new ArrayList<AssistantEventType>();
+    sut.sendMessage("noteId", conversation.getId(), "show the paragraph", authInfo, userAndRoles,
+        (type, payload) -> events.add(type)
+    ).get(5, TimeUnit.SECONDS);
+
+    assertEquals(List.of(
+        AssistantEventType.RUN_STARTED,
+        AssistantEventType.TOOL_CALL_STARTED, AssistantEventType.TOOL_CALL_DONE,
+        AssistantEventType.TOOL_CALL_STARTED, AssistantEventType.TOOL_CALL_DONE,
+        AssistantEventType.MESSAGE_DELTA, AssistantEventType.MESSAGE_DONE,
+        AssistantEventType.RUN_COMPLETED
+    ), events);
+
+    assertEquals(3, capturedHistories.size());
+    var lastHistory = capturedHistories.get(2);
+    var toolMessage = (Message.Tool) lastHistory.get(lastHistory.size() - 1);
+    assertEquals("call_2", toolMessage.getToolCallId());
+    assertTrue(toolMessage.getContent().contains("%md hello"));
+  }
 }