Skip to content

Commit 457b727

Browse files
zhuangjie1125zhuangjie1125
committed
feat: chat completion message add extra fields
1 parent 07791be commit 457b727

File tree

2 files changed

+92
-30
lines changed

2 files changed

+92
-30
lines changed

‎chat.go‎

Lines changed: 51 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"errors"
88
"io"
99
"net/http"
10+
"reflect"
1011

1112
openai "github.com/meguminnnnnnnnn/go-openai/internal"
1213

@@ -136,6 +137,8 @@ type ChatCompletionMessage struct {
136137

137138
// For Role=tool prompts this should be set to the ID given in the assistant's prior request to call a tool.
138139
ToolCallID string `json:"tool_call_id,omitempty"`
140+
141+
ExtraFields map[string]json.RawMessage `json:"-"`
139142
}
140143

141144
func (m ChatCompletionMessage) MarshalJSON() ([]byte, error) {
@@ -144,29 +147,31 @@ func (m ChatCompletionMessage) MarshalJSON() ([]byte, error) {
144147
}
145148
if len(m.MultiContent) > 0 {
146149
msg := struct {
147-
Role string `json:"role"`
148-
Content string `json:"-"`
149-
Refusal string `json:"refusal,omitempty"`
150-
MultiContent []ChatMessagePart `json:"content,omitempty"`
151-
Name string `json:"name,omitempty"`
152-
ReasoningContent string `json:"reasoning_content,omitempty"`
153-
FunctionCall *FunctionCall `json:"function_call,omitempty"`
154-
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
155-
ToolCallID string `json:"tool_call_id,omitempty"`
150+
Role string `json:"role"`
151+
Content string `json:"-"`
152+
Refusal string `json:"refusal,omitempty"`
153+
MultiContent []ChatMessagePart `json:"content,omitempty"`
154+
Name string `json:"name,omitempty"`
155+
ReasoningContent string `json:"reasoning_content,omitempty"`
156+
FunctionCall *FunctionCall `json:"function_call,omitempty"`
157+
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
158+
ToolCallID string `json:"tool_call_id,omitempty"`
159+
ExtraFields map[string]json.RawMessage `json:"-"`
156160
}(m)
157161
return json.Marshal(msg)
158162
}
159163

160164
msg := struct {
161-
Role string `json:"role"`
162-
Content string `json:"content,omitempty"`
163-
Refusal string `json:"refusal,omitempty"`
164-
MultiContent []ChatMessagePart `json:"-"`
165-
Name string `json:"name,omitempty"`
166-
ReasoningContent string `json:"reasoning_content,omitempty"`
167-
FunctionCall *FunctionCall `json:"function_call,omitempty"`
168-
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
169-
ToolCallID string `json:"tool_call_id,omitempty"`
165+
Role string `json:"role"`
166+
Content string `json:"content,omitempty"`
167+
Refusal string `json:"refusal,omitempty"`
168+
MultiContent []ChatMessagePart `json:"-"`
169+
Name string `json:"name,omitempty"`
170+
ReasoningContent string `json:"reasoning_content,omitempty"`
171+
FunctionCall *FunctionCall `json:"function_call,omitempty"`
172+
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
173+
ToolCallID string `json:"tool_call_id,omitempty"`
174+
ExtraFields map[string]json.RawMessage `json:"-"`
170175
}(m)
171176
return json.Marshal(msg)
172177
}
@@ -177,32 +182,48 @@ func (m *ChatCompletionMessage) UnmarshalJSON(bs []byte) error {
177182
Content string `json:"content"`
178183
Refusal string `json:"refusal,omitempty"`
179184
MultiContent []ChatMessagePart
180-
Name string `json:"name,omitempty"`
181-
ReasoningContent string `json:"reasoning_content,omitempty"`
182-
FunctionCall *FunctionCall `json:"function_call,omitempty"`
183-
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
184-
ToolCallID string `json:"tool_call_id,omitempty"`
185+
Name string `json:"name,omitempty"`
186+
ReasoningContent string `json:"reasoning_content,omitempty"`
187+
FunctionCall *FunctionCall `json:"function_call,omitempty"`
188+
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
189+
ToolCallID string `json:"tool_call_id,omitempty"`
190+
ExtraFields map[string]json.RawMessage `json:"-"`
185191
}{}
186192

187193
if err := json.Unmarshal(bs, &msg); err == nil {
188194
*m = ChatCompletionMessage(msg)
195+
extra, err := openai.UnmarshalExtraFields(reflect.TypeOf(m), bs)
196+
if err != nil {
197+
return err
198+
}
199+
200+
m.ExtraFields = extra
189201
return nil
190202
}
203+
191204
multiMsg := struct {
192205
Role string `json:"role"`
193206
Content string
194-
Refusal string `json:"refusal,omitempty"`
195-
MultiContent []ChatMessagePart `json:"content"`
196-
Name string `json:"name,omitempty"`
197-
ReasoningContent string `json:"reasoning_content,omitempty"`
198-
FunctionCall *FunctionCall `json:"function_call,omitempty"`
199-
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
200-
ToolCallID string `json:"tool_call_id,omitempty"`
207+
Refusal string `json:"refusal,omitempty"`
208+
MultiContent []ChatMessagePart `json:"content"`
209+
Name string `json:"name,omitempty"`
210+
ReasoningContent string `json:"reasoning_content,omitempty"`
211+
FunctionCall *FunctionCall `json:"function_call,omitempty"`
212+
ToolCalls []ToolCall `json:"tool_calls,omitempty"`
213+
ToolCallID string `json:"tool_call_id,omitempty"`
214+
ExtraFields map[string]json.RawMessage `json:"-"`
201215
}{}
202216
if err := json.Unmarshal(bs, &multiMsg); err != nil {
203217
return err
204218
}
205219
*m = ChatCompletionMessage(multiMsg)
220+
221+
extra, err := openai.UnmarshalExtraFields(reflect.TypeOf(m), bs)
222+
if err != nil {
223+
return err
224+
}
225+
226+
m.ExtraFields = extra
206227
return nil
207228
}
208229

‎chat_test.go‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
"github.com/meguminnnnnnnnn/go-openai"
1616
"github.com/meguminnnnnnnnn/go-openai/internal/test/checks"
1717
"github.com/meguminnnnnnnnn/go-openai/jsonschema"
18+
"github.com/stretchr/testify/assert"
1819
)
1920

2021
const (
@@ -1160,3 +1161,43 @@ func TestChatCompletionRequest_UnmarshalJSON(t *testing.T) {
11601161
})
11611162
}
11621163
}
1164+
1165+
func TestChatCompletionMessage_UnmarshalJSON(t *testing.T) {
1166+
bs := []byte(`{
1167+
"role": "system",
1168+
"content": "You are a helpful math tutor.",
1169+
"name": "name",
1170+
"multimodal_contents": [
1171+
{
1172+
"type": "text",
1173+
"text": "ok"
1174+
},
1175+
{
1176+
"type": "text",
1177+
"text": "Generate a picture of a Shiba Inu dog for you。"
1178+
},
1179+
{
1180+
"type": "inline_data",
1181+
"inline_data": {
1182+
"mime_type": "image/png",
1183+
"data": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAI"
1184+
}
1185+
}
1186+
]
1187+
}`)
1188+
chatMessage := &openai.ChatCompletionMessage{}
1189+
err := json.Unmarshal(bs, chatMessage)
1190+
assert.Nil(t, err)
1191+
1192+
multimodalContent := chatMessage.ExtraFields["multimodal_contents"]
1193+
mContents := make([]map[string]any, 0)
1194+
err = json.Unmarshal(multimodalContent, &mContents)
1195+
assert.Nil(t, err)
1196+
1197+
assert.Equal(t, mContents, []map[string]any{
1198+
{"type": "text", "text": "ok"},
1199+
{"type": "text", "text": "Generate a picture of a Shiba Inu dog for you。"},
1200+
{"type": "inline_data", "inline_data": map[string]any{"mime_type": "image/png", "data": "iVBORw0KGgoAAAANSUhEUgAAAgAAAAI"}},
1201+
})
1202+
1203+
}

0 commit comments

Comments
 (0)