Skip to content

feat: Adding agent support for AI Configs #893

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Jul 23, 2025
329 changes: 328 additions & 1 deletion packages/sdk/server-ai/__tests__/LDAIClientImpl.test.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,12 @@
import { LDContext } from '@launchdarkly/js-server-sdk-common';

import { LDAIDefaults } from '../src/api/config';
import { LDAIAgentDefaults } from '../src/api/agents';
import {
LDAIDefaults,
VercelAISDKConfig,
VercelAISDKMapOptions,
VercelAISDKProvider,
} from '../src/api/config';
import { LDAIClientImpl } from '../src/LDAIClientImpl';
import { LDClientMin } from '../src/LDClientMin';

Expand Down Expand Up @@ -79,6 +85,7 @@ it('includes context in variables for messages interpolation', async () => {
const result = await client.config(key, testContext, defaultValue);

expect(result.messages?.[0].content).toBe('User key: test-user');
expect(result.toVercelAISDK).toEqual(expect.any(Function));
});

it('handles missing metadata in variation', async () => {
Expand Down Expand Up @@ -132,3 +139,323 @@ it('passes the default value to the underlying client', async () => {

expect(mockLdClient.variation).toHaveBeenCalledWith(key, testContext, defaultValue);
});

// New agent-related tests
it('returns single agent config with interpolated instructions', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'test-agent';
const defaultValue: LDAIAgentDefaults = {
model: { name: 'test', parameters: { name: 'test-model' } },
instructions: 'You are a helpful assistant.',
enabled: true,
toVercelAISDK: <TMod>(
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
options?: VercelAISDKMapOptions,
): VercelAISDKConfig<TMod> => {
const modelProvider = typeof provider === 'function' ? provider : provider.test;
return {
model: modelProvider('test-model'),
messages: [],
...(options?.nonInterpolatedMessages
? {
messages: options.nonInterpolatedMessages,
}
: {}),
};
},
};

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,
mode: 'agent',
},
};

mockLdClient.variation.mockResolvedValue(mockVariation);

const variables = { name: 'John', score: 42 };
const result = await client.agent(key, testContext, defaultValue, variables);

expect(result).toEqual({
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,
toVercelAISDK: expect.any(Function),
});

// Verify tracking was called
expect(mockLdClient.track).toHaveBeenCalledWith(
'$ld:ai:agent:function:single',
testContext,
key,
1,
);
});

it('includes context in variables for agent instructions interpolation', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'test-agent';
const defaultValue: LDAIAgentDefaults = {
model: { name: 'test', parameters: { name: 'test-model' } },
instructions: 'You are a helpful assistant.',
enabled: true,
toVercelAISDK: <TMod>(
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
options?: VercelAISDKMapOptions,
): VercelAISDKConfig<TMod> => {
const modelProvider = typeof provider === 'function' ? provider : provider.test;
return {
model: modelProvider('test-model'),
messages: [],
...(options?.nonInterpolatedMessages
? {
messages: options.nonInterpolatedMessages,
}
: {}),
};
},
};

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.agent(key, testContext, defaultValue);

expect(result.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-agent';
const defaultValue: LDAIAgentDefaults = {
model: { name: 'test', parameters: { name: 'test-model' } },
instructions: 'You are a helpful assistant.',
enabled: true,
toVercelAISDK: <TMod>(
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
options?: VercelAISDKMapOptions,
): VercelAISDKConfig<TMod> => {
const modelProvider = typeof provider === 'function' ? provider : provider.test;
return {
model: modelProvider('test-model'),
messages: [],
...(options?.nonInterpolatedMessages
? {
messages: options.nonInterpolatedMessages,
}
: {}),
};
},
};

const mockVariation = {
model: { name: 'example-provider', parameters: { name: 'imagination' } },
instructions: 'Hello.',
};

mockLdClient.variation.mockResolvedValue(mockVariation);

const result = await client.agent(key, testContext, defaultValue);

expect(result).toEqual({
model: { name: 'example-provider', parameters: { name: 'imagination' } },
instructions: 'Hello.',
tracker: expect.any(Object),
enabled: false,
toVercelAISDK: expect.any(Function),
});
});

it('passes the default value to the underlying client for single agent', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'non-existent-agent';
const defaultValue: LDAIAgentDefaults = {
model: { name: 'default-model', parameters: { name: 'default' } },
provider: { name: 'default-provider' },
instructions: 'Default instructions',
enabled: true,
toVercelAISDK: <TMod>(
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
options?: VercelAISDKMapOptions,
): VercelAISDKConfig<TMod> => {
const modelProvider =
typeof provider === 'function' ? provider : provider['default-provider'];
return {
model: modelProvider('default-model'),
messages: [],
...(options?.nonInterpolatedMessages
? {
messages: options.nonInterpolatedMessages,
}
: {}),
};
},
};

mockLdClient.variation.mockResolvedValue(defaultValue);

const result = await client.agent(key, testContext, defaultValue);

expect(result).toEqual({
model: defaultValue.model,
instructions: defaultValue.instructions,
provider: defaultValue.provider,
tracker: expect.any(Object),
enabled: false,
toVercelAISDK: expect.any(Function),
});

expect(mockLdClient.variation).toHaveBeenCalledWith(key, testContext, defaultValue);
});

it('returns multiple agents config with interpolated instructions', async () => {
const client = new LDAIClientImpl(mockLdClient);

const agentConfigs = [
{
key: 'research-agent',
defaultValue: {
model: { name: 'test', parameters: { name: 'test-model' } },
instructions: 'You are a research assistant.',
enabled: true,
toVercelAISDK: <TMod>(
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
options?: VercelAISDKMapOptions,
): VercelAISDKConfig<TMod> => {
const modelProvider = typeof provider === 'function' ? provider : provider.test;
return {
model: modelProvider('test-model'),
messages: [],
...(options?.nonInterpolatedMessages
? {
messages: options.nonInterpolatedMessages,
}
: {}),
};
},
},
variables: { topic: 'climate change' },
},
{
key: 'writing-agent',
defaultValue: {
model: { name: 'test', parameters: { name: 'test-model' } },
instructions: 'You are a writing assistant.',
enabled: true,
toVercelAISDK: <TMod>(
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
options?: VercelAISDKMapOptions,
): VercelAISDKConfig<TMod> => {
const modelProvider = typeof provider === 'function' ? provider : provider.test;
return {
model: modelProvider('test-model'),
messages: [],
...(options?.nonInterpolatedMessages
? {
messages: options.nonInterpolatedMessages,
}
: {}),
};
},
},
variables: { style: 'academic' },
},
] as const;

const mockVariations = {
'research-agent': {
model: {
name: 'research-model',
parameters: { temperature: 0.3, maxTokens: 2048 },
},
provider: { name: 'openai' },
instructions: 'You are a research assistant specializing in {{topic}}.',
_ldMeta: { variationKey: 'v1', enabled: true, mode: 'agent' },
},
'writing-agent': {
model: {
name: 'writing-model',
parameters: { temperature: 0.7, maxTokens: 1024 },
},
provider: { name: 'anthropic' },
instructions: 'You are a writing assistant with {{style}} style.',
_ldMeta: { variationKey: 'v2', enabled: true, mode: 'agent' },
},
};

mockLdClient.variation.mockImplementation((key) =>
Promise.resolve(mockVariations[key as keyof typeof mockVariations]),
);

const result = await client.agents(agentConfigs, testContext);

expect(result).toEqual({
'research-agent': {
model: {
name: 'research-model',
parameters: { temperature: 0.3, maxTokens: 2048 },
},
provider: { name: 'openai' },
instructions: 'You are a research assistant specializing in climate change.',
tracker: expect.any(Object),
enabled: true,
toVercelAISDK: expect.any(Function),
},
'writing-agent': {
model: {
name: 'writing-model',
parameters: { temperature: 0.7, maxTokens: 1024 },
},
provider: { name: 'anthropic' },
instructions: 'You are a writing assistant with academic style.',
tracker: expect.any(Object),
enabled: true,
toVercelAISDK: expect.any(Function),
},
});

// Verify tracking was called
expect(mockLdClient.track).toHaveBeenCalledWith(
'$ld:ai:agent:function:multiple',
testContext,
agentConfigs.length,
agentConfigs.length,
);
});

it('handles empty agent configs array', async () => {
const client = new LDAIClientImpl(mockLdClient);

const result = await client.agents([], testContext);

expect(result).toEqual({});

// Verify tracking was called with 0 agents
expect(mockLdClient.track).toHaveBeenCalledWith(
'$ld:ai:agent:function:multiple',
testContext,
0,
0,
);
});
Loading