This is an automated email from the ASF dual-hosted git repository.
chengpan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/kyuubi.git
The following commit(s) were added to refs/heads/master by this push:
new 8ab4763ce [KYUUBI #5828] [CHAT] Kyuubi chat engine supports ernie
bot(文心一言)
8ab4763ce is described below
commit 8ab4763ce49ff19dfb34f5b2c58adb8b58b975bc
Author: hezhao2 <[email protected]>
AuthorDate: Mon Dec 11 10:18:16 2023 +0800
[KYUUBI #5828] [CHAT] Kyuubi chat engine supports ernie bot(文心一言)
# :mag: 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 :bookmark:
- [ ] 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 :coffin:
#### Behavior With This Pull Request :tada:
#### 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 <[email protected]>
Co-authored-by: Cheng Pan <[email protected]>
Signed-off-by: Cheng Pan <[email protected]>
---
docs/configuration/settings.md | 7 +-
externals/kyuubi-chat-engine/pom.xml | 5 +
.../engine/chat/ernie/enums/ChatMessageRole.java | 41 ++++++++
.../kyuubi/engine/chat/api/ApiHttpException.scala | 21 +++++
.../kyuubi/engine/chat/api/ErnieBotApi.scala | 31 ++++++
.../chat/ernie/bean/ChatCompletionRequest.scala | 36 +++++++
.../chat/ernie/bean/ChatCompletionResult.scala | 39 ++++++++
.../engine/chat/ernie/bean/ChatMessage.scala | 26 +++++
.../kyuubi/engine/chat/ernie/bean/Example.scala | 26 +++++
.../kyuubi/engine/chat/ernie/bean/Function.scala | 29 ++++++
.../engine/chat/ernie/bean/FunctionCall.scala | 25 +++++
.../engine/chat/ernie/bean/PluginUsage.scala | 29 ++++++
.../kyuubi/engine/chat/ernie/bean/SearchInfo.scala | 28 ++++++
.../engine/chat/ernie/bean/SearchResult.scala | 26 +++++
.../kyuubi/engine/chat/ernie/bean/Usage.scala | 29 ++++++
.../chat/ernie/service/ErnieBotService.scala | 90 ++++++++++++++++++
.../engine/chat/provider/ErnieBotProvider.scala | 105 +++++++++++++++++++++
.../org/apache/kyuubi/config/KyuubiConf.scala | 45 +++++++++
pom.xml | 1 +
19 files changed, 638 insertions(+), 1 deletion(-)
diff --git a/docs/configuration/settings.md b/docs/configuration/settings.md
index a3c0ca0df..6d46eeff2 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
|
[...]
|----------------------------------------------------------|---------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------
[...]
+| 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.
[...]
+| 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
[...]
+| 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.
[...]
+| 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).
[...]
+| 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
[...]
| kyuubi.engine.chat.extra.classpath | <undefined>
| The extra classpath for the Chat engine, for configuring the location
of the SDK and etc.
[...]
| 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
[...]
| 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.
[...]
@@ -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).
[...]
| kyuubi.engine.chat.java.options | <undefined>
| The extra Java options for the Chat engine
[...]
| kyuubi.engine.chat.memory | 1g
| The heap memory for the Chat engine
[...]
-| 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>
[...]
+| 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>
[...]
| 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
[...]
| 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.
[...]
| 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.
[...]
diff --git a/externals/kyuubi-chat-engine/pom.xml
b/externals/kyuubi-chat-engine/pom.xml
index 3639ceed3..0633143b9 100644
--- a/externals/kyuubi-chat-engine/pom.xml
+++ b/externals/kyuubi-chat-engine/pom.xml
@@ -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>
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 000000000..8c8921fbd
--- /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 000000000..ba54f97c8
--- /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 000000000..8593f65e5
--- /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 000000000..6cc3a6706
--- /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 000000000..e029882c5
--- /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 000000000..d2b33222d
--- /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 000000000..fe25383de
--- /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 000000000..b0ad975bf
--- /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 000000000..a6b66d78f
--- /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 000000000..dd406c775
--- /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 000000000..f97aa6c58
--- /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 000000000..76b02be92
--- /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 000000000..4696943be
--- /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 000000000..61f56421a
--- /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 000000000..967ea3332
--- /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 b655f7984..be525f4e8 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: <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")
diff --git a/pom.xml b/pom.xml
index 62791a068..36af02763 100644
--- a/pom.xml
+++ b/pom.xml
@@ -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>
