Merge pull request #64 from hotip/master

Author: phodal

docs/customize/custom-llm-server.md

@@ -41,3 +41,58 @@ class ChatInput(BaseModel):
 async def chat(msg: ChatInput):
     return StreamingResponse(fetch_chat_completion(msg.messages), media_type="text/event-stream")
 ```
+
+## Custom response format
+
+We used [JsonPathKt](https://github.com/codeniko/JsonPathKt) to parse response,
+currently we only extract the first choice and only the response message.
+If your response is this format: 
+
+```json
+{
+  "id": "chatcmpl-123",
+  "object": "chat.completion.chunk",
+  "created": 1677652288,
+  "model": "gpt-3.5-turbo",
+  "choices": [{
+    "index": 0,
+    "delta": {
+      "content": "Hello"
+    },
+    "finish_reason": "stop"
+  }]
+}
+```
+You need to set the `response format` to:
+
+```text
+$.choices[0].message.delta.content
+```
+
+## Custom request format
+
+Only support amount of request parameters like OpenAI does.
+Only support http request that don't need encryption keys(like websocket)
+
+
+### Custom Request (header and body)
+
+You can add top level field to the request body,
+And custom the origin keys for `role`, `messsage`
+
+```json
+{
+  "customHeaders": { "my header": "my value" },
+  "customFields": {"user": "userid", "date": "2012"},
+  "messageKeys": {"role": "role", "content": "message"}
+}
+```
+
+and the request body will be:
+
+```json
+{
+	"user": "userid",
+    "messages": [{"role": "user", "message": "..."}]
+  }
+```

src/main/kotlin/cc/unitmesh/devti/gui/chat/ChatCodingPanel.kt

@@ -25,6 +25,7 @@ import com.intellij.util.ui.UIUtil
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.catch
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.withContext
 import java.awt.event.ActionListener
@@ -199,13 +200,16 @@ class ChatCodingPanel(private val chatCodingService: ChatCodingService, val disp
         val startTime = System.currentTimeMillis() // 记录代码开始执行的时间
 
         var text = ""
-        runCatching {
-            content.collect {
-                text += it
-                messageView.updateSourceContent(text)
-                messageView.updateContent(text)
-                messageView.scrollToBottom()
-            }
+        content.catch {
+            it.printStackTrace()
+        }.collect {
+            text += it
+
+            // 以下两个 API 设计不合理，如果必须要同时调用，那就只提供一个就好了
+            messageView.updateSourceContent(text)
+            messageView.updateContent(text)
+
+            messageView.scrollToBottom()
         }
 
         if (delaySeconds.isNotEmpty()) {

src/main/kotlin/cc/unitmesh/devti/llms/azure/AzureOpenAIProvider.kt

@@ -26,6 +26,7 @@ import okhttp3.OkHttpClient
 import okhttp3.Request
 import okhttp3.RequestBody
 
+
 @Serializable
 data class SimpleOpenAIFormat(val role: String, val content: String) {
     companion object {

src/main/kotlin/cc/unitmesh/devti/llms/azure/ResponseBodyCallback.kt

@@ -35,6 +35,9 @@ import java.io.InputStreamReader
 import java.nio.charset.StandardCharsets
 
 class AutoDevHttpException(error: String, val statusCode: Int) : RuntimeException(error) {
+    override fun toString(): String {
+        return "AutoDevHttpException(statusCode=$statusCode, message=$message)"
+    }
 }
 
 /**
@@ -80,6 +83,11 @@ class ResponseBodyCallback(private val emitter: FlowableEmitter<SSE>, private va
                         null
                     }
 
+                    line!!.startsWith("{") -> {
+                        logger<ResponseBodyCallback>().warn("msg starts with { $line")
+                        emitter.onNext(SSE(line!!))
+                        null
+                    }
                     else -> {
                         throw SSEFormatException("Invalid sse format! $line")
                     }

src/main/kotlin/cc/unitmesh/devti/llms/custom/CustomLLMProvider.kt

@@ -22,11 +22,12 @@ import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.withContext
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.encodeToString
-import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.*
 import okhttp3.MediaType.Companion.toMediaTypeOrNull
 import okhttp3.OkHttpClient
 import okhttp3.Request
 import okhttp3.RequestBody
+import org.jetbrains.annotations.VisibleForTesting
 import java.time.Duration
 
 @Serializable
@@ -40,8 +41,9 @@ class CustomLLMProvider(val project: Project) : LLMProvider {
     private val autoDevSettingsState = AutoDevSettingsState.getInstance()
     private val url get() = autoDevSettingsState.customEngineServer
     private val key get() = autoDevSettingsState.customEngineToken
-    private val engineFormat get() = autoDevSettingsState.customEngineResponseFormat
-    private val customPromptConfig: CustomPromptConfig?
+    private val requestFormat: String get() = autoDevSettingsState.customEngineRequestFormat
+    private val responseFormat get() = autoDevSettingsState.customEngineResponseFormat
+    private val customPromptConfig: CustomPromptConfig
         get() {
             val prompts = autoDevSettingsState.customPrompts
             return CustomPromptConfig.tryParse(prompts)
@@ -73,18 +75,24 @@ class CustomLLMProvider(val project: Project) : LLMProvider {
         messages += Message("user", promptText)
 
         val customRequest = CustomRequest(messages)
-        val requestContent = Json.encodeToString<CustomRequest>(customRequest)
+        val requestContent = customRequest.updateCustomFormat(requestFormat)
 
         val body = RequestBody.create("application/json; charset=utf-8".toMediaTypeOrNull(), requestContent)
-        logger.info("Requesting from $body")
 
         val builder = Request.Builder()
         if (key.isNotEmpty()) {
             builder.addHeader("Authorization", "Bearer $key")
+            builder.addHeader("Content-Type", "application/json")
         }
+        builder.appendCustomHeaders(requestFormat)
 
-        client = client.newBuilder().readTimeout(timeout).build()
-        val request = builder.url(url).post(body).build()
+        client = client.newBuilder()
+                .readTimeout(timeout)
+                .build()
+        val request = builder
+                .url(url)
+                .post(body)
+                .build()
 
         val call = client.newCall(request)
         val emitDone = false
@@ -95,19 +103,22 @@ class CustomLLMProvider(val project: Project) : LLMProvider {
             }, BackpressureStrategy.BUFFER)
 
         try {
-            logger.info("Starting to stream:")
             return callbackFlow {
                 withContext(Dispatchers.IO) {
                     sseFlowable
-                        .doOnError(Throwable::printStackTrace)
-                        .blockingForEach { sse ->
-                            if (engineFormat.isNotEmpty()) {
-                                val chunk: String = JsonPath.parse(sse!!.data)?.read<String>(engineFormat)
-                                    ?: throw Exception("Failed to parse chunk: ${sse.data}, format: $engineFormat")
-                                trySend(chunk)
-                            } else {
-                                val result: ChatCompletionResult =
-                                    ObjectMapper().readValue(sse!!.data, ChatCompletionResult::class.java)
+                            .doOnError{
+                                it.printStackTrace()
+                                close()
+                            }
+                            .blockingForEach { sse ->
+                                if (responseFormat.isNotEmpty()) {
+                                    val chunk: String = JsonPath.parse(sse!!.data)?.read(responseFormat)
+                                            ?: throw Exception("Failed to parse chunk")
+                                    logger.warn("got msg: $chunk")
+                                    trySend(chunk)
+                                } else {
+                                    val result: ChatCompletionResult =
+                                            ObjectMapper().readValue(sse!!.data, ChatCompletionResult::class.java)
 
                                 val completion = result.choices[0].message
                                 if (completion != null && completion.content != null) {
@@ -134,7 +145,7 @@ class CustomLLMProvider(val project: Project) : LLMProvider {
 
         val body = RequestBody.create("application/json; charset=utf-8".toMediaTypeOrNull(), requestContent)
 
-        logger.info("Requesting from $body")
+        logger.info("Requesting form: $requestContent ${body.toString()}")
         val builder = Request.Builder()
         if (key.isNotEmpty()) {
             builder.addHeader("Authorization", "Bearer $key")
@@ -157,4 +168,69 @@ class CustomLLMProvider(val project: Project) : LLMProvider {
             return ""
         }
     }
-}
+}
+
+@VisibleForTesting
+fun Request.Builder.appendCustomHeaders(customRequestHeader: String): Request.Builder = apply {
+    runCatching {
+        Json.parseToJsonElement(customRequestHeader)
+                .jsonObject["customHeaders"].let { customFields ->
+            customFields?.jsonObject?.forEach { (key, value) ->
+                header(key, value.jsonPrimitive.content)
+            }
+        }
+    }.onFailure {
+        // should I warn user?
+        println("Failed to parse custom request header ${it.message}")
+    }
+}
+
+@VisibleForTesting
+fun JsonObject.updateCustomBody(customRequest: String): JsonObject {
+    return runCatching {
+        buildJsonObject {
+            // copy origin object
+            this@updateCustomBody.forEach { u, v -> put(u, v) }
+
+            val customRequestJson = Json.parseToJsonElement(customRequest).jsonObject
+
+            customRequestJson["customFields"]?.let { customFields ->
+                customFields.jsonObject.forEach { (key, value) ->
+                    put(key, value.jsonPrimitive.content)
+                }
+            }
+
+
+            // TODO clean code with magic literals
+            var roleKey = "role"
+            var contentKey = "message"
+            customRequestJson.jsonObject["messageKeys"]?.let {
+                roleKey = it.jsonObject["role"]?.jsonPrimitive?.content ?: "role"
+                contentKey = it.jsonObject["content"]?.jsonPrimitive?.content ?: "message"
+            }
+
+            val messages: JsonArray = this@updateCustomBody["messages"]?.jsonArray ?: buildJsonArray { }
+
+
+            this.put("messages", buildJsonArray {
+                messages.forEach { message ->
+                    val role: String = message.jsonObject["role"]?.jsonPrimitive?.content ?: "user"
+                    val content: String = message.jsonObject["message"]?.jsonPrimitive?.content ?: ""
+                    add(buildJsonObject {
+                        put(roleKey, role)
+                        put(contentKey, content)
+                    })
+                }
+            })
+        }
+    }.getOrElse {
+        logger<CustomLLMProvider>().error("Failed to parse custom request body", it)
+        this
+    }
+}
+
+fun CustomRequest.updateCustomFormat(format: String): String {
+    val requestContentOri = Json.encodeToString<CustomRequest>(this)
+    return Json.parseToJsonElement(requestContentOri)
+            .jsonObject.updateCustomBody(format).toString()
+}

src/main/kotlin/cc/unitmesh/devti/llms/xianghuo/XingHuoProvider.kt

@@ -4,6 +4,7 @@ package cc.unitmesh.devti.llms.xianghuo
 
 import cc.unitmesh.devti.llms.LLMProvider
 import cc.unitmesh.devti.settings.AutoDevSettingsState
+import cc.unitmesh.devti.settings.XingHuoApiVersion
 import com.intellij.openapi.components.Service
 import com.intellij.openapi.project.Project
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -26,6 +27,14 @@ class XingHuoProvider(val project: Project) : LLMProvider {
     private val secrectKey: String
         get() = autoDevSettingsState.xingHuoApiSecrect
 
+    private val apiVersion: XingHuoApiVersion
+        get() = autoDevSettingsState.xingHuoApiVersion
+    private val XingHuoApiVersion.asGeneralDomain
+        get() = when (this) {
+            XingHuoApiVersion.V1 -> ""
+            XingHuoApiVersion.V2 -> "v2"
+            else -> "v3"
+        }
 
     private val appid: String
         get() = autoDevSettingsState.xingHuoAppId
@@ -119,7 +128,7 @@ class XingHuoProvider(val project: Project) : LLMProvider {
             val header = """
             |host: spark-api.xf-yun.com
             |date: $date
-            |GET /v1.1/chat HTTP/1.1
+            |GET /v${apiVersion.value}.1/chat HTTP/1.1
         """.trimMargin()
             val signature = hmacsha256.doFinal(header.toByteArray()).encodeBase64()
             val authorization =
@@ -130,7 +139,7 @@ class XingHuoProvider(val project: Project) : LLMProvider {
                 "date" to date,
                 "host" to "spark-api.xf-yun.com"
             )
-            val urlBuilder = "https://spark-api.xf-yun.com/v1.1/chat".toHttpUrl().newBuilder()
+            val urlBuilder = "https://spark-api.xf-yun.com/v${apiVersion.value}.1/chat".toHttpUrl().newBuilder()
             params.forEach {
                 urlBuilder.addQueryParameter(it.key, it.value)
             }
@@ -147,7 +156,7 @@ class XingHuoProvider(val project: Project) : LLMProvider {
             },
             "parameter": {
                 "chat": {
-                    "domain": "general",
+                    "domain": "general${apiVersion.asGeneralDomain}",
                     "temperature": 0.5,
                     "max_tokens": 1024
                 }