Repository: softwaremill/sttp-openai
Pull request #75 (2023-05-08 08:39:06)

Edits models adt


Patch:
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 78c1c24..426ced6 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,7 +1,8 @@
 package sttp.openai.requests.completions.chat
 
-import sttp.openai.json.SnakePickle
+import sttp.openai.json.{DeserializationException, SnakePickle}
 import sttp.openai.requests.completions.Stop
+import ujson.Str
 
 object ChatRequestBody {
 
@@ -33,7 +34,7 @@ object ChatRequestBody {
     *   A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse.
     */
   case class ChatBody(
-      model: String,
+      model: ChatCompletionModel,
       messages: Seq[Message],
       temperature: Option[Double] = None,
       topP: Option[Double] = None,
@@ -50,4 +51,38 @@ object ChatRequestBody {
     implicit val chatRequestW: SnakePickle.Writer[ChatBody] = SnakePickle.macroW[ChatBody]
   }
 
+  sealed abstract class ChatCompletionModel(val value: String)
+
+  object ChatCompletionModel {
+    implicit val chatCompletionModelRW: SnakePickle.ReadWriter[ChatCompletionModel] = SnakePickle
+      .readwriter[ujson.Value]
+      .bimap[ChatCompletionModel](
+        model => SnakePickle.writeJs(model.value),
+        jsonValue =>
+          SnakePickle.read[ujson.Value](jsonValue) match {
+            case Str(value) =>
+              byChatModelValue.getOrElse(value, throw new DeserializationException(new Exception(s"Could not deserialize: $value")))
+            case e => throw new DeserializationException(new Exception(s"Could not deserialize: $e"))
+          }
+      )
+
+    case object GPT4 extends ChatCompletionModel("gpt-4")
+
+    case object GPT40314 extends ChatCompletionModel("gpt-4-0314")
+
+    case object GPT432k extends ChatCompletionModel("gpt-4-32k")
+
+    case object GPT432k0314 extends ChatCompletionModel("gpt-4-32k-0314")
+
+    case object GPT35Turbo extends ChatCompletionModel("gpt-3.5-turbo")
+
+    case object GPT35Turbo0301 extends ChatCompletionModel("gpt-3.5-turbo-0301")
+
+    case class CustomChatCompletionModel(customChatCompletionModel: String) extends ChatCompletionModel(customChatCompletionModel)
+
+    val values: Set[ChatCompletionModel] = Set(GPT4, GPT40314, GPT432k, GPT432k0314, GPT35Turbo, GPT35Turbo0301)
+
+    private val byChatModelValue = values.map(model => model.value -> model).toMap
+  }
+
 }
diff --git a/core/src/main/scala/sttp/openai/requests/completions/edit/EditRequestBody.scala b/core/src/main/scala/sttp/openai/requests/completions/edit/EditRequestBody.scala
index 88018e1..6789c86 100644
--- a/core/src/main/scala/sttp/openai/requests/completions/edit/EditRequestBody.scala
+++ b/core/src/main/scala/sttp/openai/requests/completions/edit/EditRequestBody.scala
@@ -1,11 +1,12 @@
 package sttp.openai.requests.completions.edit
 
-import sttp.openai.json.SnakePickle
+import sttp.openai.json.{DeserializationException, SnakePickle}
+import ujson.Str
 
 object EditRequestBody {
 
   /** @param model
-    *   ID of the model to use. You can use the `text-davinci-edit-001` or `code-davinci-edit-001` model with this endpoint.
+    *   ID of the [[EditModel]] to use.
     * @param input
     *   The input text to use as a starting point for the edit.
     * @param instruction
@@ -20,7 +21,7 @@ object EditRequestBody {
     *   probability mass. So 0.1 means only the tokens comprising the top 10% probability mass are considered.
     */
   case class EditBody(
-      model: String,
+      model: EditModel,
       input: Option[String] = None,
       instruction: String,
       n: Option[Int] = None,
@@ -31,4 +32,29 @@ object EditRequestBody {
   object EditBody {
     implicit val editBodyW: SnakePickle.Writer[EditBody] = SnakePickle.macroW[EditBody]
   }
+
+  sealed abstract class EditModel(val value: String)
+
+  object EditModel {
+    implicit val editModelRW: SnakePickle.ReadWriter[EditModel] = SnakePickle
+      .readwriter[ujson.Value]
+      .bimap[EditModel](
+        model => SnakePickle.writeJs(model.value),
+        jsonValue =>
+          SnakePickle.read[ujson.Value](jsonValue) match {
+            case Str(value) =>
+              byEditModelValue.getOrElse(value, throw new DeserializationException(new Exception(s"Could not deserialize: $value")))
+            case e => throw new DeserializationException(new Exception(s"Could not deserialize: $e"))
+          }
+      )
+
+    case object TextDavinciEdit001 extends EditModel("text-davinci-edit-001")
+    case object CodeDavinciEdit001 extends EditModel("code-davinci-edit-001")
+    case class CustomEditModel(customEditModel: String) extends EditModel(customEditModel)
+
+    val values: Set[EditModel] = Set(TextDavinciEdit001, CodeDavinciEdit001)
+
+    private val byEditModelValue = values.map(model => model.value -> model).toMap
+  }
+
 }


Test patch:
diff --git a/core/src/test/scala/sttp/openai/requests/completions/EditDataSpec.scala b/core/src/test/scala/sttp/openai/requests/completions/EditDataSpec.scala
index 165e117..6827f78 100644
--- a/core/src/test/scala/sttp/openai/requests/completions/EditDataSpec.scala
+++ b/core/src/test/scala/sttp/openai/requests/completions/EditDataSpec.scala
@@ -6,6 +6,7 @@ import org.scalatest.flatspec.AnyFlatSpec
 import org.scalatest.matchers.should.Matchers
 import sttp.openai.json.{SnakePickle, SttpUpickleApiExtension}
 import sttp.openai.requests.completions.edit.EditRequestBody._
+import sttp.openai.requests.completions.edit.EditRequestBody.EditModel.TextDavinciEdit001
 
 class EditDataSpec extends AnyFlatSpec with Matchers with EitherValues {
 
@@ -44,7 +45,7 @@ class EditDataSpec extends AnyFlatSpec with Matchers with EitherValues {
   "Given completions request as case class" should "be properly serialized to Json" in {
     // given
     val givenRequest: EditBody = EditBody(
-      model = "text-davinci-edit-001",
+      model = TextDavinciEdit001,
       input = Some("What day of the wek is it?"),
       instruction = "Fix the spelling mistakes"
     )
diff --git a/core/src/test/scala/sttp/openai/requests/completions/chat/ChatDataSpec.scala b/core/src/test/scala/sttp/openai/requests/completions/chat/ChatDataSpec.scala
index c677a7e..fccc479 100644
--- a/core/src/test/scala/sttp/openai/requests/completions/chat/ChatDataSpec.scala
+++ b/core/src/test/scala/sttp/openai/requests/completions/chat/ChatDataSpec.scala
@@ -7,12 +7,14 @@ import sttp.openai.fixtures
 import sttp.openai.json.{SnakePickle, SttpUpickleApiExtension}
 import sttp.openai.requests.completions.Usage
 import sttp.openai.requests.completions.Stop.SingleStop
+import sttp.openai.requests.completions.chat.ChatRequestBody.ChatCompletionModel
 
 class ChatDataSpec extends AnyFlatSpec with Matchers with EitherValues {
 
   "Given chat completions response as Json" should "be properly deserialized to case class" in {
     import sttp.openai.requests.completions.chat.ChatRequestResponseData._
     import sttp.openai.requests.completions.chat.ChatRequestResponseData.ChatResponse._
+    import sttp.openai.requests.completions.chat.ChatRequestBody.ChatCompletionModel
 
     // given
     val jsonResponse = fixtures.ChatFixture.jsonResponse
@@ -62,7 +64,7 @@ class ChatDataSpec extends AnyFlatSpec with Matchers with EitherValues {
     )
 
     val givenRequest = ChatRequestBody.ChatBody(
-      model = "gpt-3.5-turbo",
+      model = ChatCompletionModel.GPT35Turbo,
       messages = messages,
       temperature = Some(1),
       topP = Some(1),