Repository: softwaremill/sttp-openai (scala)

Issue:
JSON schema support
Would you consider adding support for the [`json_schema` response format](https://platform.openai.com/docs/guides/structured-outputs/how-to-use?context=without_parse)? If so, I would be happy to take a stab at a pull request to add it. I would model it similarly to this: https://github.com/cequence-io/openai-scala-client/blob/master/openai-core/src/main/scala/io/cequence/openaiscala/domain/JsonSchema.scala

Pull request:
Fixes #216 

@adamw re: your suggestion:

> An extension to this idea would be to derive a Tapir schema from a case class, and then [serialize it as JSON Schema](https://tapir.softwaremill.com/en/latest/docs/json-schema.html).

I chose to use the `Schema` type from [sttp-apispec](https://github.com/softwaremill/sttp-apispec) to implement this, but I didn't add explicit support for passing a Tapir schema and having it be converted automatically -- I figured downstream users could do the conversion themselves if they wanted to.

OpenAI only [supports a subset of JSON schema](https://platform.openai.com/docs/guides/structured-outputs/supported-schemas) and has some limitations, which I described in a comment in `ChatRequestBody.scala`.

I don't think it's reasonable to enforce all of these constraints in this library, users will just have to be cognizant of whether their schema violates any of them. If it does, the API returns an error like this:

```json
{
  "error": {
    "message": "Invalid schema for response_format 'schemaName': ... reason why schema was invalid ...",
    "type": "invalid_request_error",
    "param": "response_format",
    "code": null
  }
}
```

Patch:
diff --git a/README.md b/README.md
index fb23e908..ac808301 100644
--- a/README.md
+++ b/README.md
@@ -424,6 +424,121 @@ object Main extends OxApp:
 
 See also the [ChatProxy](https://github.com/softwaremill/sttp-openai/blob/master/examples/src/main/scala/examples/ChatProxy.scala) example application.
 
+#### Structured Outputs/JSON Schema support
+
+To take advantage of [OpenAI's Structured Outputs](https://platform.openai.com/docs/guides/structured-outputs/introduction)
+and support for JSON Schema, you can use `ResponseFormat.JsonSchema` when creating a completion.
+
+The example below produces a JSON object:
+
+```scala mdoc:compile-only
+//> using dep com.softwaremill.sttp.openai::core:0.2.4
+
+import scala.collection.immutable.ListMap
+import sttp.apispec.{Schema, SchemaType}
+import sttp.openai.OpenAISyncClient
+import sttp.openai.requests.completions.chat.ChatRequestResponseData.ChatResponse
+import sttp.openai.requests.completions.chat.ChatRequestBody.{ChatBody, ChatCompletionModel, ResponseFormat}
+import sttp.openai.requests.completions.chat.message._
+
+object Main extends App {
+  val apiKey = System.getenv("OPENAI_KEY")
+  val openAI = OpenAISyncClient(apiKey)
+
+  val jsonSchema: Schema =
+    Schema(SchemaType.Object).copy(properties =
+      ListMap(
+        "steps" -> Schema(SchemaType.Array).copy(items =
+          Some(Schema(SchemaType.Object).copy(properties =
+            ListMap(
+              "explanation" -> Schema(SchemaType.String),
+              "output" -> Schema(SchemaType.String)
+            )
+          ))
+        ),
+        "finalAnswer" -> Schema(SchemaType.String)
+      ),
+    )
+
+  val responseFormat: ResponseFormat.JsonSchema =
+    ResponseFormat.JsonSchema(
+      name = "mathReasoning",
+      strict = true,
+      schema = jsonSchema
+    )
+
+  val bodyMessages: Seq[Message] = Seq(
+    Message.SystemMessage(content = "You are a helpful math tutor. Guide the user through the solution step by step."),
+    Message.UserMessage(content = Content.TextContent("How can I solve 8x + 7 = -23"))
+  )
+
+  // Create body of Chat Completions Request, using our JSON Schema as the `responseFormat`
+  val chatRequestBody: ChatBody = ChatBody(
+    model = ChatCompletionModel.GPT4oMini,
+    messages = bodyMessages,
+    responseFormat = Some(responseFormat)
+  )
+
+  val chatResponse: ChatResponse = openAI.createChatCompletion(chatRequestBody)
+
+  println(chatResponse.choices)
+  /*
+    List(
+      Choices(
+        Message(
+          Assistant,
+          {
+            "steps": [
+              {"explanation": "Start with the original equation: 8x + 7 = -23", "output": "8x + 7 = -23"},
+              {"explanation": "Subtract 7 from both sides to isolate the term with x.", "output": "8x + 7 - 7 = -23 - 7"},
+              {"explanation": "This simplifies to: 8x = -30", "output": "8x = -30"},
+              {"explanation": "Now, divide both sides by 8 to solve for x.", "output": "x = -30 / 8"},
+              {"explanation": "Simplify -30 / 8 to its simplest form. Both the numerator and denominator can be divided by 2.", "output": "x = -15 / 4"}
+            ],
+            "finalAnswer": "x = -15/4"
+          },
+          List(),
+          None
+        ),
+        stop,
+        0
+      )
+    )
+  */
+}
+```
+
+##### Deriving a JSON Schema with tapir
+
+To derive the same math reasoning schema used above, you can use
+[Tapir's support for generating a JSON schema from a Tapir schema](https://tapir.softwaremill.com/en/latest/docs/json-schema.html):
+
+```scala mdoc:compile-only
+//> using dep com.softwaremill.sttp.tapir::tapir-apispec-docs:1.11.7
+
+import sttp.apispec.{Schema => ASchema}
+import sttp.tapir.Schema
+import sttp.tapir.docs.apispec.schema.TapirSchemaToJsonSchema
+import sttp.tapir.generic.auto._
+
+case class Step(
+  explanation: String,
+  output: String
+)
+
+case class MathReasoning(
+  steps: List[Step],
+  finalAnswer: String
+)
+
+val tSchema = implicitly[Schema[MathReasoning]]
+
+val jsonSchema: ASchema = TapirSchemaToJsonSchema(
+  tSchema,
+  markOptionsAsNullable = true
+)
+```
+
 ## Contributing
 
 If you have a question, or hit a problem, feel free to post on our community https://softwaremill.community/c/open-source/
diff --git a/build.sbt b/build.sbt
index aa940496..814a028a 100644
--- a/build.sbt
+++ b/build.sbt
@@ -31,7 +31,11 @@ lazy val core = (projectMatrix in file("core"))
     scalaVersions = scala2 ++ scala3
   )
   .settings(
-    libraryDependencies ++= Seq(Libraries.uPickle) ++ Libraries.sttpClient ++ Seq(Libraries.scalaTest)
+    libraryDependencies ++= Seq(
+      Libraries.tapirApispecDocs,
+      Libraries.uJsonCirce,
+      Libraries.uPickle
+    ) ++ Libraries.sttpApispec ++ Libraries.sttpClient ++ Seq(Libraries.scalaTest)
   )
   .settings(commonSettings: _*)
 
diff --git a/core/src/main/scala/sttp/openai/requests/completions/chat/ChatRequestBody.scala b/core/src/main/scala/sttp/openai/requests/completions/chat/ChatRequestBody.scala
index 1f8a8198..63a730b8 100644
--- a/core/src/main/scala/sttp/openai/requests/completions/chat/ChatRequestBody.scala
+++ b/core/src/main/scala/sttp/openai/requests/completions/chat/ChatRequestBody.scala
@@ -1,10 +1,15 @@
 package sttp.openai.requests.completions.chat
 
+import io.circe.{DecodingFailure, Json, JsonNumber, JsonObject}
+import io.circe.syntax._
+import sttp.apispec.Schema
+import sttp.apispec.circe._
 import sttp.openai.OpenAIExceptions.OpenAIException.DeserializationOpenAIException
 import sttp.openai.json.SnakePickle
 import sttp.openai.requests.completions.Stop
 import sttp.openai.requests.completions.chat.message.{Message, Tool, ToolChoice}
 import ujson._
+import ujson.circe.CirceJson
 
 object ChatRequestBody {
 
@@ -13,6 +18,110 @@ object ChatRequestBody {
   object ResponseFormat {
     case object Text extends ResponseFormat
     case object JsonObject extends ResponseFormat
+    case class JsonSchema(name: String, strict: Boolean, schema: Schema) extends ResponseFormat
+    object JsonSchema {
+      case class ParseException(circeException: DecodingFailure) extends Exception("Failed to parse JSON schema", circeException)
+
+      private case class InternalRepr(json_schema: JsonSchema)
+
+      private object InternalRepr {
+        case class FolderState(
+            fields: List[(String, Json)],
+            addAdditionalProperties: Boolean,
+            requiredProperties: List[String]
+        )
+
+        /** OpenAI's JSON schema support has a couple limitations/requirements we need to account for:
+          *
+          *   1. All fields must be `required` https://platform.openai.com/docs/guides/structured-outputs/all-fields-must-be-required
+          *
+          * 2. `additionalProperties: false` must always be set in objects
+          * https://platform.openai.com/docs/guides/structured-outputs/additionalproperties-false-must-always-be-set-in-objects
+          *
+          * We handle both of these by folding over the JSON structure, looking specifically for JSON objects.
+          *
+          * If the object has a "properties" key, then we handle both limitations listed above.
+          *
+          * If the object has a "type" key AND the "type" is "object", then we handle limitation 2.
+          */
+        private val schemaFolder: Json.Folder[Json] = new Json.Folder[Json] {
+          lazy val onNull = Json.Null
+          def onBoolean(value: Boolean): Json = Json.fromBoolean(value)
+          def onNumber(value: JsonNumber): Json = Json.fromJsonNumber(value)
+          def onString(value: String): Json = Json.fromString(value)
+          def onArray(value: Vector[Json]): Json = Json.fromValues(value.map(_.foldWith(this)))
+          def onObject(value: JsonObject): Json = {
+            val state = value.toList.foldRight(FolderState(Nil, false, Nil)) { case ((k, v), acc) =>
+              if (k == "properties")
+                acc.copy(
+                  fields = (k, v.foldWith(this)) :: acc.fields,
+                  addAdditionalProperties = true,
+                  requiredProperties = v.asObject.fold(List.empty[String])(_.keys.toList)
+                )
+              else if (k == "type")
+                acc.copy(
+                  fields = (k, v.foldWith(this)) :: acc.fields,
+                  addAdditionalProperties = acc.addAdditionalProperties || v.asString.filter(_ == "object").isDefined
+                )
+              else
+                acc.copy(fields = (k, v.foldWith(this)) :: acc.fields)
+            }
+
+            val (addlPropsRemove, addlPropsAdd) =
+              if (state.addAdditionalProperties)
+                (Set("additionalProperties"), List("additionalProperties" := false))
+              else
+                (Set(), Nil)
+
+            val (requiredRemove, requiredAdd) =
+              if (state.requiredProperties.nonEmpty)
+                (Set("required"), List("required" := state.requiredProperties))
+              else
+                (Set(), Nil)
+
+            val remove = addlPropsRemove ++ requiredRemove
+
+            val fields = addlPropsAdd ++ requiredAdd ++ state.fields.filterNot { case (k, _) => remove.contains(k) }
+
+            Json.fromFields(fields)
+          }
+        }
+
+        implicit private val schemaRW: SnakePickle.ReadWriter[Schema] = SnakePickle
+          .readwriter[Value]
+          .bimap(
+            s => CirceJson.transform(s.asJson.deepDropNullValues.foldWith(schemaFolder), upickle.default.reader[Value]),
+            v =>
+              upickle.default.transform(v).to(CirceJson).as[Schema] match {
+                case Left(e)  => throw new ParseException(e)
+                case Right(s) => s
+              }
+          )
+
+        implicit private val jsonSchemaRW: SnakePickle.ReadWriter[JsonSchema] = SnakePickle
+          .readwriter[Value]
+          .bimap(
+            s => Obj("name" -> s.name, "strict" -> s.strict, "schema" -> SnakePickle.writeJs(s.schema)),
+            v => {
+              val o = v.obj
+              JsonSchema(
+                name = o("name").str,
+                strict = o("strict").bool,
+                schema = SnakePickle.read[Schema](o("schema"))
+              )
+            }
+          )
+
+        implicit val internalReprRW: SnakePickle.ReadWriter[InternalRepr] = SnakePickle.macroRW
+      }
+
+      implicit val jsonSchemaRW: SnakePickle.ReadWriter[JsonSchema] = SnakePickle
+        .readwriter[Value]
+        .bimap[JsonSchema](
+          s => Obj(SnakePickle.writeJs(InternalRepr(s)).obj.addOne("type" -> "json_schema")),
+          v => SnakePickle.read[InternalRepr](v).json_schema
+        )
+    }
 
     implicit val textRW: SnakePickle.ReadWriter[Text.type] = SnakePickle
       .readwriter[Value]
@@ -34,11 +143,13 @@ object ChatRequestBody {
         {
           case text: Text.type             => SnakePickle.writeJs(text)
           case jsonObject: JsonObject.type => SnakePickle.writeJs(jsonObject)
+          case jsonSchema: JsonSchema      => SnakePickle.writeJs(jsonSchema)
         },
         json =>
           json("type").str match {
             case "text"        => SnakePickle.read[Text.type](json)
             case "json_object" => SnakePickle.read[JsonObject.type](json)
+            case "json_schema" => SnakePickle.read[JsonSchema](json)
           }
       )
   }
diff --git a/project/Dependencies.scala b/project/Dependencies.scala
index 27dfb486..1694acc1 100644
--- a/project/Dependencies.scala
+++ b/project/Dependencies.scala
@@ -6,9 +6,11 @@ object Dependencies {
     val scalaTest = "3.2.19"
     val scalaTestCats = "1.5.0"
 
+    val sttpApispec = "0.11.3"
     val sttpClient = "4.0.0-M19"
     val pekkoStreams = "1.1.2"
     val akkaStreams = "2.6.20"
+    val tapir = "1.11.7"
     val uPickle = "3.1.4"
   }
 
@@ -16,6 +18,11 @@ object Dependencies {
 
     val scalaTest = "org.scalatest" %% "scalatest" % V.scalaTest % Test
 
+    val sttpApispec = Seq(
+      "com.softwaremill.sttp.apispec" %% "apispec-model" % V.sttpApispec,
+      "com.softwaremill.sttp.apispec" %% "jsonschema-circe" % V.sttpApispec
+    )
+
     val sttpClient = Seq(
       "com.softwaremill.sttp.client4" %% "core" % V.sttpClient,
       "com.softwaremill.sttp.client4" %% "upickle" % V.sttpClient
@@ -42,6 +49,10 @@ object Dependencies {
       "com.softwaremill.sttp.client4" %% "ox" % V.sttpClient
     )
 
+    val tapirApispecDocs = "com.softwaremill.sttp.tapir" %% "tapir-apispec-docs" % V.tapir
+
+    val uJsonCirce = "com.lihaoyi" %% "ujson-circe" % V.uPickle
+
     val uPickle = "com.lihaoyi" %% "upickle" % V.uPickle
 
   }


Test patch:
diff --git a/core/src/test/scala/sttp/openai/fixtures/JsonSchemaFixture.scala b/core/src/test/scala/sttp/openai/fixtures/JsonSchemaFixture.scala
new file mode 100644
index 00000000..22ba900d
--- /dev/null
+++ b/core/src/test/scala/sttp/openai/fixtures/JsonSchemaFixture.scala
@@ -0,0 +1,69 @@
+package sttp.openai.fixtures
+
+object JsonSchemaFixture {
+
+  val stringSchema: String =
+    """{
+      |  "type": "json_schema",
+      |  "json_schema": {
+      |    "name": "testString",
+      |    "strict": true,
+      |    "schema": {
+      |      "type": "string"
+      |    }
+      |  }
+      |}""".stripMargin
+
+  val numberSchema: String =
+    """{
+      |  "type": "json_schema",
+      |  "json_schema": {
+      |    "name": "testNumber",
+      |    "strict": true,
+      |    "schema": {
+      |      "type": "number"
+      |    }
+      |  }
+      |}""".stripMargin
+
+  val objectSchema: String =
+    """{
+      |  "type": "json_schema",
+      |  "json_schema": {
+      |    "name": "testObject",
+      |    "strict": true,
+      |    "schema": {
+      |      "additionalProperties": false,
+      |      "required": [
+      |        "foo",
+      |        "bar"
+      |      ],
+      |      "type": "object",
+      |      "properties": {
+      |        "foo": {
+      |          "type": "string"
+      |        },
+      |        "bar": {
+      |          "type": "number"
+      |        }
+      |      }
+      |    }
+      |  }
+      |}""".stripMargin
+
+  val arraySchema: String =
+    """{
+      |  "type": "json_schema",
+      |  "json_schema": {
+      |    "name": "testArray",
+      |    "strict": true,
+      |    "schema": {
+      |      "type": "array",
+      |      "items": {
+      |        "type": "string"
+      |      }
+      |    }
+      |  }
+      |}""".stripMargin
+
+}
diff --git a/core/src/test/scala/sttp/openai/requests/completions/chat/JsonSchemaSpec.scala b/core/src/test/scala/sttp/openai/requests/completions/chat/JsonSchemaSpec.scala
new file mode 100644
index 00000000..4d31c6cb
--- /dev/null
+++ b/core/src/test/scala/sttp/openai/requests/completions/chat/JsonSchemaSpec.scala
@@ -0,0 +1,53 @@
+package sttp.openai.requests.completions.chat
+
+import org.scalatest.EitherValues
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import scala.collection.immutable.ListMap
+import sttp.apispec.{Schema, SchemaType}
+import sttp.openai.fixtures
+import sttp.openai.json.SnakePickle
+import sttp.openai.requests.completions.chat.ChatRequestBody.ResponseFormat.JsonSchema
+
+class JsonSchemaSpec extends AnyFlatSpec with Matchers with EitherValues {
+  "Given string JSON schema" should "be properly serialized to Json" in {
+    val schema = Schema(SchemaType.String)
+
+    val jsonStringSchema = ujson.read(fixtures.JsonSchemaFixture.stringSchema)
+
+    val serializedSchema = SnakePickle.writeJs(JsonSchema("testString", true, schema))
+
+    serializedSchema shouldBe jsonStringSchema
+  }
+
+  "Given number JSON schema" should "be properly serialized to Json" in {
+    val schema = Schema(SchemaType.Number)
+
+    val jsonNumberSchema = ujson.read(fixtures.JsonSchemaFixture.numberSchema)
+
+    val serializedSchema = SnakePickle.writeJs(JsonSchema("testNumber", true, schema))
+
+    serializedSchema shouldBe jsonNumberSchema
+  }
+
+  "Given object JSON schema" should "be properly serialized to Json" in {
+    val schema = Schema(SchemaType.Object)
+      .copy(properties = ListMap("foo" -> Schema(SchemaType.String), "bar" -> Schema(SchemaType.Number)))
+
+    val jsonObjectSchema = ujson.read(fixtures.JsonSchemaFixture.objectSchema)
+
+    val serializedSchema = SnakePickle.writeJs(JsonSchema("testObject", true, schema))
+
+    serializedSchema shouldBe jsonObjectSchema
+  }
+
+  "Given array JSON schema" should "be properly serialized to Json" in {
+    val schema = Schema(SchemaType.Array).copy(items = Some(Schema(SchemaType.String)))
+
+    val jsonArraySchema = ujson.read(fixtures.JsonSchemaFixture.arraySchema)
+
+    val serializedSchema = SnakePickle.writeJs(JsonSchema("testArray", true, schema))
+
+    serializedSchema shouldBe jsonArraySchema
+  }
+}