Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 85 additions & 0 deletions packages/sdk/server-ai/__tests__/LDAIClientImpl.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -846,6 +846,91 @@ describe('createTracker method', () => {
configKey: 'my-config',
variationKey: 'v1',
version: 3,
modelVersion: 1,
});
expect('modelKey' in tracker.getTrackData()).toBe(false);
});
});

describe('modelKey and modelVersion tracking', () => {
it('stamps modelKey and modelVersion on tracker from flag payload', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'test-flag';

mockLdClient.variation.mockResolvedValue({
model: {
name: 'gpt-4',
modelKey: 'my-model',
modelVersion: 2,
},
provider: { name: 'openai' },
messages: [],
_ldMeta: {
variationKey: 'v1',
enabled: true,
mode: 'completion',
version: 7,
},
});

const result = await client.completionConfig(key, testContext);
const tracker = result.createTracker();

expect(tracker.getTrackData()).toMatchObject({
modelKey: 'my-model',
modelVersion: 2,
});
});

it('defaults modelVersion to 1 and omits modelKey when absent from flag payload', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'test-flag';

mockLdClient.variation.mockResolvedValue({
model: { name: 'gpt-4' },
provider: { name: 'openai' },
messages: [],
_ldMeta: {
variationKey: 'v1',
enabled: true,
mode: 'completion',
version: 7,
},
});

const result = await client.completionConfig(key, testContext);
const tracker = result.createTracker();

const trackData = tracker.getTrackData();
expect('modelKey' in trackData).toBe(false);
expect(trackData.modelVersion).toBe(1);
});

it('exposes modelKey and modelVersion on config.model from flag payload', async () => {
const client = new LDAIClientImpl(mockLdClient);
const key = 'test-flag';

mockLdClient.variation.mockResolvedValue({
model: {
name: 'gpt-4',
modelKey: 'my-model',
modelVersion: 2,
},
provider: { name: 'openai' },
messages: [],
_ldMeta: {
variationKey: 'v1',
enabled: true,
mode: 'completion',
},
});

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

expect(result.model).toEqual({
name: 'gpt-4',
modelKey: 'my-model',
modelVersion: 2,
});
});
});
Expand Down
106 changes: 106 additions & 0 deletions packages/sdk/server-ai/__tests__/LDAIConfigTrackerImpl.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ const getExpectedTrackData = () => ({
version,
modelName,
providerName,
modelVersion: 1,
runId: testRunId,
});

Expand Down Expand Up @@ -1193,3 +1194,108 @@ describe('fromResumptionToken', () => {
expect('graphKey' in reconstructed.getTrackData()).toBe(false);
});
});

it('includes modelKey and modelVersion in getTrackData when set on constructor', () => {
const tracker = new LDAIConfigTrackerImpl(
mockLdClient,
testRunId,
configKey,
variationKey,
version,
modelName,
providerName,
testContext,
undefined,
'my-model',
2,
);

expect(tracker.getTrackData()).toEqual({
...getExpectedTrackData(),
modelKey: 'my-model',
modelVersion: 2,
});
});

it('omits modelKey from getTrackData when unset', () => {
const tracker = new LDAIConfigTrackerImpl(
mockLdClient,
testRunId,
configKey,
variationKey,
version,
modelName,
providerName,
testContext,
);

const trackData = tracker.getTrackData();
expect('modelKey' in trackData).toBe(false);
expect(trackData.modelVersion).toBe(1);
});

it('omits modelKey from getTrackData when empty string', () => {
const tracker = new LDAIConfigTrackerImpl(
mockLdClient,
testRunId,
configKey,
variationKey,
version,
modelName,
providerName,
testContext,
undefined,
'',
3,
);

const trackData = tracker.getTrackData();
expect('modelKey' in trackData).toBe(false);
expect(trackData.modelVersion).toBe(3);
});

it('excludes modelKey and modelVersion from resumption token', () => {
const tracker = new LDAIConfigTrackerImpl(
mockLdClient,
testRunId,
configKey,
variationKey,
version,
modelName,
providerName,
testContext,
undefined,
'my-model',
2,
);

const decoded = JSON.parse(Buffer.from(tracker.resumptionToken, 'base64url').toString('utf8'));
expect('modelKey' in decoded).toBe(false);
expect('modelVersion' in decoded).toBe(false);
});

it('fromResumptionToken defaults modelVersion to 1 and omits modelKey', () => {
const original = new LDAIConfigTrackerImpl(
mockLdClient,
testRunId,
configKey,
variationKey,
version,
modelName,
providerName,
testContext,
undefined,
'my-model',
2,
);

const reconstructed = LDAIConfigTrackerImpl.fromResumptionToken(
original.resumptionToken,
mockLdClient,
testContext,
);

const trackData = reconstructed.getTrackData();
expect('modelKey' in trackData).toBe(false);
expect(trackData.modelVersion).toBe(1);
});
2 changes: 2 additions & 0 deletions packages/sdk/server-ai/src/LDAIClientImpl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,8 @@ export class LDAIClientImpl implements LDAIClient {
value.provider?.name ?? '',
context,
graphKey,
value.model?.modelKey,
value.model?.modelVersion ?? 1,
);

// Validate mode match
Expand Down
6 changes: 6 additions & 0 deletions packages/sdk/server-ai/src/LDAIConfigTrackerImpl.ts
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
private _providerName: string,
private _context: LDContext,
private _graphKey?: string,
private _modelKey?: string,
private _modelVersion: number = 1,
) {
this._trackedMetrics.resumptionToken = this.resumptionToken;
}
Expand All @@ -30,6 +32,8 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
version: number;
modelName: string;
providerName: string;
modelVersion: number;
modelKey?: string;
graphKey?: string;
} {
return {
Expand All @@ -39,6 +43,8 @@ export class LDAIConfigTrackerImpl implements LDAIConfigTracker {
version: this._version,
modelName: this._modelName,
providerName: this._providerName,
modelVersion: this._modelVersion,
...(this._modelKey ? { modelKey: this._modelKey } : {}),
...(this._graphKey !== undefined ? { graphKey: this._graphKey } : {}),
};
}
Expand Down
2 changes: 2 additions & 0 deletions packages/sdk/server-ai/src/api/config/LDAIConfigTracker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@ export interface LDAIConfigTracker {
version: number;
modelName: string;
providerName: string;
modelVersion: number;
modelKey?: string;
graphKey?: string;
};

Expand Down
11 changes: 11 additions & 0 deletions packages/sdk/server-ai/src/api/config/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,17 @@ export interface LDModelConfig {
* The region for the model.
*/
region?: string;

/**
* The stable, unique key of the model (used for direct lookup; distinct from
* `name`, which is not guaranteed unique).
*/
modelKey?: string;

/**
* The pinned version of the model that this config variation references.
*/
modelVersion?: number;
}

export interface LDProviderConfig {
Expand Down
Loading