refactor(agent): decouple ai-proxy via factory injection pattern

packages/_example/src/forest/agent.ts

@@ -2,6 +2,7 @@ import type { Schema } from './typings';
 import type { AgentOptions } from '@forestadmin/agent';
 
 import { createAgent } from '@forestadmin/agent';
+import { createAiProvider } from '@forestadmin/ai-proxy';
 import { createMongoDataSource } from '@forestadmin/datasource-mongo';
 import { createMongooseDataSource } from '@forestadmin/datasource-mongoose';
 import { createSequelizeDataSource } from '@forestadmin/datasource-sequelize';
@@ -94,5 +95,11 @@ export default function makeAgent() {
     .customizeCollection('post', customizePost)
     .customizeCollection('comment', customizeComment)
     .customizeCollection('review', customizeReview)
-    .customizeCollection('sales', customizeSales);
+    .customizeCollection('sales', customizeSales)
+    .addAi(createAiProvider({
+      model: 'gpt-4o',
+      provider: 'openai',
+      name: 'test',
+      apiKey: process.env.OPENAI_API_KEY,
+    }));
 }

packages/agent/package.json

@@ -13,7 +13,6 @@
   },
   "dependencies": {
     "@fast-csv/format": "^4.3.5",
-    "@forestadmin/ai-proxy": "1.4.1",
     "@forestadmin/datasource-customizer": "1.67.3",
     "@forestadmin/datasource-toolkit": "1.50.1",
     "@forestadmin/forestadmin-client": "1.37.10",

packages/agent/src/agent.ts

@@ -1,11 +1,6 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 import type { ForestAdminHttpDriverServices } from './services';
-import type {
-  AgentOptions,
-  AgentOptionsWithDefaults,
-  AiConfiguration,
-  HttpCallback,
-} from './types';
+import type { AgentOptions, AgentOptionsWithDefaults, HttpCallback } from './types';
 import type {
   CollectionCustomizer,
   DataSourceChartDefinition,
@@ -14,7 +9,11 @@ import type {
   TCollectionName,
   TSchema,
 } from '@forestadmin/datasource-customizer';
-import type { DataSource, DataSourceFactory } from '@forestadmin/datasource-toolkit';
+import type {
+  AiProviderDefinition,
+  DataSource,
+  DataSourceFactory,
+} from '@forestadmin/datasource-toolkit';
 import type { ForestSchema } from '@forestadmin/forestadmin-client';
 
 import { DataSourceCustomizer } from '@forestadmin/datasource-customizer';
@@ -47,7 +46,7 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
   protected nocodeCustomizer: DataSourceCustomizer<S>;
   protected customizationService: CustomizationService;
   protected schemaGenerator: SchemaGenerator;
-  protected aiConfigurations: AiConfiguration[] = [];
+  protected aiProvider: AiProviderDefinition | null = null;
 
   /** Whether MCP server should be mounted */
   private mcpEnabled = false;
@@ -207,6 +206,7 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
    * Enable MCP (Model Context Protocol) server support.
    * This allows AI assistants to interact with your Forest Admin data.
    *
+   * @see {@link https://docs.forestadmin.com/developer-guide-agents-nodejs/agent-customization/ai/mcp-server}
    * @example
    * agent.mountAiMcpServer();
    */
@@ -222,42 +222,50 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
    * All AI requests from Forest Admin are forwarded to your agent and processed locally.
    * Your data and API keys never transit through Forest Admin servers, ensuring full privacy.
    *
-   * @param configuration - The AI provider configuration
-   * @param configuration.name - A unique name to identify this AI configuration
-   * @param configuration.provider - The AI provider to use ('openai')
-   * @param configuration.apiKey - Your API key for the chosen provider
-   * @param configuration.model - The model to use (e.g., 'gpt-4o')
+   * Requires the `@forestadmin/ai-proxy` package to be installed:
+   * ```bash
+   * npm install @forestadmin/ai-proxy
+   * ```
+   *
+   * @see {@link https://docs.forestadmin.com/developer-guide-agents-nodejs/agent-customization/ai/self-hosted-ai}
+   * @param provider - An AI provider definition created via `createAiProvider` from `@forestadmin/ai-proxy`
    * @returns The agent instance for chaining
    * @throws Error if addAi is called more than once
    *
    * @example
-   * agent.addAi({
+   * import { createAiProvider } from '@forestadmin/ai-proxy';
+   *
+   * agent.addAi(createAiProvider({
    *   name: 'assistant',
    *   provider: 'openai',
    *   apiKey: process.env.OPENAI_API_KEY,
    *   model: 'gpt-4o',
-   * });
+   * }));
    */
-  addAi(configuration: AiConfiguration): this {
-    if (this.aiConfigurations.length > 0) {
+  addAi(provider: AiProviderDefinition): this {
+    if (this.aiProvider) {
       throw new Error(
         'addAi can only be called once. Multiple AI configurations are not supported yet.',
       );
     }
 
-    this.options.logger(
-      'Warn',
-      `AI configuration added with model '${configuration.model}'. ` +
-        'Make sure to test Forest Admin AI features thoroughly to ensure compatibility.',
-    );
+    this.aiProvider = provider;
 
-    this.aiConfigurations.push(configuration);
+    for (const p of provider.providers) {
+      this.options.logger(
+        'Warn',
+        `AI configuration added with model '${p.model}'. ` +
+          'Make sure to test Forest Admin AI features thoroughly to ensure compatibility.',
+      );
+    }
 
     return this;
   }
 
   protected getRoutes(dataSource: DataSource, services: ForestAdminHttpDriverServices) {
-    return makeRoutes(dataSource, this.options, services, this.aiConfigurations);
+    const aiRouter = this.aiProvider?.init(this.options.logger) ?? null;
+
+    return makeRoutes(dataSource, this.options, services, aiRouter);
   }
 
   /**
@@ -380,9 +388,10 @@ export default class Agent<S extends TSchema = TSchema> extends FrameworkMounter
     let schema: Pick<ForestSchema, 'collections'>;
 
     // Get the AI configurations for schema metadata
+    const aiMeta = this.aiProvider?.providers ?? [];
     const { meta } = SchemaGenerator.buildMetadata(
       this.customizationService.buildFeatures(),
-      this.aiConfigurations,
+      aiMeta,
     );
 
     // When using experimental no-code features even in production we need to build a new schema

packages/agent/src/index.ts

@@ -9,6 +9,7 @@ export function createAgent<S extends TSchema = TSchema>(options: AgentOptions):
 
 export { Agent };
 export { AgentOptions } from './types';
+export type { AiProviderDefinition } from './types';
 export * from '@forestadmin/datasource-customizer';
 
 // export is necessary for the agent-generator package

packages/agent/src/routes/ai/ai-proxy.ts

@@ -1,74 +1,40 @@
 import type { ForestAdminHttpDriverServices } from '../../services';
-import type { AgentOptionsWithDefaults, AiConfiguration } from '../../types';
+import type { AgentOptionsWithDefaults } from '../../types';
+import type { AiRouter } from '@forestadmin/datasource-toolkit';
 import type KoaRouter from '@koa/router';
 import type { Context } from 'koa';
 
-import {
-  AIBadRequestError,
-  AIError,
-  AINotConfiguredError,
-  AINotFoundError,
-  Router as AiProxyRouter,
-  extractMcpOauthTokensFromHeaders,
-  injectOauthTokens,
-} from '@forestadmin/ai-proxy';
-import {
-  BadRequestError,
-  NotFoundError,
-  UnprocessableError,
-} from '@forestadmin/datasource-toolkit';
-
 import { HttpCode, RouteType } from '../../types';
 import BaseRoute from '../base-route';
 
 export default class AiProxyRoute extends BaseRoute {
   readonly type = RouteType.PrivateRoute;
-  private readonly aiProxyRouter: AiProxyRouter;
+  private readonly aiRouter: AiRouter;
 
   constructor(
     services: ForestAdminHttpDriverServices,
     options: AgentOptionsWithDefaults,
-    aiConfigurations: AiConfiguration[],
+    aiRouter: AiRouter,
   ) {
     super(services, options);
-    this.aiProxyRouter = new AiProxyRouter({
-      aiConfigurations,
-      logger: this.options.logger,
-    });
+    this.aiRouter = aiRouter;
   }
 
   setupRoutes(router: KoaRouter): void {
     router.post('/_internal/ai-proxy/:route', this.handleAiProxy.bind(this));
   }
 
   private async handleAiProxy(context: Context): Promise<void> {
-    try {
-      const tokensByMcpServerName = extractMcpOauthTokensFromHeaders(context.request.headers);
-
-      const mcpConfigs =
-        await this.options.forestAdminClient.mcpServerConfigService.getConfiguration();
-
-      context.response.body = await this.aiProxyRouter.route({
-        route: context.params.route,
-        body: context.request.body,
-        query: context.query,
-        mcpConfigs: injectOauthTokens({ mcpConfigs, tokensByMcpServerName }),
-      });
-      context.response.status = HttpCode.Ok;
-    } catch (error) {
-      if (error instanceof AIError) {
-        this.options.logger('Error', `AI proxy error: ${error.message}`, error);
-
-        if (error instanceof AINotConfiguredError) {
-          throw new UnprocessableError('AI is not configured. Please call addAi() on your agent.');
-        }
-
-        if (error instanceof AIBadRequestError) throw new BadRequestError(error.message);
-        if (error instanceof AINotFoundError) throw new NotFoundError(error.message);
-        throw new UnprocessableError(error.message);
-      }
-
-      throw error;
-    }
+    const mcpServerConfigs =
+      await this.options.forestAdminClient.mcpServerConfigService.getConfiguration();
+
+    context.response.body = await this.aiRouter.route({
+      route: context.params.route,
+      body: context.request.body,
+      query: context.query,
+      mcpServerConfigs,
+      headers: context.request.headers,
+    });
+    context.response.status = HttpCode.Ok;
   }
 }

packages/agent/src/routes/index.ts

@@ -1,7 +1,7 @@
 import type { ForestAdminHttpDriverServices as Services } from '../services';
-import type { AiConfiguration, AgentOptionsWithDefaults as Options } from '../types';
+import type { AgentOptionsWithDefaults as Options } from '../types';
 import type BaseRoute from './base-route';
-import type { DataSource } from '@forestadmin/datasource-toolkit';
+import type { AiRouter, DataSource } from '@forestadmin/datasource-toolkit';
 
 import CollectionApiChartRoute from './access/api-chart-collection';
 import DataSourceApiChartRoute from './access/api-chart-datasource';
@@ -165,21 +165,17 @@ function getActionRoutes(
   return routes;
 }
 
-function getAiRoutes(
-  options: Options,
-  services: Services,
-  aiConfigurations: AiConfiguration[],
-): BaseRoute[] {
-  if (aiConfigurations.length === 0) return [];
+function getAiRoutes(options: Options, services: Services, aiRouter: AiRouter | null): BaseRoute[] {
+  if (!aiRouter) return [];
 
-  return [new AiProxyRoute(services, options, aiConfigurations)];
+  return [new AiProxyRoute(services, options, aiRouter)];
 }
 
 export default function makeRoutes(
   dataSource: DataSource,
   options: Options,
   services: Services,
-  aiConfigurations: AiConfiguration[] = [],
+  aiRouter: AiRouter | null = null,
 ): BaseRoute[] {
   const routes = [
     ...getRootRoutes(options, services),
@@ -189,7 +185,7 @@ export default function makeRoutes(
     ...getApiChartRoutes(dataSource, options, services),
     ...getRelatedRoutes(dataSource, options, services),
     ...getActionRoutes(dataSource, options, services),
-    ...getAiRoutes(options, services, aiConfigurations),
+    ...getAiRoutes(options, services, aiRouter),
   ];
 
   // Ensure routes and middlewares are loaded in the right order.

packages/agent/src/types.ts

@@ -1,9 +1,13 @@
-import type { AiConfiguration, AiProvider } from '@forestadmin/ai-proxy';
-import type { CompositeId, Logger, LoggerLevel } from '@forestadmin/datasource-toolkit';
+import type {
+  AiProviderDefinition,
+  CompositeId,
+  Logger,
+  LoggerLevel,
+} from '@forestadmin/datasource-toolkit';
 import type { ForestAdminClient } from '@forestadmin/forestadmin-client';
 import type { IncomingMessage, ServerResponse } from 'http';
 
-export type { AiConfiguration, AiProvider };
+export type { AiProviderDefinition };
 
 /** Options to configure behavior of an agent's forestadmin driver */
 export type AgentOptions = {

packages/agent/src/utils/forest-schema/generator.ts

@@ -1,5 +1,5 @@
-import type { AgentOptionsWithDefaults, AiConfiguration } from '../../types';
-import type { DataSource } from '@forestadmin/datasource-toolkit';
+import type { AgentOptionsWithDefaults } from '../../types';
+import type { AiProviderMeta, DataSource } from '@forestadmin/datasource-toolkit';
 import type { ForestSchema } from '@forestadmin/forestadmin-client';
 
 import SchemaGeneratorCollection from './generator-collection';
@@ -23,7 +23,7 @@ export default class SchemaGenerator {
 
   static buildMetadata(
     features: Record<string, string> | null,
-    aiConfigurations: AiConfiguration[] = [],
+    aiProviders: AiProviderMeta[] = [],
   ): Pick<ForestSchema, 'meta'> {
     const { version } = require('../../../package.json'); // eslint-disable-line @typescript-eslint/no-var-requires,global-require
 
@@ -33,8 +33,8 @@ export default class SchemaGenerator {
         liana_version: version,
         liana_features: features,
         ai_llms:
-          aiConfigurations.length > 0
-            ? aiConfigurations.map(c => ({ name: c.name, provider: c.provider }))
+          aiProviders.length > 0
+            ? aiProviders.map(p => ({ name: p.name, provider: p.provider }))
             : null,
         stack: {
           engine: 'nodejs',