From 9a0f412fff320048b30882841ce846c1c040c4e8 Mon Sep 17 00:00:00 2001
From: Kai Xu <me@xuk.ai>
Date: Sun, 12 May 2024 06:56:04 -0600
Subject: [PATCH] fix(config.ts): add optional `config` parameter to
 `[CONFIG_KEYS.OCO_ MODEL]` validator to allow for dynamic model selection
 based on configuration (#337)

feat(engine/ollama.ts): integrate with config command to load OCO_MODEL from configuration and use it as the default AI engine model
---
 src/commands/config.ts | 4 ++--
 src/engine/ollama.ts   | 8 +++++++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/commands/config.ts b/src/commands/config.ts
index d877db3..d9ce148 100644
--- a/src/commands/config.ts
+++ b/src/commands/config.ts
@@ -180,10 +180,10 @@ export const configValidators = {
     return value;
   },
 
-  [CONFIG_KEYS.OCO_MODEL](value: any) {
+  [CONFIG_KEYS.OCO_MODEL](value: any, config: any = {}) {
     validateConfig(
       CONFIG_KEYS.OCO_MODEL,
-      [...MODEL_LIST.openai, ...MODEL_LIST.anthropic].includes(value),
+      [...MODEL_LIST.openai, ...MODEL_LIST.anthropic].includes(value) || config.OCO_AI_PROVIDER == 'ollama' || config.OCO_AI_PROVIDER == 'test',
       `${value} is not supported yet, use 'gpt-4', 'gpt-4-turbo', 'gpt-3.5-turbo' (default), 'gpt-3.5-turbo-0125', 'gpt-4-1106-preview', 'gpt-4-turbo-preview', 'gpt-4-0125-preview', 'claude-3-opus-20240229', 'claude-3-sonnet-20240229' or 'claude-3-haiku-20240307'`
     );
     return value;
diff --git a/src/engine/ollama.ts b/src/engine/ollama.ts
index f431ee1..802068a 100644
--- a/src/engine/ollama.ts
+++ b/src/engine/ollama.ts
@@ -2,11 +2,17 @@ import axios, { AxiosError } from 'axios';
 import { ChatCompletionRequestMessage } from 'openai';
 import { AiEngine } from './Engine';
 
+import {
+  getConfig
+} from '../commands/config';
+
+const config = getConfig();
+
 export class OllamaAi implements AiEngine {
   async generateCommitMessage(
     messages: Array<ChatCompletionRequestMessage>
   ): Promise<string | undefined> {
-    const model = 'mistral'; // todo: allow other models
+    const model = config?.OCO_MODEL || 'mistral';
 
     //console.log(messages);
     //process.exit()