handle user cancellation

microsoft · Jul 30, 2024 · 9ede8b6 · 9ede8b6
1 parent 30b3961
commit 9ede8b6
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 9 deletions.
diff --git a/packages/cli/src/server.ts b/packages/cli/src/server.ts
@@ -107,6 +107,11 @@ export async function startServer(options: { port: string }) {
                 // add handler
                 const chatId = randomHex(6)
                 chats[chatId] = async (chunk) => {
+                    if (!responseSoFar) {
+                        trace.itemValue("model", chunk.model)
+                        trace.appendContent("\n\n")
+                    }
+                    trace.appendToken(chunk.chunk)
                     responseSoFar += chunk.chunk ?? ""
                     tokensSoFar += chunk.tokens ?? 0
                     partialCb?.({
@@ -116,6 +121,8 @@ export async function startServer(options: { port: string }) {
                     })
                     finishReason = chunk.finishReason as any
                     if (finishReason) {
+                        trace.appendContent("\n\n")
+                        trace.itemValue(`finish reason`, finishReason)
                         delete chats[chatId]
                         resolve({ text: responseSoFar, finishReason })
                     }
@@ -129,7 +136,6 @@ export async function startServer(options: { port: string }) {
                     messages,
                 })
                 for (const ws of wss.clients) {
-                    trace.log(`chat: sending request to client`)
                     ws.send(msg)
                     break
                 }

diff --git a/packages/core/src/fence.ts b/packages/core/src/fence.ts
@@ -183,7 +183,7 @@ export function renderFencedVariables(vars: Fenced[]) {
                 validation,
                 args,
                 language,
-            }) => `-   \`${k}\` ${
+            }) => `-   ${k ? `\`${k}\`` : ""} ${
                 validation !== undefined
                     ? `schema ${args.schema}: ${validation.valid === undefined ? EMOJI_UNDEFINED : validation.valid ? EMOJI_SUCCESS : EMOJI_FAIL}`
                     : "no label"

diff --git a/packages/core/src/openai.ts b/packages/core/src/openai.ts
@@ -218,12 +218,7 @@ export const OpenAIChatCompletion: ChatCompletionHandler = async (
                 if (typeof delta?.content == "string") {
                     numTokens += estimateTokens(delta.content, encoder)
                     chatResp += delta.content
-                    if (delta.content)
-                        trace.appendContent(
-                            delta.content.includes("`")
-                                ? `\`\`\` ${delta.content.replace(/\n/g, " ")} \`\`\` `
-                                : `\`${delta.content.replace(/\n/g, " ")}\` `
-                        )
+                    trace.appendToken(delta.content)
                 } else if (Array.isArray(delta.tool_calls)) {
                     const { tool_calls } = delta
                     for (const call of tool_calls) {

diff --git a/packages/core/src/server/messages.ts b/packages/core/src/server/messages.ts
@@ -145,6 +145,7 @@ export interface ChatCancel {
 export interface ChatChunk extends RequestMessage {
     type: "chat.chunk"
     chatId: string
+    model?: string
     finishReason?: string
     chunk?: string
     tokens?: number

diff --git a/packages/core/src/trace.ts b/packages/core/src/trace.ts
@@ -59,6 +59,15 @@ export class MarkdownTrace extends EventTarget implements ToolCallTrace {
         }
     }
 
+    appendToken(content: string) {
+        if (!content) return
+        this.appendContent(
+            content.includes("`")
+                ? `\`\`\` ${content.replace(/\r?\n/g, " ")} \`\`\` `
+                : `\`${content.replace(/\r?\n/g, " ")}\` `
+        )
+    }
+
     startDetails(title: string, success?: boolean) {
         this.detailsDepth++
         title = title?.trim() || ""

diff --git a/packages/vscode/src/lmaccess.ts b/packages/vscode/src/lmaccess.ts
@@ -195,7 +195,12 @@ export function createChatModelRunner(
         const token = new vscode.CancellationTokenSource().token
         const { model, messages, modelOptions } = req
         const chatModel = await pickChatModel(state, model)
-        if (!chatModel) throw new Error("No chat model selected.")
+        if (!chatModel) {
+            onChunk({
+                finishReason: "cancel",
+            })
+            return
+        }
         const chatMessages = messagesToChatMessages(messages)
         const request = await chatModel.sendRequest(
             chatMessages,
@@ -213,6 +218,7 @@ export function createChatModelRunner(
                 chunk: fragment,
                 tokens: await chatModel.countTokens(text),
                 finishReason: undefined,
+                model: chatModel.id,
             })
         }
         onChunk({