diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantE2ETest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantE2ETest.java new file mode 100644 index 00000000000..abb12d31b8c --- /dev/null +++ b/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 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(); + 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()); + } + } +} diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantTest.java index 5befd6019f0..cc62c7fef05 100644 --- a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantTest.java +++ b/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 consumer = invocation.getArgument(3); + List 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.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>(); + 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.>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(); + 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")); + } }