feat(ldai): stamp modelKey and modelVersion on AI usage events (AIC-2850)

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)

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) {

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),
 			},

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"`
 

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))
 	}

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())
 }