@auto-engineer/ai-gateway 0.6.3 → 0.8.1

package/src/config.ts

@@ -2,6 +2,7 @@ import dotenv from 'dotenv';
 import { fileURLToPath } from 'url';
 import { dirname, resolve } from 'path';
 import createDebug from 'debug';
+import { CustomProviderConfig } from './constants';
 
 const debug = createDebug('ai-gateway:config');
 const debugEnv = createDebug('ai-gateway:config:env');
@@ -25,6 +26,7 @@ export interface AIConfig {
   xai?: {
     apiKey: string;
   };
+  custom?: CustomProviderConfig;
 }
 
 // Helper to log provider configuration
@@ -34,6 +36,34 @@ function logProviderConfig(providerName: string, apiKey: string | undefined): vo
   }
 }
 
+// Helper to build custom provider config
+function buildCustomProviderConfig(): CustomProviderConfig | undefined {
+  const name = process.env.CUSTOM_PROVIDER_NAME;
+  const baseUrl = process.env.CUSTOM_PROVIDER_BASE_URL;
+  const apiKey = process.env.CUSTOM_PROVIDER_API_KEY;
+  const defaultModel = process.env.CUSTOM_PROVIDER_DEFAULT_MODEL;
+
+  if (
+    name != null &&
+    name.length > 0 &&
+    baseUrl != null &&
+    baseUrl.length > 0 &&
+    apiKey != null &&
+    apiKey.length > 0 &&
+    defaultModel != null &&
+    defaultModel.length > 0
+  ) {
+    return {
+      name,
+      baseUrl,
+      apiKey,
+      defaultModel,
+    };
+  }
+
+  return undefined;
+}
+
 // Helper to build provider config
 function buildProviderConfig(): AIConfig {
   return {
@@ -41,6 +71,7 @@ function buildProviderConfig(): AIConfig {
     anthropic: process.env.ANTHROPIC_API_KEY != null ? { apiKey: process.env.ANTHROPIC_API_KEY } : undefined,
     google: process.env.GEMINI_API_KEY != null ? { apiKey: process.env.GEMINI_API_KEY } : undefined,
     xai: process.env.XAI_API_KEY != null ? { apiKey: process.env.XAI_API_KEY } : undefined,
+    custom: buildCustomProviderConfig(),
   };
 }
 
@@ -54,6 +85,7 @@ export function configureAIProvider(): AIConfig {
   debugEnv('Anthropic configured: %s', config.anthropic != null);
   debugEnv('Google configured: %s', config.google != null);
   debugEnv('XAI configured: %s', config.xai != null);
+  debugEnv('Custom configured: %s', config.custom != null);
 
   // Log provider configurations
   logProviderConfig('OpenAI', config.openai?.apiKey);
@@ -61,12 +93,24 @@ export function configureAIProvider(): AIConfig {
   logProviderConfig('Google', config.google?.apiKey);
   logProviderConfig('XAI', config.xai?.apiKey);
 
-  const configuredProviders = [config.openai, config.anthropic, config.google, config.xai].filter((p) => p != null);
+  if (config.custom != null) {
+    debug(
+      'Custom provider configured: %s at %s with model %s',
+      config.custom.name,
+      config.custom.baseUrl,
+      config.custom.defaultModel,
+    );
+    logProviderConfig('Custom', config.custom.apiKey);
+  }
+
+  const configuredProviders = [config.openai, config.anthropic, config.google, config.xai, config.custom].filter(
+    (p) => p != null,
+  );
 
   if (configuredProviders.length === 0) {
     debug('ERROR: No AI providers configured');
     throw new Error(
-      'At least one AI provider must be configured. Please set OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY, or XAI_API_KEY environment variables.',
+      'At least one AI provider must be configured. Please set OPENAI_API_KEY, ANTHROPIC_API_KEY, GEMINI_API_KEY, XAI_API_KEY environment variables, or configure a custom provider with CUSTOM_PROVIDER_NAME, CUSTOM_PROVIDER_BASE_URL, CUSTOM_PROVIDER_API_KEY, and CUSTOM_PROVIDER_DEFAULT_MODEL.',
     );
   }
 

package/src/constants.ts

@@ -0,0 +1,21 @@
+export enum AIProvider {
+  OpenAI = 'openai',
+  Anthropic = 'anthropic',
+  Google = 'google',
+  XAI = 'xai',
+  Custom = 'custom',
+}
+
+export interface CustomProviderConfig {
+  name: string;
+  baseUrl: string;
+  apiKey: string;
+  defaultModel: string;
+}
+
+export const DEFAULT_MODELS = {
+  [AIProvider.OpenAI]: 'gpt-4o-mini',
+  [AIProvider.Anthropic]: 'claude-sonnet-4-20250514',
+  [AIProvider.Google]: 'gemini-2.5-pro',
+  [AIProvider.XAI]: 'grok-4',
+} as const;

package/src/index-custom.specs.ts

@@ -0,0 +1,192 @@
+import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
+import { AIProvider } from './constants';
+import { getAvailableProviders, getDefaultModel, getDefaultAIProvider } from './index';
+
+// Mock the config module
+vi.mock('./config', () => ({
+  configureAIProvider: vi.fn(),
+}));
+
+// Mock the custom provider
+vi.mock('./providers/custom', () => ({
+  createCustomProvider: vi.fn(() => ({
+    languageModel: vi.fn(),
+  })),
+}));
+
+describe('Index Integration with Custom Providers', () => {
+  const originalEnv = process.env;
+
+  beforeEach(() => {
+    vi.resetModules();
+    process.env = { ...originalEnv };
+    vi.clearAllMocks();
+  });
+
+  afterEach(() => {
+    process.env = originalEnv;
+  });
+
+  describe('getAvailableProviders', () => {
+    it('should include custom provider when configured', async () => {
+      const { configureAIProvider } = await import('./config');
+      vi.mocked(configureAIProvider).mockReturnValue({
+        custom: {
+          name: 'litellm',
+          baseUrl: 'https://api.litellm.ai',
+          apiKey: 'sk-key',
+          defaultModel: 'claude-3-sonnet',
+        },
+        anthropic: { apiKey: 'sk-anthropic' },
+      });
+
+      const providers = getAvailableProviders();
+
+      expect(providers).toContain(AIProvider.Custom);
+      expect(providers).toContain(AIProvider.Anthropic);
+    });
+
+    it('should not include custom provider when not configured', async () => {
+      const { configureAIProvider } = await import('./config');
+      vi.mocked(configureAIProvider).mockReturnValue({
+        anthropic: { apiKey: 'sk-anthropic' },
+      });
+
+      const providers = getAvailableProviders();
+
+      expect(providers).not.toContain(AIProvider.Custom);
+      expect(providers).toContain(AIProvider.Anthropic);
+    });
+
+    it('should handle multiple providers including custom', async () => {
+      const { configureAIProvider } = await import('./config');
+      vi.mocked(configureAIProvider).mockReturnValue({
+        openai: { apiKey: 'sk-openai' },
+        anthropic: { apiKey: 'sk-anthropic' },
+        google: { apiKey: 'sk-google' },
+        xai: { apiKey: 'sk-xai' },
+        custom: {
+          name: 'custom',
+          baseUrl: 'https://api.custom.com',
+          apiKey: 'sk-custom',
+          defaultModel: 'custom-model',
+        },
+      });
+
+      const providers = getAvailableProviders();
+
+      expect(providers).toHaveLength(5);
+      expect(providers).toContain(AIProvider.OpenAI);
+      expect(providers).toContain(AIProvider.Anthropic);
+      expect(providers).toContain(AIProvider.Google);
+      expect(providers).toContain(AIProvider.XAI);
+      expect(providers).toContain(AIProvider.Custom);
+    });
+  });
+
+  describe('getDefaultModel', () => {
+    it('should return custom provider default model when provider is Custom', async () => {
+      const { configureAIProvider } = await import('./config');
+      vi.mocked(configureAIProvider).mockReturnValue({
+        custom: {
+          name: 'litellm',
+          baseUrl: 'https://api.litellm.ai',
+          apiKey: 'sk-key',
+          defaultModel: 'claude-3-sonnet',
+        },
+      });
+
+      const model = getDefaultModel(AIProvider.Custom);
+      expect(model).toBe('claude-3-sonnet');
+    });
+
+    it('should throw error when custom provider not configured but requested', async () => {
+      const { configureAIProvider } = await import('./config');
+      vi.mocked(configureAIProvider).mockReturnValue({
+        anthropic: { apiKey: 'sk-anthropic' },
+      });
+
+      expect(() => getDefaultModel(AIProvider.Custom)).toThrow('Custom provider not configured');
+    });
+
+    it('should use environment variable DEFAULT_AI_MODEL for custom provider', async () => {
+      process.env.DEFAULT_AI_MODEL = 'env-override-model';
+
+      const { configureAIProvider } = await import('./config');
+      vi.mocked(configureAIProvider).mockReturnValue({
+        custom: {
+          name: 'custom',
+          baseUrl: 'https://api.custom.com',
+          apiKey: 'sk-key',
+          defaultModel: 'config-model',
+        },
+      });
+
+      const model = getDefaultModel(AIProvider.Custom);
+      expect(model).toBe('env-override-model');
+    });
+
+    it('should handle different custom provider models', async () => {
+      const testCases = ['gpt-4o', 'claude-3-opus', 'llama3.1:70b', 'mistral-large', 'gemini-1.5-pro'];
+
+      const { configureAIProvider } = await import('./config');
+
+      testCases.forEach((testModel) => {
+        vi.mocked(configureAIProvider).mockReturnValue({
+          custom: {
+            name: 'test',
+            baseUrl: 'https://api.test.com',
+            apiKey: 'sk-key',
+            defaultModel: testModel,
+          },
+        });
+
+        const model = getDefaultModel(AIProvider.Custom);
+        expect(model).toBe(testModel);
+      });
+    });
+  });
+
+  describe('getDefaultAIProvider', () => {
+    it('should use custom provider as fallback when available', async () => {
+      process.env.DEFAULT_AI_PROVIDER = 'nonexistent';
+
+      const { configureAIProvider } = await import('./config');
+      vi.mocked(configureAIProvider).mockReturnValue({
+        custom: {
+          name: 'custom',
+          baseUrl: 'https://api.custom.com',
+          apiKey: 'sk-key',
+          defaultModel: 'custom-model',
+        },
+      });
+
+      // Mock getAvailableProviders to return custom
+      const mockGetAvailableProviders = vi.fn(() => [AIProvider.Custom]);
+      vi.doMock('./index', async () => ({
+        ...(await vi.importActual('./index')),
+        getAvailableProviders: mockGetAvailableProviders,
+      }));
+
+      const { getDefaultAIProvider: getDefaultAIProviderMocked } = await import('./index');
+      const provider = getDefaultAIProviderMocked();
+
+      expect(provider).toBe(AIProvider.Custom);
+    });
+
+    it('should respect explicit DEFAULT_AI_PROVIDER environment variable', () => {
+      const testCases = [
+        { env: 'openai', expected: AIProvider.OpenAI },
+        { env: 'anthropic', expected: AIProvider.Anthropic },
+        { env: 'google', expected: AIProvider.Google },
+        { env: 'xai', expected: AIProvider.XAI },
+      ];
+
+      testCases.forEach(({ env, expected }) => {
+        process.env.DEFAULT_AI_PROVIDER = env;
+        const provider = getDefaultAIProvider();
+        expect(provider).toBe(expected);
+      });
+    });
+  });
+});

package/src/index.specs.ts

@@ -1,16 +1,205 @@
-import { describe, it, expect } from 'vitest';
-import { generateTextWithAI, streamTextWithAI, AIProvider } from './index';
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import { getAvailableProviders, getDefaultAIProvider, getDefaultModel } from './index';
+import { DEFAULT_MODELS, AIProvider } from './constants';
+import type { AIConfig } from './config';
 
-describe('AI Integration', () => {
-  process.env.OPENAI_API_KEY = 'test';
+// Mock the config module
+vi.mock('./config', () => ({
+  configureAIProvider: vi.fn(
+    () =>
+      ({
+        anthropic: { apiKey: 'test-anthropic-key' },
+        openai: { apiKey: 'test-openai-key' },
+        google: { apiKey: 'test-google-key' },
+        xai: { apiKey: 'test-xai-key' },
+      }) as AIConfig,
+  ),
+}));
 
-  it('should export the correct types', () => {
-    expect(typeof generateTextWithAI).toBe('function');
-    expect(typeof streamTextWithAI).toBe('function');
+describe('Provider Selection Logic', () => {
+  describe('getDefaultAIProvider', () => {
+    let originalProvider: string | undefined;
+
+    beforeEach(() => {
+      originalProvider = process.env.DEFAULT_AI_PROVIDER;
+    });
+
+    afterEach(() => {
+      if (originalProvider !== undefined) {
+        process.env.DEFAULT_AI_PROVIDER = originalProvider;
+      } else {
+        delete process.env.DEFAULT_AI_PROVIDER;
+      }
+    });
+
+    it('should respect DEFAULT_AI_PROVIDER environment variable when set to anthropic', () => {
+      process.env.DEFAULT_AI_PROVIDER = 'anthropic';
+
+      const provider = getDefaultAIProvider();
+
+      expect(provider).toBe(AIProvider.Anthropic);
+    });
+
+    it('should respect DEFAULT_AI_PROVIDER environment variable when set to xai', () => {
+      process.env.DEFAULT_AI_PROVIDER = 'xai';
+
+      const provider = getDefaultAIProvider();
+
+      expect(provider).toBe(AIProvider.XAI);
+    });
+
+    it('should respect DEFAULT_AI_PROVIDER environment variable when set to google', () => {
+      process.env.DEFAULT_AI_PROVIDER = 'google';
+
+      const provider = getDefaultAIProvider();
+
+      expect(provider).toBe(AIProvider.Google);
+    });
+
+    it('should be case insensitive for DEFAULT_AI_PROVIDER', () => {
+      process.env.DEFAULT_AI_PROVIDER = 'ANTHROPIC';
+
+      const provider = getDefaultAIProvider();
+
+      expect(provider).toBe(AIProvider.Anthropic);
+    });
+
+    it('should handle invalid DEFAULT_AI_PROVIDER values by falling back to available providers', () => {
+      process.env.DEFAULT_AI_PROVIDER = 'invalid_provider';
+
+      const provider = getDefaultAIProvider();
+
+      expect([AIProvider.Anthropic, AIProvider.OpenAI, AIProvider.Google, AIProvider.XAI]).toContain(provider);
+    });
+
+    it('should fallback to available providers when DEFAULT_AI_PROVIDER is not set', () => {
+      delete process.env.DEFAULT_AI_PROVIDER;
+
+      const provider = getDefaultAIProvider();
+
+      expect([AIProvider.Anthropic, AIProvider.OpenAI, AIProvider.Google, AIProvider.XAI]).toContain(provider);
+    });
+  });
+
+  describe('getAvailableProviders', () => {
+    it('should return array of available providers', () => {
+      const providers = getAvailableProviders();
+
+      expect(Array.isArray(providers)).toBe(true);
+      expect(providers.length).toBeGreaterThan(0);
+
+      providers.forEach((provider) => {
+        expect([AIProvider.Anthropic, AIProvider.OpenAI, AIProvider.Google, AIProvider.XAI]).toContain(provider);
+      });
+    });
+
+    it('should maintain priority order when multiple providers are available', () => {
+      const providers = getAvailableProviders();
+      const expectedOrder = [AIProvider.Anthropic, AIProvider.OpenAI, AIProvider.Google, AIProvider.XAI];
+
+      let lastIndex = -1;
+      providers.forEach((provider) => {
+        const currentIndex = expectedOrder.indexOf(provider);
+        expect(currentIndex).toBeGreaterThan(lastIndex);
+        lastIndex = currentIndex;
+      });
+    });
+  });
+
+  describe('getDefaultModel', () => {
+    let originalModel: string | undefined;
+
+    beforeEach(() => {
+      originalModel = process.env.DEFAULT_AI_MODEL;
+    });
+
+    afterEach(() => {
+      if (originalModel !== undefined) {
+        process.env.DEFAULT_AI_MODEL = originalModel;
+      } else {
+        delete process.env.DEFAULT_AI_MODEL;
+      }
+    });
+
+    it('should use DEFAULT_AI_MODEL when set in environment', () => {
+      process.env.DEFAULT_AI_MODEL = 'custom-model-name';
+
+      const model = getDefaultModel(AIProvider.OpenAI);
+
+      expect(model).toBe('custom-model-name');
+    });
+
+    it('should use DEFAULT_AI_MODEL regardless of provider when set', () => {
+      process.env.DEFAULT_AI_MODEL = 'universal-model';
+
+      expect(getDefaultModel(AIProvider.OpenAI)).toBe('universal-model');
+      expect(getDefaultModel(AIProvider.Anthropic)).toBe('universal-model');
+      expect(getDefaultModel(AIProvider.Google)).toBe('universal-model');
+      expect(getDefaultModel(AIProvider.XAI)).toBe('universal-model');
+    });
+
+    it('should trim whitespace from DEFAULT_AI_MODEL', () => {
+      process.env.DEFAULT_AI_MODEL = '  model-with-spaces  ';
+
+      const model = getDefaultModel(AIProvider.OpenAI);
+
+      expect(model).toBe('model-with-spaces');
+    });
+
+    it('should fallback to provider-specific defaults when DEFAULT_AI_MODEL is not set', () => {
+      delete process.env.DEFAULT_AI_MODEL;
+
+      expect(getDefaultModel(AIProvider.OpenAI)).toBe(DEFAULT_MODELS[AIProvider.OpenAI]);
+      expect(getDefaultModel(AIProvider.Anthropic)).toBe(DEFAULT_MODELS[AIProvider.Anthropic]);
+      expect(getDefaultModel(AIProvider.Google)).toBe(DEFAULT_MODELS[AIProvider.Google]);
+      expect(getDefaultModel(AIProvider.XAI)).toBe(DEFAULT_MODELS[AIProvider.XAI]);
+    });
+
+    it('should fallback to provider-specific defaults when DEFAULT_AI_MODEL is empty', () => {
+      process.env.DEFAULT_AI_MODEL = '';
+
+      expect(getDefaultModel(AIProvider.OpenAI)).toBe(DEFAULT_MODELS[AIProvider.OpenAI]);
+      expect(getDefaultModel(AIProvider.Anthropic)).toBe(DEFAULT_MODELS[AIProvider.Anthropic]);
+      expect(getDefaultModel(AIProvider.Google)).toBe(DEFAULT_MODELS[AIProvider.Google]);
+      expect(getDefaultModel(AIProvider.XAI)).toBe(DEFAULT_MODELS[AIProvider.XAI]);
+    });
+
+    it('should fallback to provider-specific defaults when DEFAULT_AI_MODEL is only whitespace', () => {
+      process.env.DEFAULT_AI_MODEL = '   ';
+
+      expect(getDefaultModel(AIProvider.OpenAI)).toBe(DEFAULT_MODELS[AIProvider.OpenAI]);
+      expect(getDefaultModel(AIProvider.Anthropic)).toBe(DEFAULT_MODELS[AIProvider.Anthropic]);
+      expect(getDefaultModel(AIProvider.Google)).toBe(DEFAULT_MODELS[AIProvider.Google]);
+      expect(getDefaultModel(AIProvider.XAI)).toBe(DEFAULT_MODELS[AIProvider.XAI]);
+    });
   });
 
-  it('should have valid AIProvider type', () => {
-    const providers: AIProvider[] = [AIProvider.OpenAI, AIProvider.Anthropic, AIProvider.Google, AIProvider.XAI];
-    expect(providers).toContain(AIProvider.OpenAI);
+  describe('Integration: Provider Selection with Priority', () => {
+    it('should demonstrate overrides when keys are not set', () => {
+      const availableProviders = getAvailableProviders();
+
+      if (availableProviders.includes(AIProvider.XAI)) {
+        expect(availableProviders).toContain(AIProvider.XAI);
+      }
+    });
+
+    it('should demonstrate priority order logic', () => {
+      const availableProviders = getAvailableProviders();
+      const defaultProvider = getDefaultAIProvider();
+      const priorityOrder = [AIProvider.Anthropic, AIProvider.OpenAI, AIProvider.Google, AIProvider.XAI];
+
+      let expectedProvider = AIProvider.OpenAI;
+      for (const provider of priorityOrder) {
+        if (availableProviders.includes(provider)) {
+          expectedProvider = provider;
+          break;
+        }
+      }
+
+      const envProvider = process.env.DEFAULT_AI_PROVIDER;
+      if (envProvider === undefined || envProvider === null || envProvider.toLowerCase() === 'invalid_provider') {
+        expect(defaultProvider).toBe(expectedProvider);
+      }
+    });
   });
 });