Pass stream all the way down

Author: worksofliam

src/chat/chat.ts

@@ -138,26 +138,13 @@ async function streamModelResponse(
     return;
   }
 
-  try {
-    stream.progress(`Using ${chosenProvider}...`);
-
-    const chatResponse = await chatRequest(chosenProvider, messages, {}, token);
+  stream.progress(`Using ${chosenProvider}...`);
 
-    for await (const fragement of chatResponse.text) {
-      stream.markdown(fragement);
-    }
-  } catch (err) {
-    if (err instanceof vscode.LanguageModelError) {
-      console.log(err.message, err.code, err.stack);
-    } else {
-      console.log(err);
-    }
-
-    stream.markdown(`Failed to get a response from ${chosenProvider}.`);
-  }
+  await chatRequest(chosenProvider, messages, {}, token, stream);
 }
 
 async function selectProviderAndModel() {
+  const selected = AiConfig.getModel();
   const copilotModels = await vscode.lm.selectChatModels();
   let ollamaModels: ListResponse = {models: []};
 
@@ -167,15 +154,17 @@ async function selectProviderAndModel() {
 
   const provider = await vscode.window.showQuickPick(
     [
-      ...ollamaModels.models.map((model): ModelQuickPickItem => ({ label: model.name, family: model.name, provider: "Ollama", iconPath: new vscode.ThemeIcon("heart") })),
-      ...copilotModels.map((model): ModelQuickPickItem => ({ label: model.name, family: model.family, provider: "GitHub Copilot", iconPath: new vscode.ThemeIcon("copilot") })),
+      {kind: vscode.QuickPickItemKind.Separator, label: "Ollama Models"}, 
+      ...ollamaModels.models.map((model): ModelQuickPickItem => ({ label: model.name, family: model.name, provider: "Ollama", iconPath: new vscode.ThemeIcon("heart"), picked: model.name === selected})),
+      {kind: vscode.QuickPickItemKind.Separator, label: "GitHub Copilot Models"},
+      ...copilotModels.map((model): ModelQuickPickItem => ({ label: model.name, family: model.family, provider: "GitHub Copilot", iconPath: new vscode.ThemeIcon("copilot")})),
     ],
     {
       title: "Select the AI model",
     }
   );
 
-  if (provider) {
+  if (provider && 'provider' in provider && 'family' in provider) {
     AiConfig.setProvider(provider.provider);
     AiConfig.setModel(provider.family);
   }

src/chat/send.ts

@@ -9,40 +9,46 @@ import {
 import Configuration from "../configuration";
 import { AiConfig, AiProvider } from "./aiConfig";
 
-export async function chatRequest(
+export function chatRequest(
   provider: AiProvider,
   messages: LanguageModelChatMessage[],
   options: LanguageModelChatRequestOptions,
-  token?: CancellationToken
-): Promise<Thenable<LanguageModelChatResponse>> {
+  token: CancellationToken,
+  stream: vscode.ChatResponseStream
+): Promise<void> {
   const chosenModel = AiConfig.getModel();
 
   switch (provider) {
     case "Ollama":
-      return ollamaRequest(chosenModel, messages);
+      return ollamaRequest(chosenModel, messages, stream);
     case "GitHub Copilot":
-      return copilotRequest(chosenModel, messages, options, token);
+      return copilotRequest(chosenModel, messages, options, token, stream);
   }
 }
 
 async function copilotRequest(
   model: string,
   messages: LanguageModelChatMessage[],
   options: LanguageModelChatRequestOptions,
-  token?: CancellationToken
-): Promise<LanguageModelChatResponse> {
+  token: CancellationToken,
+  stream: vscode.ChatResponseStream
+): Promise<void> {
   const models = await vscode.lm.selectChatModels({ family: model });
   if (models.length > 0) {
     const [first] = models;
     const response = await first.sendRequest(messages, options, token);
-    return response;
+
+    for await (const fragment of response.text) {
+      stream.markdown(fragment);
+    }
   }
 }
 
 async function ollamaRequest(
   model: string,
-  messages: LanguageModelChatMessage[]
-): Promise<LanguageModelChatResponse> {
+  messages: LanguageModelChatMessage[],
+  stream: vscode.ChatResponseStream
+): Promise<void> {
   const chats = [];
   for (const message of messages) {
     chats.push({
@@ -58,11 +64,5 @@ async function ollamaRequest(
 
   console.log(response.message.content);
 
-  return {
-    text: {
-      [Symbol.asyncIterator]: async function* () {
-        yield response.message.content;
-      },
-    },
-  };
+  stream.markdown(response.message.content);
 }