refactor: fix types

chhoumann · Mar 3, 2024 · 0d14b61 · 0d14b61
1 parent 4db9991
commit 0d14b61
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 8 deletions.
diff --git a/src/gui/MacroGUIs/AIAssistantCommandSettingsModal.ts b/src/gui/MacroGUIs/AIAssistantCommandSettingsModal.ts
@@ -14,7 +14,7 @@ import {
 	DEFAULT_TOP_P,
 } from "src/ai/OpenAIModelParameters";
 import { getTokenCount } from "src/ai/AIAssistant";
-import { getModelNames } from "src/ai/aiHelpers";
+import { getModelByName, getModelNames } from "src/ai/aiHelpers";
 
 export class AIAssistantCommandSettingsModal extends Modal {
 	public waitForClose: Promise<IAIAssistantCommand>;
@@ -28,7 +28,10 @@ export class AIAssistantCommandSettingsModal extends Modal {
 	private get systemPromptTokenLength(): number {
 		if (this.settings.model === "Ask me") return Number.POSITIVE_INFINITY;
 
-		return getTokenCount(this.settings.systemPrompt, this.settings.model);
+		const model = getModelByName(this.settings.model);
+		if (!model) return Number.POSITIVE_INFINITY;
+
+		return getTokenCount(this.settings.systemPrompt, model);
 	}
 
 	constructor(settings: IAIAssistantCommand) {

diff --git a/src/gui/MacroGUIs/AIAssistantInfiniteCommandSettingsModal.ts b/src/gui/MacroGUIs/AIAssistantInfiniteCommandSettingsModal.ts
@@ -23,7 +23,9 @@ export class InfiniteAIAssistantCommandSettingsModal extends Modal {
 	private showAdvancedSettings = false;
 
 	private get systemPromptTokenLength(): number {
-		return getTokenCount(this.settings.systemPrompt, this.settings.model);
+		const model = getModelByName(this.settings.model);
+		if (!model) return Number.POSITIVE_INFINITY;
+		return getTokenCount(this.settings.systemPrompt, model);
 	}
 
 	constructor(settings: IInfiniteAIAssistantCommand) {

diff --git a/src/types/macros/QuickCommands/AIAssistantCommand.ts b/src/types/macros/QuickCommands/AIAssistantCommand.ts
@@ -1,4 +1,3 @@
-import type { Models_And_Ask_Me } from "src/ai/models";
 import { Command } from "../Command";
 import { CommandType } from "../CommandType";
 import type { IAIAssistantCommand } from "./IAIAssistantCommand";
@@ -11,7 +10,7 @@ export class AIAssistantCommand extends Command implements IAIAssistantCommand {
 	name: string;
 	type: CommandType;
 
-	model: Models_And_Ask_Me;
+	model: string;
 	systemPrompt: string;
 	outputVariableName: string;
 	promptTemplate: {

diff --git a/src/types/macros/QuickCommands/IAIAssistantCommand.ts b/src/types/macros/QuickCommands/IAIAssistantCommand.ts
@@ -1,4 +1,3 @@
-import type { Model, Models_And_Ask_Me } from "src/ai/models";
 import type { ICommand } from "../ICommand";
 import type { OpenAIModelParameters } from "src/ai/OpenAIModelParameters";
 
@@ -10,15 +9,15 @@ interface IBaseAIAssistantCommand extends ICommand {
 }
 
 export interface IAIAssistantCommand extends IBaseAIAssistantCommand {
-	model: Models_And_Ask_Me;
+	model: string;
 	promptTemplate: {
 		enable: boolean;
 		name: string;
 	};
 }
 
 export interface IInfiniteAIAssistantCommand extends IBaseAIAssistantCommand {
-	model: Model;
+	model: string;
 	resultJoiner: string;
 	chunkSeparator: string;
 	maxChunkTokens: number;