From 08f01e6261108d4712f1dd1feb81902dd3bf3316 Mon Sep 17 00:00:00 2001
From: Kazik Pogoda <morisil@xemantic.com>
Date: Mon, 14 Oct 2024 17:44:24 +0200
Subject: [PATCH] JsonSchemaGenerator update

---
 src/commonMain/kotlin/schema/JsonSchemaGenerator.kt   | 11 ++++-------
 .../kotlin/schema/JsonSchemaGeneratorTest.kt          |  8 ++++----
 2 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/src/commonMain/kotlin/schema/JsonSchemaGenerator.kt b/src/commonMain/kotlin/schema/JsonSchemaGenerator.kt
index 71ef76a..ad418cd 100644
--- a/src/commonMain/kotlin/schema/JsonSchemaGenerator.kt
+++ b/src/commonMain/kotlin/schema/JsonSchemaGenerator.kt
@@ -3,17 +3,11 @@ package com.xemantic.anthropic.schema
 import kotlinx.serialization.*
 import kotlinx.serialization.descriptors.*
 import kotlin.collections.set
-import kotlin.reflect.KClass
 
 inline fun <reified T> jsonSchemaOf(): JsonSchema = generateSchema(
   serializer<T>().descriptor
 )
 
-@OptIn(InternalSerializationApi::class)
-fun KClass<*>.toJsonSchema(): JsonSchema = generateSchema(
-  serializer().descriptor
-)
-
 @OptIn(ExperimentalSerializationApi::class)
 fun generateSchema(descriptor: SerialDescriptor): JsonSchema {
   val properties = mutableMapOf<String, JsonSchemaProperty>()
@@ -58,9 +52,12 @@ private fun generateSchemaProperty(
     )
     StructureKind.MAP -> JsonSchemaProperty("object")
     StructureKind.CLASS -> {
+      // dots are not allowed in JSON Schema name, if the @SerialName was not
+      // specified, then fully qualified class name will be used, and we need
+      // to translate it
       val refName = descriptor.serialName.replace('.', '_').trimEnd('?')
       definitions[refName] = generateSchema(descriptor)
-      JsonSchemaProperty("\$ref", ref = "#/definitions/$refName")
+      JsonSchemaProperty(ref = "#/definitions/$refName")
     }
     else -> JsonSchemaProperty("object") // Default case
   }
diff --git a/src/commonTest/kotlin/schema/JsonSchemaGeneratorTest.kt b/src/commonTest/kotlin/schema/JsonSchemaGeneratorTest.kt
index 4c824c4..411e56e 100644
--- a/src/commonTest/kotlin/schema/JsonSchemaGeneratorTest.kt
+++ b/src/commonTest/kotlin/schema/JsonSchemaGeneratorTest.kt
@@ -3,12 +3,14 @@ package com.xemantic.anthropic.schema
 import com.xemantic.anthropic.anthropicJson
 import io.kotest.assertions.json.shouldEqualJson
 import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.encodeToString
 import kotlinx.serialization.json.Json
 import kotlin.test.Test
 
 @Serializable
+@SerialName("address")
 data class Address(
   val street: String? = null,
   val city: String? = null,
@@ -35,7 +37,6 @@ class JsonSchemaGeneratorTest {
 
   @Test
   fun generateJsonSchemaForAddress() {
-
     // when
     val schema = jsonSchemaOf<Address>()
     val schemaJson = json.encodeToString(schema)
@@ -77,7 +78,7 @@ class JsonSchemaGeneratorTest {
     {
       "type": "object",
       "definitions": {
-        "com_xemantic_anthropic_schema_Address": {
+        "address": {
           "type": "object",
           "properties": {
             "street": {
@@ -116,8 +117,7 @@ class JsonSchemaGeneratorTest {
           }
         },
         "address": {
-          "type": "${'$'}ref",
-          "ref": "#/definitions/com_xemantic_anthropic_schema_Address"
+          "${'$'}ref": "#/definitions/address"
         }
       },
       "required": [