From a105e2e0ecce69f7c00e13e325c9f328c1bdfd4d Mon Sep 17 00:00:00 2001
From: Lyuben Deninski <ldeninski@gmail.com>
Date: Tue, 14 May 2024 14:39:37 +0300
Subject: [PATCH] switched to enum for `feature` param in AiClient.getModel and
 removed default param value to avoid future bugs

---
 lib/extension/src/ai/AIClient.ts                 | 10 ++++++++--
 .../src/autocomplete/AutoCompleteProvider.ts     |  2 +-
 lib/extension/src/types/util.ts                  | 16 ++++++++++++++++
 3 files changed, 25 insertions(+), 3 deletions(-)
 create mode 100644 lib/extension/src/types/util.ts

diff --git a/lib/extension/src/ai/AIClient.ts b/lib/extension/src/ai/AIClient.ts
index 10a070c..e4721ca 100644
--- a/lib/extension/src/ai/AIClient.ts
+++ b/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,
diff --git a/lib/extension/src/autocomplete/AutoCompleteProvider.ts b/lib/extension/src/autocomplete/AutoCompleteProvider.ts
index f96988d..d7a7366 100644
--- a/lib/extension/src/autocomplete/AutoCompleteProvider.ts
+++ b/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({
diff --git a/lib/extension/src/types/util.ts b/lib/extension/src/types/util.ts
new file mode 100644
index 0000000..30cc59e
--- /dev/null
+++ b/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>;