Bugfix: Wrogn template is used when two different model families used for chat and autocomplete

lib/extension/src/ai/AIClient.ts

@@ -14,6 +14,7 @@ import * as vscode from "vscode";
 import { z } from "zod";
 import { Logger } from "../logger";
 import { ApiKeyManager } from "./ApiKeyManager";
+import { EnumAsUnion } from '../types/util';
 
 function getProviderBaseUrl(): string {
   let defaultUrl = "http://localhost:8080/";
@@ -61,6 +62,11 @@ function getPromptTemplate() {
   return ollama.prompt.Llama2;
 }
 
+enum MODE_PURPOSE {
+  CHAT = "chat",
+  AUTOCOMPLETE = "autocomplete"
+};
+
 export class AIClient {
   private readonly apiKeyManager: ApiKeyManager;
   private readonly logger: Logger;
@@ -76,7 +82,7 @@ export class AIClient {
     this.logger = logger;
   }
 
-  public getModel(feature: string = "chat"): string {
+  public getModel(feature: EnumAsUnion<typeof MODE_PURPOSE>): string {
     if (feature != "chat") {
       this.logger.log(["Autocomplete Model: ", getAutoCompleteModel()]);
       return getAutoCompleteModel();
@@ -120,7 +126,7 @@ export class AIClient {
       .CompletionTextGenerator({
         api: await this.getProviderApiConfiguration(),
         promptTemplate: getPromptTemplate(),
-        model: this.getModel(),
+        model: this.getModel("chat"),
         maxGenerationTokens: maxTokens,
         stopSequences: stop,
         temperature,

lib/extension/src/autocomplete/AutoCompleteProvider.ts

@@ -141,7 +141,7 @@ export class AutoCompleteProvider
         const additionalContext = this.getAdditionalContext(document);
         const { prompt, stop } =
           this.autoCompleteTemplateProvider.getAutoCompletePrompt(
-            this.ai.getModel(),
+            this.ai.getModel("autocomplete"),
             { additionalContext, prefix, suffix }
           );
         const response = await this.ai.generateText({

lib/extension/src/types/util.ts

@@ -0,0 +1,16 @@
+export type StringValues<T> = {
+  [K in keyof T]: T[K] extends string ? T[K] : never;
+}[keyof T];
+
+type NumberValues<T> = {
+  [K in keyof T]: T[K] extends number ? T[K] : never;
+}[keyof T];
+
+type SymbolValues<T> = {
+  [K in keyof T]: T[K] extends symbol ? T[K] : never;
+}[keyof T];
+
+/**
+ * Usage : type EnumValues = EnumAsUnion<typeof anEnum>
+ */
+export type EnumAsUnion<T> = `${StringValues<T>}` | NumberValues<T> | SymbolValues<T>;