Fixing issues

zio · Nov 15, 2023 · 3177b10 · 3177b10
1 parent a6c0d37
commit 3177b10
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 12 deletions.
diff --git a/zio-openai-codegen/src/main/scala/zio/openai/codegen/generator/APIGenerator.scala b/zio-openai-codegen/src/main/scala/zio/openai/codegen/generator/APIGenerator.scala
@@ -278,9 +278,9 @@ trait APIGenerator {
 
             val url =
               if (hasQueryParams)
-                q"""baseURL.withPath(baseURL.path ++ $path).withQueryParams($queryParams)"""
+                q"""baseURL.path(baseURL.path ++ $path).queryParams($queryParams)"""
               else
-                q"""baseURL.withPath(baseURL.path ++ $path)"""
+                q"""baseURL.path(baseURL.path ++ $path)"""
 
             val body =
               endpoint.body match {
@@ -305,7 +305,7 @@ trait APIGenerator {
                 q"""${Types.zhttpHeader.term}.ContentType(${endpoint.bodyContentTypeAsMediaType})"""
 
             val request =
-              q"""${Types.zhttpRequest.term}.default(
+              q"""${Types.zhttpRequest.term}(
                method = ${endpoint.method.constructor.term},
                url = $url,
                body = body
@@ -319,6 +319,9 @@ trait APIGenerator {
                     if responseBody.contentType == ContentType.`application/json` =>
                   val responseType = responseBody.typ.scalaType(model)
                   q"""${Types.decoders.term}.tryDecodeJsonResponse[${responseType.typ}](this.codecs, req, response)"""
+                case Some(responseBody)
+                  if responseBody.contentType == ContentType.`application/octet-stream` =>
+                  q"""${Types.decoders.term}.tryDecodeBinaryResponse(req, response)"""
                 case None =>
                   q"""${Types.decoders.term}.validateEmptyResponse(req, response)"""
                 case _    =>
@@ -333,8 +336,10 @@ trait APIGenerator {
                   .typ} = {
                   $body.flatMap { body =>
                     val req = $request
-                    client.request(req).mapError(${Types.openAIFailure.term}.Unknown(_)).flatMap { response =>
-                      $mapResponse
+                    ${Types.zio_.term}.scoped {
+                      client.request(req).mapError(${Types.openAIFailure.term}.Unknown(_)).flatMap { response =>
+                        $mapResponse
+                      }
                     }
                   }
                 }
@@ -348,7 +353,7 @@ trait APIGenerator {
                   val streaming: Defn.Def =
                     q"""def ${endpoint.methodNameStreaming}(..$paramList): ${Types
                         .zstream(ScalaType.any, Types.openAIFailure, streamingResponseType)
-                        .typ} = ${Types.zstream_.term}.unwrap {
+                        .typ} = ${Types.zstream_.term}.unwrapScoped {
                           $body.flatMap { body =>
                             val req = $request
                             client.request(req).mapError(${Types.openAIFailure.term}.Unknown(_)).map { response =>

diff --git a/zio-openai-codegen/src/main/scala/zio/openai/codegen/model/API.scala b/zio-openai-codegen/src/main/scala/zio/openai/codegen/model/API.scala
@@ -10,6 +10,8 @@ final case class API(name: String, endpoints: List[Endpoint]) {
 
   def transform(f: TypeDefinition => TypeDefinition): API =
     copy(endpoints = endpoints.map(_.transform(f)))
+
+  def allTypes: Set[TypeDefinition] = endpoints.flatMap(_.allTypes).toSet
 }
 
 object API {

diff --git a/zio-openai-codegen/src/main/scala/zio/openai/codegen/model/Endpoint.scala b/zio-openai-codegen/src/main/scala/zio/openai/codegen/model/Endpoint.scala
@@ -80,4 +80,11 @@ final case class Endpoint(
       body = body.map(_.transform(f)),
       response = response.map(_.transform(f))
     )
+
+  def allTypes: Set[TypeDefinition] = {
+    val params = parameters.map(_.typ).toSet
+    val bodyTypes = body.map(_.typ).toSet
+    val responseTypes = response.map(_.typ).toSet
+    params ++ bodyTypes ++ responseTypes
+  }
 }
diff --git a/zio-openai-codegen/src/main/scala/zio/openai/codegen/model/Model.scala b/zio-openai-codegen/src/main/scala/zio/openai/codegen/model/Model.scala
@@ -9,7 +9,8 @@ final case class Model(
   initialAPIs: List[API]
 ) {
   private lazy val allTypes: Map[String, TypeDefinition] =
-    Model.collectReferencedTypes(types.values.toSeq)
+    Model.collectReferencedTypes(types.values.toSeq) ++
+      Model.collectReferencedTypes(initialAPIs.flatMap(api => api.allTypes))
 
   lazy val finalTypes: Map[String, TypeDefinition] =
     allTypes.mapValues(_.transform(unifyTypes))

diff --git a/zio-openai/src/main/scala/zio/openai/internal/Decoders.scala b/zio-openai/src/main/scala/zio/openai/internal/Decoders.scala
@@ -1,16 +1,24 @@
 package zio.openai.internal
 
-import zio.ZIO
-import zio.constraintless.{ IsElementOf, TypeList }
+import zio.{Chunk, ZIO}
+import zio.constraintless.{IsElementOf, TypeList}
 import zio.http.Status
-import zio.http.{ Request, Response }
-import zio.openai.model.{ ErrorResponse, OpenAIFailure }
-import zio.schema.codec.{ BinaryCodecs, JsonCodec }
+import zio.http.{Request, Response}
+import zio.openai.model.{ErrorResponse, OpenAIFailure}
+import zio.schema.codec.{BinaryCodecs, JsonCodec}
 
 import java.nio.charset.StandardCharsets
 
 private[openai] object Decoders {
 
+  def tryDecodeBinaryResponse(request: Request, response: Response): ZIO[Any, OpenAIFailure, Chunk[Byte]] =
+    if (response.status == Status.Ok) {
+      response.body.asChunk
+        .mapError(OpenAIFailure.Unknown(_))
+    } else {
+      failWithErrorResponse(request, response)
+    }
+
   def tryDecodeJsonResponse[T]: TryDecodeJsonResponse[T] = new TryDecodeJsonResponse[T](())
 
   def validateEmptyResponse(request: Request, response: Response): ZIO[Any, OpenAIFailure, Unit] =