use an enum for different type of bodies in invoke response

Author: sebsto

Sources/InvokeModel/InvokeModelResponse.swift

@@ -26,10 +26,13 @@ public struct InvokeModelResponse {
     let model: BedrockModel
     let contentType: ContentType
 
-    // TODO: FIXME: trasnform to an enum
-    let textCompletionBody: ContainsTextCompletion?
-    let imageGenerationBody: ContainsImageGeneration?
-    let embeddingsBody: ContainsEmbeddings?
+    let body: BodyType
+
+    public enum BodyType {
+        case textCompletion(ContainsTextCompletion)
+        case imageGeneration(ContainsImageGeneration)
+        case embeddings(ContainsEmbeddings)
+    }
 
     private init(
         model: BedrockModel,
@@ -38,9 +41,7 @@ public struct InvokeModelResponse {
     ) {
         self.model = model
         self.contentType = contentType
-        self.textCompletionBody = textCompletionBody
-        self.imageGenerationBody = nil
-        self.embeddingsBody = nil
+        self.body = .textCompletion(textCompletionBody)
     }
 
     private init(
@@ -50,9 +51,7 @@ public struct InvokeModelResponse {
     ) {
         self.model = model
         self.contentType = contentType
-        self.imageGenerationBody = imageGenerationBody
-        self.textCompletionBody = nil
-        self.embeddingsBody = nil
+        self.body = .imageGeneration(imageGenerationBody)
     }
 
     private init(
@@ -62,9 +61,7 @@ public struct InvokeModelResponse {
     ) {
         self.model = model
         self.contentType = contentType
-        self.embeddingsBody = embeddingsBody
-        self.textCompletionBody = nil
-        self.imageGenerationBody = nil
+        self.body = .embeddings(embeddingsBody)
     }
 
     /// Creates a BedrockResponse from raw response data containing text completion
@@ -115,7 +112,7 @@ public struct InvokeModelResponse {
     /// - Throws: BedrockLibraryError.decodingError if the completion cannot be extracted
     public func getTextCompletion() throws -> TextCompletion {
         do {
-            guard let textCompletionBody = textCompletionBody else {
+            guard case .textCompletion(let textCompletionBody) = self.body else {
                 throw BedrockLibraryError.decodingError("No text completion body found in the response")
             }
             return try textCompletionBody.getTextCompletion()
@@ -131,7 +128,7 @@ public struct InvokeModelResponse {
     /// - Throws: BedrockLibraryError.decodingError if the image generation cannot be extracted
     public func getGeneratedImage() throws -> ImageGenerationOutput {
         do {
-            guard let imageGenerationBody = imageGenerationBody else {
+            guard case .imageGeneration(let imageGenerationBody) = self.body else {
                 throw BedrockLibraryError.decodingError("No image generation body found in the response")
             }
             return imageGenerationBody.getGeneratedImage()
@@ -144,7 +141,7 @@ public struct InvokeModelResponse {
 
     public func getEmbeddings() throws -> Embeddings {
         do {
-            guard let embeddingsBody else {
+            guard case .embeddings(let embeddingsBody) = self.body else {
                 throw BedrockLibraryError.decodingError("No embeddings body found in the response")
             }
             return embeddingsBody.getEmbeddings()