modelMatches allows specifying modelProvider, a single item or an array.

At least one must match. Part of #5.
jkomoros · Jan 6, 2024 · 8f5b37e · 8f5b37e
1 parent 085017c
commit 8f5b37e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/src/llm.ts b/src/llm.ts
@@ -187,6 +187,13 @@ const modelMatches = (model : CompletionModelID, opts : PromptOptions = {}) : bo
 			if (contextSizeAtLeast < 0) continue;
 			if (modelInfo.maxTokens < contextSizeAtLeast) return false;
 			break;
+		case 'modelProvider':
+			let modelProvider = requirements.modelProvider;
+			if (!modelProvider) continue;
+			if (!Array.isArray(modelProvider)) modelProvider = [modelProvider];
+			const [provider] = extractModel(model);
+			if (!modelProvider.includes(provider)) return false;
+			break;
 		default:
 			assertUnreachable(key);
 		}

diff --git a/src/types.ts b/src/types.ts
@@ -160,7 +160,8 @@ export type PromptOptions = {
 	modelRequirements? : {
 		jsonResponse? : boolean,
 		imageInput?: boolean
-		contextSizeAtLeast? : number
+		contextSizeAtLeast? : number,
+		modelProvider? : ModelProvider | ModelProvider[]
 	},
 	debugLogger? : Logger
 };