From 49a5f034271638328ea5016be0a9d316a41ddc40 Mon Sep 17 00:00:00 2001
From: Richard Weerasinghe <richard@skyscrapr.io>
Date: Fri, 23 Aug 2024 16:55:14 +1200
Subject: [PATCH] fixup handling of response_format multiple type options

---
 openai/assistant.go | 23 +++++++++++++++++++++++
 openai/client.go    | 16 ++++++++++------
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/openai/assistant.go b/openai/assistant.go
index b7b5bca..6f1c2d2 100644
--- a/openai/assistant.go
+++ b/openai/assistant.go
@@ -1,6 +1,7 @@
 package openai
 
 import (
+	"encoding/json"
 	"fmt"
 	"net/url"
 	"strconv"
@@ -114,6 +115,7 @@ type AssistantToolResources struct {
 }
 
 type AssistantResponseFormat struct {
+	StringValue string `json:"-"`
 	Type       string `json:"type"`
 	JsonSchema *struct {
 		Description *string                `json:"description,omitempty"`
@@ -123,6 +125,27 @@ type AssistantResponseFormat struct {
 	} `json:"json_schema,omitempty"`
 }
 
+func (f *AssistantResponseFormat) UnmarshalJSON(data []byte) error {
+	// First try to unmarshal it as string
+	var s string
+	if err := json.Unmarshal(data, &s); err == nil {
+		// No error, fill the struct
+		f.StringValue = s
+		return nil
+	}
+
+	// Otherwise, try to unmarshal as AssistantResponseFormat
+	type Alias AssistantResponseFormat
+	var responseFormat Alias
+	if err := json.Unmarshal(data, &responseFormat); err != nil {
+		return err
+	}
+	f.Type = responseFormat.Type
+	f.JsonSchema = responseFormat.JsonSchema
+
+	return nil
+}
+
 // Create an assistant with a model and instructions.
 // [OpenAI Documentation]: https://platform.openai.com/docs/api-reference/assistants/createAssistant
 func (e *AssistantsEndpoint) CreateAssistant(req *AssistantRequest) (*Assistant, error) {
diff --git a/openai/client.go b/openai/client.go
index baf9ffc..c82a07d 100644
--- a/openai/client.go
+++ b/openai/client.go
@@ -5,7 +5,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"log"
 	"net/http"
 	"net/url"
 	"time"
@@ -96,11 +95,16 @@ func decodeResponse(body io.Reader, v any) error {
 	if v == nil {
 		return nil
 	}
-	err := json.NewDecoder(body).Decode(v)
-	if err != nil {
-		log.Fatal(err)
-	}
-	return nil
+	return json.NewDecoder(body).Decode(v)
+
+	//
+	// DEBUG - Use below if you want to see the raw response body
+	//
+	// rawBody, err := io.ReadAll(body)
+	// if err != nil {
+	// 	return err
+	// }
+	// return json.Unmarshal(rawBody, v)
 }
 
 func (c *Client) handleErrorResp(resp *http.Response) error {