From 3006ad9b50f90a1e1938293518fd6f9620725a3d Mon Sep 17 00:00:00 2001
From: Stevan Freeborn <65925598+StevanFreeborn@users.noreply.github.com>
Date: Sun, 30 Jun 2024 12:21:30 -0500
Subject: [PATCH] feat: add parameter attribute to allow customization of name,
 description, and requiredness

---
 .../Models/FunctionParameterAttribute.cs      | 42 +++++++++++++++
 .../Utils/JsonSchemaGenerator.cs              |  8 +--
 .../Unit/Utils/JsonSchemaGeneratorTests.cs    | 54 +++++++++++++++++++
 3 files changed, 101 insertions(+), 3 deletions(-)
 create mode 100644 src/AnthropicClient/Models/FunctionParameterAttribute.cs

diff --git a/src/AnthropicClient/Models/FunctionParameterAttribute.cs b/src/AnthropicClient/Models/FunctionParameterAttribute.cs
new file mode 100644
index 0000000..5fb4196
--- /dev/null
+++ b/src/AnthropicClient/Models/FunctionParameterAttribute.cs
@@ -0,0 +1,42 @@
+using AnthropicClient.Utils;
+
+namespace AnthropicClient.Models;
+
+/// <summary>
+/// Attribute to describe a function parameter.
+/// </summary>
+[AttributeUsage(AttributeTargets.Parameter)]
+public sealed class FunctionParameterAttribute : Attribute
+{ 
+  /// <summary>
+  /// The name of the parameter.
+  /// </summary>
+  public string Name { get; }
+
+  /// <summary>
+  /// The description of the parameter.
+  /// </summary>
+  public string Description { get; }
+
+  /// <summary>
+  /// Whether the parameter is required.
+  /// </summary>
+  public bool Required { get; }
+
+  /// <summary>
+  /// Initializes a new instance of the <see cref="FunctionParameterAttribute"/> class.
+  /// </summary>
+  /// <param name="name">The name of the parameter.</param>
+  /// <param name="description">The description of the parameter.</param>
+  /// <param name="required">Whether the parameter is required.</param>
+  /// <returns>A new instance of the <see cref="FunctionParameterAttribute"/> class.</returns>
+  public FunctionParameterAttribute(string description, string name = "", bool required = false)
+  {
+    ArgumentValidator.ThrowIfNullOrWhitespace(description, nameof(description));
+    ArgumentValidator.ThrowIfNull(name, nameof(name));
+
+    Name = name;
+    Description = description;
+    Required = required;
+  }
+}
\ No newline at end of file
diff --git a/src/AnthropicClient/Utils/JsonSchemaGenerator.cs b/src/AnthropicClient/Utils/JsonSchemaGenerator.cs
index 4950171..95b3c13 100644
--- a/src/AnthropicClient/Utils/JsonSchemaGenerator.cs
+++ b/src/AnthropicClient/Utils/JsonSchemaGenerator.cs
@@ -36,9 +36,11 @@ static class JsonSchemaGenerator
         continue;
       }
 
-      var paramName = parameter.Name;
-      var paramDescription = string.Empty;
-      var paramRequired = parameter.HasDefaultValue;
+      var attribute = parameter.GetCustomAttribute<FunctionParameterAttribute>();
+
+      var paramName = attribute?.Name ?? parameter.Name;
+      var paramDescription = attribute?.Description ?? string.Empty;
+      var paramRequired = attribute?.Required ?? !parameter.HasDefaultValue;
 
       var paramObject = new JsonObject();
       paramObject[DescriptionKey] = paramDescription;
diff --git a/tests/AnthropicClient.Tests/Unit/Utils/JsonSchemaGeneratorTests.cs b/tests/AnthropicClient.Tests/Unit/Utils/JsonSchemaGeneratorTests.cs
index 9478caf..7194e08 100644
--- a/tests/AnthropicClient.Tests/Unit/Utils/JsonSchemaGeneratorTests.cs
+++ b/tests/AnthropicClient.Tests/Unit/Utils/JsonSchemaGeneratorTests.cs
@@ -42,4 +42,58 @@ public class JsonSchemaGeneratorTests
 
     JsonAssert.Equal(expectedSchema, schema);
   }
+
+  [Fact]
+  public void GenerateInputSchema_GivenFunctionWithParameterThatDoesNotHaveDefaultValue_ItShouldReturnSchemaWithPropertyNameDescriptionAndRequiredProperties()
+  {
+    var expectedSchema = new JsonObject()
+    {
+      ["type"] = "object",
+      ["properties"] = new JsonObject()
+      {
+        ["name"] = new JsonObject()
+        {
+          ["description"] = string.Empty
+        }
+      },
+      ["required"] = new JsonArray()
+      {
+        "name"
+      },
+    };
+
+    var testMethod = (string name) => name;
+    var function = new AnthropicFunction(testMethod.Method);
+    
+    var schema = JsonSchemaGenerator.GenerateInputSchema(function);
+
+    JsonAssert.Equal(expectedSchema, schema);
+  }
+
+  [Fact]
+  public void GenerateInputSchema_GivenFunctionWithParameterThatHasAttribute_ItShouldReturnSchemaWithCorrectPropertyNameDescriptionAndRequiredProperty()
+  {
+    var expectedSchema = new JsonObject()
+    {
+      ["type"] = "object",
+      ["properties"] = new JsonObject()
+      {
+        ["Person's Age"] = new JsonObject()
+        {
+          ["description"] = "The age of the person."
+        }
+      },
+      ["required"] = new JsonArray()
+      {
+        "Person's Age"
+      },
+    };
+
+    var testMethod = ([FunctionParameter("The age of the person.", "Person's Age", true)] int name) => name;
+    var function = new AnthropicFunction(testMethod.Method);
+
+    var schema = JsonSchemaGenerator.GenerateInputSchema(function);
+
+    JsonAssert.Equal(expectedSchema, schema);
+  }
 }
\ No newline at end of file