diff --git a/conf/zeppelin-site.xml.template b/conf/zeppelin-site.xml.template index 19727fdec30..6cf0e071d6b 100755 --- a/conf/zeppelin-site.xml.template +++ b/conf/zeppelin-site.xml.template @@ -846,4 +846,30 @@ fields to be excluded from being saved in note files, with Paragraph prefix mean the fields in Paragraph, e.g. Paragraph.results + + zeppelin.assistant.enable + false + + + + zeppelin.assistant.openai.api.key + + + + + zeppelin.assistant.openai.endpoint + https://api.openai.com/v1 + + + + zeppelin.assistant.openai.model + gpt-6-luna + + + + zeppelin.assistant.dir + assistant + Directory for assistant conversations. + + diff --git a/zeppelin-common/src/main/java/org/apache/zeppelin/common/Message.java b/zeppelin-common/src/main/java/org/apache/zeppelin/common/Message.java index d1e1606007b..04a3f79cc2b 100644 --- a/zeppelin-common/src/main/java/org/apache/zeppelin/common/Message.java +++ b/zeppelin-common/src/main/java/org/apache/zeppelin/common/Message.java @@ -203,7 +203,9 @@ public enum OP { COLLABORATIVE_MODE_STATUS, // [s-c] collaborative mode status PATCH_PARAGRAPH, // [c-s][s-c] patch editor text NOTE_RUNNING_STATUS, // [s-c] sequential run status will be change - NOTICE // [s-c] Notice + NOTICE, // [s-c] Notice + ASSISTANT_SEND_MESSAGE, // [c-s] send a message to the notebook assistant + ASSISTANT_EVENT // [s-c] assistant run stream event (see "type" field) } // these messages will be ignored during the sequential run of the note diff --git a/zeppelin-distribution/src/bin_license/LICENSE b/zeppelin-distribution/src/bin_license/LICENSE index 3c6939d4030..09594969c90 100644 --- a/zeppelin-distribution/src/bin_license/LICENSE +++ b/zeppelin-distribution/src/bin_license/LICENSE @@ -1,5 +1,25 @@ The following components are provided under Apache License. + (Apache 2.0) ClassMate (com.fasterxml:classmate:1.7.0 - https://github.com/FasterXML/java-classmate) + (Apache 2.0) Jackson JDK8 Datatypes (com.fasterxml.jackson.datatype:jackson-datatype-jdk8:2.18.3 - https://github.com/FasterXML/jackson-modules-java8) + (Apache 2.0) Jackson Java Time Datatypes (com.fasterxml.jackson.datatype:jackson-datatype-jsr310:2.18.3 - https://github.com/FasterXML/jackson-modules-java8) + (Apache 2.0) Jackson Kotlin Module (com.fasterxml.jackson.module:jackson-module-kotlin:2.18.3 - https://github.com/FasterXML/jackson-module-kotlin) + (Apache 2.0) JSON Schema Generator (com.github.victools:jsonschema-generator:4.38.0 - https://github.com/victools/jsonschema-generator) + (Apache 2.0) JSON Schema Generator (com.github.victools:jsonschema-module-jackson:4.38.0 - https://github.com/victools/jsonschema-generator) + (Apache 2.0) JSON Schema Generator (com.github.victools:jsonschema-module-swagger-2:4.38.0 - https://github.com/victools/jsonschema-generator) + (Apache 2.0) OpenAI Java SDK (com.openai:openai-java:4.69.3 - https://github.com/openai/openai-java) + (Apache 2.0) OpenAI Java SDK (com.openai:openai-java-client-okhttp:4.69.3 - https://github.com/openai/openai-java) + (Apache 2.0) OpenAI Java SDK (com.openai:openai-java-core:4.69.3 - https://github.com/openai/openai-java) + (Apache 2.0) OkHttp (com.squareup.okhttp3:okhttp:4.12.0 - https://github.com/square/okhttp) + (Apache 2.0) Okio (com.squareup.okio:okio:3.6.0 - https://github.com/square/okio) + (Apache 2.0) Okio (com.squareup.okio:okio-jvm:3.6.0 - https://github.com/square/okio) + (Apache 2.0) Swagger Annotations (io.swagger.core.v3:swagger-annotations:2.2.31 - https://github.com/swagger-api/swagger-core) + (Apache 2.0) Kotlin (org.jetbrains.kotlin:kotlin-reflect:1.9.10 - https://github.com/JetBrains/kotlin) + (Apache 2.0) Kotlin (org.jetbrains.kotlin:kotlin-stdlib:1.9.10 - https://github.com/JetBrains/kotlin) + (Apache 2.0) Kotlin (org.jetbrains.kotlin:kotlin-stdlib-common:1.9.10 - https://github.com/JetBrains/kotlin) + (Apache 2.0) Kotlin (org.jetbrains.kotlin:kotlin-stdlib-jdk7:1.9.10 - https://github.com/JetBrains/kotlin) + (Apache 2.0) Kotlin (org.jetbrains.kotlin:kotlin-stdlib-jdk8:1.9.10 - https://github.com/JetBrains/kotlin) + (Apache 2.0) nvd3.js v1.7.1 (http://nvd3.org/) - https://github.com/novus/nvd3/blob/v1.7.1/LICENSE.md (Apache 2.0) gson v2.8.9 (com.google.code.gson:gson:jar:2.8.9 - https://github.com/google/gson) - https://github.com/google/gson/blob/gson-parent-2.8.9/LICENSE (Apache 2.0) Amazon Web Services SDK for Java v1.11.736 (https://aws.amazon.com/sdk-for-java/) - https://raw.githubusercontent.com/aws/aws-sdk-java/1.11.736/LICENSE.txt diff --git a/zeppelin-server/pom.xml b/zeppelin-server/pom.xml index 505c5ed75bb..7191c0bb017 100644 --- a/zeppelin-server/pom.xml +++ b/zeppelin-server/pom.xml @@ -432,6 +432,71 @@ gson + + com.openai + openai-java + 4.69.3 + + + + org.jetbrains.kotlin + kotlin-stdlib-jdk8 + + + org.jetbrains.kotlin + kotlin-reflect + + + com.google.errorprone + error_prone_annotations + + + org.jetbrains + annotations + + + + + + + org.jetbrains.kotlin + kotlin-stdlib-jdk8 + 1.9.10 + + + + org.jetbrains + annotations + + + + + + org.jetbrains.kotlin + kotlin-reflect + 1.9.10 + runtime + + + + org.jetbrains + annotations + + + + + + org.jetbrains + annotations + 15.0 + + + + com.google.errorprone + error_prone_annotations + 2.18.0 + + com.nimbusds nimbus-jose-jwt diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/conf/ZeppelinConfiguration.java b/zeppelin-server/src/main/java/org/apache/zeppelin/conf/ZeppelinConfiguration.java index 6a7e1b7909f..8b161d174f0 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/conf/ZeppelinConfiguration.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/conf/ZeppelinConfiguration.java @@ -888,6 +888,26 @@ public boolean isZeppelinSearchSemanticEnable() { return getBoolean(ConfVars.ZEPPELIN_SEARCH_SEMANTIC_ENABLE); } + public boolean isAssistantEnabled() { + return getBoolean(ConfVars.ZEPPELIN_ASSISTANT_ENABLE); + } + + public String getAssistantApiKey() { + return getString(ConfVars.ZEPPELIN_ASSISTANT_OPENAI_API_KEY); + } + + public String getAssistantEndpoint() { + return getString(ConfVars.ZEPPELIN_ASSISTANT_OPENAI_ENDPOINT); + } + + public String getAssistantModel() { + return getString(ConfVars.ZEPPELIN_ASSISTANT_OPENAI_MODEL); + } + + public String getAssistantDir() { + return getAbsoluteDir(ConfVars.ZEPPELIN_ASSISTANT_DIR); + } + public boolean isOnlyYarnCluster() { return getBoolean(ConfVars.ZEPPELIN_SPARK_ONLY_YARN_CLUSTER); } @@ -1193,7 +1213,16 @@ public enum ConfVars { ZEPPELIN_SPARK_ONLY_YARN_CLUSTER("zeppelin.spark.only_yarn_cluster", false), ZEPPELIN_SESSION_CHECK_INTERVAL("zeppelin.session.check_interval", 60 * 10 * 1000), ZEPPELIN_NOTE_CACHE_THRESHOLD("zeppelin.note.cache.threshold", 50), - ZEPPELIN_NOTE_FILE_EXCLUDE_FIELDS("zeppelin.note.file.exclude.fields", ""); + ZEPPELIN_NOTE_FILE_EXCLUDE_FIELDS("zeppelin.note.file.exclude.fields", ""), + ZEPPELIN_ASSISTANT_ENABLE("zeppelin.assistant.enable", false), + ZEPPELIN_ASSISTANT_DIR("zeppelin.assistant.dir", "assistant"), + ZEPPELIN_ASSISTANT_OPENAI_API_KEY("zeppelin.assistant.openai.api.key", ""), + ZEPPELIN_ASSISTANT_OPENAI_ENDPOINT( + "zeppelin.assistant.openai.endpoint", "https://api.openai.com/v1" + ), + ZEPPELIN_ASSISTANT_OPENAI_MODEL( + "zeppelin.assistant.openai.model", "gpt-6-luna" + ); private String varName; private Class varClass; diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/rest/AssistantConversationRestApi.java b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/AssistantConversationRestApi.java new file mode 100644 index 00000000000..4d3b8f52012 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/AssistantConversationRestApi.java @@ -0,0 +1,157 @@ +/* + * 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.rest; + +import com.google.gson.JsonSyntaxException; + +import jakarta.inject.Inject; +import jakarta.inject.Singleton; +import jakarta.ws.rs.BadRequestException; +import jakarta.ws.rs.DELETE; +import jakarta.ws.rs.GET; +import jakarta.ws.rs.PATCH; +import jakarta.ws.rs.POST; +import jakarta.ws.rs.Path; +import jakarta.ws.rs.PathParam; +import jakarta.ws.rs.Produces; +import jakarta.ws.rs.core.Response; + +import java.io.IOException; +import java.util.Optional; +import java.util.stream.Collectors; + +import org.apache.zeppelin.annotation.ZeppelinApi; +import org.apache.zeppelin.rest.message.ConversationRequest; +import org.apache.zeppelin.rest.message.ConversationMetadata; +import org.apache.zeppelin.rest.message.ConversationResponse; +import org.apache.zeppelin.server.JsonResponse; +import org.apache.zeppelin.service.AuthenticationService; +import org.apache.zeppelin.service.assistant.Assistant; + +@Path("/notes/{noteId}/conversations") +@Produces("application/json") +@Singleton +public class AssistantConversationRestApi extends AbstractRestApi { + + private final Assistant assistant; + + @Inject + public AssistantConversationRestApi( + AuthenticationService authenticationService, + Assistant assistant + ) { + super(authenticationService); + this.assistant = assistant; + } + + private static ConversationRequest parseBody(String body) { + try { + return GSON.fromJson(body, ConversationRequest.class); + } catch (JsonSyntaxException e) { + throw new BadRequestException(); + } + } + + @GET + @ZeppelinApi + public Response list(@PathParam("noteId") String noteId) throws IOException { + var context = getServiceContext(); + return new JsonResponse<>( + Response.Status.OK, + "", + assistant.listConversations(noteId, context.getUserAndRoles()).stream() + .map(c -> ConversationMetadata.of(c, context.getAutheInfo().getUser())) + .collect(Collectors.toList()) + ).build(); + } + + @GET + @Path("/{conversationId}") + @ZeppelinApi + public Response get( + @PathParam("noteId") String noteId, + @PathParam("conversationId") String conversationId + ) throws IOException { + var context = getServiceContext(); + var conversation = assistant.getConversation( + noteId, + conversationId, + context.getUserAndRoles() + ); + return new JsonResponse<>( + Response.Status.OK, + "", + ConversationResponse.of(conversation, context.getAutheInfo().getUser()) + ).build(); + } + + @POST + @ZeppelinApi + public Response create( + @PathParam("noteId") String noteId, + String body + ) throws IOException { + String title = Optional.ofNullable(parseBody(body)) + .map(ConversationRequest::getTitle) + .orElse(null); + var context = getServiceContext(); + var conversation = assistant.createConversation( + noteId, title, context.getAutheInfo(), context.getUserAndRoles() + ); + return new JsonResponse<>( + Response.Status.CREATED, + "", + ConversationMetadata.of(conversation, context.getAutheInfo().getUser()) + ).build(); + } + + @PATCH + @Path("/{conversationId}") + @ZeppelinApi + public Response updateTitle( + @PathParam("noteId") String noteId, + @PathParam("conversationId") String conversationId, + String body + ) throws IOException { + String title = Optional.ofNullable(parseBody(body)) + .map(ConversationRequest::getTitle) + .filter(t -> !t.isBlank()) + .orElseThrow(() -> new BadRequestException("title is required")); + var context = getServiceContext(); + var conversation = assistant.updateTitle( + noteId, conversationId, title, context.getAutheInfo().getUser(), context.getUserAndRoles()); + return new JsonResponse<>( + Response.Status.OK, + "", + ConversationMetadata.of(conversation, context.getAutheInfo().getUser()) + ).build(); + } + + @DELETE + @Path("/{conversationId}") + @ZeppelinApi + public Response delete( + @PathParam("noteId") String noteId, + @PathParam("conversationId") String conversationId + ) throws IOException { + var context = getServiceContext(); + assistant.deleteConversation( + noteId, conversationId, context.getAutheInfo().getUser(), context.getUserAndRoles()); + return Response.noContent().build(); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/rest/AssistantMessageRestApi.java b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/AssistantMessageRestApi.java new file mode 100644 index 00000000000..9d46f6ab028 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/AssistantMessageRestApi.java @@ -0,0 +1,71 @@ +/* + * 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.rest; + +import jakarta.inject.Inject; +import jakarta.inject.Singleton; +import jakarta.ws.rs.DefaultValue; +import jakarta.ws.rs.GET; +import jakarta.ws.rs.Path; +import jakarta.ws.rs.PathParam; +import jakarta.ws.rs.Produces; +import jakarta.ws.rs.QueryParam; +import jakarta.ws.rs.core.Response; + +import java.io.IOException; + +import org.apache.zeppelin.annotation.ZeppelinApi; +import org.apache.zeppelin.rest.message.MessagesResponse; +import org.apache.zeppelin.server.JsonResponse; +import org.apache.zeppelin.service.AuthenticationService; +import org.apache.zeppelin.service.assistant.Assistant; + +@Path("/notes/{noteId}/conversations/{conversationId}/messages") +@Produces("application/json") +@Singleton +public class AssistantMessageRestApi extends AbstractRestApi { + + private final Assistant assistant; + + @Inject + public AssistantMessageRestApi( + AuthenticationService authenticationService, + Assistant assistant + ) { + super(authenticationService); + this.assistant = assistant; + } + + @GET + @ZeppelinApi + public Response list( + @PathParam("noteId") String noteId, + @PathParam("conversationId") String conversationId, + @QueryParam("cursor") String cursor, + @QueryParam("limit") @DefaultValue("10") int limit + ) throws IOException { + return new JsonResponse<>( + Response.Status.OK, + "", + MessagesResponse.of( + assistant.listMessages( + noteId, conversationId, cursor, limit, getServiceContext().getUserAndRoles() + )) + ).build(); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationMetadata.java b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationMetadata.java new file mode 100644 index 00000000000..408428c7d42 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationMetadata.java @@ -0,0 +1,55 @@ +/* + * 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.rest.message; + +import org.apache.zeppelin.service.assistant.Conversation; + +public final class ConversationMetadata { + private final String id; + private final String noteId; + private final String ownerId; + private final String title; + private final String createdAt; + private final String updatedAt; + private final boolean canSendMessage; + + private ConversationMetadata( + String id, + String noteId, + String ownerId, + String title, + String createdAt, + String updatedAt, + boolean canSendMessage + ) { + this.id = id; + this.noteId = noteId; + this.ownerId = ownerId; + this.title = title; + this.createdAt = createdAt; + this.updatedAt = updatedAt; + this.canSendMessage = canSendMessage; + } + + public static ConversationMetadata of(Conversation c, String userId) { + return new ConversationMetadata( + c.getId(), c.getNoteId(), c.getOwnerId(), c.getTitle(), c.getCreatedAt(), c.getUpdatedAt(), + c.isOwner(userId) + ); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationRequest.java b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationRequest.java new file mode 100644 index 00000000000..9ecbf0dbf75 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationRequest.java @@ -0,0 +1,29 @@ +/* + * 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.rest.message; + +public class ConversationRequest { + private final String title; + + public ConversationRequest(String title) { + this.title = title; + } + + public String getTitle() { + return title; + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationResponse.java b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationResponse.java new file mode 100644 index 00000000000..e747be73b7f --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/ConversationResponse.java @@ -0,0 +1,48 @@ +/* + * 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.rest.message; + +import java.util.List; +import org.apache.zeppelin.service.assistant.Conversation; +import org.apache.zeppelin.service.assistant.Message; + +public final class ConversationResponse { + private final String id; + private final String noteId; + private final String ownerId; + private final String title; + private final String createdAt; + private final String updatedAt; + private final boolean canSendMessage; + private final List messages; + + private ConversationResponse(Conversation conversation, String userId) { + this.id = conversation.getId(); + this.noteId = conversation.getNoteId(); + this.ownerId = conversation.getOwnerId(); + this.title = conversation.getTitle(); + this.createdAt = conversation.getCreatedAt(); + this.updatedAt = conversation.getUpdatedAt(); + this.canSendMessage = conversation.isOwner(userId); + this.messages = List.copyOf(conversation.getMessages()); + } + + public static ConversationResponse of(Conversation conversation, String userId) { + return new ConversationResponse(conversation, userId); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/MessagesResponse.java b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/MessagesResponse.java new file mode 100644 index 00000000000..e026886cba0 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/rest/message/MessagesResponse.java @@ -0,0 +1,36 @@ +/* + * 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.rest.message; + +import java.util.List; +import org.apache.zeppelin.service.assistant.Message; +import org.apache.zeppelin.service.assistant.MessagePage; + +public final class MessagesResponse { + private final List messages; + private final String cursor; + + private MessagesResponse(MessagePage page) { + this.messages = List.copyOf(page.messages); + this.cursor = page.cursor; + } + + public static MessagesResponse of(MessagePage page) { + return new MessagesResponse(page); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/server/RestApiApplication.java b/zeppelin-server/src/main/java/org/apache/zeppelin/server/RestApiApplication.java index b43fe0111af..79ecf80f84d 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/server/RestApiApplication.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/server/RestApiApplication.java @@ -23,6 +23,8 @@ import jakarta.ws.rs.core.Application; +import org.apache.zeppelin.rest.AssistantConversationRestApi; +import org.apache.zeppelin.rest.AssistantMessageRestApi; import org.apache.zeppelin.rest.AdminRestApi; import org.apache.zeppelin.rest.ConfigurationsRestApi; import org.apache.zeppelin.rest.CredentialRestApi; @@ -43,6 +45,8 @@ public class RestApiApplication extends Application { @Override public Set> getClasses() { Set> s = new HashSet<>(); + s.add(AssistantConversationRestApi.class); + s.add(AssistantMessageRestApi.class); s.add(AdminRestApi.class); s.add(ConfigurationsRestApi.class); s.add(CredentialRestApi.class); diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/server/ZeppelinServer.java b/zeppelin-server/src/main/java/org/apache/zeppelin/server/ZeppelinServer.java index 480422b4a5b..af255aa2c6c 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/server/ZeppelinServer.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/server/ZeppelinServer.java @@ -92,6 +92,8 @@ import org.apache.zeppelin.search.NoSearchService; import org.apache.zeppelin.search.SearchService; import org.apache.zeppelin.service.*; +import org.apache.zeppelin.service.assistant.Assistant; +import org.apache.zeppelin.service.assistant.AssistantModule; import org.apache.zeppelin.service.AuthenticationService; import org.apache.zeppelin.service.auth.AuthenticationServiceFactory; import org.apache.zeppelin.socket.ConnectionManager; @@ -139,6 +141,7 @@ public class ZeppelinServer implements AutoCloseable { private final Server jettyWebServer; private final ServiceLocator sharedServiceLocator; private final ConfigStorage storage; + private AssistantModule assistantModule; public ZeppelinServer(ZeppelinConfiguration zConf) throws IOException { this(zConf, DEFAULT_SERVICE_LOCATOR_NAME); @@ -255,6 +258,24 @@ protected void configure() { // Lazy loading will cause paragraph recovery and cron job initialization is delayed. Notebook notebook = ServiceLocatorUtilities.getService( sharedServiceLocator, Notebook.class.getName()); + + assistantModule = new AssistantModule( + zConf.isAssistantEnabled(), + zConf.getAssistantEndpoint(), + zConf.getAssistantApiKey(), + zConf.getAssistantModel(), + zConf.getAssistantDir(), + notebook, + sharedServiceLocator.getService(NotebookService.class), + sharedServiceLocator.getService(AuthorizationService.class) + ); + ServiceLocatorUtilities.bind(sharedServiceLocator, new AbstractBinder() { + @Override + protected void configure() { + bind(assistantModule.assistant).to(Assistant.class); + } + }); + ServiceLocatorUtilities.getService( sharedServiceLocator, SearchService.class.getName()); ServiceLocatorUtilities.getService( @@ -361,6 +382,9 @@ public void shutdown(int exitCode) { if (sharedServiceLocator != null) { // Stop after Jetty so no new connection can restart the heartbeat scheduler. sharedServiceLocator.getService(NotebookServer.class).stopHeartbeatScheduler(); + if (assistantModule != null) { + assistantModule.close(); + } if (!zConf.isRecoveryEnabled()) { sharedServiceLocator.getService(InterpreterSettingManager.class).close(); } diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/ConfigurationService.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/ConfigurationService.java index a3b226770d6..001f01e8c94 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/service/ConfigurationService.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/ConfigurationService.java @@ -37,6 +37,8 @@ public Map getAllProperties(ServiceContext context, throws IOException { Map properties = zConf.dumpConfigurations(key -> !key.contains("password") && + !key.equals(ZeppelinConfiguration.ConfVars + .ZEPPELIN_ASSISTANT_OPENAI_API_KEY.getVarName()) && !key.equals(ZeppelinConfiguration.ConfVars .ZEPPELIN_NOTEBOOK_AZURE_CONNECTION_STRING.getVarName())); callback.onSuccess(properties, context); @@ -53,6 +55,8 @@ public Map getPropertiesWithPrefix(String prefix, throws IOException { Map properties = zConf.dumpConfigurations(key -> !key.contains("password") && + !key.equals(ZeppelinConfiguration.ConfVars + .ZEPPELIN_ASSISTANT_OPENAI_API_KEY.getVarName()) && !key.equals(ZeppelinConfiguration.ConfVars .ZEPPELIN_NOTEBOOK_AZURE_CONNECTION_STRING .getVarName()) && diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Assistant.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Assistant.java new file mode 100644 index 00000000000..17f0d9e3cc9 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Assistant.java @@ -0,0 +1,78 @@ +/* + * 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 java.io.IOException; +import java.util.List; +import java.util.Set; +import java.util.concurrent.Future; + +import org.apache.zeppelin.user.AuthenticationInfo; + +public interface Assistant { + + List listConversations( + String noteId, + Set userAndRoles + ) throws IOException; + + Conversation createConversation( + String noteId, + String title, + AuthenticationInfo authInfo, + Set userAndRoles + ) throws IOException; + + Conversation getConversation( + String noteId, + String conversationId, + Set userAndRoles + ) throws IOException; + + Conversation updateTitle( + String noteId, + String conversationId, + String title, + String userId, + Set userAndRoles + ) throws IOException; + + void deleteConversation( + String noteId, + String conversationId, + String userId, + Set userAndRoles + ) throws IOException; + + MessagePage listMessages( + String noteId, + String conversationId, + String cursor, + int limit, + Set userAndRoles + ) throws IOException; + + Future sendMessage( + String noteId, + String conversationId, + String userContent, + AuthenticationInfo authInfo, + Set userAndRoles, + AssistantEventListener sink + ); +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEvent.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEvent.java new file mode 100644 index 00000000000..440fc241e52 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEvent.java @@ -0,0 +1,61 @@ +/* + * 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 com.google.gson.Gson; +import com.google.gson.reflect.TypeToken; +import java.util.Collections; +import java.util.Map; +import org.apache.commons.lang3.StringUtils; + +public interface AssistantEvent { + + final class TextDelta implements AssistantEvent { + public final String delta; + + public TextDelta(String delta) { + this.delta = delta; + } + } + + final class ToolCall implements AssistantEvent { + private static final Gson GSON = new Gson(); + + public final String id; + public final String name; + public final Map arguments; + + public ToolCall(String id, String name, String arguments) { + this.id = id; + this.name = name; + this.arguments = StringUtils.isBlank(arguments) + ? Collections.emptyMap() + : GSON.fromJson(arguments, new TypeToken>() { }.getType()); + } + } + + final class Usage implements AssistantEvent { + public final int inputTokens; + public final int outputTokens; + + public Usage(int inputTokens, int outputTokens) { + this.inputTokens = inputTokens; + this.outputTokens = outputTokens; + } + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventListener.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventListener.java new file mode 100644 index 00000000000..f29b8e13733 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventListener.java @@ -0,0 +1,23 @@ +/* + * 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; + +@FunctionalInterface +public interface AssistantEventListener { + void onEvent(AssistantEventType type, AssistantEventPayload payload); +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventPayload.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventPayload.java new file mode 100644 index 00000000000..6a0ad335ff2 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventPayload.java @@ -0,0 +1,121 @@ +/* + * 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 java.util.Map; +import jakarta.ws.rs.WebApplicationException; + +public interface AssistantEventPayload { + + final class RunStarted implements AssistantEventPayload { + final String runId; + final String createdAt; + + RunStarted(String runId, String createdAt) { + this.runId = runId; + this.createdAt = createdAt; + } + } + + final class RunCompleted implements AssistantEventPayload { + final String runId; + final Usage usage; + + RunCompleted(String runId, Usage usage) { + this.runId = runId; + this.usage = usage; + } + } + + final class RunFailed implements AssistantEventPayload { + final String runId; + final Error error; + + RunFailed(String runId, Error error) { + this.runId = runId; + this.error = error; + } + } + + final class MessageDelta implements AssistantEventPayload { + final String messageId; + final String delta; + + MessageDelta(String messageId, String delta) { + this.messageId = messageId; + this.delta = delta; + } + } + + final class MessageDone implements AssistantEventPayload { + final String messageId; + final String content; + + MessageDone(String messageId, String content) { + this.messageId = messageId; + this.content = content; + } + } + + final class ToolCallStarted implements AssistantEventPayload { + final String toolCallId; + final String name; + final Map arguments; + + ToolCallStarted(String toolCallId, String name, Map arguments) { + this.toolCallId = toolCallId; + this.name = name; + this.arguments = arguments; + } + } + + final class ToolCallDone implements AssistantEventPayload { + final String toolCallId; + final ToolResult result; + + ToolCallDone(String toolCallId, ToolResult result) { + this.toolCallId = toolCallId; + this.result = result; + } + } + + final class Usage { + final int inputTokens; + final int outputTokens; + + Usage(int inputTokens, int outputTokens) { + this.inputTokens = inputTokens; + this.outputTokens = outputTokens; + } + } + + final class Error { + final int status; + + private Error(int status) { + this.status = status; + } + + public static Error of(Exception e) { + int status = e instanceof WebApplicationException + ? ((WebApplicationException) e).getResponse().getStatus() + : 500; + return new Error(status); + } + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventType.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventType.java new file mode 100644 index 00000000000..8c0137146a8 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantEventType.java @@ -0,0 +1,34 @@ +/* + * 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; + +public enum AssistantEventType { + RUN_STARTED("run.started"), + RUN_COMPLETED("run.completed"), + RUN_FAILED("run.failed"), + MESSAGE_DELTA("message.delta"), + MESSAGE_DONE("message.done"), + TOOL_CALL_STARTED("tool_call.started"), + TOOL_CALL_DONE("tool_call.done"); + + public final String wireName; + + AssistantEventType(String wireName) { + this.wireName = wireName; + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantImpl.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantImpl.java new file mode 100644 index 00000000000..9020c36625a --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantImpl.java @@ -0,0 +1,376 @@ +/* + * 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 com.google.gson.Gson; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import java.io.IOException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.stream.IntStream; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Future; +import jakarta.ws.rs.BadRequestException; +import jakarta.ws.rs.ClientErrorException; +import jakarta.ws.rs.ForbiddenException; +import jakarta.ws.rs.NotFoundException; +import jakarta.ws.rs.core.Response; +import org.apache.commons.lang3.StringUtils; +import org.apache.zeppelin.notebook.AuthorizationService; +import org.apache.zeppelin.notebook.Notebook; +import org.apache.zeppelin.rest.exception.NoteNotFoundException; +import org.apache.zeppelin.service.NotebookService; +import org.apache.zeppelin.user.AuthenticationInfo; + +public class AssistantImpl implements Assistant { + + private static final Logger LOGGER = LoggerFactory.getLogger(AssistantImpl.class); + private static final Gson GSON = new Gson(); + private static final int HISTORY_MESSAGE_LIMIT = 10; + + private static final String INSTRUCTIONS = + "You are an AI assistant helping users work with a notebook in Apache Zeppelin.\n" + + "Respond in the user's language, and keep your answers concise and clear.\n" + + "Format your responses using Markdown."; + + private final ExecutorService assistantExecutor; + + private final Set busyConversations = ConcurrentHashMap.newKeySet(); + + private final Notebook notebook; + private final ChatModel modelClient; + private final ConversationRepository conversationRepository; + private final AuthorizationService authorizationService; + private final ToolExecutor toolExecutor; + + public AssistantImpl( + Notebook notebook, + ChatModel modelClient, + NotebookService notebookService, + AuthorizationService authorizationService, + ConversationRepository conversationRepository, + ExecutorService assistantExecutor + ) { + this.assistantExecutor = assistantExecutor; + this.authorizationService = authorizationService; + this.notebook = notebook; + this.modelClient = modelClient; + this.conversationRepository = conversationRepository; + this.toolExecutor = new ToolExecutor( + List.of( + new ListParagraphsTool(notebookService) + ) + ); + } + + @Override + public List listConversations( + String noteId, + Set userAndRoles + ) throws IOException { + if (!authorizationService.isReader(noteId, userAndRoles)) { + throw new ForbiddenException(); + } + + return notebook.processNote(noteId, note -> { + if (note == null) throw new NoteNotFoundException(noteId); + return conversationRepository.findAll(noteId); + }); + } + + @Override + public Conversation createConversation( + String noteId, + String title, + AuthenticationInfo authInfo, + Set userAndRoles + ) throws IOException { + if (!authorizationService.isReader(noteId, userAndRoles)) throw new ForbiddenException(); + if (AuthenticationInfo.isAnonymous(authInfo)) throw new ForbiddenException(); + + return notebook.processNote(noteId, note -> { + if (note == null) throw new NoteNotFoundException(noteId); + var conversation = Conversation.create(noteId, title, authInfo.getUser()); + conversationRepository.create(conversation); + return conversation; + }); + } + + @Override + public Conversation getConversation( + String noteId, + String conversationId, + Set userAndRoles + ) throws IOException { + if (!authorizationService.isReader(noteId, userAndRoles)) throw new ForbiddenException(); + + return notebook.processNote(noteId, note -> { + if (note == null) throw new NoteNotFoundException(noteId); + return conversationRepository.find(noteId, conversationId) + .orElseThrow(NotFoundException::new); + }); + } + + @Override + public Conversation updateTitle( + String noteId, + String conversationId, + String title, + String userId, + Set userAndRoles + ) throws IOException { + if (!busyConversations.add(conversationId)) { + throw new ClientErrorException("", Response.Status.CONFLICT); + } + try { + var conversation = getConversation(noteId, conversationId, userAndRoles); + if (!conversation.isOwner(userId)) throw new ForbiddenException(); + return notebook.processNote(noteId, note -> { + if (note == null) throw new NoteNotFoundException(noteId); + conversationRepository.update(conversation, c -> c.setTitle(title)); + return conversation; + }); + } finally { + busyConversations.remove(conversationId); + } + } + + @Override + public void deleteConversation( + String noteId, + String conversationId, + String userId, + Set userAndRoles + ) throws IOException { + if (!busyConversations.add(conversationId)) { + throw new ClientErrorException("", Response.Status.CONFLICT); + } + try { + Conversation conversation = getConversation(noteId, conversationId, userAndRoles); + if (!conversation.isOwner(userId)) throw new ForbiddenException(); + notebook.processNote(noteId, note -> { + if (note == null) throw new NoteNotFoundException(noteId); + conversationRepository.delete(noteId, conversationId); + return null; + }); + } finally { + busyConversations.remove(conversationId); + } + } + + @Override + public MessagePage listMessages( + String noteId, + String conversationId, + String cursor, + int limit, + Set userAndRoles + ) throws IOException { + if (limit <= 0) throw new BadRequestException(); + List messages = new ArrayList<>( + getConversation(noteId, conversationId, userAndRoles).getMessages() + ); + Collections.reverse(messages); + + int startInclusive = cursor == null ? 0 : IntStream.range(0, messages.size()) + .filter(i -> messages.get(i).getId().equals(cursor)) + .findFirst() + .orElseThrow(NotFoundException::new) + 1; + int pageSize = Math.min(limit, messages.size() - startInclusive); + int endExclusive = startInclusive + pageSize; + + List page = messages.subList(startInclusive, endExclusive); + + boolean hasMore = endExclusive < messages.size(); + String nextCursor = hasMore ? page.get(page.size() - 1).getId() : null; + return new MessagePage(page, nextCursor); + } + + @Override + public Future sendMessage( + String noteId, + String conversationId, + String userContent, + AuthenticationInfo authInfo, + Set userAndRoles, + AssistantEventListener sink + ) { + return assistantExecutor.submit(() -> runMessage( + noteId, conversationId, userContent, authInfo, userAndRoles, sink) + ); + } + + private void runMessage( + String noteId, + String conversationId, + String userContent, + AuthenticationInfo authInfo, + Set userAndRoles, + AssistantEventListener sink + ) { + var runId = "run_" + UUID.randomUUID().toString().replace("-", "").substring(0, 16); + boolean acquired = false; + try { + if (StringUtils.isBlank(userContent)) throw new BadRequestException(); + if (!busyConversations.add(conversationId)) { + throw new ClientErrorException("", Response.Status.CONFLICT); + } + acquired = true; + + var conversation = getConversation(noteId, conversationId, userAndRoles); + if (!conversation.isOwner(authInfo.getUser())) throw new ForbiddenException(); + + notebook.processNote(noteId, note -> { + if (note == null) throw new NoteNotFoundException(noteId); + conversationRepository.update( + conversation, c -> c.addMessage(Message.user(Message.id(), userContent)) + ); + return null; + }); + + sink.onEvent( + AssistantEventType.RUN_STARTED, + new AssistantEventPayload.RunStarted(runId, Instant.now().toString()) + ); + + Map tokens = new HashMap<>(); + tokens.put("input", 0); + tokens.put("output", 0); + + runLoop(noteId, conversation, authInfo, userAndRoles, sink, tokens); + + sink.onEvent( + AssistantEventType.RUN_COMPLETED, + new AssistantEventPayload.RunCompleted( + runId, + new AssistantEventPayload.Usage(tokens.get("input"), tokens.get("output")) + ) + ); + } catch (Exception e) { + LOGGER.error("Error during Assistant run", e); + sink.onEvent( + AssistantEventType.RUN_FAILED, + new AssistantEventPayload.RunFailed(runId, AssistantEventPayload.Error.of(e)) + ); + } finally { + if (acquired) busyConversations.remove(conversationId); + } + } + + private void runLoop( + String noteId, + Conversation conversation, + AuthenticationInfo authInfo, + Set userAndRoles, + AssistantEventListener sink, + Map tokens + ) throws IOException { + for (int iteration = 0; iteration < 10; iteration++) { + String assistantId = Message.id(); + var textBuffer = new StringBuilder(); + List toolCalls = new ArrayList<>(); + + modelClient.stream( + INSTRUCTIONS, + conversation.getRecentMessages(HISTORY_MESSAGE_LIMIT), + toolExecutor.specs(), + event -> { + if (event instanceof AssistantEvent.TextDelta) { + String delta = ((AssistantEvent.TextDelta) event).delta; + textBuffer.append(delta); + sink.onEvent( + AssistantEventType.MESSAGE_DELTA, + new AssistantEventPayload.MessageDelta(assistantId, delta) + ); + } else if (event instanceof AssistantEvent.ToolCall) { + AssistantEvent.ToolCall tc = (AssistantEvent.ToolCall) event; + toolCalls.add(new ToolCall(tc.id, tc.name, tc.arguments)); + } else if (event instanceof AssistantEvent.Usage) { + AssistantEvent.Usage u = (AssistantEvent.Usage) event; + tokens.put("input", tokens.get("input") + u.inputTokens); + tokens.put("output", tokens.get("output") + u.outputTokens); + } + } + ); + + if (!toolCalls.isEmpty()) { + Message.Assistant assistantMsg = Message.assistant(assistantId, textBuffer.toString()); + toolCalls.forEach(assistantMsg::addToolCall); + + List turn = new ArrayList<>(); + turn.add(assistantMsg); + + if (textBuffer.length() > 0) { + sink.onEvent( + AssistantEventType.MESSAGE_DONE, + new AssistantEventPayload.MessageDone(assistantId, textBuffer.toString()) + ); + } + + for (ToolCall tc : assistantMsg.getToolCalls()) { + sink.onEvent( + AssistantEventType.TOOL_CALL_STARTED, + new AssistantEventPayload.ToolCallStarted(tc.getId(), tc.getName(), tc.getArguments()) + ); + ToolResult result = toolExecutor.callTool( + noteId, tc.getName(), tc.getArguments(), authInfo, userAndRoles + ); + tc.setResult(result); + turn.add(Message.tool(Message.id(), tc.getId(), GSON.toJson(result))); + sink.onEvent( + AssistantEventType.TOOL_CALL_DONE, + new AssistantEventPayload.ToolCallDone(tc.getId(), result) + ); + } + + notebook.processNote(noteId, note -> { + if (note == null) throw new NoteNotFoundException(noteId); + conversationRepository.update(conversation, c -> turn.forEach(c::addMessage)); + return null; + }); + } else { + // No tool calls; finish the run. + String assistantText = textBuffer.toString(); + notebook.processNote(noteId, note -> { + if (note == null) throw new NoteNotFoundException(noteId); + conversationRepository.update( + conversation, + c -> c.addMessage(Message.assistant(assistantId, assistantText)) + ); + return null; + }); + sink.onEvent( + AssistantEventType.MESSAGE_DONE, + new AssistantEventPayload.MessageDone(assistantId, assistantText) + ); + return; + } + } + throw new IllegalStateException("Tool iteration limit exceeded"); + } + +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantModule.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantModule.java new file mode 100644 index 00000000000..cc0a4361021 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/AssistantModule.java @@ -0,0 +1,83 @@ +/* + * 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 com.google.common.util.concurrent.ThreadFactoryBuilder; + +import java.io.File; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; + +import org.apache.commons.lang3.StringUtils; +import org.apache.zeppelin.notebook.AuthorizationService; +import org.apache.zeppelin.notebook.Notebook; +import org.apache.zeppelin.service.NotebookService; + +public class AssistantModule implements AutoCloseable { + + public final Assistant assistant; + + public final OpenAiChatModel model; + + public final ConversationRepository repository; + + private final ExecutorService assistantExecutor; + + public AssistantModule( + boolean enabled, + String endpoint, + String apiKey, + String modelName, + String conversationDir, + Notebook notebook, + NotebookService notebookService, + AuthorizationService authorizationService + ) { + if (!enabled || StringUtils.isBlank(apiKey)) { + assistant = new DisabledAssistant(); + model = null; + repository = null; + assistantExecutor = null; + return; + } + repository = new FileConversationRepository(new File(conversationDir)); + model = new OpenAiChatModel(endpoint, apiKey, modelName); + assistantExecutor = new ThreadPoolExecutor( + 10, 10, 0L, TimeUnit.MILLISECONDS, new ArrayBlockingQueue<>(10), + new ThreadFactoryBuilder().setNameFormat("assistant-run-%d").setDaemon(true).build() + ); + assistant = new AssistantImpl( + notebook, model, notebookService, authorizationService, repository, assistantExecutor + ); + } + + @Override + public void close() { + try { + if (assistantExecutor != null) { + assistantExecutor.shutdownNow(); + } + } finally { + if (model != null) { + model.close(); + } + } + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ChatModel.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ChatModel.java new file mode 100644 index 00000000000..7c4ac2134ae --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ChatModel.java @@ -0,0 +1,31 @@ +/* + * 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 java.util.List; +import java.util.function.Consumer; + +public interface ChatModel { + + void stream( + String systemPrompt, + List history, + List tools, + Consumer consumer + ); +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Conversation.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Conversation.java new file mode 100644 index 00000000000..bb156ff7f22 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Conversation.java @@ -0,0 +1,88 @@ +/* + * 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 java.time.Instant; +import java.time.LocalDateTime; +import java.time.format.DateTimeFormatter; +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; + +public class Conversation { + + // yy-MM-dd HH:mm, e.g. "26-10-02 14:30" + private static final DateTimeFormatter TITLE_TIME = DateTimeFormatter.ofPattern("yy-MM-dd HH:mm"); + + private String id; + private String noteId; + private String ownerId; + private String title; + private String createdAt; + private String updatedAt; + private List messages; + + Conversation() { + this.messages = new ArrayList<>(); + } + + public static Conversation create(String noteId, String title, String ownerId) { + Conversation c = new Conversation(); + c.id = "conv_" + UUID.randomUUID().toString().replace("-", "").substring(0, 16); + c.noteId = noteId; + c.ownerId = ownerId; + c.title = title != null ? title : LocalDateTime.now().format(TITLE_TIME); + c.createdAt = Instant.now().toString(); + c.updatedAt = c.createdAt; + return c; + } + + public String getId() { return id; } + public String getNoteId() { return noteId; } + public String getOwnerId() { return ownerId; } + public String getTitle() { return title; } + public String getCreatedAt() { return createdAt; } + public String getUpdatedAt() { return updatedAt; } + public List getMessages() { return messages; } + + public List getRecentMessages(int limit) { + if (limit <= 0) throw new IllegalArgumentException(); + int start = Math.max(0, messages.size() - limit); + // Keep tool results together with the assistant message that requested them. + while (start > 0 && messages.get(start) instanceof Message.Tool) start--; + return messages.subList(start, messages.size()); + } + + public boolean isOwner(String user) { + return ownerId.equals(user); + } + + public void setTitle(String title) { + this.title = title; + touch(); + } + + public void addMessage(Message message) { + messages.add(message); + updatedAt = Instant.now().toString(); + } + + public void touch() { + updatedAt = Instant.now().toString(); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationJsonCodec.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationJsonCodec.java new file mode 100644 index 00000000000..c4d3776f1ac --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationJsonCodec.java @@ -0,0 +1,89 @@ +/* + * 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 com.google.gson.Gson; +import com.google.gson.GsonBuilder; +import com.google.gson.JsonDeserializationContext; +import com.google.gson.JsonDeserializer; +import com.google.gson.JsonElement; +import com.google.gson.JsonObject; +import com.google.gson.JsonParseException; +import com.google.gson.JsonSerializationContext; +import com.google.gson.JsonSerializer; + +import java.lang.reflect.Type; + +/** + * JSON wire format for an assistant conversation. + */ +final class ConversationJsonCodec { + + private static final Gson GSON = new GsonBuilder() + .registerTypeAdapter(Message.class, new MessageAdapter()) + .create(); + + private ConversationJsonCodec() { + } + + static String serialize(Conversation conversation) { + return GSON.toJson(conversation); + } + + static Conversation deserialize(String json) { + Conversation conversation = GSON.fromJson(json, Conversation.class); + if (conversation == null || conversation.getId() == null || conversation.getNoteId() == null + || conversation.getCreatedAt() == null || conversation.getMessages() == null) { + throw new JsonParseException("Invalid assistant conversation storage"); + } + return conversation; + } + + /** + * Routes deserialization to the concrete Message subtype based on {@code role}. + */ + private static class MessageAdapter + implements JsonDeserializer, JsonSerializer { + private static final Gson SUBTYPE = new Gson(); + + @Override + public Message deserialize(JsonElement json, Type type, JsonDeserializationContext ctx) + throws JsonParseException { + JsonObject obj = json.getAsJsonObject(); + if (!obj.has("role")) { + throw new JsonParseException("Message missing 'role' field"); + } + Message.Role role = Message.Role.fromValue(obj.get("role").getAsString()); + switch (role) { + case USER: + return SUBTYPE.fromJson(json, Message.User.class); + case ASSISTANT: + return SUBTYPE.fromJson(json, Message.Assistant.class); + case TOOL: + return SUBTYPE.fromJson(json, Message.Tool.class); + default: + throw new JsonParseException("Unsupported role: " + role); + } + } + + @Override + public JsonElement serialize(Message src, Type type, JsonSerializationContext ctx) { + return SUBTYPE.toJsonTree(src); + } + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationRepository.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationRepository.java new file mode 100644 index 00000000000..a88db48edb9 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ConversationRepository.java @@ -0,0 +1,35 @@ +/* + * 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 java.io.IOException; +import java.util.List; +import java.util.Optional; +import java.util.function.Consumer; + +public interface ConversationRepository { + List findAll(String noteId) throws IOException; + + Optional find(String noteId, String conversationId) throws IOException; + + void create(Conversation conversation) throws IOException; + + void update(Conversation entity, Consumer change) throws IOException; + + void delete(String noteId, String conversationId) throws IOException; +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/DisabledAssistant.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/DisabledAssistant.java new file mode 100644 index 00000000000..bcc46330e3f --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/DisabledAssistant.java @@ -0,0 +1,101 @@ +/* + * 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 java.io.IOException; +import java.util.List; +import java.util.Set; +import java.util.concurrent.Future; + +import org.apache.zeppelin.user.AuthenticationInfo; + +import jakarta.ws.rs.ServiceUnavailableException; + +public class DisabledAssistant implements Assistant { + + @Override + public List listConversations( + String noteId, + Set userAndRoles + ) throws IOException { + throw new ServiceUnavailableException(); + } + + @Override + public Conversation createConversation( + String noteId, + String title, + AuthenticationInfo authInfo, + Set userAndRoles + ) throws IOException { + throw new ServiceUnavailableException(); + } + + @Override + public Conversation getConversation( + String noteId, + String conversationId, + Set userAndRoles + ) throws IOException { + throw new ServiceUnavailableException(); + } + + @Override + public Conversation updateTitle( + String noteId, + String conversationId, + String title, + String userId, + Set userAndRoles + ) throws IOException { + throw new ServiceUnavailableException(); + } + + @Override + public void deleteConversation( + String noteId, + String conversationId, + String userId, + Set userAndRoles + ) throws IOException { + throw new ServiceUnavailableException(); + } + + @Override + public MessagePage listMessages( + String noteId, + String conversationId, + String cursor, + int limit, + Set userAndRoles + ) throws IOException { + throw new ServiceUnavailableException(); + } + + @Override + public Future sendMessage( + String noteId, + String conversationId, + String userContent, + AuthenticationInfo authInfo, + Set userAndRoles, + AssistantEventListener sink + ) { + throw new ServiceUnavailableException(); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/FileConversationRepository.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/FileConversationRepository.java new file mode 100644 index 00000000000..b9f371b33ec --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/FileConversationRepository.java @@ -0,0 +1,102 @@ +/* + * 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 java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.NoSuchFileException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.Optional; +import java.util.function.Consumer; +import org.apache.zeppelin.util.FileUtils; + +public final class FileConversationRepository implements ConversationRepository { + private final File directory; + + public FileConversationRepository(File directory) { + this.directory = directory; + } + + @Override + public List findAll(String noteId) throws IOException { + File noteDirectory = noteDirectory(noteId); + if (!noteDirectory.exists()) return Collections.emptyList(); + File[] files = noteDirectory.listFiles((dir, name) -> name.endsWith(".json")); + if (files == null) throw new IOException(); + + List conversations = new ArrayList<>(); + for (File file : files) read(file).ifPresent(conversations::add); + conversations.sort(Comparator.comparing(c -> Instant.parse(c.getCreatedAt()))); + return Collections.unmodifiableList(conversations); + } + + @Override + public Optional find(String noteId, String conversationId) throws IOException { + return read(file(noteId, conversationId)); + } + + @Override + public void create(Conversation conversation) throws IOException { + File file = file(conversation.getNoteId(), conversation.getId()); + if (file.exists()) throw new IllegalStateException(); + commit(file, conversation); + } + + @Override + public void update(Conversation entity, Consumer change) throws IOException { + File file = file(entity.getNoteId(), entity.getId()); + if (!file.exists()) throw new IOException(); + change.accept(entity); + commit(file, entity); + } + + @Override + public void delete(String noteId, String conversationId) throws IOException { + Files.delete(file(noteId, conversationId).toPath()); + } + + private File noteDirectory(String noteId) throws IOException { + File noteDirectory = new File(directory, noteId).getCanonicalFile(); + if (!directory.getCanonicalFile().equals(noteDirectory.getParentFile())) throw new IOException(); + return noteDirectory; + } + + private File file(String noteId, String conversationId) throws IOException { + File noteDirectory = noteDirectory(noteId); + File file = new File(noteDirectory, conversationId + ".json").getCanonicalFile(); + if (!noteDirectory.equals(file.getParentFile())) throw new IOException(); + return file; + } + + private static Optional read(File file) throws IOException { + try { + return Optional.of(ConversationJsonCodec.deserialize(Files.readString(file.toPath()))); + } catch (NoSuchFileException e) { + return Optional.empty(); + } + } + + private static void commit(File file, Conversation conversation) throws IOException { + FileUtils.atomicWriteToFile(ConversationJsonCodec.serialize(conversation), file); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java new file mode 100644 index 00000000000..9104ae45f3a --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ListParagraphsTool.java @@ -0,0 +1,91 @@ +/* + * 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 java.io.IOException; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.commons.lang3.StringUtils; +import org.apache.zeppelin.user.AuthenticationInfo; + +import org.apache.zeppelin.notebook.Paragraph; +import org.apache.zeppelin.service.NotebookService; +import org.apache.zeppelin.service.ServiceContext; +import org.apache.zeppelin.service.SimpleServiceCallback; + +public class ListParagraphsTool implements Tool { + + private final NotebookService notebookService; + + public ListParagraphsTool(NotebookService notebookService) { + this.notebookService = notebookService; + } + + @Override + public String name() { + return "list_paragraphs"; + } + + @Override + public String description() { + return "List notebook paragraphs."; + } + + @Override + public Map parameters() { + return Map.of("type", "object", "properties", Map.of()); + } + + @Override + public List> call( + String noteId, + Map args, + AuthenticationInfo authInfo, + Set userAndRoles + ) throws IOException { + ServiceContext context = new ServiceContext(authInfo, userAndRoles); + return notebookService.getNote(noteId, context, callback(), note -> { + List> result = new ArrayList<>(); + for (int i = 0; i < note.getParagraphCount(); i++) { + Paragraph p = note.getParagraph(i); + Map info = new HashMap<>(); + info.put("id", p.getId()); + info.put("title", p.getTitle()); + info.put("text", StringUtils.abbreviate(p.getText(), 500)); + info.put("index", i); + result.add(info); + } + return result; + }); + } + + private static SimpleServiceCallback callback() { + return new SimpleServiceCallback() { + @Override + public void onFailure(Exception ex, ServiceContext context) throws IOException { + if (ex instanceof RuntimeException) { + throw (RuntimeException) ex; + } + throw new IOException(ex); + } + }; + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Message.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Message.java new file mode 100644 index 00000000000..b0e1b3d4cd4 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Message.java @@ -0,0 +1,156 @@ +/* + * 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 com.google.gson.annotations.SerializedName; + +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.UUID; + +public abstract class Message { + + public enum Role { + @SerializedName("system") SYSTEM, + @SerializedName("user") USER, + @SerializedName("assistant") ASSISTANT, + @SerializedName("tool") TOOL; + + public static Role fromValue(String value) { + for (Role r : values()) { + if (r.name().equalsIgnoreCase(value)) { + return r; + } + } + throw new IllegalArgumentException(); + } + } + + private String id; + private String createdAt; + + private Message() {} + + Message(String id, String createdAt) { + this.id = id; + this.createdAt = createdAt; + } + + public String getId() { + return id; + } + + public String getCreatedAt() { return createdAt; } + + public abstract Role getRole(); + + public static String id() { + return "msg_" + UUID.randomUUID().toString().replace("-", "").substring(0, 16); + } + + public static Message.User user(String id, String content) { + return new Message.User(id, content, Instant.now().toString()); + } + + public static Message.Assistant assistant(String id, String content) { + return new Message.Assistant(id, content, Instant.now().toString()); + } + + public static Message.Tool tool(String id, String toolCallId, String content) { + return new Message.Tool(id, toolCallId, content, Instant.now().toString()); + } + + public static class User extends Message { + private Role role = Role.USER; + private String content; + + private User() {} + + User(String id, String content, String createdAt) { + super(id, createdAt); + this.content = content; + } + + @Override + public Role getRole() { + return role; + } + + public String getContent() { + return content; + } + } + + public static class Assistant extends Message { + private Role role = Role.ASSISTANT; + private String content; + private List toolCalls = new ArrayList<>(); + + private Assistant() {} + + Assistant(String id, String content, String createdAt) { + super(id, createdAt); + this.content = content; + } + + @Override + public Role getRole() { + return role; + } + + public String getContent() { + return content; + } + + public List getToolCalls() { + return toolCalls; + } + + public void addToolCall(ToolCall toolCall) { + toolCalls.add(toolCall); + } + } + + public static class Tool extends Message { + private Role role = Role.TOOL; + private String toolCallId; + private String content; + + private Tool() {} + + Tool(String id, String toolCallId, String content, String createdAt) { + super(id, createdAt); + this.toolCallId = toolCallId; + this.content = content; + } + + @Override + public Role getRole() { + return role; + } + + public String getToolCallId() { + return toolCallId; + } + + public String getContent() { + return content; + } + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/MessagePage.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/MessagePage.java new file mode 100644 index 00000000000..6fee457130e --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/MessagePage.java @@ -0,0 +1,30 @@ +/* + * 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 java.util.List; + +public final class MessagePage { + public final List messages; + public final String cursor; + + public MessagePage(List messages, String cursor) { + this.messages = messages; + this.cursor = cursor; + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/OpenAiChatModel.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/OpenAiChatModel.java new file mode 100644 index 00000000000..88ab7a27db4 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/OpenAiChatModel.java @@ -0,0 +1,201 @@ +/* + * 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 com.google.gson.Gson; +import com.openai.client.OpenAIClient; +import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.core.JsonValue; +import com.openai.core.http.StreamResponse; +import com.openai.models.responses.EasyInputMessage; +import com.openai.models.responses.FunctionTool; +import com.openai.models.responses.ResponseCreateParams; +import com.openai.models.responses.ResponseFunctionToolCall; +import com.openai.models.responses.ResponseInputItem; +import com.openai.models.responses.ResponseStreamEvent; + +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CancellationException; +import java.util.function.Consumer; + +/** + * {@link ChatModel} backed by the official {@code com.openai:openai-java} SDK, + */ +public class OpenAiChatModel implements ChatModel, AutoCloseable { + + private static final Gson GSON = new Gson(); + + private final String baseUrl; + private final String apiKey; + private final String model; + private OpenAIClient cachedClient; + private boolean closed; + + public OpenAiChatModel(String baseUrl, String apiKey, String model) { + this.baseUrl = baseUrl; + this.apiKey = apiKey; + this.model = model; + } + + private synchronized OpenAIClient client() { + if (closed) throw new CancellationException(); + if (cachedClient == null) { + cachedClient = OpenAIOkHttpClient.builder() + .baseUrl(baseUrl) + .apiKey(apiKey) + .timeout(Duration.ofMinutes(1)) + .build(); + } + return cachedClient; + } + + @Override + public synchronized void close() { + if (closed) return; + closed = true; + if (cachedClient != null) { + cachedClient.close(); + } + } + + @Override + public void stream( + String instruction, + List history, + List tools, + Consumer consumer + ) { + List input = new ArrayList<>(); + for (Message m : history) { + addHistoryItems(input, m); + } + + ResponseCreateParams.Builder params = ResponseCreateParams.builder() + .model(model) + .instructions(instruction) + .store(false) + .inputOfResponse(input); + + for (ToolSpec t : tools) { + params.addTool(FunctionTool.builder() + .name(t.name) + .description(t.description) + .parameters(toParameters(t.parameters)) + .strict(false) + .build()); + } + + boolean[] completed = {false}; + try ( + StreamResponse stream = client() + .responses() + .createStreaming(params.build()) + ) { + stream.stream().forEach(event -> { + if (event.completed().isPresent()) completed[0] = true; + handleEvent(event, consumer); + }); + } + // A stream that ends without a completed response + if (!completed[0]) { + throw new IllegalStateException(); + } + } + + private static FunctionTool.Parameters toParameters(Map schema) { + FunctionTool.Parameters.Builder builder = FunctionTool.Parameters.builder(); + for (Map.Entry e : schema.entrySet()) { + builder.putAdditionalProperty(e.getKey(), JsonValue.from(e.getValue())); + } + return builder.build(); + } + + // TODO(seung-00) replace with pattern matching (Java 16+) + private void addHistoryItems(List input, Message m) { + if (m instanceof Message.User) { + String content = ((Message.User) m).getContent(); + if (content != null) { + input.add(easyMessage(EasyInputMessage.Role.USER, content)); + } + } else if (m instanceof Message.Assistant) { + Message.Assistant assistant = (Message.Assistant) m; + if (assistant.getContent() != null && !assistant.getContent().isEmpty()) { + input.add(easyMessage(EasyInputMessage.Role.ASSISTANT, assistant.getContent())); + } + for (ToolCall tc : assistant.getToolCalls()) { + input.add(ResponseInputItem.ofFunctionCall(ResponseFunctionToolCall.builder() + .callId(tc.getId()) + .name(tc.getName()) + .arguments(GSON.toJson(tc.getArguments())) + .build())); + } + } else if (m instanceof Message.Tool) { + Message.Tool tool = (Message.Tool) m; + input.add(ResponseInputItem.ofFunctionCallOutput( + ResponseInputItem.FunctionCallOutput.builder() + .callId(tool.getToolCallId()) + .output(tool.getContent() != null ? tool.getContent() : "") + .build() + ) + ); + } + } + + private static ResponseInputItem easyMessage(EasyInputMessage.Role role, String content) { + return ResponseInputItem.ofEasyInputMessage( + EasyInputMessage.builder().role(role).content(content).build()); + } + + private void handleEvent( + ResponseStreamEvent event, + Consumer consumer + ) { + // A failure event ends the stream without a completion; surface it instead of silently + // saving an empty/partial reply and reporting success. + event.error().ifPresent(e -> { + throw new IllegalStateException("OpenAI stream error: " + e.message()); + }); + event.failed().ifPresent(f -> { + throw new IllegalStateException("OpenAI response failed"); + }); + event.incomplete().ifPresent(i -> { + throw new IllegalStateException("OpenAI response incomplete"); + }); + + event.outputTextDelta().ifPresent( + d -> consumer.accept(new AssistantEvent.TextDelta(d.delta()))); + + // A refusal is a real reply the user must see; stream it as text so it is delivered and saved. + event.refusalDelta().ifPresent( + r -> consumer.accept(new AssistantEvent.TextDelta(r.delta()))); + + event.outputItemDone().flatMap(done -> done.item().functionCall()) + .ifPresent(call -> + consumer.accept( + new AssistantEvent.ToolCall(call.callId(), call.name(), call.arguments()) + )); + + event.completed().flatMap(done -> done.response().usage()) + .ifPresent(u -> consumer.accept( + new AssistantEvent.Usage((int) u.inputTokens(), (int) u.outputTokens()) + )); + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Tool.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Tool.java new file mode 100644 index 00000000000..cb390eb4163 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/Tool.java @@ -0,0 +1,40 @@ +/* + * 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 java.io.IOException; +import java.util.Map; +import java.util.Set; +import org.apache.zeppelin.user.AuthenticationInfo; + +public interface Tool { + + String name(); + + String description(); + + /** JSON Schema of the tool's arguments. */ + Map parameters(); + + Object call( + String noteId, + Map args, + AuthenticationInfo authInfo, + Set userAndRoles + ) throws IOException; +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolCall.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolCall.java new file mode 100644 index 00000000000..dd5c404ac46 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolCall.java @@ -0,0 +1,43 @@ +/* + * 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 java.util.Map; + +public class ToolCall { + + private String id; + private String name; + private Map arguments; + private ToolResult result; + + private ToolCall() {} + + public ToolCall(String id, String name, Map arguments) { + this.id = id; + this.name = name; + this.arguments = arguments; + } + + public String getId() { return id; } + public String getName() { return name; } + public Map getArguments() { return arguments; } + public ToolResult getResult() { return result; } + + public void setResult(ToolResult result) { this.result = result; } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolExecutor.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolExecutor.java new file mode 100644 index 00000000000..c77db4977fb --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolExecutor.java @@ -0,0 +1,61 @@ +/* + * 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 java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.zeppelin.user.AuthenticationInfo; +public class ToolExecutor { + + private final Map byName; + + public ToolExecutor(List tools) { + Map map = new LinkedHashMap<>(); + for (Tool t : tools) { + map.put(t.name(), t); + } + this.byName = map; + } + + public List specs() { + return byName.values().stream() + .map(t -> new ToolSpec(t.name(), t.description(), t.parameters())) + .collect(Collectors.toList()); + } + + public ToolResult callTool( + String noteId, + String name, + Map args, + AuthenticationInfo authInfo, + Set userAndRoles + ) { + Tool tool = byName.get(name); + + if (tool == null) return new ToolResult(null, "Unknown tool: " + name); + + try { + return new ToolResult(tool.call(noteId, args, authInfo, userAndRoles), null); + } catch (Exception ex) { + return new ToolResult(null, String.valueOf(ex.getMessage())); + } + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolResult.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolResult.java new file mode 100644 index 00000000000..7e76c93dcb2 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolResult.java @@ -0,0 +1,28 @@ +/* + * 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; + +public final class ToolResult { + public final Object value; + public final String error; + + public ToolResult(Object value, String error) { + this.value = value; + this.error = error; + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolSpec.java b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolSpec.java new file mode 100644 index 00000000000..f5f0a419e00 --- /dev/null +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/service/assistant/ToolSpec.java @@ -0,0 +1,32 @@ +/* + * 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 java.util.Map; + +public final class ToolSpec { + public final String name; + public final String description; + public final Map parameters; + + public ToolSpec(String name, String description, Map parameters) { + this.name = name; + this.description = description; + this.parameters = parameters; + } +} diff --git a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java index ba5bffc76e5..75eb183e9b7 100644 --- a/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java +++ b/zeppelin-server/src/main/java/org/apache/zeppelin/socket/NotebookServer.java @@ -91,6 +91,8 @@ import org.apache.zeppelin.scheduler.Job.Status; import org.apache.zeppelin.service.JobManagerService; import org.apache.zeppelin.service.NotebookService; +import org.apache.zeppelin.service.assistant.AssistantEventListener; +import org.apache.zeppelin.service.assistant.Assistant; import org.apache.zeppelin.service.ServiceContext; import org.apache.zeppelin.service.SimpleServiceCallback; import org.apache.zeppelin.service.exception.JobManagerForbiddenException; @@ -162,6 +164,7 @@ String getKey() { private Provider notebookServiceProvider; private AuthorizationService authorizationService; private Provider jobManagerServiceProvider; + private Provider assistantProvider; public NotebookServer() { NotebookServer.self.set(this); @@ -220,6 +223,16 @@ public void setJobManagerService( this.jobManagerServiceProvider = jobManagerServiceProvider; } + @Inject + public void setAssistant( + Provider assistantProvider) { + this.assistantProvider = assistantProvider; + } + + public Assistant getAssistant() { + return assistantProvider.get(); + } + public Notebook getNotebook() { return notebookProvider.get(); } @@ -568,6 +581,9 @@ public void onMessage(NotebookSocket conn, String msg) { case PATCH_PARAGRAPH: patchParagraph(conn, context, receivedMessage); break; + case ASSISTANT_SEND_MESSAGE: + sendAssistantMessage(conn, context, receivedMessage); + break; default: break; } @@ -1245,6 +1261,34 @@ public void onSuccess(String result, ServiceContext context) throws IOException }); } + private void sendAssistantMessage( + NotebookSocket conn, + ServiceContext context, + Message fromMessage + ) { + String noteId = (String) fromMessage.get("noteId"); + String conversationId = (String) fromMessage.get("conversationId"); + String content = (String) fromMessage.get("content"); + AssistantEventListener sink = (type, payload) -> { + try { + conn.send(serializeMessage(new Message(OP.ASSISTANT_EVENT) + .put("conversationId", conversationId) + .put("type", type.wireName) + .put("payload", payload))); + } catch (IOException e) { + LOGGER.warn("Failed to send assistant event to connection", e); + } + }; + getAssistant().sendMessage( + noteId, + conversationId, + content, + context.getAutheInfo(), + context.getUserAndRoles(), + sink + ); + } + private void cloneNote(NotebookSocket conn, ServiceContext context, Message fromMessage) throws IOException { diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantEventTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantEventTest.java new file mode 100644 index 00000000000..3d8a25c40b6 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantEventTest.java @@ -0,0 +1,45 @@ +/* + * 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.*; + +import com.google.gson.JsonParseException; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class AssistantEventTest { + @Test + void parsesArguments() { + var call = new AssistantEvent.ToolCall( + "call", + "tool", + "{\"name\":\"value\",\"options\":{\"enabled\":true}}" + ); + assertEquals(Map.of("name", "value", "options", Map.of("enabled", true)), call.arguments); + + for (String arguments : new String[]{null, "", " \t", "{}"}) { + assertTrue(new AssistantEvent.ToolCall("call", "tool", arguments).arguments.isEmpty()); + } + + assertThrows( + JsonParseException.class, + () -> new AssistantEvent.ToolCall("call", "tool", "{") + ); + } +} 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 new file mode 100644 index 00000000000..5befd6019f0 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/AssistantTest.java @@ -0,0 +1,456 @@ +/* + * 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.*; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; + +import java.io.File; +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import jakarta.ws.rs.BadRequestException; +import jakarta.ws.rs.ClientErrorException; +import jakarta.ws.rs.ForbiddenException; +import jakarta.ws.rs.ServiceUnavailableException; +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.rest.exception.NoteNotFoundException; +import org.apache.zeppelin.service.NotebookService; +import org.apache.zeppelin.user.AuthenticationInfo; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class AssistantTest { + + private final ExecutorService executor = Executors.newFixedThreadPool(2); + + @AfterEach + void tearDown() { + executor.shutdownNow(); + } + + private final AuthenticationInfo authInfo = new AuthenticationInfo("holden"); + private final AuthenticationInfo otherAuthInfo = new AuthenticationInfo("phoebe"); + private final Set userAndRoles = Set.of("user"); + + @Test + void rejectsUnavailable() { + var sut = new DisabledAssistant(); + + assertThrows(ServiceUnavailableException.class, + () -> sut.listConversations("noteId", userAndRoles)); + assertThrows(ServiceUnavailableException.class, + () -> sut.createConversation("noteId", "title", authInfo, userAndRoles)); + assertThrows(ServiceUnavailableException.class, + () -> sut.getConversation("noteId", "conv", userAndRoles)); + assertThrows(ServiceUnavailableException.class, + () -> sut.updateTitle("noteId", "conv", "title", authInfo.getUser(), userAndRoles)); + assertThrows(ServiceUnavailableException.class, + () -> sut.deleteConversation("noteId", "conv", authInfo.getUser(), userAndRoles)); + assertThrows(ServiceUnavailableException.class, + () -> sut.listMessages("noteId", "conv", null, 10, userAndRoles)); + assertThrows(ServiceUnavailableException.class, + () -> sut.sendMessage("noteId", "conv", "hi", authInfo, userAndRoles, + (type, payload) -> fail("Disabled service must not emit events"))); + } + + @Test + void rejectsNonReaders() { + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(false); + + var sut = new AssistantImpl( + null, mock(ChatModel.class), null, authorization, null, executor + ); + + assertThrows(ForbiddenException.class, () -> sut.listConversations("noteId", userAndRoles)); + assertThrows(ForbiddenException.class, + () -> sut.createConversation("noteId", "title", authInfo, userAndRoles) + ); + assertThrows(ForbiddenException.class, + () -> sut.deleteConversation("noteId", "conv", authInfo.getUser(), userAndRoles) + ); + } + + @Test + void rejectsAnonymousCreate() { + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + var sut = new AssistantImpl( + null, mock(ChatModel.class), null, authorization, null, executor); + + assertThrows( + ForbiddenException.class, + () -> sut.createConversation("noteId", "title", AuthenticationInfo.ANONYMOUS, userAndRoles) + ); + } + + @Test + void createsConversation(@TempDir File tempDir) throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())) + .thenAnswer(invocation -> ((NoteProcessor) invocation.getArgument(1)).process(new Note())); + + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + + var repository = new FileConversationRepository(tempDir); + var sut = new AssistantImpl( + notebook, mock(ChatModel.class), null, authorization, repository, executor + ); + + var created = sut.createConversation("noteId", "test", authInfo, userAndRoles); + assertTrue(new File(tempDir, "noteId" + "/" + created.getId() + ".json").isFile()); + } + + @Test + void rejectsMissingNote() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())) + .thenAnswer(invocation -> ((NoteProcessor) invocation.getArgument(1)).process(null)); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + var repository = mock(ConversationRepository.class); + var sut = new AssistantImpl( + notebook, mock(ChatModel.class), null, authorization, repository, executor + ); + + assertThrows( + NoteNotFoundException.class, + () -> sut.createConversation("noteId", "test", authInfo, userAndRoles) + ); + verifyNoInteractions(repository); + } + + @Test + void releasesSlotAfterFailure() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())) + .thenAnswer(invocation -> ((NoteProcessor) invocation.getArgument(1)).process(new Note())); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + + var conversation = Conversation.create("noteId", "test", authInfo.getUser()); + var repository = mock(ConversationRepository.class); + when(repository.find("noteId", conversation.getId())) + .thenReturn(Optional.of(conversation)); + var sut = new AssistantImpl( + notebook, mock(ChatModel.class), null, authorization, repository, executor + ); + + doThrow(new IOException()).when(repository).update(any(), any()); + assertThrows( + IOException.class, + () -> sut.updateTitle("noteId", conversation.getId(), "discarded", authInfo.getUser(), userAndRoles) + ); + + doAnswer(invocation -> { + invocation.>getArgument(1).accept(invocation.getArgument(0)); + return null; + }).when(repository).update(any(), any()); + var updated = sut.updateTitle( + "noteId", conversation.getId(), "after-failure", authInfo.getUser(), userAndRoles + ); + assertEquals("after-failure", updated.getTitle()); + } + + @Test + void enforcesOwnership() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())) + .thenAnswer(invocation -> ((NoteProcessor) invocation.getArgument(1)) + .process(new Note())); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + + var conversation = Conversation.create("noteId", "test", authInfo.getUser()); + var repository = mock(ConversationRepository.class); + when(repository.find("noteId", conversation.getId())).thenReturn(Optional.of(conversation)); + when(repository.findAll("noteId")).thenReturn(List.of(conversation)); + var sut = new AssistantImpl( + notebook, mock(ChatModel.class), null, authorization, repository, executor + ); + + // A different reader can read. + assertEquals(1, sut.listConversations("noteId", userAndRoles).size()); + assertNotNull(sut.getConversation("noteId", conversation.getId(), userAndRoles)); + assertNotNull(sut.listMessages("noteId", conversation.getId(), null, 10, userAndRoles)); + + // But cannot mutate. + assertThrows(ForbiddenException.class, + () -> sut.deleteConversation("noteId", conversation.getId(), otherAuthInfo.getUser(), userAndRoles) + ); + assertThrows(ForbiddenException.class, + () -> sut.updateTitle("noteId", conversation.getId(), "x", otherAuthInfo.getUser(), userAndRoles)); + var events = new ArrayList(); + sut.sendMessage("noteId", conversation.getId(), "hi", otherAuthInfo, userAndRoles, + (type, payload) -> events.add(payload) + ).get(5, TimeUnit.SECONDS); + assertEquals(1, events.size()); + assertEquals(403, ((AssistantEventPayload.RunFailed) events.get(0)).error.status); + } + + @Test + void paginatesMessages() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())) + .thenAnswer(invocation -> ((NoteProcessor) invocation.getArgument(1)) + .process(new Note())); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + var conversation = Conversation.create("noteId", "test", authInfo.getUser()); + var repository = mock(ConversationRepository.class); + when(repository.find("noteId", conversation.getId())).thenReturn(Optional.of(conversation)); + var sut = new AssistantImpl( + notebook, mock(ChatModel.class), null, authorization, repository, executor + ); + + var messages = List.of( + Message.user("m1", "first"), + Message.user("m2", "second"), + Message.user("m3", "third")); + messages.forEach(conversation::addMessage); + + var first = sut.listMessages("noteId", conversation.getId(), null, 2, userAndRoles); + assertEquals(List.of(messages.get(2), messages.get(1)), first.messages); + assertEquals("m2", first.cursor); + + var last = sut.listMessages("noteId", conversation.getId(), first.cursor, 2, userAndRoles); + assertEquals(List.of(messages.get(0)), last.messages); + assertNull(last.cursor); + } + + @Test + void rejectsInvalidLimit() { + var sut = new AssistantImpl( + null, mock(ChatModel.class), null, null, null, executor); + assertThrows( + BadRequestException.class, + () -> sut.listMessages("noteId", "conv", null, 0, userAndRoles) + ); + assertThrows( + BadRequestException.class, + () -> sut.listMessages("noteId", "conv", null, -1, userAndRoles) + ); + } + + @Test + void limitsModelHistory() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())).thenAnswer(invocation -> + ((NoteProcessor) invocation.getArgument(1)).process(new Note())); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + var chatModel = mock(ChatModel.class); + var repository = mock(ConversationRepository.class); + var conversation = Conversation.create("noteId", "test", authInfo.getUser()); + var sut = new AssistantImpl( + notebook, chatModel, null, authorization, repository, executor + ); + var limit = 10; + + for (int i = 0; i < limit + 2; i++) conversation.addMessage(Message.user("msg_" + i, "m" + i)); + 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()); + doAnswer(invocation -> { + List history = invocation.getArgument(1); + assertEquals(limit, history.size()); + assertEquals("m3", ((Message.User) history.get(0)).getContent()); + assertEquals("hi", ((Message.User) history.get(9)).getContent()); + return null; + }).when(chatModel).stream(any(), any(), any(), any()); + + sut.sendMessage("noteId", conversation.getId(), "hi", authInfo, userAndRoles, + (type, payload) -> { } + ).get(5, TimeUnit.SECONDS); + } + + @Test + void streamsReply() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())).thenAnswer(invocation -> + ((NoteProcessor) invocation.getArgument(1)).process(new Note())); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + var repository = mock(ConversationRepository.class); + var conversation = Conversation.create("noteId", "test", authInfo.getUser()); + 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 chatModel = mock(ChatModel.class); + var sut = new AssistantImpl( + notebook, chatModel, null, authorization, repository, executor + ); + + doAnswer(invocation -> { + Consumer consumer = invocation.getArgument(3); + consumer.accept(new AssistantEvent.TextDelta("Hel")); + consumer.accept(new AssistantEvent.TextDelta("lo")); + return null; + }).when(chatModel).stream(any(), any(), any(), any()); + + var deltas = new ArrayList(); + var dones = new ArrayList(); + sut.sendMessage("noteId", conversation.getId(), "hi", authInfo, userAndRoles, + (type, payload) -> { + if (type == AssistantEventType.MESSAGE_DELTA) { + deltas.add(((AssistantEventPayload.MessageDelta) payload).delta); + } else if (type == AssistantEventType.MESSAGE_DONE) { + dones.add(((AssistantEventPayload.MessageDone) payload).content); + } + }).get(5, TimeUnit.SECONDS); + + assertEquals(List.of("Hel", "lo"), deltas); + assertEquals(List.of("Hello"), dones); + assertEquals("Hello", ((Message.Assistant) conversation.getMessages().get(1)).getContent()); + } + + @Test + void completesToolTurn() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())) + .thenAnswer(invocation -> ((NoteProcessor) invocation.getArgument(1)) + .process(new Note())); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + var repository = mock(ConversationRepository.class); + + var conversation = Conversation.create("noteId", "test", authInfo.getUser()); + 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 chatModel = mock(ChatModel.class); + var notebookService = mock(NotebookService.class); + var sut = new AssistantImpl( + notebook, chatModel, notebookService, authorization, repository, executor + ); + + doAnswer(invocation -> { + Consumer consumer = invocation.getArgument(3); + consumer.accept(new AssistantEvent.ToolCall("call_1", "list_paragraphs", "{}")); + return null; + }).doAnswer(invocation -> { + List history = invocation.getArgument(1); + assertEquals("call_1", ((Message.Tool) history.get(2)).getToolCallId()); + Consumer consumer = invocation.getArgument(3); + consumer.accept(new AssistantEvent.TextDelta("final")); + return null; + }).when(chatModel).stream(any(), any(), any(), any()); + + sut.sendMessage("noteId", conversation.getId(), "hi", authInfo, userAndRoles, + (type, payload) -> { } + ).get(5, TimeUnit.SECONDS); + + var stored = conversation.getMessages(); + assertEquals(4, stored.size()); + + Message.Assistant toolTurn = (Message.Assistant) stored.get(1); + assertEquals("list_paragraphs", toolTurn.getToolCalls().get(0).getName()); + assertEquals("call_1", ((Message.Tool) stored.get(2)).getToolCallId()); + assertEquals("final", ((Message.Assistant) stored.get(3)).getContent()); + } + + @Test + void guardsRunningConversation() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), any())).thenAnswer(invocation -> + ((NoteProcessor) invocation.getArgument(1)).process(new Note())); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(true); + var repository = mock(ConversationRepository.class); + var conversation = Conversation.create("noteId", "test", authInfo.getUser()); + 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 chatModel = mock(ChatModel.class); + var sut = new AssistantImpl( + notebook, chatModel, null, authorization, repository, executor); + + var streaming = new CountDownLatch(1); + var finish = new CountDownLatch(1); + + doAnswer(invocation -> { + streaming.countDown(); // signal streaming has started + assertTrue(finish.await(5, TimeUnit.SECONDS)); // wait for finish + return null; + }).when(chatModel).stream(any(), any(), any(), any()); + + try { + var running = sut.sendMessage( + "noteId", conversation.getId(), "hi", authInfo, userAndRoles, + (type, payload) -> { } + ); + assertTrue(streaming.await(5, TimeUnit.SECONDS)); + + // other conversations can be mutated. + var second = Conversation.create("noteId", "independent", authInfo.getUser()); + when(repository.find("noteId", second.getId())).thenReturn(Optional.of(second)); + assertEquals("updated", + sut.updateTitle("noteId", second.getId(), "updated", authInfo.getUser(), userAndRoles) + .getTitle() + ); + + // running conversation cannot be mutated. + assertThrows(ClientErrorException.class, + () -> sut.updateTitle("noteId", conversation.getId(), "renamed", authInfo.getUser(), + userAndRoles) + ); + assertThrows(ClientErrorException.class, + () -> sut.deleteConversation("noteId", conversation.getId(), authInfo.getUser(), + userAndRoles) + ); + + var events = new ArrayList(); + sut.sendMessage("noteId", conversation.getId(), "duplicated", authInfo, userAndRoles, + (type, payload) -> events.add(payload) + ).get(5, TimeUnit.SECONDS); + assertEquals(409, ((AssistantEventPayload.RunFailed) events.get(0)).error.status); + + finish.countDown(); + running.get(5, TimeUnit.SECONDS); // wait for message processing to complete. + + assertEquals("released", + sut.updateTitle("noteId", conversation.getId(), "released", authInfo.getUser(), userAndRoles) + .getTitle() + ); + } finally { + finish.countDown(); + } + } +} diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ConversationTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ConversationTest.java new file mode 100644 index 00000000000..5c274899707 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ConversationTest.java @@ -0,0 +1,49 @@ +/* + * 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.*; + +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class ConversationTest { + + @Test + void validatesHistoryWindow() { + var conversation = Conversation.create("noteId", "test", "holden"); + assertTrue(conversation.getRecentMessages(10).isEmpty()); + assertThrows(IllegalArgumentException.class, () -> conversation.getRecentMessages(0)); + assertThrows(IllegalArgumentException.class, () -> conversation.getRecentMessages(-1)); + } + + @Test + void keepsToolCallsWithResults() { + var conversation = Conversation.create("noteId", "test", "holden"); + conversation.addMessage(Message.user("old", "old")); + var assistant = Message.assistant("assistant", ""); + assistant.addToolCall(new ToolCall("call_1", "list_paragraphs", Map.of())); + conversation.addMessage(assistant); + var result = Message.tool("tool", "call_1", "[]"); + conversation.addMessage(result); + + assertEquals(List.of(assistant, result), conversation.getRecentMessages(1)); + assertEquals(3, conversation.getMessages().size()); + } +} diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/FileConversationRepositoryTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/FileConversationRepositoryTest.java new file mode 100644 index 00000000000..efeaa94b323 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/FileConversationRepositoryTest.java @@ -0,0 +1,114 @@ +/* + * 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.*; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Files; +import java.util.Set; +import java.util.stream.Collectors; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class FileConversationRepositoryTest { + @Test + void storesSeparateFiles(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + + var first = Conversation.create("noteId", "first", "holden"); + sut.create(first); + var second = Conversation.create("noteId", "second", "holden"); + sut.create(second); + + var loaded = sut.findAll("noteId"); + + assertEquals( + Set.of(first.getId(), second.getId()), loaded.stream().map(Conversation::getId) + .collect(Collectors.toSet()) + ); + assertTrue(new File(directory, "noteId/" + first.getId() + ".json").isFile()); + assertTrue(new File(directory, "noteId/" + second.getId() + ".json").isFile()); + } + + @Test + void propagatesReadFailure(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + Files.createDirectories(new File(directory, "noteId/unreadable.json").toPath()); + assertThrows(IOException.class, () -> sut.findAll("noteId")); + } + + @Test + void rejectsDuplicate(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + var conversation = Conversation.create("noteId", "title", "holden"); + sut.create(conversation); + + conversation.setTitle("discarded"); + assertThrows(IllegalStateException.class, () -> sut.create(conversation)); + assertEquals("title", sut.find("noteId", conversation.getId()).orElseThrow().getTitle()); + } + + @Test + void persistsUpdate(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + var conversation = Conversation.create("noteId", "title", "holden"); + sut.create(conversation); + + sut.update(conversation, c -> { + c.setTitle("renamed"); + c.addMessage(Message.user("msg_1", "hello")); + }); + var stored = sut.find("noteId", conversation.getId()).orElseThrow(); + + assertEquals("renamed", stored.getTitle()); + assertEquals(1, stored.getMessages().size()); + assertEquals("hello", ((Message.User) stored.getMessages().get(0)).getContent()); + } + + @Test + void preservesFailedUpdate(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + var conversation = Conversation.create("noteId", "title", "holden"); + sut.create(conversation); + + assertThrows(IllegalStateException.class, () -> sut.update(conversation, c -> { + c.setTitle("discarded"); + throw new IllegalStateException(); + })); + assertEquals("title", sut.find("noteId", conversation.getId()).orElseThrow().getTitle()); + } + + @Test + void rejectsDeletedConversation(@TempDir File directory) throws IOException { + var sut = new FileConversationRepository(directory); + var conversation = Conversation.create("noteId", "title", "holden"); + sut.create(conversation); + + sut.delete("noteId", conversation.getId()); + + assertTrue(sut.find("noteId", conversation.getId()).isEmpty()); + assertTrue(sut.findAll("noteId").isEmpty()); + + assertThrows(IOException.class, () -> sut.delete("noteId", conversation.getId())); + assertThrows(IOException.class, + () -> sut.update(conversation, c -> fail("update callback shouldn't run")) + ); + } +} diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java new file mode 100644 index 00000000000..4cdab27b636 --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ListParagraphsToolTest.java @@ -0,0 +1,87 @@ +/* + * 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.*; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; + +import java.util.Map; +import java.util.Set; +import org.apache.zeppelin.interpreter.InterpreterFactory; +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.rest.exception.ForbiddenException; +import org.apache.zeppelin.service.NotebookService; +import org.apache.zeppelin.user.AuthenticationInfo; +import org.junit.jupiter.api.Test; + +class ListParagraphsToolTest { + private final AuthenticationInfo authInfo = AuthenticationInfo.ANONYMOUS; + private final Set userAndRoles = Set.of("user"); + + @Test + void listsParagraphs() throws Exception { + var note = new Note(); + note.setInterpreterFactory(mock(InterpreterFactory.class)); + var notebook = mock(Notebook.class); + 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 notebookService = new NotebookService(notebook, authorization, null, null); + var sut = new ListParagraphsTool(notebookService); + + var shortParagraph = note.addNewParagraph(AuthenticationInfo.ANONYMOUS); + shortParagraph.setText("%md hello"); + var longParagraph = note.addNewParagraph(AuthenticationInfo.ANONYMOUS); + var longText = "%python\n" + "x".repeat(500); + longParagraph.setText(longText); + + var paragraphs = sut.call("noteId", Map.of(), authInfo, userAndRoles); + assertEquals(2, paragraphs.size()); + assertEquals(shortParagraph.getId(), paragraphs.get(0).get("id")); + assertEquals(longParagraph.getId(), paragraphs.get(1).get("id")); + assertEquals("%md hello", paragraphs.get(0).get("text")); + assertEquals(longText.substring(0, 497) + "...", paragraphs.get(1).get("text")); + assertFalse(paragraphs.get(1).containsKey("interpreter")); + assertEquals(longText, longParagraph.getText()); + assertEquals(1, paragraphs.get(1).get("index")); + } + + @Test + void rejectsNonReaders() throws Exception { + var notebook = mock(Notebook.class); + when(notebook.processNote(eq("noteId"), eq(false), any())) + .thenAnswer(invocation -> ((NoteProcessor) invocation.getArgument(2)) + .process(new Note())); + var authorization = mock(AuthorizationService.class); + when(authorization.isReader("noteId", userAndRoles)).thenReturn(false); + + var notebookService = new NotebookService(notebook, authorization, null, null); + + var sut = new ListParagraphsTool(notebookService); + + assertThrows( + ForbiddenException.class, + () -> sut.call("noteId", Map.of(), authInfo, userAndRoles) + ); + } +} diff --git a/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ToolExecutorTest.java b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ToolExecutorTest.java new file mode 100644 index 00000000000..5f4ce1691bb --- /dev/null +++ b/zeppelin-server/src/test/java/org/apache/zeppelin/service/assistant/ToolExecutorTest.java @@ -0,0 +1,97 @@ +/* + * 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.*; + +import java.io.IOException; +import java.util.List; +import java.util.Map; +import java.util.Set; +import org.apache.zeppelin.user.AuthenticationInfo; +import org.junit.jupiter.api.Test; + +class ToolExecutorTest { + private final AuthenticationInfo authInfo = AuthenticationInfo.ANONYMOUS; + private final Set userAndRoles = Set.of("user"); + + @Test + void exposesTools() { + var tool = new EchoTool(); + var sut = new ToolExecutor(List.of(tool)); + var specs = sut.specs(); + assertEquals(1, specs.size()); + assertEquals(tool.name(), specs.get(0).name); + assertEquals(tool.description(), specs.get(0).description); + assertEquals(tool.parameters(), specs.get(0).parameters); + } + + @Test + void callsTool() { + var sut = new ToolExecutor(List.of(new EchoTool())); + var args = Map.of("text", "hello"); + var result = sut.callTool("noteId", "echo", args, authInfo, userAndRoles); + assertNull(result.error); + assertEquals(args, result.value); + } + + @Test + void rejectsUnknownTool() { + var sut = new ToolExecutor(List.of(new EchoTool())); + var result = sut.callTool("noteId", "unknown_tool", Map.of(), authInfo, userAndRoles); + assertNotNull(result.error); + } + + @Test + void returnsToolError() { + var tool = new EchoTool() { + @Override + public Object call(String noteId, Map args, AuthenticationInfo authInfo, + Set userAndRoles) throws IOException { + throw new IOException("failed"); + } + }; + var sut = new ToolExecutor(List.of(tool)); + var result = sut.callTool("noteId", "echo", Map.of(), authInfo, userAndRoles); + assertEquals("failed", result.error); + assertNull(result.value); + } + + private static class EchoTool implements Tool { + @Override + public String name() { + return "echo"; + } + + @Override + public String description() { + return "Return the arguments."; + } + + @Override + public Map parameters() { + return Map.of("type", "object"); + } + + @Override + public Object call(String noteId, Map args, AuthenticationInfo authInfo, + Set userAndRoles) throws IOException { + return args; + } + } +} diff --git a/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-assistant.interface.ts b/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-assistant.interface.ts new file mode 100644 index 00000000000..e3d08061301 --- /dev/null +++ b/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-assistant.interface.ts @@ -0,0 +1,32 @@ +/* + * Licensed 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. + */ + +export interface AssistantSendMessage { + noteId: string; + conversationId: string; + content: string; +} + +export type AssistantEventType = + | 'run.started' + | 'run.completed' + | 'run.failed' + | 'message.delta' + | 'message.done' + | 'tool_call.started' + | 'tool_call.done'; + +export interface AssistantEvent { + conversationId: string; + type: AssistantEventType; + payload: unknown; +} diff --git a/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-data-type-map.interface.ts b/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-data-type-map.interface.ts index 6c023a09231..9fcfe7205c9 100644 --- a/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-data-type-map.interface.ts +++ b/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-data-type-map.interface.ts @@ -11,6 +11,7 @@ */ import { AuthInfo, ErrorInfo } from './message-common.interface'; +import { AssistantEvent, AssistantSendMessage } from './message-assistant.interface'; import { CheckpointNote, CloneNote, @@ -91,6 +92,7 @@ import { OP } from './message-operator.interface'; export type MessageDataTypeMap = MessageSendDataTypeMap | MessageReceiveDataTypeMap; export interface MessageReceiveDataTypeMap { + [OP.ASSISTANT_EVENT]: AssistantEvent; [OP.COMPLETION_LIST]: CompletionReceived; [OP.NOTES_INFO]: NotesInfo; [OP.NOTE]: Note; @@ -127,6 +129,7 @@ export interface MessageReceiveDataTypeMap { } export interface MessageSendDataTypeMap { + [OP.ASSISTANT_SEND_MESSAGE]: AssistantSendMessage; [OP.PING]: undefined; [OP.LIST_NOTES]: undefined; [OP.GET_HOME_NOTE]: undefined; diff --git a/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-operator.interface.ts b/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-operator.interface.ts index e1812877dae..a87527623ab 100644 --- a/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-operator.interface.ts +++ b/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/message-operator.interface.ts @@ -511,6 +511,18 @@ export enum OP { */ NOTE_RUNNING_STATUS = 'NOTE_RUNNING_STATUS', + /** + * [c-s] + * send a message to the notebook assistant + */ + ASSISTANT_SEND_MESSAGE = 'ASSISTANT_SEND_MESSAGE', + + /** + * [s-c] + * assistant run stream event + */ + ASSISTANT_EVENT = 'ASSISTANT_EVENT', + /** * [s-c] * Notice diff --git a/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/public-api.ts b/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/public-api.ts index fca7b31a317..230906eba2d 100644 --- a/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/public-api.ts +++ b/zeppelin-web-angular/projects/zeppelin-sdk/src/interfaces/public-api.ts @@ -10,6 +10,7 @@ * limitations under the License. */ +export * from './message-assistant.interface'; export * from './message-common.interface'; export * from './message-data-type-map.interface'; export * from './message-interpreter.interface';