feat: Adding agent support for AI Configs

packages/sdk/server-ai/__tests__/LDAIClientImpl.test.ts

@@ -1,5 +1,6 @@
 import { LDContext } from '@launchdarkly/js-server-sdk-common';
 
+import { LDAIAgentDefaults } from '../src/api/agents';
 import { LDAIDefaults } from '../src/api/config';
 import { LDAIClientImpl } from '../src/LDAIClientImpl';
 import { LDClientMin } from '../src/LDClientMin';
@@ -129,3 +130,122 @@ it('passes the default value to the underlying client', async () => {
 
   expect(mockLdClient.variation).toHaveBeenCalledWith(key, testContext, defaultValue);
 });
+
+it('returns agent config with interpolated instructions', async () => {
+  const client = new LDAIClientImpl(mockLdClient);
+  const key = 'test-flag';
+  const defaultValue: LDAIAgentDefaults = {
+    model: { name: 'test', parameters: { name: 'test-model' } },
+    instructions: 'You are a helpful assistant.',
+    enabled: true,
+  };
+
+  const mockVariation = {
+    model: {
+      name: 'example-model',
+      parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
+    },
+    provider: {
+      name: 'example-provider',
+    },
+    instructions: 'You are a helpful assistant. your name is {{name}} and your score is {{score}}',
+    _ldMeta: {
+      variationKey: 'v1',
+      enabled: true,
+    },
+  };
+
+  mockLdClient.variation.mockResolvedValue(mockVariation);
+
+  const variables = { name: 'John', score: 42 };
+  const result = await client.agents([key], testContext, defaultValue, variables);
+
+  expect(result).toEqual({
+    'test-flag': {
+      model: {
+        name: 'example-model',
+        parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
+      },
+      provider: {
+        name: 'example-provider',
+      },
+      instructions: 'You are a helpful assistant. your name is John and your score is 42',
+      tracker: expect.any(Object),
+      enabled: true,
+    },
+  });
+});
+
+it('includes context in variables for agent instructions interpolation', async () => {
+  const client = new LDAIClientImpl(mockLdClient);
+  const key = 'test-flag';
+  const defaultValue: LDAIAgentDefaults = {
+    model: { name: 'test', parameters: { name: 'test-model' } },
+    instructions: 'You are a helpful assistant.',
+  };
+
+  const mockVariation = {
+    instructions: 'You are a helpful assistant. your user key is {{ldctx.key}}',
+    _ldMeta: { variationKey: 'v1', enabled: true, mode: 'agent' },
+  };
+
+  mockLdClient.variation.mockResolvedValue(mockVariation);
+
+  const result = await client.agents([key], testContext, defaultValue);
+
+  expect(result[key].instructions).toBe('You are a helpful assistant. your user key is test-user');
+});
+
+it('handles missing metadata in agent variation', async () => {
+  const client = new LDAIClientImpl(mockLdClient);
+  const key = 'test-flag';
+  const defaultValue: LDAIAgentDefaults = {
+    model: { name: 'test', parameters: { name: 'test-model' } },
+    instructions: 'You are a helpful assistant.',
+  };
+
+  const mockVariation = {
+    model: { name: 'example-provider', parameters: { name: 'imagination' } },
+    instructions: 'Hello.',
+  };
+
+  mockLdClient.variation.mockResolvedValue(mockVariation);
+
+  const result = await client.agents([key], testContext, defaultValue);
+
+  expect(result).toEqual({
+    'test-flag': {
+      model: { name: 'example-provider', parameters: { name: 'imagination' } },
+      instructions: 'Hello.',
+      tracker: expect.any(Object),
+      enabled: false,
+    },
+  });
+});
+
+it('passes the default value to the underlying client for agents', async () => {
+  const client = new LDAIClientImpl(mockLdClient);
+  const key = 'non-existent-flag';
+  const defaultValue: LDAIAgentDefaults = {
+    model: { name: 'default-model', parameters: { name: 'default' } },
+    provider: { name: 'default-provider' },
+    instructions: 'Default instructions',
+    enabled: true,
+  };
+
+  mockLdClient.variation.mockResolvedValue(defaultValue);
+
+  const result = await client.agents([key], testContext, defaultValue);
+
+  expect(result).toEqual({
+    'non-existent-flag': {
+      model: defaultValue.model,
+      instructions: defaultValue.instructions,
+      provider: defaultValue.provider,
+      tracker: expect.any(Object),
+      enabled: false,
+    },
+  });
+
+  expect(mockLdClient.variation).toHaveBeenCalledWith(key, testContext, defaultValue);
+});

packages/sdk/server-ai/src/LDAIClientImpl.ts

@@ -2,18 +2,29 @@ import * as Mustache from 'mustache';
 
 import { LDContext } from '@launchdarkly/js-server-sdk-common';
 
-import { LDAIConfig, LDAIDefaults, LDMessage, LDModelConfig, LDProviderConfig } from './api/config';
+import { LDAIAgent, LDAIAgentDefaults, LDAIAgents } from './api/agents';
+import {
+  LDAIConfig,
+  LDAIConfigTracker,
+  LDAIDefaults,
+  LDMessage,
+  LDModelConfig,
+  LDProviderConfig,
+} from './api/config';
 import { LDAIClient } from './api/LDAIClient';
 import { LDAIConfigTrackerImpl } from './LDAIConfigTrackerImpl';
 import { LDClientMin } from './LDClientMin';
 
+type Mode = 'completion' | 'agent';
+
 /**
  * Metadata assorted with a model configuration variation.
  */
 interface LDMeta {
   variationKey: string;
   enabled: boolean;
   version?: number;
+  mode?: Mode;
 }
 
 /**
@@ -23,24 +34,38 @@ interface LDMeta {
 interface VariationContent {
   model?: LDModelConfig;
   messages?: LDMessage[];
+  instructions?: string;
   provider?: LDProviderConfig;
   _ldMeta?: LDMeta;
 }
 
+/**
+ * The result of evaluating a configuration.
+ */
+interface EvaluationResult {
+  tracker: LDAIConfigTracker;
+  enabled: boolean;
+  model?: LDModelConfig;
+  provider?: LDProviderConfig;
+  messages?: LDMessage[];
+  instructions?: string;
+  mode?: string;
+}
+
 export class LDAIClientImpl implements LDAIClient {
   constructor(private _ldClient: LDClientMin) {}
 
   private _interpolateTemplate(template: string, variables: Record<string, unknown>): string {
     return Mustache.render(template, variables, undefined, { escape: (item: any) => item });
   }
 
-  async config(
+  private async _evaluate(
     key: string,
     context: LDContext,
     defaultValue: LDAIDefaults,
-    variables?: Record<string, unknown>,
-  ): Promise<LDAIConfig> {
+  ): Promise<EvaluationResult> {
     const value: VariationContent = await this._ldClient.variation(key, context, defaultValue);
+
     const tracker = new LDAIConfigTrackerImpl(
       this._ldClient,
       key,
@@ -50,29 +75,108 @@ export class LDAIClientImpl implements LDAIClient {
       value._ldMeta?.version ?? 1,
       context,
     );
+
     // eslint-disable-next-line no-underscore-dangle
     const enabled = !!value._ldMeta?.enabled;
+
+    return {
+      tracker,
+      enabled,
+      model: value.model,
+      provider: value.provider,
+      messages: value.messages,
+      instructions: value.instructions,
+      // eslint-disable-next-line no-underscore-dangle
+      mode: value._ldMeta?.mode ?? 'completion',
+    };
+  }
+
+  private async _evaluateAgent(
+    key: string,
+    context: LDContext,
+    defaultValue: LDAIAgentDefaults,
+    variables?: Record<string, unknown>,
+  ): Promise<LDAIAgent> {
+    const { tracker, enabled, model, provider, instructions } = await this._evaluate(
+      key,
+      context,
+      defaultValue,
+    );
+
+    const agent: LDAIAgent = {
+      tracker,
+      enabled,
+    };
+    // We are going to modify the contents before returning them, so we make a copy.
+    // This isn't a deep copy and the application developer should not modify the returned content.
+    if (model) {
+      agent.model = { ...model };
+    }
+
+    if (provider) {
+      agent.provider = { ...provider };
+    }
+
+    const allVariables = { ...variables, ldctx: context };
+
+    if (instructions) {
+      agent.instructions = this._interpolateTemplate(instructions, allVariables);
+    }
+
+    return agent;
+  }
+
+  async config(
+    key: string,
+    context: LDContext,
+    defaultValue: LDAIDefaults,
+    variables?: Record<string, unknown>,
+  ): Promise<LDAIConfig> {
+    const { tracker, enabled, model, provider, messages } = await this._evaluate(
+      key,
+      context,
+      defaultValue,
+    );
+
     const config: LDAIConfig = {
       tracker,
       enabled,
     };
     // We are going to modify the contents before returning them, so we make a copy.
     // This isn't a deep copy and the application developer should not modify the returned content.
-    if (value.model) {
-      config.model = { ...value.model };
+    if (model) {
+      config.model = { ...model };
     }
-    if (value.provider) {
-      config.provider = { ...value.provider };
+    if (provider) {
+      config.provider = { ...provider };
     }
     const allVariables = { ...variables, ldctx: context };
 
-    if (value.messages) {
-      config.messages = value.messages.map((entry: any) => ({
+    if (messages) {
+      config.messages = messages.map((entry: any) => ({
         ...entry,
         content: this._interpolateTemplate(entry.content, allVariables),
       }));
     }
 
     return config;
   }
+
+  async agents<TKey extends string>(
+    agentKeys: readonly TKey[],
+    context: LDContext,
+    defaultValue: LDAIAgentDefaults,
+    variables?: Record<string, unknown>,
+  ): Promise<LDAIAgents<TKey>> {
+    const agents = {} as LDAIAgents<TKey>;
+
+    await Promise.all(
+      agentKeys.map(async (agentKey) => {
+        const result = await this._evaluateAgent(agentKey, context, defaultValue, variables);
+        agents[agentKey] = result;
+      }),
+    );
+
+    return agents;
+  }
 }

packages/sdk/server-ai/src/api/LDAIClient.ts

@@ -1,5 +1,6 @@
 import { LDContext } from '@launchdarkly/js-server-sdk-common';
 
+import { LDAIAgentDefaults, LDAIAgents } from './agents';
 import { LDAIConfig, LDAIDefaults } from './config/LDAIConfig';
 
 /**
@@ -63,4 +64,56 @@ export interface LDAIClient {
     defaultValue: LDAIDefaults,
     variables?: Record<string, unknown>,
   ): Promise<LDAIConfig>;
+
+  /**
+   * Retrieves and processes an AI Config agents based on the provided keys, LaunchDarkly context,
+   * and variables. This includes the model configuration and the customized instructions.
+   *
+   * @param agentKeys The keys of the AI Config Agents.
+   * @param context The LaunchDarkly context object that contains relevant information about the
+   * current environment, user, or session. This context may influence how the configuration is
+   * processed or personalized.
+   * @param defaultValue A fallback value containing model configuration and messages. This will
+   * be used if the configuration is not available from LaunchDarkly.
+   * @param variables A map of key-value pairs representing dynamic variables to be injected into
+   * the instruction. The keys correspond to placeholders within the template, and the values
+   * are the corresponding replacements.
+   *
+   * @returns Map of AI `config` agent keys to `agent`, customized `instructions`, and a `tracker`. If the configuration cannot be accessed from
+   * LaunchDarkly, then the return value will include information from the `defaultValue`. The returned `tracker` can
+   * be used to track AI operation metrics (latency, token usage, etc.).
+   *
+   * @example
+   * ```
+   * const agentKeys = ["agent-key-1", "agent-key-2"];
+   * const context = {...};
+   * const variables = {username: 'john'};
+   * const defaultValue = {
+   *  enabled: false,
+   * };
+   *
+   * const result = agents(agentKeys, context, defaultValue, variables);
+   * // Output:
+   * {
+   *   'agent-key-1': {
+   *      enabled: true,
+   *      config: {
+   *        modelId: "gpt-4o",
+   *        temperature: 0.2,
+   *        maxTokens: 4096,
+   *        userDefinedKey: "myValue",
+   *     },
+   *     instructions: "You are an amazing GPT.",
+   *     tracker: ...
+   *   },
+   *   'agent-key-2': {...},
+   * }
+   * ```
+   */
+  agents<TKey extends string>(
+    agentKeys: readonly TKey[],
+    context: LDContext,
+    defaultValue: LDAIAgentDefaults,
+    variables?: Record<string, unknown>,
+  ): Promise<LDAIAgents<TKey>>;
 }