refactoring_v1 (#391)

src/utils/engine.ts

@@ -1,37 +1,65 @@
+import { getConfig, OCO_AI_PROVIDER_ENUM } from '../commands/config';
+import { AnthropicEngine } from '../engine/anthropic';
+import { AzureEngine } from '../engine/azure';
 import { AiEngine } from '../engine/Engine';
-import { OpenAi } from '../engine/openAi';
+import { FlowiseAi } from '../engine/flowise';
 import { Gemini } from '../engine/gemini';
-import { getConfig } from '../commands/config';
 import { OllamaAi } from '../engine/ollama';
-import { AnthropicAi } from '../engine/anthropic'
-import { TestAi } from '../engine/testAi';
-import { Azure } from '../engine/azure';
-import { FlowiseAi } from '../engine/flowise'
+import { OpenAiEngine } from '../engine/openAi';
+import { TestAi, TestMockType } from '../engine/testAi';
 
 export function getEngine(): AiEngine {
   const config = getConfig();
-  const provider = config?.OCO_AI_PROVIDER;
-  
-  if (provider?.startsWith('ollama')) {
-    const ollamaAi = new OllamaAi();
-    const model = provider.substring('ollama/'.length);
-    if (model) {
-      ollamaAi.setModel(model);
-      ollamaAi.setUrl(config?.OCO_OLLAMA_API_URL);
-    }
-    return ollamaAi;
-  } else if (provider == 'anthropic') {
-    return new AnthropicAi();
-  } else if (provider == 'test') {
-    return new TestAi();
-  } else if (provider == 'gemini') {
-    return new Gemini();  
-  } else if (provider == 'azure') {
-  	return new Azure();
-  } else if( provider == 'flowise'){
-    return new FlowiseAi();
+  const provider = config.OCO_AI_PROVIDER;
+
+  const DEFAULT_CONFIG = {
+    model: config.OCO_MODEL!,
+    maxTokensOutput: config.OCO_TOKENS_MAX_OUTPUT!,
+    maxTokensInput: config.OCO_TOKENS_MAX_INPUT!,
+    baseURL: config.OCO_OPENAI_BASE_PATH!
+  };
+
+  switch (provider) {
+    case OCO_AI_PROVIDER_ENUM.OLLAMA:
+      return new OllamaAi({
+        ...DEFAULT_CONFIG,
+        apiKey: '',
+        baseURL: config.OCO_OLLAMA_API_URL!
+      });
+
+    case OCO_AI_PROVIDER_ENUM.ANTHROPIC:
+      return new AnthropicEngine({
+        ...DEFAULT_CONFIG,
+        apiKey: config.OCO_ANTHROPIC_API_KEY!
+      });
+
+    case OCO_AI_PROVIDER_ENUM.TEST:
+      return new TestAi(config.OCO_TEST_MOCK_TYPE as TestMockType);
+
+    case OCO_AI_PROVIDER_ENUM.GEMINI:
+      return new Gemini({
+        ...DEFAULT_CONFIG,
+        apiKey: config.OCO_GEMINI_API_KEY!,
+        baseURL: config.OCO_GEMINI_BASE_PATH!
+      });
+
+    case OCO_AI_PROVIDER_ENUM.AZURE:
+      return new AzureEngine({
+        ...DEFAULT_CONFIG,
+        apiKey: config.OCO_AZURE_API_KEY!
+      });
+
+    case OCO_AI_PROVIDER_ENUM.FLOWISE:
+      return new FlowiseAi({
+        ...DEFAULT_CONFIG,
+        baseURL: config.OCO_FLOWISE_ENDPOINT || DEFAULT_CONFIG.baseURL,
+        apiKey: config.OCO_FLOWISE_API_KEY!
+      });
+
+    default:
+      return new OpenAiEngine({
+        ...DEFAULT_CONFIG,
+        apiKey: config.OCO_OPENAI_API_KEY!
+      });
   }
-  
-  //open ai gpt by default
-  return new OpenAi();
 }