From 2f128cb1e3b6b3ad9abb82c74059711788755cc0 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Wed, 8 Jul 2026 03:13:57 +0000 Subject: [PATCH] feat(ldai): stamp modelKey and modelVersion on AI usage events (AIC-2850) Read modelKey and modelVersion from the AI Config variation payload and expose them on Config. Stamp both fields on every metric event alongside existing modelName and providerName. modelVersion defaults to 1 when absent; modelKey is omitted when empty. Neither field is included in the resumption token. Additive and backward compatible with older payloads. Co-authored-by: Anthony Torns II --- ldai/client.go | 4 +++ ldai/client_test.go | 68 ++++++++++++++++++++++++++++++++++++ ldai/config.go | 30 ++++++++++++++++ ldai/datamodel/datamodel.go | 6 ++++ ldai/tracker.go | 15 +++++--- ldai/tracker_test.go | 69 +++++++++++++++++++++++++++++++++++-- 6 files changed, 184 insertions(+), 8 deletions(-) diff --git a/ldai/client.go b/ldai/client.go index 76eddc5e..afebf6be 100644 --- a/ldai/client.go +++ b/ldai/client.go @@ -156,12 +156,16 @@ func (c *Client) evaluateConfig( builder := NewConfig(). WithModelName(parsed.Model.Name). + WithModelKey(parsed.Model.Key). WithProviderName(parsed.Provider.Name). WithEnabled(parsed.Meta.Enabled). WithMode(parsed.Mode). WithEvaluationMetricKey(parsed.EvaluationMetricKey). WithEvaluationMetricKeys(parsed.EvaluationMetricKeys). WithJudgeConfiguration(parsed.JudgeConfiguration) + if parsed.Model.Version != nil { + builder.WithModelVersion(*parsed.Model.Version) + } for k, v := range parsed.Model.Parameters { builder.WithModelParam(k, v) diff --git a/ldai/client_test.go b/ldai/client_test.go index 01805f1e..30847178 100644 --- a/ldai/client_test.go +++ b/ldai/client_test.go @@ -150,6 +150,72 @@ func TestParseModelName(t *testing.T) { } } +func TestParseModelKeyAndVersion(t *testing.T) { + tests := []struct { + name string + json []byte + expectedKey string + expectedVersion int + }{ + { + name: "missing", + json: []byte(`{"model": {"name": "gpt-4"}}`), + expectedKey: "", + expectedVersion: 1, + }, + { + name: "modelKey and modelVersion set", + json: []byte(`{"model": {"name": "gpt-4", "modelKey": "my-model", "modelVersion": 2}}`), + expectedKey: "my-model", + expectedVersion: 2, + }, + { + name: "modelVersion only", + json: []byte(`{"model": {"name": "gpt-4", "modelVersion": 3}}`), + expectedKey: "", + expectedVersion: 3, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + client, err := NewClient(newMockSDK(test.json, nil)) + require.NoError(t, err) + require.NotNil(t, client) + + defaultVal := NewConfig().Enable().WithMessage("hello", datamodel.User).Build() + cfg := client.CompletionConfig("key", ldcontext.New("user"), defaultVal, nil) + + assert.Equal(t, test.expectedKey, cfg.ModelKey()) + assert.Equal(t, test.expectedVersion, cfg.ModelVersion()) + }) + } +} + +func TestCreateTrackerStampsModelKeyAndVersionOnTrackData(t *testing.T) { + configJSON := []byte(`{ + "_ldMeta": {"variationKey": "var-1", "enabled": true, "version": 1}, + "model": {"name": "gpt-4", "modelKey": "my-model", "modelVersion": 2}, + "provider": {"name": "openai"}, + "messages": [{"content": "hello", "role": "user"}] + }`) + + mockSDK := newMockSDK(configJSON, nil) + client, err := NewClient(mockSDK) + require.NoError(t, err) + mockSDK.events = nil + + cfg := client.CompletionConfig("my-config", ldcontext.New("user"), Disabled(), nil) + tracker := cfg.CreateTracker() + require.NotNil(t, tracker) + assert.NoError(t, tracker.TrackSuccess()) + + require.NotEmpty(t, mockSDK.events) + data := mockSDK.events[len(mockSDK.events)-1].data + assert.Equal(t, "my-model", data.GetByKey("modelKey").StringValue()) + assert.Equal(t, 2, data.GetByKey("modelVersion").IntValue()) +} + func TestParseProviderName(t *testing.T) { tests := []struct { name string @@ -1009,6 +1075,8 @@ func TestClient_CreateTracker_RoundTrip(t *testing.T) { // modelName and providerName should be empty on reconstructed tracker assert.Equal(t, "", feedbackEvent.data.GetByKey("modelName").StringValue()) assert.Equal(t, "", feedbackEvent.data.GetByKey("providerName").StringValue()) + assert.False(t, feedbackEvent.data.GetByKey("modelKey").IsDefined()) + assert.Equal(t, 1, feedbackEvent.data.GetByKey("modelVersion").IntValue()) } func TestClient_CreateTracker_InvalidToken(t *testing.T) { diff --git a/ldai/config.go b/ldai/config.go index 8fa58bfd..39e455c6 100644 --- a/ldai/config.go +++ b/ldai/config.go @@ -49,6 +49,20 @@ func (c *Config) ModelName() string { return c.c.Model.Name } +// ModelKey returns the stable, unique key of the model (used for direct lookup; distinct from +// ModelName, which is not guaranteed unique). +func (c *Config) ModelKey() string { + return c.c.Model.Key +} + +// ModelVersion returns the pinned version of the model that this config variation references. +func (c *Config) ModelVersion() int { + if c.c.Model.Version == nil { + return 1 + } + return *c.c.Model.Version +} + // ModelParam returns the model parameter named by key. The second parameter is true if the key exists. func (c *Config) ModelParam(key string) (ldvalue.Value, bool) { val, ok := c.c.Model.Parameters[key] @@ -113,6 +127,8 @@ type ConfigBuilder struct { enabled bool providerName string modelName string + modelKey string + modelVersion *int modelParams map[string]ldvalue.Value modelCustomParams map[string]ldvalue.Value mode string @@ -165,6 +181,18 @@ func (cb *ConfigBuilder) WithModelName(modelName string) *ConfigBuilder { return cb } +// WithModelKey sets the stable, unique key of the model associated with the config. +func (cb *ConfigBuilder) WithModelKey(modelKey string) *ConfigBuilder { + cb.modelKey = modelKey + return cb +} + +// WithModelVersion sets the pinned version of the model associated with the config. +func (cb *ConfigBuilder) WithModelVersion(modelVersion int) *ConfigBuilder { + cb.modelVersion = &modelVersion + return cb +} + // WithProviderName sets the provider name associated with the config. func (cb *ConfigBuilder) WithProviderName(providerName string) *ConfigBuilder { cb.providerName = providerName @@ -235,6 +263,8 @@ func (cb *ConfigBuilder) Build() Config { }, Model: datamodel.Model{ Name: cb.modelName, + Key: cb.modelKey, + Version: cb.modelVersion, Parameters: maps.Clone(cb.modelParams), Custom: maps.Clone(cb.modelCustomParams), }, diff --git a/ldai/datamodel/datamodel.go b/ldai/datamodel/datamodel.go index b5de1961..1f27aa6f 100644 --- a/ldai/datamodel/datamodel.go +++ b/ldai/datamodel/datamodel.go @@ -19,6 +19,12 @@ type Model struct { // Name identifies the model. Name string `json:"name"` + // Key is the model's stable, unique key (distinct from Name, which is not guaranteed unique). + Key string `json:"modelKey,omitempty"` + + // Version is the pinned version of the model that the variation references. + Version *int `json:"modelVersion,omitempty"` + // Parameters are the model parameters, generally provided by LaunchDarkly. Parameters map[string]ldvalue.Value `json:"parameters,omitempty"` diff --git a/ldai/tracker.go b/ldai/tracker.go index 66473315..caa825ff 100644 --- a/ldai/tracker.go +++ b/ldai/tracker.go @@ -203,7 +203,11 @@ func newTrackerWithStopwatch( Set("configKey", ldvalue.String(key)). Set("version", ldvalue.Int(version)). Set("providerName", ldvalue.String(config.ProviderName())). - Set("modelName", ldvalue.String(config.ModelName())) + Set("modelName", ldvalue.String(config.ModelName())). + Set("modelVersion", ldvalue.Int(config.ModelVersion())) + if config.ModelKey() != "" { + builder.Set("modelKey", ldvalue.String(config.ModelKey())) + } if variationKey != "" { builder.Set("variationKey", ldvalue.String(variationKey)) } @@ -230,7 +234,7 @@ func (t *Tracker) logWarning(format string, args ...interface{}) { // ResumptionToken returns a URL-safe Base64-encoded token that can be used to reconstruct a tracker // in a different process (e.g., for deferred feedback). The token contains the runId, configKey, -// variationKey, and version. It does not contain modelName or providerName. +// variationKey, and version. It does not contain modelName, providerName, modelKey, or modelVersion. func (t *Tracker) ResumptionToken() string { payload := resumptionPayload{ RunID: t.runID, @@ -245,8 +249,8 @@ func (t *Tracker) ResumptionToken() string { // TrackerFromResumptionToken reconstructs a Tracker from a resumption token and the given context. // This is used for cross-process scenarios (e.g., deferred feedback) where the original tracker // is no longer available but its runId must be reused. The token is obtained from Tracker.ResumptionToken(). -// The reconstructed tracker will have empty modelName and providerName since these are not included -// in the token. +// The reconstructed tracker will have empty modelName, providerName, and modelKey, and modelVersion +// defaults to 1, since these are not included in the token. func TrackerFromResumptionToken(token string, sdk ServerSDK, context ldcontext.Context) (*Tracker, error) { decoded, err := base64.RawURLEncoding.DecodeString(token) if err != nil { @@ -262,7 +266,8 @@ func TrackerFromResumptionToken(token string, sdk ServerSDK, context ldcontext.C Set("configKey", ldvalue.String(payload.ConfigKey)). Set("version", ldvalue.Int(payload.Version)). Set("providerName", ldvalue.String("")). - Set("modelName", ldvalue.String("")) + Set("modelName", ldvalue.String("")). + Set("modelVersion", ldvalue.Int(1)) if payload.VariationKey != "" { builder.Set("variationKey", ldvalue.String(payload.VariationKey)) } diff --git a/ldai/tracker_test.go b/ldai/tracker_test.go index ef7eac15..f1b42fe3 100644 --- a/ldai/tracker_test.go +++ b/ldai/tracker_test.go @@ -55,7 +55,11 @@ func makeTrackData(configKey, variationKey string, version int, config *Config, Set("configKey", ldvalue.String(configKey)). Set("version", ldvalue.Int(version)). Set("providerName", ldvalue.String(config.ProviderName())). - Set("modelName", ldvalue.String(config.ModelName())) + Set("modelName", ldvalue.String(config.ModelName())). + Set("modelVersion", ldvalue.Int(config.ModelVersion())) + if config.ModelKey() != "" { + builder.Set("modelKey", ldvalue.String(config.ModelKey())) + } if variationKey != "" { builder.Set("variationKey", ldvalue.String(variationKey)) } @@ -579,9 +583,14 @@ func TestTracker_ResumptionToken(t *testing.T) { assert.Equal(t, 3, payload.Version) }) - t.Run("does not include modelName or providerName", func(t *testing.T) { + t.Run("does not include modelName, providerName, modelKey, or modelVersion", func(t *testing.T) { events := newMockEvents() - config := NewConfig().WithModelName("gpt-4").WithProviderName("openai").Build() + config := NewConfig(). + WithModelName("gpt-4"). + WithProviderName("openai"). + WithModelKey("my-model"). + WithModelVersion(2). + Build() tracker := newTracker(events, newRunID(), "key", "var", 1, ldcontext.New("key"), &config, nil) token := tracker.ResumptionToken() @@ -593,7 +602,61 @@ func TestTracker_ResumptionToken(t *testing.T) { _, hasModel := raw["modelName"] _, hasProvider := raw["providerName"] + _, hasModelKey := raw["modelKey"] + _, hasModelVersion := raw["modelVersion"] assert.False(t, hasModel, "token should not contain modelName") assert.False(t, hasProvider, "token should not contain providerName") + assert.False(t, hasModelKey, "token should not contain modelKey") + assert.False(t, hasModelVersion, "token should not contain modelVersion") + }) +} + +func TestTracker_TrackDataIncludesModelKeyAndVersion(t *testing.T) { + t.Run("includes modelKey and modelVersion when set on config", func(t *testing.T) { + events := newMockEvents() + config := NewConfig(). + WithModelName("gpt-4"). + WithModelKey("my-model"). + WithModelVersion(2). + Build() + tracker := newTracker(events, newRunID(), "key", "var", 1, ldcontext.New("key"), &config, nil) + assert.NoError(t, tracker.TrackSuccess()) + + require.Len(t, events.events, 1) + data := events.events[0].data + assert.Equal(t, "my-model", data.GetByKey("modelKey").StringValue()) + assert.Equal(t, 2, data.GetByKey("modelVersion").IntValue()) }) + + t.Run("omits modelKey when empty but still includes modelVersion", func(t *testing.T) { + events := newMockEvents() + config := NewConfig().WithModelName("gpt-4").Build() + tracker := newTracker(events, newRunID(), "key", "var", 1, ldcontext.New("key"), &config, nil) + assert.NoError(t, tracker.TrackSuccess()) + + require.Len(t, events.events, 1) + data := events.events[0].data + assert.False(t, data.GetByKey("modelKey").IsDefined()) + assert.Equal(t, 1, data.GetByKey("modelVersion").IntValue()) + }) +} + +func TestTrackerFromResumptionToken_ModelKeyAndVersionDefaults(t *testing.T) { + mockSDK := newMockSDK(nil, nil) + config := NewConfig(). + WithModelName("gpt-4"). + WithModelKey("my-model"). + WithModelVersion(2). + Build() + tracker := newTracker(mockSDK, newRunID(), "key", "var", 1, ldcontext.New("key"), &config, mockSDK.log.Loggers) + + token := tracker.ResumptionToken() + reconstructed, err := TrackerFromResumptionToken(token, mockSDK, ldcontext.New("key")) + require.NoError(t, err) + assert.NoError(t, reconstructed.TrackSuccess()) + + require.Len(t, mockSDK.events, 1) + data := mockSDK.events[0].data + assert.False(t, data.GetByKey("modelKey").IsDefined()) + assert.Equal(t, 1, data.GetByKey("modelVersion").IntValue()) }