Skip to content

Commit

Permalink
Merge branch 'main' into rlamb/add-http-collectors
Browse files Browse the repository at this point in the history
  • Loading branch information
kinyoklion authored Nov 8, 2024
2 parents 4717551 + 89ce6db commit 537fe98
Show file tree
Hide file tree
Showing 6 changed files with 496 additions and 2 deletions.
136 changes: 136 additions & 0 deletions packages/sdk/server-ai/__tests__/LDAIClientImpl.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
import { LDContext } from '@launchdarkly/js-server-sdk-common';

import { LDGenerationConfig } from '../src/api/config';
import { LDAIClientImpl } from '../src/LDAIClientImpl';
import { LDClientMin } from '../src/LDClientMin';

const mockLdClient: jest.Mocked<LDClientMin> = {
variation: jest.fn(),
track: jest.fn(),
};

const testContext: LDContext = { kind: 'user', key: 'test-user' };

it('interpolates template variables', () => {
const client = new LDAIClientImpl(mockLdClient);
const template = 'Hello {{name}}, your score is {{score}}';
const variables = { name: 'John', score: 42 };

const result = client.interpolateTemplate(template, variables);
expect(result).toBe('Hello John, your score is 42');
});

it('handles empty variables in template interpolation', () => {
const client = new LDAIClientImpl(mockLdClient);
const template = 'Hello {{name}}';
const variables = {};

const result = client.interpolateTemplate(template, variables);
expect(result).toBe('Hello ');
});

it('returns model config with interpolated prompts', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'test-flag';
const defaultValue: LDGenerationConfig = {
model: { modelId: 'test', name: 'test-model' },
prompt: [],
};

const mockVariation = {
model: { modelId: 'example-provider', name: 'imagination' },
prompt: [
{ role: 'system', content: 'Hello {{name}}' },
{ role: 'user', content: 'Score: {{score}}' },
],
_ldMeta: {
versionKey: 'v1',
enabled: true,
},
};

mockLdClient.variation.mockResolvedValue(mockVariation);

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

expect(result).toEqual({
config: {
model: { modelId: 'example-provider', name: 'imagination' },
prompt: [
{ role: 'system', content: 'Hello John' },
{ role: 'user', content: 'Score: 42' },
],
},
tracker: expect.any(Object),
enabled: true,
});
});

it('includes context in variables for prompt interpolation', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'test-flag';
const defaultValue: LDGenerationConfig = {
model: { modelId: 'test', name: 'test-model' },
prompt: [],
};

const mockVariation = {
prompt: [{ role: 'system', content: 'User key: {{ldctx.key}}' }],
_ldMeta: { versionKey: 'v1', enabled: true },
};

mockLdClient.variation.mockResolvedValue(mockVariation);

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

expect(result.config.prompt?.[0].content).toBe('User key: test-user');
});

it('handles missing metadata in variation', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'test-flag';
const defaultValue: LDGenerationConfig = {
model: { modelId: 'test', name: 'test-model' },
prompt: [],
};

const mockVariation = {
model: { modelId: 'example-provider', name: 'imagination' },
prompt: [{ role: 'system', content: 'Hello' }],
};

mockLdClient.variation.mockResolvedValue(mockVariation);

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

expect(result).toEqual({
config: {
model: { modelId: 'example-provider', name: 'imagination' },
prompt: [{ role: 'system', content: 'Hello' }],
},
tracker: expect.any(Object),
enabled: false,
});
});

it('passes the default value to the underlying client', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'non-existent-flag';
const defaultValue: LDGenerationConfig = {
model: { modelId: 'default-model', name: 'default' },
prompt: [{ role: 'system', content: 'Default prompt' }],
};

mockLdClient.variation.mockResolvedValue(defaultValue);

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

expect(result).toEqual({
config: defaultValue,
tracker: expect.any(Object),
enabled: false,
});

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

0 comments on commit 537fe98

Please sign in to comment.