fix: more fixes related to model mapping

Author: dosco

src/ax/ai/anthropic/api.ts

@@ -5,6 +5,7 @@ import type {
   AxChatRequest,
   AxChatResponse,
   AxChatResponseResult,
+  AxInternalChatRequest,
   AxModelConfig
 } from '../types.js';
 
@@ -88,9 +89,9 @@ export class AxAIAnthropic extends AxBaseAI<
   }
 
   override generateChatReq = (
-    req: Readonly<AxChatRequest>
+    req: Readonly<AxInternalChatRequest>
   ): [API, AxAIAnthropicChatRequest] => {
-    const model = this.config.model;
+    const model = req.model;
 
     const apiConfig = {
       name: '/messages'

src/ax/ai/base.ts

@@ -15,6 +15,8 @@ import type {
   AxChatResponse,
   AxEmbedRequest,
   AxEmbedResponse,
+  AxInternalChatRequest,
+  AxInternalEmbedRequest,
   AxModelConfig,
   AxModelInfo,
   AxModelInfoWithProvider,
@@ -65,10 +67,12 @@ export class AxBaseAI<
 > implements AxAIService
 {
   generateChatReq?: (
-    req: Readonly<AxChatRequest>,
+    req: Readonly<AxInternalChatRequest>,
     config: Readonly<AxAIPromptConfig>
   ) => [API, TChatRequest];
-  generateEmbedReq?: (req: Readonly<AxEmbedRequest>) => [API, TEmbedRequest];
+  generateEmbedReq?: (
+    req: Readonly<AxInternalEmbedRequest>
+  ) => [API, TEmbedRequest];
   generateChatResp?: (resp: Readonly<TChatResponse>) => AxChatResponse;
   generateChatStreamResp?: (
     resp: Readonly<TChatResponseDelta>,
@@ -86,10 +90,10 @@ export class AxBaseAI<
   private modelInfo: readonly AxModelInfo[];
   private modelUsage?: AxTokenUsage;
   private embedModelUsage?: AxTokenUsage;
+  private models: AxBaseAIArgs['models'];
 
   protected apiURL: string;
   protected name: string;
-  protected models: AxBaseAIArgs['models'];
   protected headers: Record<string, string>;
   protected supportFor: AxBaseAIFeatures;
 
@@ -266,7 +270,7 @@ export class AxBaseAI<
       model,
       functions,
       modelConfig: { ...chatReq.modelConfig, stream }
-    } as Readonly<AxChatRequest>;
+    };
 
     const fn = async () => {
       const [apiConfig, reqValue] = reqFn(req, options as AxAIPromptConfig);
@@ -406,7 +410,7 @@ export class AxBaseAI<
     const req = {
       ...embedReq,
       embedModel
-    } as Readonly<AxEmbedRequest>;
+    };
 
     const fn = async () => {
       const [apiConfig, reqValue] = this.generateEmbedReq!(req);

src/ax/ai/cohere/api.ts

@@ -9,8 +9,9 @@ import type {
   AxAIServiceOptions,
   AxChatRequest,
   AxChatResponse,
-  AxEmbedRequest,
   AxEmbedResponse,
+  AxInternalChatRequest,
+  AxInternalEmbedRequest,
   AxModelConfig
 } from '../types.js';
 
@@ -98,11 +99,11 @@ export class AxAICohere extends AxBaseAI<
   }
 
   override generateChatReq = (
-    req: Readonly<AxChatRequest>,
+    req: Readonly<AxInternalChatRequest>,
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     _config: Readonly<AxAIPromptConfig>
   ): [API, AxAICohereChatRequest] => {
-    const model = this.config.model;
+    const model = req.model;
 
     const lastChatMsg = req.chatPrompt.at(-1);
     const restOfChat = req.chatPrompt.slice(0, -1);
@@ -187,9 +188,9 @@ export class AxAICohere extends AxBaseAI<
   };
 
   override generateEmbedReq = (
-    req: Readonly<AxEmbedRequest>
+    req: Readonly<AxInternalEmbedRequest>
   ): [API, AxAICohereEmbedRequest] => {
-    const model = this.config.embedModel;
+    const model = req.embedModel;
 
     if (!model) {
       throw new Error('Embed model not set');

src/ax/ai/google-gemini/api.ts

@@ -6,11 +6,11 @@ import {
 } from '../base.js';
 import type {
   AxAIServiceOptions,
-  AxChatRequest,
   AxChatResponse,
   AxChatResponseResult,
-  AxEmbedRequest,
   AxEmbedResponse,
+  AxInternalChatRequest,
+  AxInternalEmbedRequest,
   AxModelConfig,
   AxTokenUsage
 } from '../types.js';
@@ -150,9 +150,9 @@ export class AxAIGoogleGemini extends AxBaseAI<
   }
 
   override generateChatReq = (
-    req: Readonly<AxChatRequest>
+    req: Readonly<AxInternalChatRequest>
   ): [API, AxAIGoogleGeminiChatRequest] => {
-    const model = this.config.model;
+    const model = req.model;
     const stream = req.modelConfig?.stream ?? this.config.stream;
 
     if (!req.chatPrompt || req.chatPrompt.length === 0) {
@@ -332,9 +332,9 @@ export class AxAIGoogleGemini extends AxBaseAI<
   };
 
   override generateEmbedReq = (
-    req: Readonly<AxEmbedRequest>
+    req: Readonly<AxInternalEmbedRequest>
   ): [API, AxAIGoogleGeminiBatchEmbedRequest] => {
-    const model = this.config.embedModel;
+    const model = req.embedModel;
 
     if (!model) {
       throw new Error('Embed model not set');

src/ax/ai/huggingface/api.ts

@@ -7,8 +7,8 @@ import {
 import type {
   AxAIPromptConfig,
   AxAIServiceOptions,
-  AxChatRequest,
   AxChatResponse,
+  AxInternalChatRequest,
   AxModelConfig
 } from '../types.js';
 
@@ -88,11 +88,11 @@ export class AxAIHuggingFace extends AxBaseAI<
   }
 
   override generateChatReq = (
-    req: Readonly<AxChatRequest>,
+    req: Readonly<AxInternalChatRequest>,
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     _config: Readonly<AxAIPromptConfig>
   ): [API, AxAIHuggingFaceRequest] => {
-    const model = this.config.model;
+    const model = req.model;
 
     const functionsList = req.functions
       ? `Functions:\n${JSON.stringify(req.functions, null, 2)}\n`

src/ax/ai/openai/api.ts

@@ -10,8 +10,9 @@ import type {
   AxChatRequest,
   AxChatResponse,
   AxChatResponseResult,
-  AxEmbedRequest,
   AxEmbedResponse,
+  AxInternalChatRequest,
+  AxInternalEmbedRequest,
   AxModelConfig,
   AxModelInfo
 } from '../types.js';
@@ -121,11 +122,11 @@ export class AxAIOpenAI extends AxBaseAI<
   }
 
   override generateChatReq = (
-    req: Readonly<AxChatRequest>,
+    req: Readonly<AxInternalChatRequest>,
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
     _config: Readonly<AxAIPromptConfig>
   ): [API, AxAIOpenAIChatRequest] => {
-    const model = this.config.model;
+    const model = req.model;
 
     if (!req.chatPrompt || req.chatPrompt.length === 0) {
       throw new Error('Chat prompt is empty');
@@ -182,9 +183,9 @@ export class AxAIOpenAI extends AxBaseAI<
   };
 
   override generateEmbedReq = (
-    req: Readonly<AxEmbedRequest>
+    req: Readonly<AxInternalEmbedRequest>
   ): [API, AxAIOpenAIEmbedRequest] => {
-    const model = this.config.embedModel;
+    const model = req.embedModel;
 
     if (!model) {
       throw new Error('Embed model not set');

src/ax/ai/types.ts

@@ -143,11 +143,17 @@ export type AxChatRequest = {
   model?: string;
 };
 
+export type AxInternalChatRequest = Omit<AxChatRequest, 'model'> &
+  Required<Pick<AxChatRequest, 'model'>>;
+
 export type AxEmbedRequest = {
   texts?: readonly string[];
   embedModel?: string;
 };
 
+export type AxInternalEmbedRequest = Omit<AxEmbedRequest, 'embedModel'> &
+  Required<Pick<AxEmbedRequest, 'embedModel'>>;
+
 export type AxRateLimiterFunction = <T = unknown>(
   reqFunc: () => Promise<T | ReadableStream<T>>,
   info: Readonly<{ modelUsage?: AxTokenUsage; embedModelUsage?: AxTokenUsage }>

src/examples/summarize.ts

@@ -2,14 +2,16 @@ import { AxAI, AxAIOpenAIModel, AxChainOfThought } from '@ax-llm/ax';
 
 const noteText = `The technological singularity—or simply the singularity[1]—is a hypothetical future point in time at which technological growth becomes uncontrollable and irreversible, resulting in unforeseeable changes to human civilization.[2][3] According to the most popular version of the singularity hypothesis, I.J. Good's intelligence explosion model, an upgradable intelligent agent will eventually enter a "runaway reaction" of self-improvement cycles, each new and more intelligent generation appearing more and more rapidly, causing an "explosion" in intelligence and resulting in a powerful superintelligence that qualitatively far surpasses all human intelligence.[4]`;
 
+// Example with OpenAI using custom labels in place of model names
 const ai = new AxAI({
   name: 'openai',
   apiKey: process.env.OPENAI_APIKEY as string,
-  config: { model: AxAIOpenAIModel.GPT35Turbo }
+  config: { model: 'model-a' },
+  modelMap: {
+    'model-a': AxAIOpenAIModel.GPT35Turbo
+  }
 });
 
-ai.setOptions({ debug: true });
-
 // const ai = new AxAI({ name: 'ollama', model: 'nous-hermes2' });
 
 const gen = new AxChainOfThought(