diff --git a/docs/configuration/settings.md b/docs/configuration/settings.md index a3c0ca0df98..6d46eeff29b 100644 --- a/docs/configuration/settings.md +++ b/docs/configuration/settings.md @@ -122,6 +122,11 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co | Key | Default | Meaning | Type | Since | |----------------------------------------------------------|---------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------|-------| +| kyuubi.engine.chat.ernie.http.connect.timeout | PT2M | The timeout[ms] for establishing the connection with the ernie bot server. A timeout value of zero is interpreted as an infinite timeout. | duration | 1.9.0 | +| kyuubi.engine.chat.ernie.http.proxy | <undefined> | HTTP proxy url for API calling in ernie bot engine. e.g. http://127.0.0.1:1088 | string | 1.9.0 | +| kyuubi.engine.chat.ernie.http.socket.timeout | PT2M | The timeout[ms] for waiting for data packets after ernie bot server connection is established. A timeout value of zero is interpreted as an infinite timeout. | duration | 1.9.0 | +| kyuubi.engine.chat.ernie.model | completions | ID of the model used in ernie bot. Available models are completions_pro, ernie_bot_8k, completions and eb-instant[Model overview](https://cloud.baidu.com/doc/WENXINWORKSHOP/s/6lp69is2a). | string | 1.9.0 | +| kyuubi.engine.chat.ernie.token | <undefined> | The token to access ernie bot open API, which could be got at https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Ilkkrb0i5 | string | 1.9.0 | | kyuubi.engine.chat.extra.classpath | <undefined> | The extra classpath for the Chat engine, for configuring the location of the SDK and etc. | string | 1.8.0 | | kyuubi.engine.chat.gpt.apiKey | <undefined> | The key to access OpenAI open API, which could be got at https://platform.openai.com/account/api-keys | string | 1.8.0 | | kyuubi.engine.chat.gpt.http.connect.timeout | PT2M | The timeout[ms] for establishing the connection with the Chat GPT server. A timeout value of zero is interpreted as an infinite timeout. | duration | 1.8.0 | @@ -130,7 +135,7 @@ You can configure the Kyuubi properties in `$KYUUBI_HOME/conf/kyuubi-defaults.co | kyuubi.engine.chat.gpt.model | gpt-3.5-turbo | ID of the model used in ChatGPT. Available models refer to OpenAI's [Model overview](https://platform.openai.com/docs/models/overview). | string | 1.8.0 | | kyuubi.engine.chat.java.options | <undefined> | The extra Java options for the Chat engine | string | 1.8.0 | | kyuubi.engine.chat.memory | 1g | The heap memory for the Chat engine | string | 1.8.0 | -| kyuubi.engine.chat.provider | ECHO | The provider for the Chat engine. Candidates: | string | 1.8.0 | +| kyuubi.engine.chat.provider | ECHO | The provider for the Chat engine. Candidates: | string | 1.8.0 | | kyuubi.engine.connection.url.use.hostname | true | (deprecated) When true, the engine registers with hostname to zookeeper. When Spark runs on K8s with cluster mode, set to false to ensure that server can connect to engine | boolean | 1.3.0 | | kyuubi.engine.deregister.exception.classes || A comma-separated list of exception classes. If there is any exception thrown, whose class matches the specified classes, the engine would deregister itself. | set | 1.2.0 | | kyuubi.engine.deregister.exception.messages || A comma-separated list of exception messages. If there is any exception thrown, whose message or stacktrace matches the specified message list, the engine would deregister itself. | set | 1.2.0 | diff --git a/externals/kyuubi-chat-engine/pom.xml b/externals/kyuubi-chat-engine/pom.xml index 3639ceed329..0633143b9fa 100644 --- a/externals/kyuubi-chat-engine/pom.xml +++ b/externals/kyuubi-chat-engine/pom.xml @@ -65,6 +65,11 @@ test + + com.squareup.retrofit2 + converter-jackson + ${retrofit.version} + diff --git a/externals/kyuubi-chat-engine/src/main/java/org/apache/kyuubi/engine/chat/ernie/enums/ChatMessageRole.java b/externals/kyuubi-chat-engine/src/main/java/org/apache/kyuubi/engine/chat/ernie/enums/ChatMessageRole.java new file mode 100644 index 00000000000..8c8921fbdf0 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/java/org/apache/kyuubi/engine/chat/ernie/enums/ChatMessageRole.java @@ -0,0 +1,41 @@ +/* + * 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.kyuubi.engine.chat.ernie.enums; + +public enum ChatMessageRole { + FUNCTION("function"), + + USER("user"), + + ASSISTANT("assistant"); + + private final String value; + + private ChatMessageRole(String value) { + this.value = value; + } + + public String value() { + return this.value; + } + + @Override + public String toString() { + return "ChatMessageRole{" + "value='" + value + '\'' + '}'; + } +} diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/api/ApiHttpException.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/api/ApiHttpException.scala new file mode 100644 index 00000000000..ba54f97c840 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/api/ApiHttpException.scala @@ -0,0 +1,21 @@ +/* + * 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.kyuubi.engine.chat.api + +class ApiHttpException(statusCode: Int, message: String, exception: Exception) + extends Exception(message, exception) {} diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/api/ErnieBotApi.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/api/ErnieBotApi.scala new file mode 100644 index 00000000000..8593f65e51a --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/api/ErnieBotApi.scala @@ -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.kyuubi.engine.chat.api + +import io.reactivex.Single +import retrofit2.http.{Body, Path, POST, Query} + +import org.apache.kyuubi.engine.chat.ernie.bean.{ChatCompletionRequest, ChatCompletionResult} + +trait ErnieBotApi { + @POST("/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/{model}") + def createChatCompletion( + @Path("model") model: String, + @Query("access_token") accessToken: String, + @Body chatCompletionRequest: ChatCompletionRequest): Single[ChatCompletionResult] +} diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatCompletionRequest.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatCompletionRequest.scala new file mode 100644 index 00000000000..6cc3a6706bf --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatCompletionRequest.scala @@ -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.kyuubi.engine.chat.ernie.bean + +import java.lang.{Double => JDouble} +import java.util.{List => JList} + +import com.fasterxml.jackson.annotation.JsonProperty + +case class ChatCompletionRequest( + @JsonProperty("messages") messages: JList[ChatMessage], + @JsonProperty("functions") functions: JList[Function] = null, + @JsonProperty("temperature") temperature: JDouble = null, + @JsonProperty("top_p") topP: JDouble = null, + @JsonProperty("penalty_score") presenceScore: JDouble = null, + @JsonProperty("stream") stream: Boolean = false, + @JsonProperty("system") system: String = null, + @JsonProperty("stop") stop: JList[String] = null, + @JsonProperty("disable_search") disableSearch: Boolean = false, + @JsonProperty("enable_citation") enableCitation: Boolean = false, + @JsonProperty("user_id") userId: String = null) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatCompletionResult.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatCompletionResult.scala new file mode 100644 index 00000000000..e029882c5e4 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatCompletionResult.scala @@ -0,0 +1,39 @@ +/* + * 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.kyuubi.engine.chat.ernie.bean + +import java.lang.{Long => JLong} + +import com.fasterxml.jackson.annotation.JsonProperty + +case class ChatCompletionResult( + @JsonProperty("id") id: String, + @JsonProperty("object") obj: String, + @JsonProperty("created") created: JLong, + @JsonProperty("sentence_id") sentenceId: JLong, + @JsonProperty("is_end") isEnd: Boolean, + @JsonProperty("is_truncated") isTruncated: Boolean, + @JsonProperty("finish_reason") finishReason: String, + @JsonProperty("search_info") searchInfo: SearchInfo, + @JsonProperty("result") result: String, + @JsonProperty("need_clear_history") needClearHistory: Boolean, + @JsonProperty("ban_round") banRound: JLong, + @JsonProperty("usage") usage: Usage, + @JsonProperty("function_call") functionCall: FunctionCall, + @JsonProperty("error_msg") errorMsg: String, + @JsonProperty("error_code") errorCode: JLong) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatMessage.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatMessage.scala new file mode 100644 index 00000000000..d2b33222d94 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/ChatMessage.scala @@ -0,0 +1,26 @@ +/* + * 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.kyuubi.engine.chat.ernie.bean + +import com.fasterxml.jackson.annotation.JsonProperty + +case class ChatMessage( + @JsonProperty("role") role: String, + @JsonProperty("content") content: String, + @JsonProperty("name") name: String, + @JsonProperty("function_call") functionCall: FunctionCall = null) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Example.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Example.scala new file mode 100644 index 00000000000..fe25383dec4 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Example.scala @@ -0,0 +1,26 @@ +/* + * 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.kyuubi.engine.chat.ernie.bean + +import com.fasterxml.jackson.annotation.JsonProperty + +case class Example( + @JsonProperty("role") role: String, + @JsonProperty("name") name: String, + @JsonProperty("content") content: String = null, + @JsonProperty("function_call") functionCall: FunctionCall = null) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Function.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Function.scala new file mode 100644 index 00000000000..b0ad975bfa1 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Function.scala @@ -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.kyuubi.engine.chat.ernie.bean + +import java.util.{List => JList} + +import com.fasterxml.jackson.annotation.JsonProperty + +case class Function( + @JsonProperty("name") name: String, + @JsonProperty("description") description: String, + @JsonProperty("parameters") parameters: Object, + @JsonProperty("responses") responses: Object = null, + @JsonProperty("examples") examples: JList[Example] = null) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/FunctionCall.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/FunctionCall.scala new file mode 100644 index 00000000000..a6b66d78f8b --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/FunctionCall.scala @@ -0,0 +1,25 @@ +/* + * 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.kyuubi.engine.chat.ernie.bean + +import com.fasterxml.jackson.annotation.JsonProperty + +case class FunctionCall( + @JsonProperty("name") name: String, + @JsonProperty("thoughts") thoughts: String, + @JsonProperty("arguments") arguments: String = null) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/PluginUsage.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/PluginUsage.scala new file mode 100644 index 00000000000..dd406c775a8 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/PluginUsage.scala @@ -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.kyuubi.engine.chat.ernie.bean + +import java.lang.{Long => JLong} + +import com.fasterxml.jackson.annotation.JsonProperty + +case class PluginUsage( + @JsonProperty("name") name: String, + @JsonProperty("parse_tokens") parseTokens: JLong, + @JsonProperty("abstract_tokens") abstractTokens: JLong, + @JsonProperty("search_tokens") searchTokens: JLong, + @JsonProperty("total_tokens") totalTokens: JLong) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/SearchInfo.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/SearchInfo.scala new file mode 100644 index 00000000000..f97aa6c5863 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/SearchInfo.scala @@ -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.kyuubi.engine.chat.ernie.bean + +import java.lang.{Long => JLong} +import java.util.{List => JList} + +import com.fasterxml.jackson.annotation.JsonProperty + +case class SearchInfo( + @JsonProperty("is_beset") isBeset: JLong, + @JsonProperty("rewrite_query") rewriteQuery: String, + @JsonProperty("search_results") searchResults: JList[SearchResult]) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/SearchResult.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/SearchResult.scala new file mode 100644 index 00000000000..76b02be9243 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/SearchResult.scala @@ -0,0 +1,26 @@ +/* + * 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.kyuubi.engine.chat.ernie.bean + +import com.fasterxml.jackson.annotation.JsonProperty + +case class SearchResult( + @JsonProperty("index") index: java.lang.Long, + @JsonProperty("url") url: String, + @JsonProperty("title") title: String, + @JsonProperty("datasource_id") datasourceId: String) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Usage.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Usage.scala new file mode 100644 index 00000000000..4696943be6b --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Usage.scala @@ -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.kyuubi.engine.chat.ernie.bean + +import java.lang.{Long => JLong} +import java.util.{List => JList} + +import com.fasterxml.jackson.annotation.JsonProperty + +case class Usage( + @JsonProperty("prompt_tokens") promptTokens: JLong, + @JsonProperty("completion_tokens") completionTokens: JLong, + @JsonProperty("total_tokens") totalTokens: JLong, + @JsonProperty("plugins") plugins: JList[PluginUsage]) diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/service/ErnieBotService.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/service/ErnieBotService.scala new file mode 100644 index 00000000000..61f56421abe --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/service/ErnieBotService.scala @@ -0,0 +1,90 @@ +/* + * 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.kyuubi.engine.chat.ernie.service + +import java.io.IOException +import java.time.Duration +import java.util.concurrent.TimeUnit + +import com.fasterxml.jackson.annotation.JsonInclude +import com.fasterxml.jackson.databind.{DeserializationFeature, ObjectMapper, PropertyNamingStrategy} +import io.reactivex.Single +import okhttp3.{ConnectionPool, OkHttpClient} +import retrofit2.{HttpException, Retrofit} +import retrofit2.adapter.rxjava2.RxJava2CallAdapterFactory +import retrofit2.converter.jackson.JacksonConverterFactory + +import org.apache.kyuubi.engine.chat.api.{ApiHttpException, ErnieBotApi} +import org.apache.kyuubi.engine.chat.ernie.bean.{ChatCompletionRequest, ChatCompletionResult} + +class ErnieBotService(api: ErnieBotApi) { + + def execute[T](apiCall: Single[T]): T = { + try apiCall.blockingGet + catch { + case httpException: HttpException => + try if (httpException.response != null && httpException.response.errorBody != null) { + val errorBody: String = httpException.response.errorBody.string + val statusCode: Int = httpException.response.code + throw new ApiHttpException(statusCode, errorBody, httpException) + } else { + throw httpException + } + catch { + case ioException: IOException => + throw httpException + } + } + } + + def createChatCompletion( + request: ChatCompletionRequest, + model: String, + accessToken: String): ChatCompletionResult = { + execute(this.api.createChatCompletion(model, accessToken, request)) + } +} + +object ErnieBotService { + final private val BASE_URL = "https://aip.baidubce.com/" + + def apply(api: ErnieBotApi): ErnieBotService = new ErnieBotService(api) + + def defaultObjectMapper: ObjectMapper = { + val mapper: ObjectMapper = new ObjectMapper + mapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false) + mapper.setSerializationInclusion(JsonInclude.Include.NON_NULL) + mapper.setPropertyNamingStrategy(PropertyNamingStrategy.SNAKE_CASE) + mapper + } + + def defaultClient(timeout: Duration): OkHttpClient = { + new OkHttpClient.Builder() + .connectionPool(new ConnectionPool(5, 1, TimeUnit.SECONDS)) + .readTimeout(timeout.toMillis, TimeUnit.MILLISECONDS) + .build + } + + def defaultRetrofit(client: OkHttpClient, mapper: ObjectMapper): Retrofit = { + new Retrofit.Builder().baseUrl(BASE_URL).client(client) + .addConverterFactory(JacksonConverterFactory.create(mapper)) + .addCallAdapterFactory(RxJava2CallAdapterFactory.create) + .build + } + +} diff --git a/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ErnieBotProvider.scala b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ErnieBotProvider.scala new file mode 100644 index 00000000000..967ea333223 --- /dev/null +++ b/externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/provider/ErnieBotProvider.scala @@ -0,0 +1,105 @@ +/* + * 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.kyuubi.engine.chat.provider + +import java.net.{InetSocketAddress, Proxy, URL} +import java.time.Duration +import java.util +import java.util.concurrent.TimeUnit + +import scala.collection.JavaConverters._ + +import com.google.common.cache.{CacheBuilder, CacheLoader, LoadingCache} +import com.theokanning.openai.service.OpenAiService.defaultObjectMapper + +import org.apache.kyuubi.config.KyuubiConf +import org.apache.kyuubi.engine.chat.api.ErnieBotApi +import org.apache.kyuubi.engine.chat.ernie.bean.{ChatCompletionRequest, ChatMessage} +import org.apache.kyuubi.engine.chat.ernie.enums.ChatMessageRole +import org.apache.kyuubi.engine.chat.ernie.service.ErnieBotService +import org.apache.kyuubi.engine.chat.ernie.service.ErnieBotService.{defaultClient, defaultRetrofit} + +class ErnieBotProvider(conf: KyuubiConf) extends ChatProvider { + + private val accessToken = conf.get(KyuubiConf.ENGINE_ERNIE_BOT_ACCESS_TOKEN).getOrElse { + throw new IllegalArgumentException( + s"'${KyuubiConf.ENGINE_ERNIE_BOT_ACCESS_TOKEN.key}' must be configured, " + + s"which could be got at https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Ilkkrb0i5") + } + + private val model = conf.get(KyuubiConf.ENGINE_ERNIE_BOT_MODEL) + + private val ernieBotService: ErnieBotService = { + val builder = defaultClient( + Duration.ofMillis(conf.get(KyuubiConf.ENGINE_ERNIE_HTTP_SOCKET_TIMEOUT))) + .newBuilder + .connectTimeout(Duration.ofMillis(conf.get(KyuubiConf.ENGINE_ERNIE_HTTP_CONNECT_TIMEOUT))) + + conf.get(KyuubiConf.ENGINE_CHAT_GPT_HTTP_PROXY) match { + case Some(httpProxyUrl) => + val url = new URL(httpProxyUrl) + val proxy = new Proxy(Proxy.Type.HTTP, new InetSocketAddress(url.getHost, url.getPort)) + builder.proxy(proxy) + case _ => + } + + val retrofit = defaultRetrofit(builder.build(), defaultObjectMapper) + val ernieBotApi = retrofit.create(classOf[ErnieBotApi]) + new ErnieBotService(ernieBotApi) + } + + private var sessionUser: Option[String] = None + + private val chatHistory: LoadingCache[String, util.ArrayDeque[ChatMessage]] = + CacheBuilder.newBuilder() + .expireAfterWrite(10, TimeUnit.MINUTES) + .build(new CacheLoader[String, util.ArrayDeque[ChatMessage]] { + override def load(sessionId: String): util.ArrayDeque[ChatMessage] = + new util.ArrayDeque[ChatMessage] + }) + + override def open(sessionId: String, user: Option[String]): Unit = { + sessionUser = user + chatHistory.getIfPresent(sessionId) + } + + override def ask(sessionId: String, q: String): String = { + val messages = chatHistory.get(sessionId) + try { + messages.addLast(ChatMessage(ChatMessageRole.USER.value(), q, null)) + val completionRequest = ChatCompletionRequest( + messages = messages.asScala.toList.asJava, + userId = sessionUser.orNull) + val chatCompletionResult = ernieBotService + .createChatCompletion(completionRequest, model, accessToken) + if (chatCompletionResult.errorMsg != null) { + throw new RuntimeException(chatCompletionResult.errorMsg) + } + val responseText = chatCompletionResult.result + responseText + } catch { + case e: Throwable => + messages.removeLast() + s"Chat failed. Error: ${e.getMessage}" + } + } + + override def close(sessionId: String): Unit = { + chatHistory.invalidate(sessionId) + } +} diff --git a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala index b655f79840b..be525f4e874 100644 --- a/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala +++ b/kyuubi-common/src/main/scala/org/apache/kyuubi/config/KyuubiConf.scala @@ -3059,12 +3059,15 @@ object KyuubiConf { .doc("The provider for the Chat engine. Candidates: ") .version("1.8.0") .stringConf .transform { case "ECHO" | "echo" => "org.apache.kyuubi.engine.chat.provider.EchoProvider" case "GPT" | "gpt" | "ChatGPT" => "org.apache.kyuubi.engine.chat.provider.ChatGPTProvider" + case "ERNIE" | "ernie" | "ErnieBot" => + "org.apache.kyuubi.engine.chat.provider.ErnieBotProvider" case other => other } .createWithDefault("ECHO") @@ -3085,6 +3088,23 @@ object KyuubiConf { .stringConf .createWithDefault("gpt-3.5-turbo") + val ENGINE_ERNIE_BOT_ACCESS_TOKEN: OptionalConfigEntry[String] = + buildConf("kyuubi.engine.chat.ernie.token") + .doc("The token to access ernie bot open API, which could be got at " + + "https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Ilkkrb0i5") + .version("1.9.0") + .stringConf + .createOptional + + val ENGINE_ERNIE_BOT_MODEL: ConfigEntry[String] = + buildConf("kyuubi.engine.chat.ernie.model") + .doc("ID of the model used in ernie bot. " + + "Available models are completions_pro, ernie_bot_8k, completions and eb-instant" + + "[Model overview](https://cloud.baidu.com/doc/WENXINWORKSHOP/s/6lp69is2a).") + .version("1.9.0") + .stringConf + .createWithDefault("completions") + val ENGINE_CHAT_EXTRA_CLASSPATH: OptionalConfigEntry[String] = buildConf("kyuubi.engine.chat.extra.classpath") .doc("The extra classpath for the Chat engine, for configuring the location " + @@ -3100,6 +3120,13 @@ object KyuubiConf { .stringConf .createOptional + val ENGINE_ERNIE_BOT_HTTP_PROXY: OptionalConfigEntry[String] = + buildConf("kyuubi.engine.chat.ernie.http.proxy") + .doc("HTTP proxy url for API calling in ernie bot engine. e.g. http://127.0.0.1:1088") + .version("1.9.0") + .stringConf + .createOptional + val ENGINE_CHAT_GPT_HTTP_CONNECT_TIMEOUT: ConfigEntry[Long] = buildConf("kyuubi.engine.chat.gpt.http.connect.timeout") .doc("The timeout[ms] for establishing the connection with the Chat GPT server. " + @@ -3109,6 +3136,15 @@ object KyuubiConf { .checkValue(_ >= 0, "must be 0 or positive number") .createWithDefault(Duration.ofSeconds(120).toMillis) + val ENGINE_ERNIE_HTTP_CONNECT_TIMEOUT: ConfigEntry[Long] = + buildConf("kyuubi.engine.chat.ernie.http.connect.timeout") + .doc("The timeout[ms] for establishing the connection with the ernie bot server. " + + "A timeout value of zero is interpreted as an infinite timeout.") + .version("1.9.0") + .timeConf + .checkValue(_ >= 0, "must be 0 or positive number") + .createWithDefault(Duration.ofSeconds(120).toMillis) + val ENGINE_CHAT_GPT_HTTP_SOCKET_TIMEOUT: ConfigEntry[Long] = buildConf("kyuubi.engine.chat.gpt.http.socket.timeout") .doc("The timeout[ms] for waiting for data packets after Chat GPT server " + @@ -3118,6 +3154,15 @@ object KyuubiConf { .checkValue(_ >= 0, "must be 0 or positive number") .createWithDefault(Duration.ofSeconds(120).toMillis) + val ENGINE_ERNIE_HTTP_SOCKET_TIMEOUT: ConfigEntry[Long] = + buildConf("kyuubi.engine.chat.ernie.http.socket.timeout") + .doc("The timeout[ms] for waiting for data packets after ernie bot server " + + "connection is established. A timeout value of zero is interpreted as an infinite timeout.") + .version("1.9.0") + .timeConf + .checkValue(_ >= 0, "must be 0 or positive number") + .createWithDefault(Duration.ofSeconds(120).toMillis) + val ENGINE_JDBC_MEMORY: ConfigEntry[String] = buildConf("kyuubi.engine.jdbc.memory") .doc("The heap memory for the JDBC query engine") diff --git a/pom.xml b/pom.xml index 62791a0687a..36af02763d1 100644 --- a/pom.xml +++ b/pom.xml @@ -181,6 +181,7 @@ 4.11.0 4.1.100.Final 0.12.0 + 2.9.0 0.5.0-incubating ${spark.binary.version} 1.10.1