@ax-llm/ax 13.0.10 → 13.0.11

package/index.d.ts

@@ -73,18 +73,18 @@ declare class AxAIRefusalError extends Error {
     toString(): string;
 }
 
-type AxAIInputModelList<TModel, TEmbedModel> = (AxAIModelListBase & {
+type AxAIInputModelList<TModel, TEmbedModel, TModelKey> = (AxAIModelListBase<TModelKey> & {
     isInternal?: boolean;
 } & ({
     model: TModel;
 } | {
     embedModel: TEmbedModel;
 }))[];
-type AxAIModelListBase = {
-    key: string;
+type AxAIModelListBase<TModelKey> = {
+    key: TModelKey;
     description: string;
 };
-type AxAIModelList = (AxAIModelListBase & ({
+type AxAIModelList<TModelKey> = (AxAIModelListBase<TModelKey> & ({
     model: string;
 } | {
     embedModel: string;
@@ -364,8 +364,8 @@ type AxAIServiceOptions = {
     abortSignal?: AbortSignal;
     logger?: AxLoggerFunction;
 };
-type AxAIServiceActionOptions<TModel = unknown, TEmbedModel = unknown> = {
-    ai?: Readonly<AxAIService<TModel, TEmbedModel>>;
+type AxAIServiceActionOptions<TModel = unknown, TEmbedModel = unknown, TModelKey = string> = {
+    ai?: Readonly<AxAIService<TModel, TEmbedModel, TModelKey>>;
     sessionId?: string;
     traceId?: string | undefined;
     timeout?: number;
@@ -378,18 +378,18 @@ type AxAIServiceActionOptions<TModel = unknown, TEmbedModel = unknown> = {
     stepIndex?: number;
     functionResultFormatter?: (result: unknown) => string;
 };
-interface AxAIService<TModel = unknown, TEmbedModel = unknown> {
+interface AxAIService<TModel = unknown, TEmbedModel = unknown, TModelKey = string> {
     getId(): string;
     getName(): string;
     getFeatures(model?: TModel): AxAIFeatures;
-    getModelList(): AxAIModelList | undefined;
+    getModelList(): AxAIModelList<TModelKey> | undefined;
     getMetrics(): AxAIServiceMetrics;
     getLogger(): AxLoggerFunction;
     getLastUsedChatModel(): TModel | undefined;
     getLastUsedEmbedModel(): TEmbedModel | undefined;
     getLastUsedModelConfig(): AxModelConfig | undefined;
-    chat(req: Readonly<AxChatRequest<TModel>>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<TModel, TEmbedModel>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
-    embed(req: Readonly<AxEmbedRequest<TEmbedModel>>, options?: Readonly<AxAIServiceActionOptions<TModel, TEmbedModel>>): Promise<AxEmbedResponse>;
+    chat(req: Readonly<AxChatRequest<TModel>>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<TModel, TEmbedModel, TModelKey>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
+    embed(req: Readonly<AxEmbedRequest<TEmbedModel>>, options?: Readonly<AxAIServiceActionOptions<TModel, TEmbedModel, TModelKey>>): Promise<AxEmbedResponse>;
     setOptions(options: Readonly<AxAIServiceOptions>): void;
     getOptions(): Readonly<AxAIServiceOptions>;
 }
@@ -410,7 +410,7 @@ interface AxAIFeatures {
     hasThinkingBudget?: boolean;
     hasShowThoughts?: boolean;
 }
-interface AxBaseAIArgs<TModel, TEmbedModel> {
+interface AxBaseAIArgs<TModel, TEmbedModel, TModelKey> {
     name: string;
     apiURL: string;
     headers: () => Promise<Record<string, string>>;
@@ -421,11 +421,11 @@ interface AxBaseAIArgs<TModel, TEmbedModel> {
     }>;
     options?: Readonly<AxAIServiceOptions>;
     supportFor: AxAIFeatures | ((model: TModel) => AxAIFeatures);
-    models?: AxAIInputModelList<TModel, TEmbedModel>;
+    models?: AxAIInputModelList<TModel, TEmbedModel, TModelKey>;
 }
 declare const axBaseAIDefaultConfig: () => AxModelConfig;
 declare const axBaseAIDefaultCreativeConfig: () => AxModelConfig;
-declare class AxBaseAI<TModel, TEmbedModel, TChatRequest, TEmbedRequest, TChatResponse, TChatResponseDelta, TEmbedResponse> implements AxAIService<TModel, TEmbedModel> {
+declare class AxBaseAI<TModel, TEmbedModel, TChatRequest, TEmbedRequest, TChatResponse, TChatResponseDelta, TEmbedResponse, TModelKey> implements AxAIService<TModel, TEmbedModel, TModelKey> {
     private readonly aiImpl;
     private debug;
     private rt?;
@@ -450,7 +450,7 @@ declare class AxBaseAI<TModel, TEmbedModel, TChatRequest, TEmbedRequest, TChatRe
     protected headers: () => Promise<Record<string, string>>;
     protected supportFor: AxAIFeatures | ((model: TModel) => AxAIFeatures);
     private metrics;
-    constructor(aiImpl: Readonly<AxAIServiceImpl<TModel, TEmbedModel, TChatRequest, TEmbedRequest, TChatResponse, TChatResponseDelta, TEmbedResponse>>, { name, apiURL, headers, modelInfo, defaults, options, supportFor, models, }: Readonly<AxBaseAIArgs<TModel, TEmbedModel>>);
+    constructor(aiImpl: Readonly<AxAIServiceImpl<TModel, TEmbedModel, TChatRequest, TEmbedRequest, TChatResponse, TChatResponseDelta, TEmbedResponse>>, { name, apiURL, headers, modelInfo, defaults, options, supportFor, models, }: Readonly<AxBaseAIArgs<TModel, TEmbedModel, TModelKey>>);
     private getMetricsInstruments;
     setName(name: string): void;
     getId(): string;
@@ -459,7 +459,7 @@ declare class AxBaseAI<TModel, TEmbedModel, TChatRequest, TEmbedRequest, TChatRe
     setOptions(options: Readonly<AxAIServiceOptions>): void;
     getOptions(): Readonly<AxAIServiceOptions>;
     getLogger(): AxLoggerFunction;
-    getModelList(): AxAIModelList | undefined;
+    getModelList(): AxAIModelList<TModelKey>;
     getName(): string;
     getFeatures(model?: TModel): AxAIFeatures;
     getLastUsedChatModel(): TModel | undefined;
@@ -482,11 +482,11 @@ declare class AxBaseAI<TModel, TEmbedModel, TChatRequest, TEmbedRequest, TChatRe
     private recordChatMetrics;
     private recordEmbedMetrics;
     getMetrics(): AxAIServiceMetrics;
-    chat(req: Readonly<AxChatRequest<TModel>>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<TModel, TEmbedModel>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
+    chat(req: Readonly<AxChatRequest<TModel>>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<TModel, TEmbedModel, TModelKey>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
     private _chat1;
     private cleanupFunctionSchema;
     private _chat2;
-    embed(req: Readonly<AxEmbedRequest<TEmbedModel>>, options?: Readonly<AxAIServiceActionOptions<TModel, TEmbedModel>>): Promise<AxEmbedResponse>;
+    embed(req: Readonly<AxEmbedRequest<TEmbedModel>>, options?: Readonly<AxAIServiceActionOptions<TModel, TEmbedModel, TModelKey>>): Promise<AxEmbedResponse>;
     private _embed1;
     private _embed2;
     private buildHeaders;
@@ -729,17 +729,17 @@ type AxAIAnthropicChatResponseDelta = AxAIAnthropicMessageStartEvent | AxAIAnthr
 
 declare const axAIAnthropicDefaultConfig: () => AxAIAnthropicConfig;
 declare const axAIAnthropicVertexDefaultConfig: () => AxAIAnthropicConfig;
-interface AxAIAnthropicArgs {
+interface AxAIAnthropicArgs<TModelKey = string> {
     name: 'anthropic';
     apiKey?: string | (() => Promise<string>);
     projectId?: string;
     region?: string;
     config?: Readonly<Partial<AxAIAnthropicConfig>>;
     options?: Readonly<AxAIServiceOptions>;
-    models?: AxAIInputModelList<AxAIAnthropicModel | AxAIAnthropicVertexModel, undefined>;
+    models?: AxAIInputModelList<AxAIAnthropicModel | AxAIAnthropicVertexModel, undefined, TModelKey>;
 }
-declare class AxAIAnthropic extends AxBaseAI<AxAIAnthropicModel | AxAIAnthropicVertexModel, unknown, AxAIAnthropicChatRequest, unknown, AxAIAnthropicChatResponse, AxAIAnthropicChatResponseDelta, unknown> {
-    constructor({ apiKey, projectId, region, config, options, models, }: Readonly<Omit<AxAIAnthropicArgs, 'name'>>);
+declare class AxAIAnthropic<TModelKey = string> extends AxBaseAI<AxAIAnthropicModel | AxAIAnthropicVertexModel, unknown, AxAIAnthropicChatRequest, unknown, AxAIAnthropicChatResponse, AxAIAnthropicChatResponseDelta, unknown, TModelKey> {
+    constructor({ apiKey, projectId, region, config, options, models, }: Readonly<Omit<AxAIAnthropicArgs<TModelKey>, 'name'>>);
 }
 
 declare const axModelInfoAnthropic: AxModelInfo[];
@@ -990,13 +990,13 @@ declare const axAIOpenAIDefaultConfig: () => AxAIOpenAIConfig<AxAIOpenAIModel, A
 declare const axAIOpenAIBestConfig: () => AxAIOpenAIConfig<AxAIOpenAIModel, AxAIOpenAIEmbedModel>;
 declare const axAIOpenAICreativeConfig: () => AxAIOpenAIConfig<AxAIOpenAIModel, AxAIOpenAIEmbedModel>;
 declare const axAIOpenAIFastConfig: () => AxAIOpenAIConfig<AxAIOpenAIModel, AxAIOpenAIEmbedModel>;
-interface AxAIOpenAIArgs<TName = 'openai', TModel = AxAIOpenAIModel, TEmbedModel = AxAIOpenAIEmbedModel, TChatReq extends AxAIOpenAIChatRequest<TModel> = AxAIOpenAIChatRequest<TModel>> extends Omit<AxAIOpenAIBaseArgs<TModel, TEmbedModel, TChatReq>, 'config' | 'supportFor' | 'modelInfo'> {
+interface AxAIOpenAIArgs<TName = 'openai', TModel = AxAIOpenAIModel, TEmbedModel = AxAIOpenAIEmbedModel, TModelKey = string, TChatReq extends AxAIOpenAIChatRequest<TModel> = AxAIOpenAIChatRequest<TModel>> extends Omit<AxAIOpenAIBaseArgs<TModel, TEmbedModel, TModelKey, TChatReq>, 'config' | 'supportFor' | 'modelInfo'> {
     name: TName;
     modelInfo?: AxModelInfo[];
-    config?: Partial<AxAIOpenAIBaseArgs<TModel, TEmbedModel, TChatReq>['config']>;
+    config?: Partial<AxAIOpenAIBaseArgs<TModel, TEmbedModel, TModelKey, TChatReq>['config']>;
 }
 type ChatReqUpdater<TModel, TChatReq extends AxAIOpenAIChatRequest<TModel>> = (req: Readonly<TChatReq>) => TChatReq;
-interface AxAIOpenAIBaseArgs<TModel, TEmbedModel, TChatReq extends AxAIOpenAIChatRequest<TModel>> {
+interface AxAIOpenAIBaseArgs<TModel, TEmbedModel, TModelKey, TChatReq extends AxAIOpenAIChatRequest<TModel>> {
     apiKey: string;
     apiURL?: string;
     config: Readonly<AxAIOpenAIConfig<TModel, TEmbedModel>>;
@@ -1004,15 +1004,15 @@ interface AxAIOpenAIBaseArgs<TModel, TEmbedModel, TChatReq extends AxAIOpenAICha
         streamingUsage?: boolean;
     }>;
     modelInfo: Readonly<AxModelInfo[]>;
-    models?: AxAIInputModelList<TModel, TEmbedModel>;
+    models?: AxAIInputModelList<TModel, TEmbedModel, TModelKey>;
     chatReqUpdater?: ChatReqUpdater<TModel, TChatReq>;
     supportFor: AxAIFeatures | ((model: TModel) => AxAIFeatures);
 }
-declare class AxAIOpenAIBase<TModel, TEmbedModel, TChatReq extends AxAIOpenAIChatRequest<TModel> = AxAIOpenAIChatRequest<TModel>> extends AxBaseAI<TModel, TEmbedModel, AxAIOpenAIChatRequest<TModel>, AxAIOpenAIEmbedRequest<TEmbedModel>, AxAIOpenAIChatResponse, AxAIOpenAIChatResponseDelta, AxAIOpenAIEmbedResponse> {
-    constructor({ apiKey, config, options, apiURL, modelInfo, models, chatReqUpdater, supportFor, }: Readonly<Omit<AxAIOpenAIBaseArgs<TModel, TEmbedModel, TChatReq>, 'name'>>);
+declare class AxAIOpenAIBase<TModel, TEmbedModel, TModelKey, TChatReq extends AxAIOpenAIChatRequest<TModel> = AxAIOpenAIChatRequest<TModel>> extends AxBaseAI<TModel, TEmbedModel, AxAIOpenAIChatRequest<TModel>, AxAIOpenAIEmbedRequest<TEmbedModel>, AxAIOpenAIChatResponse, AxAIOpenAIChatResponseDelta, AxAIOpenAIEmbedResponse, TModelKey> {
+    constructor({ apiKey, config, options, apiURL, modelInfo, models, chatReqUpdater, supportFor, }: Readonly<Omit<AxAIOpenAIBaseArgs<TModel, TEmbedModel, TModelKey, TChatReq>, 'name'>>);
 }
-declare class AxAIOpenAI extends AxAIOpenAIBase<AxAIOpenAIModel, AxAIOpenAIEmbedModel> {
-    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIOpenAIArgs, 'name'>>);
+declare class AxAIOpenAI<TModelKey = string> extends AxAIOpenAIBase<AxAIOpenAIModel, AxAIOpenAIEmbedModel, TModelKey> {
+    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIOpenAIArgs<'openai', AxAIOpenAIModel, AxAIOpenAIEmbedModel, TModelKey>, 'name'>>);
 }
 
 declare const axAIAzureOpenAIDefaultConfig: () => AxAIOpenAIConfig<AxAIOpenAIModel, AxAIOpenAIEmbedModel>;
@@ -1020,20 +1020,20 @@ declare const axAIAzureOpenAICreativeConfig: () => AxAIOpenAIConfig<AxAIOpenAIMo
 declare const axAIAzureOpenAIFastConfig: () => AxAIOpenAIConfig<AxAIOpenAIModel, AxAIOpenAIEmbedModel>;
 declare const axAIAzureOpenAIBestConfig: () => AxAIOpenAIConfig<AxAIOpenAIModel, AxAIOpenAIEmbedModel>;
 type AxAIAzureOpenAIConfig = AxAIOpenAIConfig<AxAIOpenAIModel, AxAIOpenAIEmbedModel>;
-type AxAIAzureOpenAIArgs = AxAIOpenAIArgs<'azure-openai', AxAIOpenAIModel, AxAIOpenAIEmbedModel> & {
+type AxAIAzureOpenAIArgs<TModelKey = string> = AxAIOpenAIArgs<'azure-openai', AxAIOpenAIModel, AxAIOpenAIEmbedModel, TModelKey> & {
     resourceName: string;
     deploymentName: string;
     version?: string;
 };
-declare class AxAIAzureOpenAI extends AxAIOpenAIBase<AxAIOpenAIModel, AxAIOpenAIEmbedModel> {
-    constructor({ apiKey, resourceName, deploymentName, version, config, options, models, modelInfo, }: Readonly<Omit<AxAIAzureOpenAIArgs, 'name'>>);
+declare class AxAIAzureOpenAI<TModelKey = string> extends AxAIOpenAIBase<AxAIOpenAIModel, AxAIOpenAIEmbedModel, TModelKey> {
+    constructor({ apiKey, resourceName, deploymentName, version, config, options, models, modelInfo, }: Readonly<Omit<AxAIAzureOpenAIArgs<TModelKey>, 'name'>>);
 }
 
 /**
  * Options for the balancer.
  */
-type AxBalancerOptions = {
-    comparator?: (a: AxAIService, b: AxAIService) => number;
+type AxBalancerOptions<TModelKey = string> = {
+    comparator?: (a: AxAIService<unknown, unknown, TModelKey>, b: AxAIService<unknown, unknown, TModelKey>) => number;
     debug?: boolean;
     initialBackoffMs?: number;
     maxBackoffMs?: number;
@@ -1042,7 +1042,7 @@ type AxBalancerOptions = {
 /**
  * Balancer that rotates through services.
  */
-declare class AxBalancer implements AxAIService<unknown, unknown> {
+declare class AxBalancer<TModelKey = string> implements AxAIService<unknown, unknown, TModelKey> {
     private services;
     private currentServiceIndex;
     private currentService;
@@ -1051,7 +1051,7 @@ declare class AxBalancer implements AxAIService<unknown, unknown> {
     private maxBackoffMs;
     private maxRetries;
     private serviceFailures;
-    constructor(services: readonly AxAIService[], options?: AxBalancerOptions);
+    constructor(services: readonly AxAIService<unknown, unknown, TModelKey>[], options?: AxBalancerOptions<TModelKey>);
     getLastUsedChatModel(): unknown;
     getLastUsedEmbedModel(): unknown;
     getLastUsedModelConfig(): AxModelConfig | undefined;
@@ -1062,8 +1062,8 @@ declare class AxBalancer implements AxAIService<unknown, unknown> {
     /**
      * Service comparator that sorts services by cost.
      */
-    static metricComparator: (a: AxAIService, b: AxAIService) => number;
-    getModelList(): AxAIModelList | undefined;
+    static metricComparator: <TModelKey_1 = string>(a: AxAIService<unknown, unknown, TModelKey_1>, b: AxAIService<unknown, unknown, TModelKey_1>) => number;
+    getModelList(): AxAIModelList<TModelKey> | undefined;
     private getNextService;
     private reset;
     getName(): string;
@@ -1073,8 +1073,8 @@ declare class AxBalancer implements AxAIService<unknown, unknown> {
     private canRetryService;
     private handleFailure;
     private handleSuccess;
-    chat(req: Readonly<AxChatRequest>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions> | undefined): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
-    embed(req: Readonly<AxEmbedRequest>, options?: Readonly<AxAIServiceActionOptions> | undefined): Promise<AxEmbedResponse>;
+    chat(req: Readonly<AxChatRequest>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<unknown, unknown, TModelKey>> | undefined): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
+    embed(req: Readonly<AxEmbedRequest>, options?: Readonly<AxAIServiceActionOptions<unknown, unknown, TModelKey>> | undefined): Promise<AxEmbedResponse>;
     setOptions(options: Readonly<AxAIServiceOptions>): void;
     getOptions(): Readonly<AxAIServiceOptions>;
     getLogger(): AxLoggerFunction;
@@ -1181,15 +1181,15 @@ type AxAICohereEmbedResponse = {
 
 declare const axAICohereDefaultConfig: () => AxAICohereConfig;
 declare const axAICohereCreativeConfig: () => AxAICohereConfig;
-interface AxAICohereArgs {
+interface AxAICohereArgs<TModelKey = string> {
     name: 'cohere';
     apiKey: string;
     config?: Readonly<Partial<AxAICohereConfig>>;
     options?: Readonly<AxAIServiceOptions>;
-    models?: AxAIInputModelList<AxAICohereModel, AxAICohereEmbedModel>;
+    models?: AxAIInputModelList<AxAICohereModel, AxAICohereEmbedModel, TModelKey>;
 }
-declare class AxAICohere extends AxBaseAI<AxAICohereModel, AxAICohereEmbedModel, AxAICohereChatRequest, AxAICohereEmbedRequest, AxAICohereChatResponse, AxAICohereChatResponseDelta, AxAICohereEmbedResponse> {
-    constructor({ apiKey, config, options, models, }: Readonly<Omit<AxAICohereArgs, 'name'>>);
+declare class AxAICohere<TModelKey = string> extends AxBaseAI<AxAICohereModel, AxAICohereEmbedModel, AxAICohereChatRequest, AxAICohereEmbedRequest, AxAICohereChatResponse, AxAICohereChatResponseDelta, AxAICohereEmbedResponse, TModelKey> {
+    constructor({ apiKey, config, options, models, }: Readonly<Omit<AxAICohereArgs<TModelKey>, 'name'>>);
 }
 
 declare const axModelInfoCohere: AxModelInfo[];
@@ -1206,9 +1206,9 @@ declare enum AxAIDeepSeekModel {
 type DeepSeekConfig = AxAIOpenAIConfig<AxAIDeepSeekModel, undefined>;
 declare const axAIDeepSeekDefaultConfig: () => DeepSeekConfig;
 declare const axAIDeepSeekCodeConfig: () => DeepSeekConfig;
-type AxAIDeepSeekArgs = AxAIOpenAIArgs<'deepseek', AxAIDeepSeekModel, undefined>;
-declare class AxAIDeepSeek extends AxAIOpenAIBase<AxAIDeepSeekModel, undefined> {
-    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIDeepSeekArgs, 'name'>>);
+type AxAIDeepSeekArgs<TModelKey = string> = AxAIOpenAIArgs<'deepseek', AxAIDeepSeekModel, undefined, TModelKey>;
+declare class AxAIDeepSeek<TModelKey = string> extends AxAIOpenAIBase<AxAIDeepSeekModel, undefined, TModelKey> {
+    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIDeepSeekArgs<TModelKey>, 'name'>>);
 }
 
 declare const axModelInfoDeepSeek: AxModelInfo[];
@@ -1451,7 +1451,7 @@ interface AxAIGoogleGeminiOptionsTools {
     googleSearch?: boolean;
     urlContext?: boolean;
 }
-interface AxAIGoogleGeminiArgs {
+interface AxAIGoogleGeminiArgs<TModelKey = string> {
     name: 'google-gemini';
     apiKey?: string | (() => Promise<string>);
     projectId?: string;
@@ -1459,14 +1459,14 @@ interface AxAIGoogleGeminiArgs {
     endpointId?: string;
     config?: Readonly<Partial<AxAIGoogleGeminiConfig>>;
     options?: Readonly<AxAIServiceOptions & AxAIGoogleGeminiOptionsTools>;
-    models?: AxAIInputModelList<AxAIGoogleGeminiModel, AxAIGoogleGeminiEmbedModel>;
+    models?: AxAIInputModelList<AxAIGoogleGeminiModel, AxAIGoogleGeminiEmbedModel, TModelKey>;
     modelInfo?: AxModelInfo[];
 }
 /**
  * AxAIGoogleGemini: AI Service
  */
-declare class AxAIGoogleGemini extends AxBaseAI<AxAIGoogleGeminiModel, AxAIGoogleGeminiEmbedModel, AxAIGoogleGeminiChatRequest, AxAIGoogleGeminiBatchEmbedRequest | AxAIGoogleVertexBatchEmbedRequest, AxAIGoogleGeminiChatResponse, AxAIGoogleGeminiChatResponseDelta, AxAIGoogleGeminiBatchEmbedResponse | AxAIGoogleVertexBatchEmbedResponse> {
-    constructor({ apiKey, projectId, region, endpointId, config, options, models, modelInfo, }: Readonly<Omit<AxAIGoogleGeminiArgs, 'name'>>);
+declare class AxAIGoogleGemini<TModelKey = string> extends AxBaseAI<AxAIGoogleGeminiModel, AxAIGoogleGeminiEmbedModel, AxAIGoogleGeminiChatRequest, AxAIGoogleGeminiBatchEmbedRequest | AxAIGoogleVertexBatchEmbedRequest, AxAIGoogleGeminiChatResponse, AxAIGoogleGeminiChatResponseDelta, AxAIGoogleGeminiBatchEmbedResponse | AxAIGoogleVertexBatchEmbedResponse, TModelKey> {
+    constructor({ apiKey, projectId, region, endpointId, config, options, models, modelInfo, }: Readonly<Omit<AxAIGoogleGeminiArgs<TModelKey>, 'name'>>);
 }
 
 /**
@@ -1481,14 +1481,14 @@ declare enum AxAIGroqModel {
     Gemma2_9B = "gemma2-9b-it"
 }
 
-type AxAIGroqArgs = AxAIOpenAIArgs<'groq', AxAIGroqModel, undefined> & {
+type AxAIGroqArgs<TModelKey = string> = AxAIOpenAIArgs<'groq', AxAIGroqModel, undefined, TModelKey> & {
     options?: Readonly<AxAIServiceOptions> & {
         tokensPerMinute?: number;
     };
     modelInfo?: AxModelInfo[];
 };
-declare class AxAIGroq extends AxAIOpenAIBase<AxAIGroqModel, undefined> {
-    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIGroqArgs, 'name'>>);
+declare class AxAIGroq<TModelKey = string> extends AxAIOpenAIBase<AxAIGroqModel, undefined, TModelKey> {
+    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIGroqArgs<TModelKey>, 'name'>>);
     setOptions: (options: Readonly<AxAIServiceOptions>) => void;
     private newRateLimiter;
 }
@@ -1534,15 +1534,15 @@ type AxAIHuggingFaceResponse = {
 
 declare const axAIHuggingFaceDefaultConfig: () => AxAIHuggingFaceConfig;
 declare const axAIHuggingFaceCreativeConfig: () => AxAIHuggingFaceConfig;
-interface AxAIHuggingFaceArgs {
+interface AxAIHuggingFaceArgs<TModelKey = string> {
     name: 'huggingface';
     apiKey: string;
     config?: Readonly<Partial<AxAIHuggingFaceConfig>>;
     options?: Readonly<AxAIServiceOptions>;
-    models?: AxAIInputModelList<AxAIHuggingFaceModel, undefined>;
+    models?: AxAIInputModelList<AxAIHuggingFaceModel, undefined, TModelKey>;
 }
-declare class AxAIHuggingFace extends AxBaseAI<AxAIHuggingFaceModel, unknown, AxAIHuggingFaceRequest, unknown, AxAIHuggingFaceResponse, unknown, unknown> {
-    constructor({ apiKey, config, options, models, }: Readonly<Omit<AxAIHuggingFaceArgs, 'name'>>);
+declare class AxAIHuggingFace<TModelKey = string> extends AxBaseAI<AxAIHuggingFaceModel, unknown, AxAIHuggingFaceRequest, unknown, AxAIHuggingFaceResponse, unknown, unknown, TModelKey> {
+    constructor({ apiKey, config, options, models, }: Readonly<Omit<AxAIHuggingFaceArgs<TModelKey>, 'name'>>);
 }
 
 /**
@@ -1626,20 +1626,20 @@ type AxAIMistralChatRequest = Omit<AxAIOpenAIChatRequest<AxAIMistralModel>, 'max
         tool_call_id: string;
     })[];
 };
-type AxAIMistralArgs = AxAIOpenAIArgs<'mistral', AxAIMistralModel, AxAIMistralEmbedModels> & {
+type AxAIMistralArgs<TModelKey = string> = AxAIOpenAIArgs<'mistral', AxAIMistralModel, AxAIMistralEmbedModels, TModelKey> & {
     options?: Readonly<AxAIServiceOptions> & {
         tokensPerMinute?: number;
     };
     modelInfo?: AxModelInfo[];
 };
-declare class AxAIMistral extends AxAIOpenAIBase<AxAIMistralModel, AxAIMistralEmbedModels> {
-    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIMistralArgs, 'name'>>);
+declare class AxAIMistral<TModelKey = string> extends AxAIOpenAIBase<AxAIMistralModel, AxAIMistralEmbedModels, TModelKey> {
+    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIMistralArgs<TModelKey>, 'name'>>);
     private updateMessages;
 }
 
 declare const axModelInfoMistral: AxModelInfo[];
 
-type AxMockAIServiceConfig = {
+type AxMockAIServiceConfig<TModelKey = string> = {
     name?: string;
     id?: string;
     modelInfo?: Partial<AxModelInfoWithProvider>;
@@ -1648,7 +1648,7 @@ type AxMockAIServiceConfig = {
         functions?: boolean;
         streaming?: boolean;
     };
-    models?: AxAIModelList;
+    models?: AxAIModelList<TModelKey>;
     options?: AxAIServiceOptions;
     chatResponse?: AxChatResponse | ReadableStream<AxChatResponse> | (() => Promise<AxChatResponse | ReadableStream<AxChatResponse>>) | ((req: Readonly<AxChatRequest<unknown>>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<unknown, unknown>>) => Promise<AxChatResponse | ReadableStream<AxChatResponse>>);
     embedResponse?: AxEmbedResponse | ((req: Readonly<AxEmbedRequest>) => AxEmbedResponse | Promise<AxEmbedResponse>);
@@ -1656,10 +1656,10 @@ type AxMockAIServiceConfig = {
     errorMessage?: string;
     latencyMs?: number;
 };
-declare class AxMockAIService implements AxAIService {
+declare class AxMockAIService<TModelKey = string> implements AxAIService<unknown, unknown, TModelKey> {
     private readonly config;
     private metrics;
-    constructor(config?: AxMockAIServiceConfig);
+    constructor(config?: AxMockAIServiceConfig<TModelKey>);
     getLastUsedChatModel(): unknown;
     getLastUsedEmbedModel(): unknown;
     getLastUsedModelConfig(): AxModelConfig | undefined;
@@ -1669,23 +1669,23 @@ declare class AxMockAIService implements AxAIService {
         functions: boolean;
         streaming: boolean;
     };
-    getModelList(): AxAIModelList | undefined;
+    getModelList(): AxAIModelList<TModelKey> | undefined;
     getMetrics(): AxAIServiceMetrics;
-    chat(req: Readonly<AxChatRequest<unknown>>, _options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<unknown, unknown>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
-    embed(req: Readonly<AxEmbedRequest>, _options?: Readonly<AxAIServiceActionOptions>): Promise<AxEmbedResponse>;
+    chat(req: Readonly<AxChatRequest<unknown>>, _options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<unknown, unknown, TModelKey>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
+    embed(req: Readonly<AxEmbedRequest>, _options?: Readonly<AxAIServiceActionOptions<unknown, unknown, TModelKey>>): Promise<AxEmbedResponse>;
     setOptions(options: Readonly<AxAIServiceOptions>): void;
     getOptions(): Readonly<AxAIServiceOptions>;
     getLogger(): AxLoggerFunction;
     private updateMetrics;
 }
 
-type AxAIServiceListItem<TModel = unknown, TEmbedModel = unknown> = {
-    key: string;
-    service: AxAIService<TModel, TEmbedModel>;
+type AxAIServiceListItem<TModel = unknown, TEmbedModel = unknown, TModelKey = string> = {
+    key: TModelKey;
+    service: AxAIService<TModel, TEmbedModel, TModelKey>;
     description: string;
     isInternal?: boolean;
 };
-declare class AxMultiServiceRouter implements AxAIService<string, string> {
+declare class AxMultiServiceRouter<TModelKey = string> implements AxAIService<unknown, unknown, TModelKey> {
     private options?;
     private lastUsedService?;
     private services;
@@ -1694,18 +1694,18 @@ declare class AxMultiServiceRouter implements AxAIService<string, string> {
      * It validates that each service provides a unique set of model keys,
      * then builds a lookup (map) for routing the chat/embed requests.
      */
-    constructor(services: (AxAIServiceListItem<string, string> | AxAIService<string, string>)[]);
-    getLastUsedChatModel(): string | undefined;
-    getLastUsedEmbedModel(): string | undefined;
+    constructor(services: (AxAIServiceListItem<unknown, unknown, TModelKey> | AxAIService<unknown, unknown, TModelKey>)[]);
+    getLastUsedChatModel(): unknown | undefined;
+    getLastUsedEmbedModel(): unknown | undefined;
     getLastUsedModelConfig(): AxModelConfig | undefined;
     /**
      * Delegates the chat call to the service matching the provided model key.
      */
-    chat(req: Readonly<AxChatRequest<string>>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<string, string>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
+    chat(req: Readonly<AxChatRequest<string>>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<unknown, unknown, TModelKey>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
     /**
      * Delegates the embed call to the service matching the provided embed model key.
      */
-    embed(req: Readonly<AxEmbedRequest<string>>, options?: Readonly<AxAIServiceActionOptions<string, string>>): Promise<AxEmbedResponse>;
+    embed(req: Readonly<AxEmbedRequest<string>>, options?: Readonly<AxAIServiceActionOptions<unknown, unknown, TModelKey>>): Promise<AxEmbedResponse>;
     /**
      * Returns a composite ID built from the IDs of the underlying services.
      */
@@ -1717,12 +1717,12 @@ declare class AxMultiServiceRouter implements AxAIService<string, string> {
     /**
      * Aggregates all available models across the underlying services.
      */
-    getModelList(): AxAIModelList;
+    getModelList(): AxAIModelList<TModelKey>;
     /**
      * If a model key is provided, delegate to the corresponding service's features.
      * Otherwise, returns a default feature set.
      */
-    getFeatures(model?: string): {
+    getFeatures(model?: TModelKey): {
         functions: boolean;
         streaming: boolean;
         functionCot?: boolean;
@@ -1752,19 +1752,19 @@ declare class AxMultiServiceRouter implements AxAIService<string, string> {
      * @param key - The model key
      * @param entry - The service entry to set
      */
-    setServiceEntry(key: string, entry: {
+    setServiceEntry(key: TModelKey, entry: {
         isInternal?: boolean;
         description: string;
         model?: string;
         embedModel?: string;
-        service: AxAIService<string, string>;
+        service: AxAIService<unknown, unknown, TModelKey>;
     }): void;
 }
 
 type AxAIOllamaAIConfig = AxAIOpenAIConfig<string, string>;
 declare const axAIOllamaDefaultConfig: () => AxAIOllamaAIConfig;
 declare const axAIOllamaDefaultCreativeConfig: () => AxAIOllamaAIConfig;
-type AxAIOllamaArgs = AxAIOpenAIArgs<'ollama', string, string> & {
+type AxAIOllamaArgs<TModelKey = string> = AxAIOpenAIArgs<'ollama', string, string, TModelKey> & {
     model?: string;
     embedModel?: string;
     url?: string;
@@ -1772,8 +1772,8 @@ type AxAIOllamaArgs = AxAIOpenAIArgs<'ollama', string, string> & {
 /**
  * OllamaAI: AI Service
  */
-declare class AxAIOllama extends AxAIOpenAIBase<string, string> {
-    constructor({ apiKey, url, config, options, models, }: Readonly<Omit<AxAIOllamaArgs, 'name'>>);
+declare class AxAIOllama<TModelKey = string> extends AxAIOpenAIBase<string, string, TModelKey> {
+    constructor({ apiKey, url, config, options, models, }: Readonly<Omit<AxAIOllamaArgs<TModelKey>, 'name'>>);
 }
 
 /**
@@ -2302,7 +2302,7 @@ Readonly<AxAIOpenAIEmbedResponse>> {
 declare const axAIOpenAIResponsesDefaultConfig: () => AxAIOpenAIResponsesConfig<AxAIOpenAIResponsesModel, AxAIOpenAIEmbedModel>;
 declare const axAIOpenAIResponsesBestConfig: () => AxAIOpenAIResponsesConfig<AxAIOpenAIResponsesModel, AxAIOpenAIEmbedModel>;
 declare const axAIOpenAIResponsesCreativeConfig: () => AxAIOpenAIResponsesConfig<AxAIOpenAIResponsesModel, AxAIOpenAIEmbedModel>;
-interface AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TResponsesReq extends AxAIOpenAIResponsesRequest<TModel>> {
+interface AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TModelKey, TResponsesReq extends AxAIOpenAIResponsesRequest<TModel>> {
     apiKey: string;
     config: AxAIOpenAIResponsesConfig<TModel, TEmbedModel>;
     options?: {
@@ -2310,27 +2310,27 @@ interface AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TResponsesReq extends
     } & AxAIServiceOptions;
     apiURL?: string;
     modelInfo?: ReadonlyArray<AxModelInfo>;
-    models?: AxAIInputModelList<TModel, TEmbedModel>;
+    models?: AxAIInputModelList<TModel, TEmbedModel, TModelKey>;
     responsesReqUpdater?: (req: Readonly<TResponsesReq>) => Readonly<TResponsesReq>;
     supportFor?: AxAIFeatures | ((model: TModel) => AxAIFeatures);
 }
 /**
  * Base class for OpenAI AI services using the /v1/responses API endpoint
  */
-declare class AxAIOpenAIResponsesBase<TModel, TEmbedModel, TResponsesReq extends AxAIOpenAIResponsesRequest<TModel>> extends AxBaseAI<TModel, TEmbedModel, AxAIOpenAIResponsesRequest<TModel>, AxAIOpenAIEmbedRequest<TEmbedModel>, AxAIOpenAIResponsesResponse, AxAIOpenAIResponsesResponseDelta, AxAIOpenAIEmbedResponse> {
-    constructor({ apiKey, config, options, apiURL, modelInfo, models, responsesReqUpdater, supportFor, }: Readonly<AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TResponsesReq>>);
+declare class AxAIOpenAIResponsesBase<TModel, TEmbedModel, TModelKey, TResponsesReq extends AxAIOpenAIResponsesRequest<TModel>> extends AxBaseAI<TModel, TEmbedModel, AxAIOpenAIResponsesRequest<TModel>, AxAIOpenAIEmbedRequest<TEmbedModel>, AxAIOpenAIResponsesResponse, AxAIOpenAIResponsesResponseDelta, AxAIOpenAIEmbedResponse, TModelKey> {
+    constructor({ apiKey, config, options, apiURL, modelInfo, models, responsesReqUpdater, supportFor, }: Readonly<AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TModelKey, TResponsesReq>>);
 }
 /**
  * Ready-to-use implementation of the OpenAI Responses API client
  * This class uses OpenAI's /v1/responses API endpoint which supports text, image, and audio inputs
  */
-interface AxAIOpenAIResponsesArgs<TName = 'openai-responses', TModel = AxAIOpenAIResponsesModel, TEmbedModel = AxAIOpenAIEmbedModel, TChatReq extends AxAIOpenAIResponsesRequest<TModel> = AxAIOpenAIResponsesRequest<TModel>> extends Omit<AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TChatReq>, 'config' | 'supportFor' | 'modelInfo'> {
+interface AxAIOpenAIResponsesArgs<TName = 'openai-responses', TModel = AxAIOpenAIResponsesModel, TEmbedModel = AxAIOpenAIEmbedModel, TModelKey = string, TChatReq extends AxAIOpenAIResponsesRequest<TModel> = AxAIOpenAIResponsesRequest<TModel>> extends Omit<AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TModelKey, TChatReq>, 'config' | 'supportFor' | 'modelInfo'> {
     name: TName;
     modelInfo?: AxModelInfo[];
-    config?: Partial<AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TChatReq>['config']>;
+    config?: Partial<AxAIOpenAIResponsesBaseArgs<TModel, TEmbedModel, TModelKey, TChatReq>['config']>;
 }
-declare class AxAIOpenAIResponses extends AxAIOpenAIResponsesBase<AxAIOpenAIResponsesModel, AxAIOpenAIEmbedModel, AxAIOpenAIResponsesRequest<AxAIOpenAIResponsesModel>> {
-    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIOpenAIResponsesArgs, 'name'>>);
+declare class AxAIOpenAIResponses<TModelKey = string> extends AxAIOpenAIResponsesBase<AxAIOpenAIResponsesModel, AxAIOpenAIEmbedModel, TModelKey, AxAIOpenAIResponsesRequest<AxAIOpenAIResponsesModel>> {
+    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIOpenAIResponsesArgs<'openai-responses', AxAIOpenAIResponsesModel, AxAIOpenAIEmbedModel, TModelKey>, 'name'>>);
 }
 
 declare enum AxAIRekaModel {
@@ -2404,7 +2404,7 @@ declare const axAIRekaDefaultConfig: () => AxAIRekaConfig;
 declare const axAIRekaBestConfig: () => AxAIRekaConfig;
 declare const axAIRekaCreativeConfig: () => AxAIRekaConfig;
 declare const axAIRekaFastConfig: () => AxAIRekaConfig;
-interface AxAIRekaArgs {
+interface AxAIRekaArgs<TModelKey = string> {
     name: 'reka';
     apiKey: string;
     apiURL?: string;
@@ -2413,10 +2413,10 @@ interface AxAIRekaArgs {
         streamingUsage?: boolean;
     }>;
     modelInfo?: Readonly<AxModelInfo[]>;
-    models?: AxAIInputModelList<AxAIRekaModel, undefined>;
+    models?: AxAIInputModelList<AxAIRekaModel, undefined, TModelKey>;
 }
-declare class AxAIReka extends AxBaseAI<AxAIRekaModel, undefined, AxAIRekaChatRequest, unknown, AxAIRekaChatResponse, AxAIRekaChatResponseDelta, unknown> {
-    constructor({ apiKey, config, options, apiURL, modelInfo, models, }: Readonly<Omit<AxAIRekaArgs, 'name'>>);
+declare class AxAIReka<TModelKey = string> extends AxBaseAI<AxAIRekaModel, undefined, AxAIRekaChatRequest, unknown, AxAIRekaChatResponse, AxAIRekaChatResponseDelta, unknown, TModelKey> {
+    constructor({ apiKey, config, options, apiURL, modelInfo, models, }: Readonly<Omit<AxAIRekaArgs<TModelKey>, 'name'>>);
 }
 
 /**
@@ -2426,9 +2426,9 @@ declare const axModelInfoReka: AxModelInfo[];
 
 type TogetherAIConfig = AxAIOpenAIConfig<string, unknown>;
 declare const axAITogetherDefaultConfig: () => TogetherAIConfig;
-type AxAITogetherArgs = AxAIOpenAIArgs<'together', string, unknown>;
-declare class AxAITogether extends AxAIOpenAIBase<string, unknown> {
-    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAITogetherArgs, 'name'>>);
+type AxAITogetherArgs<TModelKey> = AxAIOpenAIArgs<'together', string, unknown, TModelKey>;
+declare class AxAITogether<TModelKey = string> extends AxAIOpenAIBase<string, unknown, TModelKey> {
+    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAITogetherArgs<TModelKey>, 'name'>>);
 }
 
 declare const axModelInfoTogether: AxModelInfo[];
@@ -2447,30 +2447,6 @@ declare function axValidateChatRequestMessage(item: AxChatRequestMessage): void;
  */
 declare function axValidateChatResponseResult(results: Readonly<AxChatResponseResult[]> | Readonly<AxChatResponseResult>): void;
 
-type AxAIArgs = AxAIOpenAIArgs | AxAIOpenAIResponsesArgs | AxAIAzureOpenAIArgs | AxAITogetherArgs | AxAIAnthropicArgs | AxAIGroqArgs | AxAIGoogleGeminiArgs | AxAICohereArgs | AxAIHuggingFaceArgs | AxAIMistralArgs | AxAIDeepSeekArgs | AxAIOllamaArgs | AxAIRekaArgs;
-type AxAIModels = AxAIOpenAIModel | AxAIAnthropicModel | AxAIGroqModel | AxAIGoogleGeminiModel | AxAICohereModel | AxAIHuggingFaceModel | AxAIMistralModel | AxAIDeepSeekModel;
-type AxAIEmbedModels = AxAIOpenAIEmbedModel | AxAIGoogleGeminiEmbedModel | AxAICohereEmbedModel;
-declare class AxAI implements AxAIService {
-    private ai;
-    constructor(options: Readonly<AxAIArgs>);
-    getName(): string;
-    getId(): string;
-    getFeatures(model?: string): {
-        functions: boolean;
-        streaming: boolean;
-    };
-    getModelList(): AxAIModelList | undefined;
-    getLastUsedChatModel(): unknown;
-    getLastUsedEmbedModel(): unknown;
-    getLastUsedModelConfig(): AxModelConfig | undefined;
-    getMetrics(): AxAIServiceMetrics;
-    chat(req: Readonly<AxChatRequest>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
-    embed(req: Readonly<AxEmbedRequest>, options?: Readonly<AxAIServiceActionOptions & AxAIServiceActionOptions>): Promise<AxEmbedResponse>;
-    setOptions(options: Readonly<AxAIServiceOptions>): void;
-    getOptions(): Readonly<AxAIServiceOptions>;
-    getLogger(): AxLoggerFunction;
-}
-
 declare enum AxAIGrokModel {
     Grok3 = "grok-3",
     Grok3Mini = "grok-3-mini",
@@ -2512,14 +2488,38 @@ type AxAIGrokChatRequest = AxAIOpenAIChatRequest<AxAIGrokModel> & {
         sources?: AxAIGrokSearchSource[];
     };
 };
-type AxAIGrokArgs = AxAIOpenAIArgs<'grok', AxAIGrokModel, AxAIGrokEmbedModels, AxAIGrokChatRequest> & {
+type AxAIGrokArgs<TModelKey = string> = AxAIOpenAIArgs<'grok', AxAIGrokModel, AxAIGrokEmbedModels, TModelKey, AxAIGrokChatRequest> & {
     options?: Readonly<AxAIServiceOptions & AxAIGrokOptionsTools> & {
         tokensPerMinute?: number;
     };
     modelInfo?: AxModelInfo[];
 };
-declare class AxAIGrok extends AxAIOpenAIBase<AxAIGrokModel, AxAIGrokEmbedModels, AxAIGrokChatRequest> {
-    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIGrokArgs, 'name'>>);
+declare class AxAIGrok<TModelKey = string> extends AxAIOpenAIBase<AxAIGrokModel, AxAIGrokEmbedModels, TModelKey, AxAIGrokChatRequest> {
+    constructor({ apiKey, config, options, models, modelInfo, }: Readonly<Omit<AxAIGrokArgs<TModelKey>, 'name'>>);
+}
+
+type AxAIArgs<TModelKey> = AxAIOpenAIArgs<'openai', AxAIOpenAIModel, AxAIOpenAIEmbedModel, TModelKey> | AxAIOpenAIResponsesArgs<'openai-responses', AxAIOpenAIResponsesModel, AxAIOpenAIEmbedModel, TModelKey> | AxAIAzureOpenAIArgs<TModelKey> | AxAITogetherArgs<TModelKey> | AxAIAnthropicArgs<TModelKey> | AxAIGroqArgs<TModelKey> | AxAIGoogleGeminiArgs<TModelKey> | AxAICohereArgs<TModelKey> | AxAIHuggingFaceArgs<TModelKey> | AxAIMistralArgs<TModelKey> | AxAIDeepSeekArgs<TModelKey> | AxAIOllamaArgs<TModelKey> | AxAIRekaArgs<TModelKey> | AxAIGrokArgs<TModelKey>;
+type AxAIModels = AxAIOpenAIModel | AxAIAnthropicModel | AxAIGroqModel | AxAIGoogleGeminiModel | AxAICohereModel | AxAIHuggingFaceModel | AxAIMistralModel | AxAIDeepSeekModel | AxAIGrokModel;
+type AxAIEmbedModels = AxAIOpenAIEmbedModel | AxAIGoogleGeminiEmbedModel | AxAICohereEmbedModel;
+declare class AxAI<TModelKey = string> implements AxAIService<unknown, unknown, TModelKey> {
+    private ai;
+    constructor(options: Readonly<AxAIArgs<TModelKey>>);
+    getName(): string;
+    getId(): string;
+    getFeatures(model?: string): {
+        functions: boolean;
+        streaming: boolean;
+    };
+    getModelList(): AxAIModelList<TModelKey> | undefined;
+    getLastUsedChatModel(): unknown;
+    getLastUsedEmbedModel(): unknown;
+    getLastUsedModelConfig(): AxModelConfig | undefined;
+    getMetrics(): AxAIServiceMetrics;
+    chat(req: Readonly<AxChatRequest>, options?: Readonly<AxAIPromptConfig & AxAIServiceActionOptions<unknown, unknown, TModelKey>>): Promise<AxChatResponse | ReadableStream<AxChatResponse>>;
+    embed(req: Readonly<AxEmbedRequest>, options?: Readonly<AxAIServiceActionOptions<unknown, unknown, TModelKey>>): Promise<AxEmbedResponse>;
+    setOptions(options: Readonly<AxAIServiceOptions>): void;
+    getOptions(): Readonly<AxAIServiceOptions>;
+    getLogger(): AxLoggerFunction;
 }
 
 declare const axModelInfoGrok: AxModelInfo[];
@@ -2945,13 +2945,13 @@ type AxResultPickerFunctionFunctionResults = {
     }[];
 };
 type AxResultPickerFunction<OUT extends AxGenOut> = (data: AxResultPickerFunctionFieldResults<OUT> | AxResultPickerFunctionFunctionResults) => number | Promise<number>;
-type AxProgramForwardOptions = {
+type AxProgramForwardOptions<MODEL> = {
     maxRetries?: number;
     maxSteps?: number;
     mem?: AxAIMemory;
     ai?: AxAIService;
     modelConfig?: AxModelConfig;
-    model?: string;
+    model?: MODEL;
     sessionId?: string;
     traceId?: string | undefined;
     tracer?: Tracer;
@@ -2978,7 +2978,7 @@ type AxProgramForwardOptions = {
     excludeContentFromTrace?: boolean;
     strictMode?: boolean;
 };
-type AxProgramStreamingForwardOptions = Omit<AxProgramForwardOptions, 'stream'>;
+type AxProgramStreamingForwardOptions<MODEL> = Omit<AxProgramForwardOptions<MODEL>, 'stream'>;
 type AxGenDeltaOut<OUT extends AxGenOut> = {
     version: number;
     index: number;
@@ -2986,9 +2986,9 @@ type AxGenDeltaOut<OUT extends AxGenOut> = {
 };
 type AxGenStreamingOut<OUT extends AxGenOut> = AsyncGenerator<AxGenDeltaOut<OUT>, void, unknown>;
 type AxSetExamplesOptions = {};
-interface AxForwardable<IN extends AxGenIn, OUT extends AxGenOut> {
-    forward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+interface AxForwardable<IN extends AxGenIn, OUT extends AxGenOut, TModelKey> {
+    forward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions<TModelKey>>): Promise<OUT>;
+    streamingForward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions<TModelKey>>): AxGenStreamingOut<OUT>;
 }
 interface AxTunable<IN extends AxGenIn, OUT extends AxGenOut> {
     setExamples: (examples: Readonly<AxProgramExamples<IN, OUT>>, options?: Readonly<AxSetExamplesOptions>) => void;
@@ -3001,7 +3001,7 @@ interface AxUsable {
     getUsage: () => AxProgramUsage[];
     resetUsage: () => void;
 }
-interface AxProgrammable<IN extends AxGenIn, OUT extends AxGenOut> extends AxForwardable<IN, OUT>, AxTunable<IN, OUT>, AxUsable {
+interface AxProgrammable<IN extends AxGenIn, OUT extends AxGenOut, TModelKey = string> extends AxForwardable<IN, OUT, TModelKey>, AxTunable<IN, OUT>, AxUsable {
     getSignature: () => AxSignature;
 }
 type AxProgramUsage = AxChatResponse['modelUsage'] & {
@@ -3115,7 +3115,7 @@ declare class AxGen<IN extends AxGenIn = AxGenIn, OUT extends AxGenOut = AxGenOu
     private streamingFieldProcessors;
     private excludeContentFromTrace;
     private thoughtFieldName;
-    constructor(signature: NonNullable<ConstructorParameters<typeof AxSignature>[0]>, options?: Readonly<AxProgramForwardOptions>);
+    constructor(signature: NonNullable<ConstructorParameters<typeof AxSignature>[0]>, options?: Readonly<AxProgramForwardOptions<any>>);
     private getSignatureName;
     private getMetricsInstruments;
     updateMeter(meter?: Meter): void;
@@ -3128,9 +3128,9 @@ declare class AxGen<IN extends AxGenIn = AxGenIn, OUT extends AxGenOut = AxGenOu
     private forwardSendRequest;
     private forwardCore;
     private _forward2;
-    _forward1(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options: Readonly<AxProgramForwardOptions>): AxGenStreamingOut<OUT>;
-    forward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    _forward1(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options: Readonly<AxProgramForwardOptions<any>>): AxGenStreamingOut<OUT>;
+    forward<T extends Readonly<AxAIService>>(ai: T, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions<NonNullable<ReturnType<T['getModelList']>>[number]['key']>>): Promise<OUT>;
+    streamingForward<T extends Readonly<AxAIService>>(ai: T, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions<NonNullable<ReturnType<T['getModelList']>>[number]['key']>>): AxGenStreamingOut<OUT>;
     setExamples(examples: Readonly<AxProgramExamples<IN, OUT>>, options?: Readonly<AxSetExamplesOptions>): void;
     private isDebug;
     private getLogger;
@@ -3201,8 +3201,8 @@ declare class AxDBManager {
 }
 
 declare class AxDefaultResultReranker extends AxGen<AxRerankerIn, AxRerankerOut> {
-    constructor(options?: Readonly<AxProgramForwardOptions>);
-    forward: (ai: Readonly<AxAIService>, input: Readonly<AxRerankerIn>, options?: Readonly<AxProgramForwardOptions>) => Promise<AxRerankerOut>;
+    constructor(options?: Readonly<AxProgramForwardOptions<string>>);
+    forward: <T extends Readonly<AxAIService>>(ai: T, input: Readonly<AxRerankerIn>, options?: Readonly<AxProgramForwardOptions<NonNullable<ReturnType<T["getModelList"]>>[number]["key"]>>) => Promise<AxRerankerOut>;
 }
 
 interface AxApacheTikaArgs {
@@ -4293,11 +4293,11 @@ interface AxFlowNodeDefinition {
 }
 type AxFlowStepFunction = (state: AxFlowState, context: Readonly<{
     mainAi: AxAIService;
-    mainOptions?: AxProgramForwardOptions;
+    mainOptions?: AxProgramForwardOptions<string>;
 }>) => Promise<AxFlowState> | AxFlowState;
-interface AxFlowDynamicContext {
-    ai?: AxAIService;
-    options?: AxProgramForwardOptions;
+interface AxFlowDynamicContext<T extends Readonly<AxAIService>> {
+    ai?: T;
+    options?: AxProgramForwardOptions<NonNullable<ReturnType<T['getModelList']>>[number]['key']>;
 }
 type GetGenIn<T extends AxProgrammable<AxGenIn, AxGenOut>> = T extends AxProgrammable<infer IN, AxGenOut> ? IN : never;
 type GetGenOut<T extends AxProgrammable<AxGenIn, AxGenOut>> = T extends AxProgrammable<AxGenIn, infer OUT> ? OUT : never;
@@ -4314,20 +4314,20 @@ interface AxFlowable<IN extends AxGenIn, OUT extends AxGenOut> extends AxProgram
 }
 type AxFlowTypedParallelBranch<TNodes extends Record<string, AxProgrammable<any, any>>, TState extends AxFlowState> = (subFlow: AxFlowTypedSubContext<TNodes, TState>) => AxFlowTypedSubContext<TNodes, AxFlowState>;
 interface AxFlowTypedSubContext<TNodes extends Record<string, AxProgrammable<any, any>>, TState extends AxFlowState> {
-    execute<TNodeName extends keyof TNodes & string>(nodeName: TNodeName, mapping: (state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext): AxFlowTypedSubContext<TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
+    execute<TNodeName extends keyof TNodes & string, TAI extends Readonly<AxAIService>>(nodeName: TNodeName, mapping: (state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext<TAI>): AxFlowTypedSubContext<TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
     map<TNewState extends AxFlowState>(transform: (state: TState) => TNewState): AxFlowTypedSubContext<TNodes, TNewState>;
     executeSteps(initialState: TState, context: Readonly<{
         mainAi: AxAIService;
-        mainOptions?: AxProgramForwardOptions;
+        mainOptions?: AxProgramForwardOptions<string>;
     }>): Promise<AxFlowState>;
 }
 type AxFlowParallelBranch = (subFlow: AxFlowSubContext) => AxFlowSubContext;
 interface AxFlowSubContext {
-    execute(nodeName: string, mapping: (state: AxFlowState) => Record<string, AxFieldValue>, dynamicContext?: AxFlowDynamicContext): this;
+    execute<TAI extends Readonly<AxAIService>>(nodeName: string, mapping: (state: AxFlowState) => Record<string, AxFieldValue>, dynamicContext?: AxFlowDynamicContext<TAI>): this;
     map(transform: (state: AxFlowState) => AxFlowState): this;
-    executeSteps(initialState: AxFlowState, context: Readonly<{
-        mainAi: AxAIService;
-        mainOptions?: AxProgramForwardOptions;
+    executeSteps<TAI extends Readonly<AxAIService>>(initialState: AxFlowState, context: Readonly<{
+        mainAi: TAI;
+        mainOptions?: AxProgramForwardOptions<NonNullable<ReturnType<TAI['getModelList']>>[number]['key']>;
     }>): Promise<AxFlowState>;
 }
 interface AxFlowBranchContext {
@@ -4601,6 +4601,8 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
     private readonly autoParallelConfig;
     private readonly executionPlanner;
     private program?;
+    private nodeUsage;
+    private nodeTraces;
     /**
      * Converts a string to camelCase for valid field names
      */
@@ -4640,7 +4642,26 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
     setDemos(demos: readonly AxProgramDemos<IN, OUT>[]): void;
     getUsage(): AxProgramUsage[];
     resetUsage(): void;
-    streamingForward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    /**
+     * Resets trace tracking for the flow.
+     * This is called automatically on each forward/streamingForward call.
+     */
+    resetTraces(): void;
+    /**
+     * Gets a detailed usage report broken down by node name.
+     * This provides visibility into which nodes are consuming the most tokens.
+     *
+     * @returns Object mapping node names to their usage statistics
+     */
+    getUsageReport(): Record<string, AxProgramUsage[]>;
+    /**
+     * Gets a detailed trace report broken down by node name.
+     * This provides visibility into the execution traces for each node.
+     *
+     * @returns Object mapping node names to their trace data
+     */
+    getTracesReport(): Record<string, AxProgramTrace<any, any>[]>;
+    streamingForward<T extends Readonly<AxAIService>>(ai: T, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions<NonNullable<ReturnType<T['getModelList']>>[number]['key']>>): AxGenStreamingOut<OUT>;
     /**
      * Executes the flow with the given AI service and input values.
      *
@@ -4671,7 +4692,7 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
      * @param options - Optional forward options to use as defaults (includes autoParallel override)
      * @returns Promise that resolves to the final output
      */
-    forward(ai: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions & {
+    forward<T extends Readonly<AxAIService>>(ai: T, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions<NonNullable<ReturnType<T['getModelList']>>[number]['key']> & {
         autoParallel?: boolean;
     }>): Promise<OUT>;
     /**
@@ -4680,7 +4701,6 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
      *
      * @param name - The name of the node
      * @param signature - Signature string in the same format as AxSignature
-     * @param options - Optional program forward options (same as AxGen)
      * @returns New AxFlow instance with updated TNodes type
      *
      * @example
@@ -4689,7 +4709,7 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
      * flow.node('analyzer', 'text:string -> analysis:string, confidence:number', { debug: true })
      * ```
      */
-    node<TName extends string, TSig extends string>(name: TName, signature: TSig, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
+    node<TName extends string, TSig extends string>(name: TName, signature: TSig): AxFlow<IN, OUT, TNodes & {
         [K in TName]: InferAxGen<TSig>;
     }, // Add new node to registry
     TState>;
@@ -4699,7 +4719,6 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
      *
      * @param name - The name of the node
      * @param signature - AxSignature instance to use for this node
-     * @param options - Optional program forward options (same as AxGen)
      * @returns New AxFlow instance with updated TNodes type
      *
      * @example
@@ -4708,7 +4727,7 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
      * flow.node('summarizer', sig, { temperature: 0.1 })
      * ```
      */
-    node<TName extends string>(name: TName, signature: AxSignature, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
+    node<TName extends string>(name: TName, signature: AxSignature): AxFlow<IN, OUT, TNodes & {
         [K in TName]: AxGen<AxGenIn, AxGenOut>;
     }, // Add new node to registry
     TState>;
@@ -4747,10 +4766,10 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
     /**
      * Short alias for node() - supports signature strings, AxSignature instances, AxGen instances, and program classes
      */
-    n<TName extends string, TSig extends string>(name: TName, signature: TSig, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
+    n<TName extends string, TSig extends string>(name: TName, signature: TSig): AxFlow<IN, OUT, TNodes & {
         [K in TName]: InferAxGen<TSig>;
     }, TState>;
-    n<TName extends string>(name: TName, signature: AxSignature, options?: Readonly<AxProgramForwardOptions>): AxFlow<IN, OUT, TNodes & {
+    n<TName extends string>(name: TName, signature: AxSignature): AxFlow<IN, OUT, TNodes & {
         [K in TName]: AxGen<AxGenIn, AxGenOut>;
     }, TState>;
     n<TName extends string, TProgram extends new () => AxProgrammable<any, any>>(name: TName, programClass: TProgram): AxFlow<IN, OUT, TNodes & {
@@ -4835,11 +4854,11 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
      * flow.execute('summarizer', state => ({ text: state.originalText }), { ai: cheapAI })
      * ```
      */
-    execute<TNodeName extends keyof TNodes & string>(nodeName: TNodeName, mapping: (_state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext): AxFlow<IN, OUT, TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
+    execute<TNodeName extends keyof TNodes & string, TAI extends Readonly<AxAIService>>(nodeName: TNodeName, mapping: (_state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext<TAI>): AxFlow<IN, OUT, TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
     /**
      * Short alias for execute()
      */
-    e<TNodeName extends keyof TNodes & string>(nodeName: TNodeName, mapping: (_state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext): AxFlow<IN, OUT, TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
+    e<TNodeName extends keyof TNodes & string, TAI extends Readonly<AxAIService>>(nodeName: TNodeName, mapping: (_state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext<TAI>): AxFlow<IN, OUT, TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
     /**
      * Starts a conditional branch based on a predicate function.
      *
@@ -5050,12 +5069,12 @@ TState extends AxFlowState = IN> implements AxFlowable<IN, OUT> {
 declare class AxFlowSubContextImpl implements AxFlowSubContext {
     private readonly nodeGenerators;
     private readonly steps;
-    constructor(nodeGenerators: Map<string, AxForwardable<AxGenIn, AxGenOut> & AxTunable<AxGenIn, AxGenOut> & AxUsable>);
-    execute(nodeName: string, mapping: (state: AxFlowState) => Record<string, AxFieldValue>, dynamicContext?: AxFlowDynamicContext): this;
+    constructor(nodeGenerators: Map<string, AxForwardable<AxGenIn, AxGenOut, string> & AxTunable<AxGenIn, AxGenOut> & AxUsable>);
+    execute<TAI extends Readonly<AxAIService>>(nodeName: string, mapping: (state: AxFlowState) => Record<string, AxFieldValue>, dynamicContext?: AxFlowDynamicContext<TAI>): this;
     map(transform: (state: AxFlowState) => AxFlowState): this;
     executeSteps(initialState: AxFlowState, context: Readonly<{
         mainAi: AxAIService;
-        mainOptions?: AxProgramForwardOptions;
+        mainOptions?: AxProgramForwardOptions<string>;
     }>): Promise<AxFlowState>;
 }
 /**
@@ -5065,11 +5084,11 @@ declare class AxFlowTypedSubContextImpl<TNodes extends Record<string, AxGen<any,
     private readonly nodeGenerators;
     private readonly steps;
     constructor(nodeGenerators: Map<string, AxGen<AxGenIn, AxGenOut> | AxProgram<AxGenIn, AxGenOut>>);
-    execute<TNodeName extends keyof TNodes & string>(nodeName: TNodeName, mapping: (state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext): AxFlowTypedSubContext<TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
+    execute<TNodeName extends keyof TNodes & string, TAI extends Readonly<AxAIService>>(nodeName: TNodeName, mapping: (state: TState) => GetGenIn<TNodes[TNodeName]>, dynamicContext?: AxFlowDynamicContext<TAI>): AxFlowTypedSubContext<TNodes, AddNodeResult<TState, TNodeName, GetGenOut<TNodes[TNodeName]>>>;
     map<TNewState extends AxFlowState>(transform: (state: TState) => TNewState): AxFlowTypedSubContext<TNodes, TNewState>;
     executeSteps(initialState: TState, context: Readonly<{
         mainAi: AxAIService;
-        mainOptions?: AxProgramForwardOptions;
+        mainOptions?: AxProgramForwardOptions<string>;
     }>): Promise<AxFlowState>;
 }
 
@@ -5409,7 +5428,7 @@ interface AxAgentic<IN extends AxGenIn, OUT extends AxGenOut> extends AxProgramm
     getFunction(): AxFunction;
     getFeatures(): AxAgentFeatures;
 }
-type AxAgentOptions = Omit<AxProgramForwardOptions, 'functions'> & {
+type AxAgentOptions = Omit<AxProgramForwardOptions<string>, 'functions'> & {
     disableSmartModelRouting?: boolean;
     /** List of field names that should not be automatically passed from parent to child agents */
     excludeFieldsFromPassthrough?: string[];
@@ -5460,8 +5479,8 @@ declare class AxAgent<IN extends AxGenIn, OUT extends AxGenOut> implements AxAge
      * Initializes the agent's execution context, processing child agents and their functions.
      */
     private init;
-    forward(parentAi: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions>): Promise<OUT>;
-    streamingForward(parentAi: Readonly<AxAIService>, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions>): AxGenStreamingOut<OUT>;
+    forward<T extends Readonly<AxAIService>>(parentAi: T, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramForwardOptions<NonNullable<ReturnType<T['getModelList']>>[number]['key']>>): Promise<OUT>;
+    streamingForward<T extends Readonly<AxAIService>>(parentAi: T, values: IN | AxMessage<IN>[], options?: Readonly<AxProgramStreamingForwardOptions<NonNullable<ReturnType<T['getModelList']>>[number]['key']>>): AxGenStreamingOut<OUT>;
     /**
      * Updates the agent's description.
      * This updates both the stored description and the function's description.
@@ -5477,7 +5496,7 @@ declare class AxAgent<IN extends AxGenIn, OUT extends AxGenOut> implements AxAge
 }
 
 declare class AxChainOfThought<IN extends AxGenIn = AxGenIn, OUT extends AxGenOut = AxGenOut> extends AxGen<IN, OUT> {
-    constructor(signature: Readonly<AxSignature | string>, options?: Readonly<AxProgramForwardOptions & {
+    constructor(signature: Readonly<AxSignature | string>, options?: Readonly<AxProgramForwardOptions<string> & {
         setVisibleReasoning?: boolean;
     }>);
 }
@@ -5491,16 +5510,16 @@ declare class AxRAG extends AxChainOfThought<{
     private genQuery;
     private queryFn;
     private maxHops;
-    constructor(queryFn: (query: string) => Promise<string>, options: Readonly<AxProgramForwardOptions & {
+    constructor(queryFn: (query: string) => Promise<string>, options: Readonly<AxProgramForwardOptions<string> & {
         maxHops?: number;
     }>);
-    forward(ai: Readonly<AxAIService>, values: {
+    forward<T extends Readonly<AxAIService>>(ai: T, values: {
         context: string[];
         question: string;
     } | AxMessage<{
         context: string[];
         question: string;
-    }>[], options?: Readonly<AxProgramForwardOptions>): Promise<{
+    }>[], options?: Readonly<AxProgramForwardOptions<NonNullable<ReturnType<T['getModelList']>>[number]['key']>>): Promise<{
         answer: string;
     }>;
 }