[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:
parent
f9d27de9bb
commit
8ab4763ce4
@ -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: <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 |
|
||||
|
||||
5
externals/kyuubi-chat-engine/pom.xml
vendored
5
externals/kyuubi-chat-engine/pom.xml
vendored
@ -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>
|
||||
|
||||
@ -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 + '\'' + '}';
|
||||
}
|
||||
}
|
||||
@ -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) {}
|
||||
31
externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/api/ErnieBotApi.scala
vendored
Normal file
31
externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/api/ErnieBotApi.scala
vendored
Normal 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]
|
||||
}
|
||||
@ -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)
|
||||
@ -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)
|
||||
@ -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)
|
||||
@ -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)
|
||||
@ -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)
|
||||
@ -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)
|
||||
@ -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)
|
||||
@ -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])
|
||||
@ -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)
|
||||
29
externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Usage.scala
vendored
Normal file
29
externals/kyuubi-chat-engine/src/main/scala/org/apache/kyuubi/engine/chat/ernie/bean/Usage.scala
vendored
Normal 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])
|
||||
@ -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
|
||||
}
|
||||
|
||||
}
|
||||
@ -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)
|
||||
}
|
||||
}
|
||||
@ -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")
|
||||
|
||||
1
pom.xml
1
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>
|
||||
|
||||
Loading…
Reference in New Issue
Block a user