From b594927df072d1bbe3169abfe925e05e7d5755f9 Mon Sep 17 00:00:00 2001
From: Stevan Freeborn <65925598+StevanFreeborn@users.noreply.github.com>
Date: Sun, 30 Jun 2024 15:25:53 -0500
Subject: [PATCH] fix: deal with nested complex types

- make sure the definitions object is the same
  when making recursive calls to generate schema.
  this way all definitions are captured and can
  be referenced correctly.
- add tests for nested complex types
---
 .../Utils/JsonSchemaGenerator.cs              |  6 +-
 .../Unit/Utils/JsonSchemaGeneratorTestData.cs | 71 +++++++++++++++++++
 2 files changed, 75 insertions(+), 2 deletions(-)

diff --git a/src/AnthropicClient/Utils/JsonSchemaGenerator.cs b/src/AnthropicClient/Utils/JsonSchemaGenerator.cs
index abd23a7..e270464 100644
--- a/src/AnthropicClient/Utils/JsonSchemaGenerator.cs
+++ b/src/AnthropicClient/Utils/JsonSchemaGenerator.cs
@@ -68,12 +68,12 @@ static class JsonSchemaGenerator
 
   private static JsonObject GenerateParameterTypeSchema(Type type, JsonObject inputSchema)
   {
-    var definitions = inputSchema[DefinitionsKey];
+    var definitions = inputSchema[DefinitionsKey] ?? new JsonObject();
 
     // Check if a definition for the type already exists
     // If it does, return a reference to the definition
     // no need to evaluate the type further
-    if (definitions is not null && definitions.AsObject().ContainsKey(type.FullName))
+    if (definitions.AsObject().ContainsKey(type.FullName))
     {
       return new JsonObject()
       {
@@ -150,6 +150,8 @@ static class JsonSchemaGenerator
   private static JsonObject GenerateTypeDefinitionSchema(Type type, JsonObject inputSchema)
   {
     var definitions = inputSchema[DefinitionsKey] ?? new JsonObject();
+    inputSchema[DefinitionsKey] = definitions;
+
     var typeSchema = new JsonObject()
     {
       [TypeKey] = ObjectType
diff --git a/tests/AnthropicClient.Tests/Unit/Utils/JsonSchemaGeneratorTestData.cs b/tests/AnthropicClient.Tests/Unit/Utils/JsonSchemaGeneratorTestData.cs
index b4ee0b2..2eecd02 100644
--- a/tests/AnthropicClient.Tests/Unit/Utils/JsonSchemaGeneratorTestData.cs
+++ b/tests/AnthropicClient.Tests/Unit/Utils/JsonSchemaGeneratorTestData.cs
@@ -494,11 +494,82 @@ public class JsonSchemaGeneratorTestData : IEnumerable<object[]>
         },
       }
     };
+
+    // nested class parameter type with no attributes
+    yield return new object[]
+    {
+      Tool.CreateFromFunction(TestToolName,TestToolDescription,(Family family) => family),
+      new JsonObject()
+      {
+        ["type"] = "object",
+        ["definitions"] = new JsonObject()
+        {
+          [$"{typeof(Person).FullName}"] = new JsonObject()
+          {
+            ["type"] = "object",
+            ["properties"] = new JsonObject()
+            {
+              ["Name"] = new JsonObject()
+              {
+                ["type"] = "string",
+                ["description"] = string.Empty
+              },
+              ["Age"] = new JsonObject()
+              {
+                ["type"] = "integer",
+                ["description"] = string.Empty
+              }
+            },
+            ["required"] = new JsonArray()
+            {
+              "Name",
+              "Age"
+            }
+          },
+          [$"{typeof(Family).FullName}"] = new JsonObject()
+          {
+            ["type"] = "object",
+            ["properties"] = new JsonObject()
+            {
+              ["Members"] = new JsonObject()
+              {
+                ["type"] = "array",
+                ["items"] = new JsonObject()
+                {
+                  ["$ref"] = $"#/definitions/{typeof(Person).FullName}"
+                },
+                ["description"] = string.Empty
+              }
+            },
+            ["required"] = new JsonArray()
+            {
+              "Members"
+            }
+          }
+        },
+        ["properties"] = new JsonObject()
+        {
+          ["family"] = new JsonObject()
+          {
+            ["$ref"] = $"#/definitions/{typeof(Family).FullName}",
+            ["description"] = string.Empty
+          }
+        },
+        ["required"] = new JsonArray()
+        {
+          "family",
+        },
+      }
+    };
   }
 
   IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 }
 
+class Family 
+{
+  public List<Person> Members { get; } = [];
+}
 
 class Person
 {