Skip to content

Commit bf95b92

Browse files
ctawiahctawiah
authored
feat: Adding agent support for AI Configs (#893)
**Requirements** - [x] I have added test coverage for new or changed functionality - [x] I have followed the repository's [pull request submission guidelines](../blob/main/CONTRIBUTING.md#submitting-pull-requests) - [x] I have validated my changes against all supported platform versions **Related issues** Provide links to any issues in this repository or elsewhere relating to this pull request. **Describe the solution you've provided** Provide a clear and concise description of what you expect to happen. **Describe alternatives you've considered** Provide a clear and concise description of any alternative solutions or features you've considered. **Additional context** Add any other context about the pull request here.
1 parent f53315c commit bf95b92

File tree

6 files changed

+590
-13
lines changed

6 files changed

+590
-13
lines changed

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

Lines changed: 328 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
11
import { LDContext } from '@launchdarkly/js-server-sdk-common';
22

3-
import { LDAIDefaults } from '../src/api/config';
3+
import { LDAIAgentDefaults } from '../src/api/agents';
4+
import {
5+
LDAIDefaults,
6+
VercelAISDKConfig,
7+
VercelAISDKMapOptions,
8+
VercelAISDKProvider,
9+
} from '../src/api/config';
410
import { LDAIClientImpl } from '../src/LDAIClientImpl';
511
import { LDClientMin } from '../src/LDClientMin';
612

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

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

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

133140
expect(mockLdClient.variation).toHaveBeenCalledWith(key, testContext, defaultValue);
134141
});
142+
143+
// New agent-related tests
144+
it('returns single agent config with interpolated instructions', async () => {
145+
const client = new LDAIClientImpl(mockLdClient);
146+
const key = 'test-agent';
147+
const defaultValue: LDAIAgentDefaults = {
148+
model: { name: 'test', parameters: { name: 'test-model' } },
149+
instructions: 'You are a helpful assistant.',
150+
enabled: true,
151+
toVercelAISDK: <TMod>(
152+
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
153+
options?: VercelAISDKMapOptions,
154+
): VercelAISDKConfig<TMod> => {
155+
const modelProvider = typeof provider === 'function' ? provider : provider.test;
156+
return {
157+
model: modelProvider('test-model'),
158+
messages: [],
159+
...(options?.nonInterpolatedMessages
160+
? {
161+
messages: options.nonInterpolatedMessages,
162+
}
163+
: {}),
164+
};
165+
},
166+
};
167+
168+
const mockVariation = {
169+
model: {
170+
name: 'example-model',
171+
parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
172+
},
173+
provider: {
174+
name: 'example-provider',
175+
},
176+
instructions: 'You are a helpful assistant. Your name is {{name}} and your score is {{score}}',
177+
_ldMeta: {
178+
variationKey: 'v1',
179+
enabled: true,
180+
mode: 'agent',
181+
},
182+
};
183+
184+
mockLdClient.variation.mockResolvedValue(mockVariation);
185+
186+
const variables = { name: 'John', score: 42 };
187+
const result = await client.agent(key, testContext, defaultValue, variables);
188+
189+
expect(result).toEqual({
190+
model: {
191+
name: 'example-model',
192+
parameters: { name: 'imagination', temperature: 0.7, maxTokens: 4096 },
193+
},
194+
provider: {
195+
name: 'example-provider',
196+
},
197+
instructions: 'You are a helpful assistant. Your name is John and your score is 42',
198+
tracker: expect.any(Object),
199+
enabled: true,
200+
toVercelAISDK: expect.any(Function),
201+
});
202+
203+
// Verify tracking was called
204+
expect(mockLdClient.track).toHaveBeenCalledWith(
205+
'$ld:ai:agent:function:single',
206+
testContext,
207+
key,
208+
1,
209+
);
210+
});
211+
212+
it('includes context in variables for agent instructions interpolation', async () => {
213+
const client = new LDAIClientImpl(mockLdClient);
214+
const key = 'test-agent';
215+
const defaultValue: LDAIAgentDefaults = {
216+
model: { name: 'test', parameters: { name: 'test-model' } },
217+
instructions: 'You are a helpful assistant.',
218+
enabled: true,
219+
toVercelAISDK: <TMod>(
220+
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
221+
options?: VercelAISDKMapOptions,
222+
): VercelAISDKConfig<TMod> => {
223+
const modelProvider = typeof provider === 'function' ? provider : provider.test;
224+
return {
225+
model: modelProvider('test-model'),
226+
messages: [],
227+
...(options?.nonInterpolatedMessages
228+
? {
229+
messages: options.nonInterpolatedMessages,
230+
}
231+
: {}),
232+
};
233+
},
234+
};
235+
236+
const mockVariation = {
237+
instructions: 'You are a helpful assistant. Your user key is {{ldctx.key}}',
238+
_ldMeta: { variationKey: 'v1', enabled: true, mode: 'agent' },
239+
};
240+
241+
mockLdClient.variation.mockResolvedValue(mockVariation);
242+
243+
const result = await client.agent(key, testContext, defaultValue);
244+
245+
expect(result.instructions).toBe('You are a helpful assistant. Your user key is test-user');
246+
});
247+
248+
it('handles missing metadata in agent variation', async () => {
249+
const client = new LDAIClientImpl(mockLdClient);
250+
const key = 'test-agent';
251+
const defaultValue: LDAIAgentDefaults = {
252+
model: { name: 'test', parameters: { name: 'test-model' } },
253+
instructions: 'You are a helpful assistant.',
254+
enabled: true,
255+
toVercelAISDK: <TMod>(
256+
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
257+
options?: VercelAISDKMapOptions,
258+
): VercelAISDKConfig<TMod> => {
259+
const modelProvider = typeof provider === 'function' ? provider : provider.test;
260+
return {
261+
model: modelProvider('test-model'),
262+
messages: [],
263+
...(options?.nonInterpolatedMessages
264+
? {
265+
messages: options.nonInterpolatedMessages,
266+
}
267+
: {}),
268+
};
269+
},
270+
};
271+
272+
const mockVariation = {
273+
model: { name: 'example-provider', parameters: { name: 'imagination' } },
274+
instructions: 'Hello.',
275+
};
276+
277+
mockLdClient.variation.mockResolvedValue(mockVariation);
278+
279+
const result = await client.agent(key, testContext, defaultValue);
280+
281+
expect(result).toEqual({
282+
model: { name: 'example-provider', parameters: { name: 'imagination' } },
283+
instructions: 'Hello.',
284+
tracker: expect.any(Object),
285+
enabled: false,
286+
toVercelAISDK: expect.any(Function),
287+
});
288+
});
289+
290+
it('passes the default value to the underlying client for single agent', async () => {
291+
const client = new LDAIClientImpl(mockLdClient);
292+
const key = 'non-existent-agent';
293+
const defaultValue: LDAIAgentDefaults = {
294+
model: { name: 'default-model', parameters: { name: 'default' } },
295+
provider: { name: 'default-provider' },
296+
instructions: 'Default instructions',
297+
enabled: true,
298+
toVercelAISDK: <TMod>(
299+
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
300+
options?: VercelAISDKMapOptions,
301+
): VercelAISDKConfig<TMod> => {
302+
const modelProvider =
303+
typeof provider === 'function' ? provider : provider['default-provider'];
304+
return {
305+
model: modelProvider('default-model'),
306+
messages: [],
307+
...(options?.nonInterpolatedMessages
308+
? {
309+
messages: options.nonInterpolatedMessages,
310+
}
311+
: {}),
312+
};
313+
},
314+
};
315+
316+
mockLdClient.variation.mockResolvedValue(defaultValue);
317+
318+
const result = await client.agent(key, testContext, defaultValue);
319+
320+
expect(result).toEqual({
321+
model: defaultValue.model,
322+
instructions: defaultValue.instructions,
323+
provider: defaultValue.provider,
324+
tracker: expect.any(Object),
325+
enabled: false,
326+
toVercelAISDK: expect.any(Function),
327+
});
328+
329+
expect(mockLdClient.variation).toHaveBeenCalledWith(key, testContext, defaultValue);
330+
});
331+
332+
it('returns multiple agents config with interpolated instructions', async () => {
333+
const client = new LDAIClientImpl(mockLdClient);
334+
335+
const agentConfigs = [
336+
{
337+
key: 'research-agent',
338+
defaultValue: {
339+
model: { name: 'test', parameters: { name: 'test-model' } },
340+
instructions: 'You are a research assistant.',
341+
enabled: true,
342+
toVercelAISDK: <TMod>(
343+
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
344+
options?: VercelAISDKMapOptions,
345+
): VercelAISDKConfig<TMod> => {
346+
const modelProvider = typeof provider === 'function' ? provider : provider.test;
347+
return {
348+
model: modelProvider('test-model'),
349+
messages: [],
350+
...(options?.nonInterpolatedMessages
351+
? {
352+
messages: options.nonInterpolatedMessages,
353+
}
354+
: {}),
355+
};
356+
},
357+
},
358+
variables: { topic: 'climate change' },
359+
},
360+
{
361+
key: 'writing-agent',
362+
defaultValue: {
363+
model: { name: 'test', parameters: { name: 'test-model' } },
364+
instructions: 'You are a writing assistant.',
365+
enabled: true,
366+
toVercelAISDK: <TMod>(
367+
provider: VercelAISDKProvider<TMod> | Record<string, VercelAISDKProvider<TMod>>,
368+
options?: VercelAISDKMapOptions,
369+
): VercelAISDKConfig<TMod> => {
370+
const modelProvider = typeof provider === 'function' ? provider : provider.test;
371+
return {
372+
model: modelProvider('test-model'),
373+
messages: [],
374+
...(options?.nonInterpolatedMessages
375+
? {
376+
messages: options.nonInterpolatedMessages,
377+
}
378+
: {}),
379+
};
380+
},
381+
},
382+
variables: { style: 'academic' },
383+
},
384+
] as const;
385+
386+
const mockVariations = {
387+
'research-agent': {
388+
model: {
389+
name: 'research-model',
390+
parameters: { temperature: 0.3, maxTokens: 2048 },
391+
},
392+
provider: { name: 'openai' },
393+
instructions: 'You are a research assistant specializing in {{topic}}.',
394+
_ldMeta: { variationKey: 'v1', enabled: true, mode: 'agent' },
395+
},
396+
'writing-agent': {
397+
model: {
398+
name: 'writing-model',
399+
parameters: { temperature: 0.7, maxTokens: 1024 },
400+
},
401+
provider: { name: 'anthropic' },
402+
instructions: 'You are a writing assistant with {{style}} style.',
403+
_ldMeta: { variationKey: 'v2', enabled: true, mode: 'agent' },
404+
},
405+
};
406+
407+
mockLdClient.variation.mockImplementation((key) =>
408+
Promise.resolve(mockVariations[key as keyof typeof mockVariations]),
409+
);
410+
411+
const result = await client.agents(agentConfigs, testContext);
412+
413+
expect(result).toEqual({
414+
'research-agent': {
415+
model: {
416+
name: 'research-model',
417+
parameters: { temperature: 0.3, maxTokens: 2048 },
418+
},
419+
provider: { name: 'openai' },
420+
instructions: 'You are a research assistant specializing in climate change.',
421+
tracker: expect.any(Object),
422+
enabled: true,
423+
toVercelAISDK: expect.any(Function),
424+
},
425+
'writing-agent': {
426+
model: {
427+
name: 'writing-model',
428+
parameters: { temperature: 0.7, maxTokens: 1024 },
429+
},
430+
provider: { name: 'anthropic' },
431+
instructions: 'You are a writing assistant with academic style.',
432+
tracker: expect.any(Object),
433+
enabled: true,
434+
toVercelAISDK: expect.any(Function),
435+
},
436+
});
437+
438+
// Verify tracking was called
439+
expect(mockLdClient.track).toHaveBeenCalledWith(
440+
'$ld:ai:agent:function:multiple',
441+
testContext,
442+
agentConfigs.length,
443+
agentConfigs.length,
444+
);
445+
});
446+
447+
it('handles empty agent configs array', async () => {
448+
const client = new LDAIClientImpl(mockLdClient);
449+
450+
const result = await client.agents([], testContext);
451+
452+
expect(result).toEqual({});
453+
454+
// Verify tracking was called with 0 agents
455+
expect(mockLdClient.track).toHaveBeenCalledWith(
456+
'$ld:ai:agent:function:multiple',
457+
testContext,
458+
0,
459+
0,
460+
);
461+
});

0 commit comments

Comments
 (0)