@auto-engineer/ai-gateway 0.7.0 → 0.8.2

package/src/index.ts

@@ -4,6 +4,10 @@ import { anthropic } from '@ai-sdk/anthropic';
 import { google } from '@ai-sdk/google';
 import { xai } from '@ai-sdk/xai';
 import { configureAIProvider } from './config';
+import { DEFAULT_MODELS, AIProvider } from './constants';
+import { createCustomProvider } from './providers/custom';
+
+export { AIProvider } from './constants';
 import { z } from 'zod';
 import { getRegisteredToolsForAI } from './mcp-server';
 import { startServer } from './mcp-server';
@@ -137,14 +141,8 @@ interface AIToolValidationError extends Error {
   [key: string]: unknown;
 }
 
-export enum AIProvider {
-  OpenAI = 'openai',
-  Anthropic = 'anthropic',
-  Google = 'google',
-  XAI = 'xai',
-}
-
 export interface AIOptions {
+  provider?: AIProvider;
   model?: string;
   temperature?: number;
   maxTokens?: number;
@@ -167,22 +165,47 @@ const defaultOptions: AIOptions = {
   maxTokens: 1000,
 };
 
-function getDefaultModel(provider: AIProvider): string {
-  const model = (() => {
-    switch (provider) {
-      case AIProvider.OpenAI:
-        return 'gpt-4o-mini'; // maybe 5
-      case AIProvider.Anthropic:
-        return 'claude-sonnet-4-20250514'; // 4
-      case AIProvider.Google:
-        return 'gemini-2.5-pro';
-      case AIProvider.XAI:
-        return 'grok-4';
-      default:
-        // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
-        throw new Error(`Unknown provider: ${provider}`);
+export function getDefaultAIProvider(): AIProvider {
+  const envProvider = process.env.DEFAULT_AI_PROVIDER?.toLowerCase();
+  switch (envProvider) {
+    case 'openai':
+      return AIProvider.OpenAI;
+    case 'anthropic':
+      return AIProvider.Anthropic;
+    case 'google':
+      return AIProvider.Google;
+    case 'xai':
+      return AIProvider.XAI;
+    default: {
+      // Fallback to the first available provider
+      const available = getAvailableProviders();
+      return available.length > 0 ? available[0] : AIProvider.Anthropic;
+    }
+  }
+}
+
+export function getDefaultModel(provider: AIProvider): string {
+  // Check if DEFAULT_AI_MODEL is set in environment
+  const envModel = process.env.DEFAULT_AI_MODEL;
+  if (envModel !== undefined && envModel !== null && envModel.trim().length > 0) {
+    debugConfig('Using DEFAULT_AI_MODEL from environment: %s for provider %s', envModel, provider);
+    return envModel.trim();
+  }
+
+  if (provider === AIProvider.Custom) {
+    const config = configureAIProvider();
+    if (config.custom != null) {
+      debugConfig('Selected custom provider default model %s for provider %s', config.custom.defaultModel, provider);
+      return config.custom.defaultModel;
     }
-  })();
+    throw new Error('Custom provider not configured');
+  }
+
+  const model =
+    DEFAULT_MODELS[provider] ??
+    (() => {
+      throw new Error(`Unknown provider: ${provider}`);
+    })();
   debugConfig('Selected default model %s for provider %s', model, provider);
   return model;
 }
@@ -200,30 +223,35 @@ function getModel(provider: AIProvider, model?: string) {
       return google(modelName);
     case AIProvider.XAI:
       return xai(modelName);
+    case AIProvider.Custom: {
+      const config = configureAIProvider();
+      if (config.custom == null) {
+        throw new Error('Custom provider not configured');
+      }
+      const customProvider = createCustomProvider(config.custom);
+      return customProvider.languageModel(modelName);
+    }
     default:
       throw new Error(`Unknown provider: ${provider as string}`);
   }
 }
 
-export async function generateTextWithAI(
-  prompt: string,
-  provider: AIProvider,
-  options: AIOptions = {},
-): Promise<string> {
-  debugAPI('generateTextWithAI called - provider: %s, promptLength: %d', provider, prompt.length);
+export async function generateTextWithAI(prompt: string, options: AIOptions = {}): Promise<string> {
+  const resolvedProvider = options.provider ?? getDefaultAIProvider();
+  debugAPI('generateTextWithAI called - provider: %s, promptLength: %d', resolvedProvider, prompt.length);
   const finalOptions = { ...defaultOptions, ...options };
-  const model = finalOptions.model ?? getDefaultModel(provider);
-  const modelInstance = getModel(provider, model);
+  const model = finalOptions.model ?? getDefaultModel(resolvedProvider);
+  const modelInstance = getModel(resolvedProvider, model);
 
   if (finalOptions.includeTools === true) {
     debugTools('Tools requested, starting MCP server');
     await startServer();
-    const result = await generateTextWithToolsAI(prompt, provider, options);
+    const result = await generateTextWithToolsAI(prompt, options);
     return result.text;
   }
 
   try {
-    debugAPI('Making API call to %s with model %s', provider, model);
+    debugAPI('Making API call to %s with model %s', resolvedProvider, model);
     debugAPI('Request params - temperature: %d, maxTokens: %d', finalOptions.temperature, finalOptions.maxTokens);
 
     const result = await generateText({
@@ -236,22 +264,19 @@ export async function generateTextWithAI(
     debugAPI('API call successful - response length: %d, usage: %o', result.text.length, result.usage);
     return result.text;
   } catch (error) {
-    extractAndLogError(error, provider, 'generateTextWithAI');
+    extractAndLogError(error, resolvedProvider, 'generateTextWithAI');
     throw error;
   }
 }
 
-export async function* streamTextWithAI(
-  prompt: string,
-  provider: AIProvider,
-  options: AIOptions = {},
-): AsyncGenerator<string> {
-  debugStream('streamTextWithAI called - provider: %s, promptLength: %d', provider, prompt.length);
+export async function* streamTextWithAI(prompt: string, options: AIOptions = {}): AsyncGenerator<string> {
+  const resolvedProvider = options.provider ?? getDefaultAIProvider();
+  debugStream('streamTextWithAI called - provider: %s, promptLength: %d', resolvedProvider, prompt.length);
   const finalOptions = { ...defaultOptions, ...options };
-  const model = getModel(provider, finalOptions.model);
+  const model = getModel(resolvedProvider, finalOptions.model);
 
   try {
-    debugStream('Starting stream from %s', provider);
+    debugStream('Starting stream from %s', resolvedProvider);
     const stream = await streamText({
       model,
       prompt,
@@ -269,7 +294,7 @@ export async function* streamTextWithAI(
     }
     debugStream('Stream completed - total chunks: %d, total length: %d', totalChunks, totalLength);
   } catch (error) {
-    extractAndLogError(error, provider, 'streamTextWithAI');
+    extractAndLogError(error, resolvedProvider, 'streamTextWithAI');
     throw error;
   }
 }
@@ -279,16 +304,13 @@ export async function* streamTextWithAI(
  * Optionally calls a stream callback for each token if provided.
  * Always returns the complete collected response.
  */
-export async function generateTextStreamingWithAI(
-  prompt: string,
-  provider: AIProvider,
-  options: AIOptions = {},
-): Promise<string> {
-  debugStream('generateTextStreamingWithAI called - provider: %s', provider);
+export async function generateTextStreamingWithAI(prompt: string, options: AIOptions = {}): Promise<string> {
+  const resolvedProvider = options.provider ?? getDefaultAIProvider();
+  debugStream('generateTextStreamingWithAI called - provider: %s', resolvedProvider);
   const finalOptions = { ...defaultOptions, ...options };
   let collectedResult = '';
 
-  const stream = streamTextWithAI(prompt, provider, finalOptions);
+  const stream = streamTextWithAI(prompt, finalOptions);
 
   let tokenCount = 0;
   for await (const token of stream) {
@@ -386,13 +408,13 @@ async function executeToolConversation(
 
 export async function generateTextWithToolsAI(
   prompt: string,
-  provider: AIProvider,
   options: AIOptions = {},
 ): Promise<{ text: string; toolCalls?: unknown[] }> {
-  debugTools('generateTextWithToolsAI called - provider: %s', provider);
+  const resolvedProvider = options.provider ?? getDefaultAIProvider();
+  debugTools('generateTextWithToolsAI called - provider: %s', resolvedProvider);
   const finalOptions = { ...defaultOptions, ...options };
-  const model = finalOptions.model ?? getDefaultModel(provider);
-  const modelInstance = getModel(provider, model);
+  const model = finalOptions.model ?? getDefaultModel(resolvedProvider);
+  const modelInstance = getModel(resolvedProvider, model);
 
   const registeredTools = finalOptions.includeTools === true ? getRegisteredToolsForAI() : {};
   debugTools('Registered tools: %o', Object.keys(registeredTools));
@@ -408,7 +430,7 @@ export async function generateTextWithToolsAI(
     registeredTools,
     hasTools,
     finalOptions,
-    provider,
+    resolvedProvider,
   );
 
   return {
@@ -450,26 +472,26 @@ async function executeToolCalls(
 export async function generateTextWithImageAI(
   text: string,
   imageBase64: string,
-  provider: AIProvider,
   options: AIOptions = {},
 ): Promise<string> {
+  const resolvedProvider = options.provider ?? getDefaultAIProvider();
   debugAPI(
     'generateTextWithImageAI called - provider: %s, textLength: %d, imageSize: %d',
-    provider,
+    resolvedProvider,
     text.length,
     imageBase64.length,
   );
   const finalOptions = { ...defaultOptions, ...options };
-  const model = finalOptions.model ?? getDefaultModel(provider);
-  const modelInstance = getModel(provider, model);
+  const model = finalOptions.model ?? getDefaultModel(resolvedProvider);
+  const modelInstance = getModel(resolvedProvider, model);
 
-  if (provider !== AIProvider.OpenAI && provider !== AIProvider.XAI) {
-    debugError('Provider %s does not support image inputs', provider);
-    throw new Error(`Provider ${provider} does not support image inputs`);
+  if (resolvedProvider !== AIProvider.OpenAI && resolvedProvider !== AIProvider.XAI) {
+    debugError('Provider %s does not support image inputs', resolvedProvider);
+    throw new Error(`Provider ${resolvedProvider} does not support image inputs`);
   }
 
   try {
-    debugAPI('Sending image+text to %s', provider);
+    debugAPI('Sending image+text to %s', resolvedProvider);
     const result = await generateText({
       model: modelInstance,
       messages: [
@@ -488,7 +510,7 @@ export async function generateTextWithImageAI(
     debugAPI('Image API call successful - response length: %d', result.text.length);
     return result.text;
   } catch (error) {
-    extractAndLogError(error, provider, 'generateTextWithImageAI');
+    extractAndLogError(error, resolvedProvider, 'generateTextWithImageAI');
     throw error;
   }
 }
@@ -496,10 +518,11 @@ export async function generateTextWithImageAI(
 export function getAvailableProviders(): AIProvider[] {
   const config = configureAIProvider();
   const providers: AIProvider[] = [];
-  if (config.openai != null) providers.push(AIProvider.OpenAI);
   if (config.anthropic != null) providers.push(AIProvider.Anthropic);
+  if (config.openai != null) providers.push(AIProvider.OpenAI);
   if (config.google != null) providers.push(AIProvider.Google);
   if (config.xai != null) providers.push(AIProvider.XAI);
+  if (config.custom != null) providers.push(AIProvider.Custom);
   debugConfig('Available providers: %o', providers);
   return providers;
 }
@@ -626,12 +649,13 @@ async function attemptStructuredGeneration<T>(
   throw lastError;
 }
 
-export async function generateStructuredDataWithAI<T>(
-  prompt: string,
-  provider: AIProvider,
-  options: StructuredAIOptions<T>,
-): Promise<T> {
-  debugAPI('generateStructuredDataWithAI called - provider: %s, schema: %s', provider, options.schemaName ?? 'unnamed');
+export async function generateStructuredDataWithAI<T>(prompt: string, options: StructuredAIOptions<T>): Promise<T> {
+  const resolvedProvider = options.provider ?? getDefaultAIProvider();
+  debugAPI(
+    'generateStructuredDataWithAI called - provider: %s, schema: %s',
+    resolvedProvider,
+    options.schemaName ?? 'unnamed',
+  );
 
   if (options.includeTools === true) {
     debugTools('Tools requested, starting MCP server');
@@ -641,17 +665,14 @@ export async function generateStructuredDataWithAI<T>(
   debugTools('Registered tools for structured data: %o', Object.keys(registeredTools));
   const hasTools = Object.keys(registeredTools).length > 0;
 
-  return attemptStructuredGeneration(prompt, provider, options, registeredTools, hasTools);
+  return attemptStructuredGeneration(prompt, resolvedProvider, options, registeredTools, hasTools);
 }
 
-export async function streamStructuredDataWithAI<T>(
-  prompt: string,
-  provider: AIProvider,
-  options: StreamStructuredAIOptions<T>,
-): Promise<T> {
+export async function streamStructuredDataWithAI<T>(prompt: string, options: StreamStructuredAIOptions<T>): Promise<T> {
+  const resolvedProvider = options.provider ?? getDefaultAIProvider();
   debugStream(
     'streamStructuredDataWithAI called - provider: %s, schema: %s',
-    provider,
+    resolvedProvider,
     options.schemaName ?? 'unnamed',
   );
   const maxSchemaRetries = 3;
@@ -660,7 +681,7 @@ export async function streamStructuredDataWithAI<T>(
   for (let attempt = 0; attempt < maxSchemaRetries; attempt++) {
     try {
       debugValidation('Stream structured data attempt %d/%d', attempt + 1, maxSchemaRetries);
-      const model = getModel(provider, options.model);
+      const model = getModel(resolvedProvider, options.model);
 
       const enhancedPrompt = attempt > 0 && lastError ? getEnhancedPrompt(prompt, lastError) : prompt;
 

package/src/providers/custom.specs.ts

@@ -0,0 +1,161 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import { createCustomProvider } from './custom';
+import { CustomProviderConfig } from '../constants';
+
+interface MockConfig {
+  name: string;
+  baseURL: string;
+  apiKey: string;
+}
+
+// Mock the createOpenAI function
+vi.mock('@ai-sdk/openai', () => ({
+  createOpenAI: vi.fn((config: MockConfig) => ({
+    languageModel: vi.fn((modelId: string) => ({
+      modelId,
+      provider: config.name,
+      specificationVersion: 'v1' as const,
+      defaultObjectGenerationMode: 'json' as const,
+    })),
+  })),
+}));
+
+describe('Custom Provider', () => {
+  const mockConfig: CustomProviderConfig = {
+    name: 'test-provider',
+    baseUrl: 'https://api.example.com/v1',
+    apiKey: 'test-api-key',
+    defaultModel: 'test-model',
+  };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe('createCustomProvider', () => {
+    it('should create a custom provider with the correct configuration', () => {
+      const provider = createCustomProvider(mockConfig);
+
+      expect(provider).toBeDefined();
+      // Provider creation should succeed with valid config
+    });
+
+    it('should call createOpenAI with the correct parameters', async () => {
+      const { createOpenAI } = await import('@ai-sdk/openai');
+
+      createCustomProvider(mockConfig);
+
+      expect(createOpenAI).toHaveBeenCalledWith({
+        name: mockConfig.name,
+        baseURL: mockConfig.baseUrl,
+        apiKey: mockConfig.apiKey,
+      });
+    });
+
+    it('should create a provider that can create language models', () => {
+      const provider = createCustomProvider(mockConfig);
+      const model = provider.languageModel('test-model');
+
+      expect(model).toBeDefined();
+      expect(model.modelId).toBe('test-model');
+      expect(model.provider).toBe(mockConfig.name);
+    });
+
+    it('should handle different base URLs correctly', () => {
+      const configs = [
+        { ...mockConfig, baseUrl: 'https://api.litellm.ai' },
+        { ...mockConfig, baseUrl: 'http://localhost:8000' },
+        { ...mockConfig, baseUrl: 'https://custom-llm.company.com/api' },
+      ];
+
+      configs.forEach((config) => {
+        const provider = createCustomProvider(config);
+        expect(provider).toBeDefined();
+      });
+    });
+
+    it('should handle different provider names', () => {
+      const configs = [
+        { ...mockConfig, name: 'litellm' },
+        { ...mockConfig, name: 'local-llm' },
+        { ...mockConfig, name: 'company-custom-llm' },
+      ];
+
+      configs.forEach((config) => {
+        const provider = createCustomProvider(config);
+        expect(provider).toBeDefined();
+      });
+    });
+  });
+
+  describe('provider compatibility', () => {
+    it('should be compatible with OpenAI-style endpoints', () => {
+      const litellmConfig: CustomProviderConfig = {
+        name: 'litellm',
+        baseUrl: 'https://api.litellm.ai/chat/completions',
+        apiKey: 'sk-litellm-key',
+        defaultModel: 'claude-3-sonnet',
+      };
+
+      const provider = createCustomProvider(litellmConfig);
+      expect(provider).toBeDefined();
+    });
+
+    it('should work with localhost endpoints for development', () => {
+      const localConfig: CustomProviderConfig = {
+        name: 'local-dev',
+        baseUrl: 'http://localhost:8000',
+        apiKey: 'local-key',
+        defaultModel: 'local-model',
+      };
+
+      const provider = createCustomProvider(localConfig);
+      expect(provider).toBeDefined();
+    });
+  });
+
+  describe('error handling', () => {
+    it('should pass through any errors from createOpenAI', async () => {
+      const { createOpenAI } = await import('@ai-sdk/openai');
+      const mockError = new Error('Invalid API key');
+
+      vi.mocked(createOpenAI).mockImplementationOnce(() => {
+        throw mockError;
+      });
+
+      expect(() => createCustomProvider(mockConfig)).toThrow('Invalid API key');
+    });
+  });
+
+  describe('configuration validation', () => {
+    it('should handle minimal required configuration', () => {
+      const minimalConfig: CustomProviderConfig = {
+        name: 'minimal',
+        baseUrl: 'https://api.example.com',
+        apiKey: 'key',
+        defaultModel: 'model',
+      };
+
+      const provider = createCustomProvider(minimalConfig);
+      expect(provider).toBeDefined();
+    });
+
+    it('should preserve all configuration properties', async () => {
+      const fullConfig: CustomProviderConfig = {
+        name: 'full-config',
+        baseUrl: 'https://api.example.com/v1',
+        apiKey: 'sk-test-key-123',
+        defaultModel: 'gpt-4o',
+      };
+
+      createCustomProvider(fullConfig);
+
+      const { createOpenAI } = await import('@ai-sdk/openai');
+      expect(createOpenAI).toHaveBeenCalledWith({
+        name: fullConfig.name,
+        baseURL: fullConfig.baseUrl,
+        apiKey: fullConfig.apiKey,
+      });
+    });
+  });
+});

package/src/providers/custom.ts

@@ -0,0 +1,24 @@
+import { createOpenAI } from '@ai-sdk/openai';
+import { CustomProviderConfig } from '../constants';
+import createDebug from 'debug';
+
+const debug = createDebug('ai-gateway:custom');
+
+export interface CustomProviderOptions {
+  config: CustomProviderConfig;
+}
+
+export function createCustomProvider(config: CustomProviderConfig) {
+  debug('Creating custom provider: %s with baseUrl: %s', config.name, config.baseUrl);
+
+  // Use OpenAI's provider implementation but with custom baseUrl
+  // This leverages the existing, battle-tested OpenAI provider while allowing custom endpoints
+  const customProvider = createOpenAI({
+    name: config.name,
+    baseURL: config.baseUrl,
+    apiKey: config.apiKey,
+  });
+
+  debug('Custom provider created successfully: %s', config.name);
+  return customProvider;
+}