[KYUUBI #5828] [CHAT] Kyuubi chat engine supports ernie bot(文心一言)

# 🔍 Description
## Issue References 🔗

This pull request fixes #5386

## Describe Your Solution 🔧

add a new backend(ernie bot) for the Chat engine.
<img width="1672" alt="Screenshot 2023-12-07 at 16 20 56" src="https://github.com/apache/kyuubi/assets/32693629/9850916c-3c4a-433c-8278-3068e7b37314">

## Types of changes 🔖

- [ ] Bugfix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)

## Test Plan 🧪

#### Behavior Without This Pull Request ⚰️

#### Behavior With This Pull Request 🎉

#### Related Unit Tests

---

# Checklists
## 📝 Author Self Checklist

- [ ] My code follows the [style guidelines](https://kyuubi.readthedocs.io/en/master/contributing/code/style.html) of this project
- [ ] I have performed a self-review
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] This patch was not authored or co-authored using [Generative Tooling](https://www.apache.org/legal/generative-tooling.html)

## 📝 Committer Pre-Merge Checklist

- [x] Pull request title is okay.
- [x] No license issues.
- [x] Milestone correctly set?
- [ ] Test coverage is ok
- [x] Assignees are selected.
- [x] Minimum number of approvals
- [x] No changes are requested

**Be nice. Be informative.**

Closes #5828 from zhaohehuhu/dev-1207.

Closes #5828

c7314d4a8 [Cheng Pan] Update externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/service/ErnieBotService.scala
0c4e01007 [hezhao2] update doc
78e51b3d1 [hezhao2] add ernie into doc
2f4a63845 [hezhao2] refactor
832bc0453 [hezhao2] delete enum model
67214c575 [hezhao2] get rid of some java code
567772679 [hezhao2] java bean to scale case class
4d5c5940d [hezhao2] refactor some params
7c44eb83f [hezhao2] refactor some params
56b9ad13a [hezhao2] refactor
a8e3d6cf6 [hezhao2] refactor
7376d800d [hezhao2] Kyuubi chat engine supports ernie bot
4b72a09c0 [hezhao2] Kyuubi chat engine supports ernie bot

Lead-authored-by: hezhao2 <hezhao2@cisco.com>
Co-authored-by: Cheng Pan <pan3793@gmail.com>
Signed-off-by: Cheng Pan <chengpan@apache.org>
This commit is contained in:
hezhao2 2023-12-11 10:18:16 +08:00 committed by Cheng Pan
parent f9d27de9bb
commit 8ab4763ce4
19 changed files with 638 additions and 1 deletions

View File

@ -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 | &lt;undefined&gt; | 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 | &lt;undefined&gt; | 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 | &lt;undefined&gt; | 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 | &lt;undefined&gt; | 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 | &lt;undefined&gt; | 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: <ul> <li>ECHO: simply replies a welcome message.</li> <li>GPT: a.k.a ChatGPT, powered by OpenAI.</li></ul> | string | 1.8.0 |
| kyuubi.engine.chat.provider | ECHO | The provider for the Chat engine. Candidates: <ul> <li>ECHO: simply replies a welcome message.</li> <li>GPT: a.k.a ChatGPT, powered by OpenAI.</li> <li>ERNIE: ErnieBot, powered by Baidu.</li></ul> | 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 |

View File

@ -65,6 +65,11 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>com.squareup.retrofit2</groupId>
<artifactId>converter-jackson</artifactId>
<version>${retrofit.version}</version>
</dependency>
</dependencies>
<build>

View File

@ -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 + '\'' + '}';
}
}

View File

@ -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) {}

View File

@ -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]
}

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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])

View File

@ -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)

View File

@ -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])

View File

@ -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
}
}

View File

@ -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)
}
}

View File

@ -3059,12 +3059,15 @@ object KyuubiConf {
.doc("The provider for the Chat engine. Candidates: <ul>" +
" <li>ECHO: simply replies a welcome message.</li>" +
" <li>GPT: a.k.a ChatGPT, powered by OpenAI.</li>" +
" <li>ERNIE: ErnieBot, powered by Baidu.</li>" +
"</ul>")
.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")

View File

@ -181,6 +181,7 @@
<mockito.version>4.11.0</mockito.version>
<netty.version>4.1.100.Final</netty.version>
<openai.java.version>0.12.0</openai.java.version>
<retrofit.version>2.9.0</retrofit.version>
<paimon.version>0.5.0-incubating</paimon.version>
<paimon.spark.binary.version>${spark.binary.version}</paimon.spark.binary.version>
<parquet.version>1.10.1</parquet.version>