Add completion model

devchat-ai · Jun 15, 2024 · 82cbe74 · 82cbe74
1 parent 64b19f7
commit 82cbe74
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/src/main/kotlin/ai/devchat/plugin/completion/agent/Agent.kt b/src/main/kotlin/ai/devchat/plugin/completion/agent/Agent.kt
@@ -53,6 +53,7 @@ class Agent(val scope: CoroutineScope) {
 
   data class CompletionResponse(
     val id: String,
+    val model: String?,
     val choices: List<Choice>,
     val promptBuildingElapse: Long,
     val llmRequestElapse: Long
@@ -71,6 +72,7 @@ class Agent(val scope: CoroutineScope) {
     @SerializedName("language") val language: String,
     @SerializedName("prompt_time") val promptBuildingElapse: Long,
     @SerializedName("llm_time") val llmRequestElapse: Long,
+    @SerializedName("model") val model: String? = null,
     @SerializedName("cache_hit") val cacheHit: Boolean = false
   ) {
     enum class EventType {
@@ -269,12 +271,13 @@ class Agent(val scope: CoroutineScope) {
     completionRequest: CompletionRequest
   ): CompletionResponse? = suspendCancellableCoroutine { continuation ->
     currentRequest = RequestInfo.fromCompletionRequest(completionRequest)
+    val model = CONFIG["complete_model"] as? String
     var startTime = System.currentTimeMillis()
     val prompt = ContextBuilder(
       completionRequest.filepath,
       completionRequest.text,
       completionRequest.position
-    ).createPrompt(CONFIG["complete_model"] as? String)
+    ).createPrompt(model)
     val promptBuildingElapse = System.currentTimeMillis() - startTime
 
     scope.launch {
@@ -289,7 +292,7 @@ class Agent(val scope: CoroutineScope) {
       val offset = completionRequest.position
       val replaceRange = CompletionResponse.Choice.Range(start = offset, end = offset)
       val choice = CompletionResponse.Choice(index = 0, text = completion.text, replaceRange = replaceRange)
-      val response = CompletionResponse(completion.id, listOf(choice), promptBuildingElapse, llmRequestElapse)
+      val response = CompletionResponse(completion.id, model, listOf(choice), promptBuildingElapse, llmRequestElapse)
       continuation.resumeWith(Result.success(response))
     }
 

diff --git a/src/main/kotlin/ai/devchat/plugin/completion/editor/InlineCompletionService.kt b/src/main/kotlin/ai/devchat/plugin/completion/editor/InlineCompletionService.kt
@@ -166,7 +166,8 @@ class InlineCompletionService {
             ide = "intellij",
             language = virtualFile?.extension ?: "",
             promptBuildingElapse = completion.promptBuildingElapse,
-            llmRequestElapse = completion.llmRequestElapse
+            llmRequestElapse = completion.llmRequestElapse,
+            model = completion.model,
           )
         )
       }
@@ -234,7 +235,8 @@ class InlineCompletionService {
           ide = "intellij",
           language = virtualFile?.extension ?: "",
           promptBuildingElapse = currentCompletion.completion.promptBuildingElapse,
-          llmRequestElapse = currentCompletion.completion.llmRequestElapse
+          llmRequestElapse = currentCompletion.completion.llmRequestElapse,
+          model = currentCompletion.completion.model,
         )
       )
     }
@@ -253,6 +255,7 @@ class InlineCompletionService {
                 )),
                 promptBuildingElapse = 0,
                 llmRequestElapse = 0,
+                model = currentCompletion.completion.model,
               )
             )
           }