From e639bf3395d60730612409b937e44b2c7fb391ff Mon Sep 17 00:00:00 2001
From: Casey Waldren <cwaldren@launchdarkly.com>
Date: Fri, 22 Nov 2024 16:21:04 -0800
Subject: [PATCH] fix: add setter for ModelId

---
 pkgs/sdk/server-ai/src/Config/LdAiConfig.cs | 29 +++++++++++++++++++--
 pkgs/sdk/server-ai/test/LdAiClientTest.cs   |  2 ++
 pkgs/sdk/server-ai/test/LdAiConfigTest.cs   |  7 +++++
 3 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/pkgs/sdk/server-ai/src/Config/LdAiConfig.cs b/pkgs/sdk/server-ai/src/Config/LdAiConfig.cs
index 97193a0..0baf031 100644
--- a/pkgs/sdk/server-ai/src/Config/LdAiConfig.cs
+++ b/pkgs/sdk/server-ai/src/Config/LdAiConfig.cs
@@ -88,6 +88,7 @@ public class Builder
         private readonly Dictionary<string, LdValue> _modelParams;
         private readonly Dictionary<string, LdValue> _customModelParams;
         private string _providerId;
+        private string _modelId;
 
         internal Builder()
         {
@@ -96,6 +97,7 @@ internal Builder()
             _modelParams = new Dictionary<string, LdValue>();
             _customModelParams = new Dictionary<string, LdValue>();
             _providerId = "";
+            _modelId = "";
         }
 
         /// <summary>
@@ -157,11 +159,22 @@ public Builder SetCustomModelParam(string name, LdValue value)
             return this;
         }
 
+        /// <summary>
+        /// Sets the model's ID. By default, this will be the empty string.
+        /// </summary>
+        /// <param name="id">the model ID</param>
+        /// <returns>the builder</returns>
+        public Builder SetModelId(string id)
+        {
+            _modelId = id;
+            return this;
+        }
+
         /// <summary>
         /// Sets the model provider's ID. By default, this will be the empty string.
         /// </summary>
         /// <param name="id">the ID</param>
-        /// <returns></returns>
+        /// <returns>the builder</returns>
         public Builder SetModelProviderId(string id)
         {
             _providerId = id;
@@ -174,7 +187,18 @@ public Builder SetModelProviderId(string id)
         /// <returns>a new LdAiConfig</returns>
         public LdAiConfig Build()
         {
-            return new LdAiConfig(_enabled, _messages, new Meta(), new Model {Parameters = _modelParams, Custom = _customModelParams}, new Provider{ Id = _providerId });
+            return new LdAiConfig(
+                _enabled,
+                _messages,
+                new Meta(),
+                new Model
+                {
+                    Id = _modelId,
+                    Parameters = _modelParams,
+                    Custom = _customModelParams
+                },
+                new Provider{ Id = _providerId }
+            );
         }
     }
 
@@ -219,6 +243,7 @@ internal LdValue ToLdValue()
             }))) },
             { "model", LdValue.ObjectFrom(new Dictionary<string, LdValue>
             {
+                { "id", LdValue.Of(Model.Id) },
                 { "parameters", LdValue.ObjectFrom(Model.Parameters) },
                 { "custom", LdValue.ObjectFrom(Model.Custom) }
             }) },
diff --git a/pkgs/sdk/server-ai/test/LdAiClientTest.cs b/pkgs/sdk/server-ai/test/LdAiClientTest.cs
index 58634a5..a338e6e 100644
--- a/pkgs/sdk/server-ai/test/LdAiClientTest.cs
+++ b/pkgs/sdk/server-ai/test/LdAiClientTest.cs
@@ -119,6 +119,7 @@ public void CanSetAllDefaultValueFields()
             LdAiConfig.New().
                 AddMessage("foo").
                 SetModelParam("foo", LdValue.Of("bar")).
+                SetModelId("awesome-model").
                 SetCustomModelParam("foo", LdValue.Of("baz")).
                 SetModelProviderId("amazing-provider").
                 SetEnabled(true).Build());
@@ -133,6 +134,7 @@ public void CanSetAllDefaultValueFields()
         Assert.Equal("amazing-provider", tracker.Config.Provider.Id);
         Assert.Equal("bar", tracker.Config.Model.Parameters["foo"].AsString);
         Assert.Equal("baz", tracker.Config.Model.Custom["foo"].AsString);
+        Assert.Equal("awesome-model", tracker.Config.Model.Id);
     }
 
     [Fact]
diff --git a/pkgs/sdk/server-ai/test/LdAiConfigTest.cs b/pkgs/sdk/server-ai/test/LdAiConfigTest.cs
index abb53dd..872d506 100644
--- a/pkgs/sdk/server-ai/test/LdAiConfigTest.cs
+++ b/pkgs/sdk/server-ai/test/LdAiConfigTest.cs
@@ -72,6 +72,13 @@ public void CanSetModelParams()
         Assert.Equal(LdValue.Of(43), config.Model.Custom["baz"]);
     }
 
+    [Fact]
+    public void CanSetModelId()
+    {
+        var config = LdAiConfig.New().SetModelId("awesome-model").Build();
+        Assert.Equal("awesome-model", config.Model.Id);
+    }
+
     [Fact]
     public void CanSetModelProviderId()
     {