feat(engine): add support for MLX AI provider (#437)

* docs(CONTRIBUTING.md): update `TODO.md` reference (#435)

Signed-off-by: Emmanuel Ferdman <emmanuelferdman@gmail.com>

* feat(engine): add support for MLX AI provider
docs/engine: update documentation to include new engine providers

* fix(mlx.ts): add repetition_penalty option to generateCommitMessage method for improved model behavior

---------

Signed-off-by: Emmanuel Ferdman <emmanuelferdman@gmail.com>
Co-authored-by: Emmanuel Ferdman <emmanuelferdman@gmail.com>

src/commands/config.ts

@@ -93,6 +93,8 @@ const getDefaultModel = (provider: string | undefined): string => {
   switch (provider) {
     case 'ollama':
       return '';
+    case 'mlx':
+      return '';
     case 'anthropic':
       return MODEL_LIST.anthropic[0];
     case 'gemini':
@@ -138,7 +140,7 @@ export const configValidators = {
     validateConfig(
       'OCO_API_KEY',
       value,
-      'You need to provide the OCO_API_KEY when OCO_AI_PROVIDER set to "openai" (default) or "ollama" or "azure" or "gemini" or "flowise" or "anthropic". Run `oco config set OCO_API_KEY=your_key OCO_AI_PROVIDER=openai`'
+      'You need to provide the OCO_API_KEY when OCO_AI_PROVIDER set to "openai" (default) or "ollama" or "mlx" or "azure" or "gemini" or "flowise" or "anthropic". Run `oco config set OCO_API_KEY=your_key OCO_AI_PROVIDER=openai`'
     );
 
     return value;
@@ -261,8 +263,8 @@ export const configValidators = {
         'test',
         'flowise',
         'groq'
-      ].includes(value) || value.startsWith('ollama'),
-      `${value} is not supported yet, use 'ollama', 'anthropic', 'azure', 'gemini', 'flowise' or 'openai' (default)`
+      ].includes(value) || value.startsWith('ollama') || value.startsWith('mlx'),
+      `${value} is not supported yet, use 'ollama', 'mlx', anthropic', 'azure', 'gemini', 'flowise' or 'openai' (default)`
     );
 
     return value;
@@ -307,7 +309,8 @@ export enum OCO_AI_PROVIDER_ENUM {
   AZURE = 'azure',
   TEST = 'test',
   FLOWISE = 'flowise',
-  GROQ = 'groq'
+  GROQ = 'groq',
+  MLX = 'mlx'
 }
 
 export type ConfigType = {

src/engine/mlx.ts

@@ -0,0 +1,47 @@
+import axios, { AxiosInstance } from 'axios';
+import { OpenAI } from 'openai';
+import { AiEngine, AiEngineConfig } from './Engine';
+import { chown } from 'fs';
+
+interface MLXConfig extends AiEngineConfig {}
+
+export class MLXEngine implements AiEngine {
+    config: MLXConfig;
+    client: AxiosInstance;
+
+    constructor(config) {
+        this.config = config;
+        this.client = axios.create({
+            url: config.baseURL
+            ? `${config.baseURL}/${config.apiKey}`
+            : 'http://localhost:8080/v1/chat/completions',
+            headers: { 'Content-Type': 'application/json' }
+        });
+    }
+
+    async generateCommitMessage(
+        messages: Array<OpenAI.Chat.Completions.ChatCompletionMessageParam>): 
+        Promise<string | undefined> {
+            const params = {
+                messages,
+                temperature: 0, 
+                top_p: 0.1,
+                repetition_penalty: 1.5,
+                stream: false
+            };
+            try {
+                const response = await this.client.post(
+                    this.client.getUri(this.config),
+                    params
+                );
+                
+                const choices = response.data.choices;
+                const message = choices[0].message;
+                
+                return message?.content;
+            } catch (err: any) {
+                const message = err.response?.data?.error ?? err.message;
+                throw new Error(`MLX provider error: ${message}`);
+            }
+        }
+}

src/utils/engine.ts

@@ -8,6 +8,7 @@ import { OllamaEngine } from '../engine/ollama';
 import { OpenAiEngine } from '../engine/openAi';
 import { TestAi, TestMockType } from '../engine/testAi';
 import { GroqEngine } from '../engine/groq';
+import { MLXEngine } from '../engine/mlx';
 
 export function getEngine(): AiEngine {
   const config = getConfig();
@@ -43,6 +44,9 @@ export function getEngine(): AiEngine {
     case OCO_AI_PROVIDER_ENUM.GROQ:
       return new GroqEngine(DEFAULT_CONFIG);
 
+    case OCO_AI_PROVIDER_ENUM.MLX:
+      return new MLXEngine(DEFAULT_CONFIG);
+
     default:
       return new OpenAiEngine(DEFAULT_CONFIG);
   }