move raw schema

sashabaranov · Aug 10, 2024 · c2367ef · c2367ef
1 parent 194fad1
commit c2367ef
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 26 deletions.
diff --git a/api_integration_test.go b/api_integration_test.go
@@ -250,7 +250,7 @@ func TestChatCompletionResponseFormat_JSONSchemaRaw(t *testing.T) {
 	c := openai.NewClient(apiToken)
 	ctx := context.Background()
 
-	schema := []byte(`{"name":"cases","schema":{"type":"object","properties":{"CamelCase":{"type":"string"},"KebabCase":{"type":"string"},"PascalCase":{"type":"string"},"SnakeCase":{"type":"string"}},"required":["PascalCase","CamelCase","KebabCase","SnakeCase"],"additionalProperties":false},"strict":true}`)
+	schema := []byte(`{"type":"object","properties":{"CamelCase":{"type":"string"},"KebabCase":{"type":"string"},"PascalCase":{"type":"string"},"SnakeCase":{"type":"string"}},"required":["PascalCase","CamelCase","KebabCase","SnakeCase"],"additionalProperties":false}`)
 
 	resp, err := c.CreateChatCompletion(
 		ctx,
@@ -272,8 +272,12 @@ func TestChatCompletionResponseFormat_JSONSchemaRaw(t *testing.T) {
 				},
 			},
 			ResponseFormat: &openai.ChatCompletionResponseFormat{
-				Type:          openai.ChatCompletionResponseFormatTypeJSONSchema,
-				JSONSchemaRaw: &schema,
+				Type: openai.ChatCompletionResponseFormatTypeJSONSchema,
+				JSONSchema: &openai.ChatCompletionResponseFormatJSONSchema{
+					Name:      "cases",
+					SchemaRaw: &schema,
+					Strict:    true,
+				},
 			},
 		},
 	)

diff --git a/chat.go b/chat.go
@@ -182,45 +182,39 @@ const (
 )
 
 type ChatCompletionResponseFormat struct {
-	Type          ChatCompletionResponseFormatType        `json:"type,omitempty"`
-	JSONSchema    *ChatCompletionResponseFormatJSONSchema `json:"-"`
-	JSONSchemaRaw *[]byte                                 `json:"-"`
+	Type       ChatCompletionResponseFormatType        `json:"type,omitempty"`
+	JSONSchema *ChatCompletionResponseFormatJSONSchema `json:"json_schema,omitempty"`
 }
 
-func (c *ChatCompletionResponseFormat) MarshalJSON() ([]byte, error) {
-	if c.JSONSchema != nil && c.JSONSchemaRaw != nil {
-		return nil, errors.New("cannot define both JSONSchema and JSONSchemaRaw")
-	}
-
-	type Alias ChatCompletionResponseFormat
+type ChatCompletionResponseFormatJSONSchema struct {
+	Name        string                `json:"name"`
+	Description string                `json:"description,omitempty"`
+	Schema      jsonschema.Definition `json:"-"`
+	SchemaRaw   *[]byte               `json:"-"`
+	Strict      bool                  `json:"strict"`
+}
 
+func (c *ChatCompletionResponseFormatJSONSchema) MarshalJSON() ([]byte, error) {
+	type Alias ChatCompletionResponseFormatJSONSchema
 	var data struct {
 		*Alias
-		JSONSchema interface{} `json:"json_schema,omitempty"`
+		Schema interface{} `json:"schema,omitempty"`
 	}
 
 	data.Alias = (*Alias)(c)
 
-	if c.JSONSchema != nil {
-		data.JSONSchema = c.JSONSchema
-	} else if c.JSONSchemaRaw != nil {
-		var rawJSON interface{}
-		if err := json.Unmarshal(*c.JSONSchemaRaw, &rawJSON); err != nil {
+    data.Schema = c.Schema
+	if c.SchemaRaw != nil {
+		var rawSchema interface{}
+		if err := json.Unmarshal(*c.SchemaRaw, &rawSchema); err != nil {
 			return nil, err
 		}
-		data.JSONSchema = rawJSON
+		data.Schema = rawSchema
 	}
 
 	return json.Marshal(data)
 }
 
-type ChatCompletionResponseFormatJSONSchema struct {
-	Name        string                `json:"name"`
-	Description string                `json:"description,omitempty"`
-	Schema      jsonschema.Definition `json:"schema"`
-	Strict      bool                  `json:"strict"`
-}
-
 // ChatCompletionRequest represents a request structure for chat completion API.
 type ChatCompletionRequest struct {
 	Model            string                        `json:"model"`