fix: redesigned model map feature

Author: dosco

src/ax/ai/anthropic/api.ts

@@ -34,11 +34,10 @@ export interface AxAIAnthropicArgs {
   apiKey: string;
   config?: Readonly<Partial<AxAIAnthropicConfig>>;
   options?: Readonly<AxAIServiceOptions>;
+  modelMap?: Record<string, AxAIAnthropicModel | string>;
 }
 
 export class AxAIAnthropic extends AxBaseAI<
-  AxAIAnthropicModel,
-  string,
   AxAIAnthropicChatRequest,
   unknown,
   AxAIAnthropicChatResponse,
@@ -50,7 +49,8 @@ export class AxAIAnthropic extends AxBaseAI<
   constructor({
     apiKey,
     config,
-    options
+    options,
+    modelMap
   }: Readonly<Omit<AxAIAnthropicArgs, 'name'>>) {
     if (!apiKey || apiKey === '') {
       throw new Error('Anthropic API key not set');
@@ -69,7 +69,8 @@ export class AxAIAnthropic extends AxBaseAI<
       modelInfo: axModelInfoAnthropic,
       models: { model: _config.model },
       options,
-      supportFor: { functions: true, streaming: true }
+      supportFor: { functions: true, streaming: true },
+      modelMap
     });
 
     this.config = _config;

src/ax/ai/azure-openai/api.ts

@@ -5,7 +5,11 @@ import {
   axAIOpenAIDefaultConfig,
   axAIOpenAIFastConfig
 } from '../openai/api.js';
-import type { AxAIOpenAIConfig } from '../openai/types.js';
+import type {
+  AxAIOpenAIConfig,
+  AxAIOpenAIEmbedModel,
+  AxAIOpenAIModel
+} from '../openai/types.js';
 import type { AxAIServiceOptions } from '../types.js';
 
 export const axAIAzureOpenAIDefaultConfig = axAIOpenAIDefaultConfig;
@@ -24,6 +28,7 @@ export interface AxAIAzureOpenAIArgs {
   version?: string;
   config?: Readonly<Partial<AxAIOpenAIConfig>>;
   options?: Readonly<AxAIServiceOptions>;
+  modelMap?: Record<string, AxAIOpenAIModel | AxAIOpenAIEmbedModel>;
 }
 
 export class AxAIAzureOpenAI extends AxAIOpenAI {
@@ -33,7 +38,8 @@ export class AxAIAzureOpenAI extends AxAIOpenAI {
     deploymentName,
     version = 'api-version=2024-02-15-preview',
     config,
-    options
+    options,
+    modelMap
   }: Readonly<Omit<AxAIAzureOpenAIArgs, 'name'>>) {
     if (!apiKey || apiKey === '') {
       throw new Error('Azure OpenAPI API key not set');
@@ -48,7 +54,7 @@ export class AxAIAzureOpenAI extends AxAIOpenAI {
       ...axAIAzureOpenAIDefaultConfig(),
       ...config
     };
-    super({ apiKey, config: _config, options });
+    super({ apiKey, config: _config, options, modelMap });
 
     const host = resourceName.includes('://')
       ? resourceName

src/ax/ai/balance.ts

@@ -1,7 +1,7 @@
 import type { ReadableStream } from 'stream/web';
 
-import type { AxAIModelMap } from './base.js';
 import type {
+  AxAIModelMap,
   AxAIPromptConfig,
   AxAIService,
   AxAIServiceActionOptions,
@@ -11,18 +11,20 @@ import type {
   AxEmbedRequest,
   AxEmbedResponse,
   AxModelConfig,
-  AxModelInfo
+  AxModelInfoWithProvider
 } from './types.js';
 
 export class AxBalancer implements AxAIService {
   private services: AxAIService[];
   private currentServiceIndex: number = 0;
   private currentService: AxAIService;
+  private serviceMap: Record<string, AxAIService> = {};
 
   constructor(services: readonly AxAIService[]) {
     if (services.length === 0) {
       throw new Error('No AI services provided.');
     }
+
     this.services = [...services].sort((a, b) => {
       const aInfo = a.getModelInfo();
       const bInfo = b.getModelInfo();
@@ -42,12 +44,8 @@ export class AxBalancer implements AxAIService {
     this.currentService = cs;
   }
 
-  setModelMap(modelMap: AxAIModelMap<string>): void {
-    this.currentService.setModelMap(modelMap);
-  }
-
-  setEmbedModelMap(embedModelMap: AxAIModelMap<string>): void {
-    this.currentService.setEmbedModelMap(embedModelMap);
+  getModelMap(): AxAIModelMap | undefined {
+    throw new Error('Method not implemented.');
   }
 
   private getNextService(): boolean {
@@ -72,11 +70,11 @@ export class AxBalancer implements AxAIService {
     return this.currentService.getName();
   }
 
-  getModelInfo(): Readonly<AxModelInfo & { provider: string }> {
+  getModelInfo(): Readonly<AxModelInfoWithProvider> {
     return this.currentService.getModelInfo();
   }
 
-  getEmbedModelInfo(): Readonly<AxModelInfo> | undefined {
+  getEmbedModelInfo(): Readonly<AxModelInfoWithProvider> | undefined {
     return this.currentService.getEmbedModelInfo();
   }