Repository: softwaremill/sttp-openai
Pull request #35 (2023-04-19 10:10:03)

Create image


Patch:
diff --git a/core/src/main/scala/sttp/openai/OpenAi.scala b/core/src/main/scala/sttp/openai/OpenAi.scala
index 2ed940f..4ef67cb 100644
--- a/core/src/main/scala/sttp/openai/OpenAi.scala
+++ b/core/src/main/scala/sttp/openai/OpenAi.scala
@@ -10,6 +10,8 @@ import sttp.openai.requests.completions.chat.ChatRequestResponseData.ChatRespons
 import sttp.openai.requests.completions.edit.EditRequestBody.EditBody
 import sttp.openai.requests.completions.edit.EditRequestResponseData.EditResponse
 import sttp.openai.requests.files.FilesResponseData._
+import sttp.openai.requests.images.ImageCreationRequestBody.ImageCreationBody
+import sttp.openai.requests.images.ImageCreationResponseData.ImageCreationResponse
 import sttp.openai.requests.models.ModelsResponseData.{ModelData, ModelsResponse}
 
 class OpenAi(authToken: String) {
@@ -48,16 +50,27 @@ class OpenAi(authToken: String) {
       .get(OpenAIEndpoints.FilesEndpoint)
       .response(asJsonSnake[FilesResponse])
 
-  /** @param fileId
-    *   The ID of the file to use for this request.
-    * @return
-    *   Information about deleted file.
-    */
+    /** @param fileId
+      *   The ID of the file to use for this request.
+      * @return
+      *   Information about deleted file.
+      */
   def deleteFile(fileId: String): Request[Either[ResponseException[String, Exception], DeletedFileData]] =
     openApiAuthRequest
       .delete(OpenAIEndpoints.deleteFileEndpoint(fileId))
       .response(asJsonSnake[DeletedFileData])
 
+  /** @param imageCreationBody
+    *   Create image request body
+    *
+    * Creates an image given a prompt in request body and send it over to [[https://api.openai.com/v1/images/generations]]
+    */
+  def createImage(imageCreationBody: ImageCreationBody): Request[Either[ResponseException[String, Exception], ImageCreationResponse]] =
+    openApiAuthRequest
+      .post(OpenAIEndpoints.CreateImageEndpoint)
+      .body(imageCreationBody)
+      .response(asJsonSnake[ImageCreationResponse])
+
   /** @param editRequestBody
     *   Edit request body
     *
@@ -95,11 +108,17 @@ class OpenAi(authToken: String) {
 }
 
 private object OpenAIEndpoints {
+  private val ImageEndpointBase: Uri = uri"https://api.openai.com/v1/images"
+
   val ChatEndpoint: Uri = uri"https://api.openai.com/v1/chat/completions"
   val CompletionsEndpoint: Uri = uri"https://api.openai.com/v1/completions"
+  val CreateImageEndpoint: Uri = ImageEndpointBase.addPath("generations")
   val EditEndpoint: Uri = uri"https://api.openai.com/v1/edits"
+  val EditImageEndpoint: Uri = ImageEndpointBase.addPath("edits")
   val FilesEndpoint: Uri = uri"https://api.openai.com/v1/files"
   val ModelEndpoint: Uri = uri"https://api.openai.com/v1/models"
+  val VariationsImageEndpoint: Uri = ImageEndpointBase.addPath("variations")
+
   def deleteFileEndpoint(fileId: String): Uri = FilesEndpoint.addPath(fileId)
   def retrieveFileEndpoint(fileId: String): Uri = FilesEndpoint.addPath(fileId)
   def retrieveModelEndpoint(modelId: String): Uri = ModelEndpoint.addPath(modelId)
diff --git a/core/src/main/scala/sttp/openai/requests/images/ImageCreationRequestBody.scala b/core/src/main/scala/sttp/openai/requests/images/ImageCreationRequestBody.scala
new file mode 100644
index 0000000..547463a
--- /dev/null
+++ b/core/src/main/scala/sttp/openai/requests/images/ImageCreationRequestBody.scala
@@ -0,0 +1,94 @@
+package sttp.openai.requests.images
+
+import sttp.client4.DeserializationException
+import sttp.openai.json.SnakePickle
+import ujson.Str
+
+object ImageCreationRequestBody {
+
+  case class ImageCreationBody(
+      prompt: String,
+      n: Option[Int] = None,
+      size: Option[Size] = None,
+      responseFormat: Option[ResponseFormat] = None,
+      user: Option[String] = None
+  )
+
+  object ImageCreationBody {
+    implicit val imageCreationBodyRW: SnakePickle.ReadWriter[ImageCreationBody] = SnakePickle.macroRW[ImageCreationBody]
+  }
+
+  sealed abstract class ResponseFormat(val value: String)
+
+  object ResponseFormat {
+    case object URL extends ResponseFormat("url")
+
+    case object B64Json extends ResponseFormat("b64_json")
+
+    /** Use only as a workaround if API supports a format that's not yet predefined as a case object of Response Format. Otherwise, a custom
+      * format would be rejected. See [[https://platform.openai.com/docs/api-reference/images/create-edit]] for current list of supported
+      * formats
+      */
+    case class Custom(customResponseFormat: String) extends ResponseFormat(customResponseFormat)
+
+    val values: Set[ResponseFormat] = Set(URL, B64Json)
+
+    private val byValue =
+      values.map(s => (s.value, s)).toMap
+
+    private def withValue(responseFormat: String): Option[ResponseFormat] = byValue.get(responseFormat)
+
+    implicit val responseFormatRW: SnakePickle.ReadWriter[ResponseFormat] = SnakePickle
+      .readwriter[ujson.Value]
+      .bimap[ResponseFormat](
+        _.value,
+        jsonValue =>
+          SnakePickle.read[ujson.Value](jsonValue) match {
+            case Str(value) =>
+              withValue(value) match {
+                case Some(responseFormat) => responseFormat
+                case None                 => Custom(value)
+              }
+            case e => throw DeserializationException(e.str, new Exception(s"Could not deserialize: $e"))
+          }
+      )
+  }
+
+  sealed abstract class Size(val value: String)
+
+  object Size {
+    case object Small extends Size("256x256")
+
+    case object Medium extends Size("512x512")
+
+    case object Large extends Size("1024x1024")
+
+    /** Use only as a workaround if API supports a format that's not yet predefined as a case object of Size. Otherwise, a custom format
+      * would be rejected. See [[https://platform.openai.com/docs/api-reference/images/create-edit]] for current list of supported formats
+      */
+    case class Custom(customSize: String) extends Size(customSize)
+
+    val values: Set[Size] = Set(Small, Medium, Large)
+
+    private val byValue =
+      values.map(s => (s.value, s)).toMap
+
+    private def withValue(size: String): Option[Size] = byValue.get(size)
+
+    implicit val sizeRW: SnakePickle.ReadWriter[Size] = SnakePickle
+      .readwriter[ujson.Value]
+      .bimap[Size](
+        _.value,
+        jsonValue =>
+          SnakePickle.read[ujson.Value](jsonValue) match {
+            case Str(value) =>
+              withValue(value) match {
+                case Some(size) => size
+                case None       => Custom(value)
+              }
+            case e => throw DeserializationException(e.str, new Exception(s"Could not deserialize: $e"))
+          }
+      )
+  }
+
+}
diff --git a/core/src/main/scala/sttp/openai/requests/images/ImageCreationResponseData.scala b/core/src/main/scala/sttp/openai/requests/images/ImageCreationResponseData.scala
new file mode 100644
index 0000000..6c62298
--- /dev/null
+++ b/core/src/main/scala/sttp/openai/requests/images/ImageCreationResponseData.scala
@@ -0,0 +1,22 @@
+package sttp.openai.requests.images
+
+import sttp.openai.json.SnakePickle
+
+object ImageCreationResponseData {
+
+  case class ImageCreationResponse(
+      created: Int,
+      data: Seq[GeneratedImageData]
+  )
+
+  object ImageCreationResponse {
+    implicit val imageCreationResponseRW: SnakePickle.ReadWriter[ImageCreationResponse] = SnakePickle.macroRW[ImageCreationResponse]
+  }
+
+  case class GeneratedImageData(url: String)
+
+  object GeneratedImageData {
+    implicit val generatedImageDataRW: SnakePickle.ReadWriter[GeneratedImageData] = SnakePickle.macroRW[GeneratedImageData]
+  }
+
+}


Test patch:
diff --git a/core/src/test/scala/sttp/openai/fixtures/ImageCreationFixture.scala b/core/src/test/scala/sttp/openai/fixtures/ImageCreationFixture.scala
new file mode 100644
index 0000000..54197ad
--- /dev/null
+++ b/core/src/test/scala/sttp/openai/fixtures/ImageCreationFixture.scala
@@ -0,0 +1,20 @@
+package sttp.openai.fixtures
+
+object ImageCreationFixture {
+  val jsonRequest: String = """{
+      |"prompt": "cute fish",
+      |"n": 1,
+      |"size": "1024x1024",
+      |"response_format": "url",
+      |"user": "user1"
+      |}""".stripMargin
+
+  val jsonResponse: String = """{
+      |"created": 1681893694,
+      |"data": [
+      |  {
+      |    "url": "https://generated.image.url"
+      |  }
+      |]
+      |}""".stripMargin
+}
diff --git a/core/src/test/scala/sttp/openai/requests/completions/CompletionsDataSpec.scala b/core/src/test/scala/sttp/openai/requests/completions/CompletionsDataSpec.scala
index 1c5f688..53e2c57 100644
--- a/core/src/test/scala/sttp/openai/requests/completions/CompletionsDataSpec.scala
+++ b/core/src/test/scala/sttp/openai/requests/completions/CompletionsDataSpec.scala
@@ -132,4 +132,52 @@ class CompletionsDataSpec extends AnyFlatSpec with Matchers with EitherValues {
     serializedJson shouldBe jsonRequest
   }
 
+  "Given request of single prompt as a JSON" should "be properly serialize to case class" in {
+    import sttp.openai.requests.completions.CompletionsRequestBody._
+    import sttp.openai.requests.completions.CompletionsRequestBody.CompletionsBody._
+
+    // given
+    val request = CompletionsRequestBody.CompletionsBody(
+      model = "text-davinci-003",
+      prompt = Some(SinglePrompt("Say this is a test")),
+      maxTokens = Some(7),
+      temperature = Some(0),
+      topP = Some(1),
+      n = Some(1),
+      stream = Some(false),
+      logprobs = None,
+      stop = Some(SingleStop("\n"))
+    )
+
+    // when
+    val givenRequest: CompletionsBody = SnakePickle.read[CompletionsBody](fixtures.CompletionsFixture.jsonSinglePromptRequest)
+
+    // then
+    givenRequest shouldBe request
+  }
+
+  "Given request of multiple prompt as a JSON" should "be properly serialize to case class" in {
+    import sttp.openai.requests.completions.CompletionsRequestBody._
+    import sttp.openai.requests.completions.CompletionsRequestBody.CompletionsBody._
+
+    // given
+    val request = CompletionsRequestBody.CompletionsBody(
+      model = "text-davinci-003",
+      prompt = Some(MultiplePrompt(Seq("Say this is a test", "Say this is also a test"))),
+      maxTokens = Some(7),
+      temperature = Some(0),
+      topP = Some(1),
+      n = Some(1),
+      stream = Some(false),
+      logprobs = None,
+      stop = Some(SingleStop("\n"))
+    )
+
+    // when
+    val givenRequest: CompletionsBody = SnakePickle.read[CompletionsBody](fixtures.CompletionsFixture.jsonMultiplePromptRequest)
+
+    // then
+    givenRequest shouldBe request
+  }
+
 }
diff --git a/core/src/test/scala/sttp/openai/requests/images/ImageCreationDataSpec.scala b/core/src/test/scala/sttp/openai/requests/images/ImageCreationDataSpec.scala
new file mode 100644
index 0000000..3aae8a6
--- /dev/null
+++ b/core/src/test/scala/sttp/openai/requests/images/ImageCreationDataSpec.scala
@@ -0,0 +1,76 @@
+package sttp.openai.requests.images
+
+import org.scalatest.EitherValues
+import org.scalatest.flatspec.AnyFlatSpec
+import org.scalatest.matchers.should.Matchers
+import sttp.openai.fixtures
+import sttp.openai.json.{SnakePickle, SttpUpickleApiExtension}
+class ImageCreationDataSpec extends AnyFlatSpec with Matchers with EitherValues {
+
+  "Given image generation response as Json" should "be properly deserialized to case class" in {
+    import sttp.openai.requests.images.ImageCreationResponseData._
+    import sttp.openai.requests.images.ImageCreationResponseData.ImageCreationResponse._
+
+    // given
+    val jsonResponse = fixtures.ImageCreationFixture.jsonResponse
+
+    val generatedImageData = Seq(
+      GeneratedImageData("https://generated.image.url")
+    )
+
+    val expectedResponse: ImageCreationResponse = ImageCreationResponse(
+      created = 1681893694,
+      data = generatedImageData
+    )
+    // when
+    val givenResponse = SttpUpickleApiExtension.deserializeJsonSnake.apply(jsonResponse)
+
+    // then
+    givenResponse.value shouldBe expectedResponse
+  }
+
+  "Given create image request as case class" should "be properly serialized to Json" in {
+    import sttp.openai.requests.images.ImageCreationRequestBody._
+    import sttp.openai.requests.images.ImageCreationRequestBody.ImageCreationBody._
+
+    // given
+    val givenRequest: ImageCreationBody = ImageCreationBody(
+      prompt = "cute fish",
+      Some(1),
+      size = Some(Size.Custom("1024x1024")),
+      Some(ResponseFormat.Custom("url")),
+      Some("user1")
+    )
+
+    val jsonRequest = ujson.read(fixtures.ImageCreationFixture.jsonRequest)
+
+    // when
+    val serializedJson: ujson.Value = SnakePickle.writeJs(givenRequest)
+
+    // then
+    serializedJson shouldBe jsonRequest
+  }
+
+  "Given create image request as case class created with enum values" should "be properly serialized to Json" in {
+    import sttp.openai.requests.images.ImageCreationRequestBody._
+    import sttp.openai.requests.images.ImageCreationRequestBody.ImageCreationBody._
+
+    // given
+    val givenRequest: ImageCreationBody = ImageCreationBody(
+      "cute fish",
+      Some(1),
+      Some(Size.Large),
+      Some(ResponseFormat.URL),
+      Some("user1")
+    )
+
+    val jsonRequest = ujson.read(fixtures.ImageCreationFixture.jsonRequest)
+
+    // when
+    val serializedJson: ujson.Value = SnakePickle.writeJs(givenRequest)
+
+    // then
+    serializedJson shouldBe jsonRequest
+  }
+
+}