Skip to content

Commit d8cd419

Browse files
learner0810BenjaminBraunDev
learner0810
authored and
BenjaminBraunDev
committed
support multi modal inputs (kubernetes-sigs#1617)
1 parent a00349f commit d8cd419

File tree

3 files changed

+137
-25
lines changed

3 files changed

+137
-25
lines changed

‎pkg/epp/scheduling/framework/plugins/multi/prefix/plugin_test.go‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -217,8 +217,8 @@ func TestPrefixPluginChatCompletions(t *testing.T) {
217217
Body: &types.LLMRequestBody{
218218
ChatCompletions: &types.ChatCompletionsRequest{
219219
Messages: []types.Message{
220-
{Role: "user", Content: "hello world"},
221-
{Role: "assistant", Content: "hi there"},
220+
{Role: "user", Content: types.Content{Raw: "hello world"}},
221+
{Role: "assistant", Content: types.Content{Raw: "hi there"}},
222222
},
223223
},
224224
},
@@ -252,8 +252,8 @@ func TestPrefixPluginChatCompletionsGrowth(t *testing.T) {
252252
Body: &types.LLMRequestBody{
253253
ChatCompletions: &types.ChatCompletionsRequest{
254254
Messages: []types.Message{
255-
{Role: "system", Content: "You are a helpful assistant"},
256-
{Role: "user", Content: "Hello, how are you?"},
255+
{Role: "system", Content: types.Content{Raw: "You are a helpful assistant"}},
256+
{Role: "user", Content: types.Content{Raw: "Hello, how are you?"}},
257257
},
258258
},
259259
},
@@ -285,10 +285,10 @@ func TestPrefixPluginChatCompletionsGrowth(t *testing.T) {
285285
Body: &types.LLMRequestBody{
286286
ChatCompletions: &types.ChatCompletionsRequest{
287287
Messages: []types.Message{
288-
{Role: "system", Content: "You are a helpful assistant"},
289-
{Role: "user", Content: "Hello, how are you?"},
290-
{Role: "assistant", Content: "I'm doing well, thank you! How can I help you today?"},
291-
{Role: "user", Content: "Can you explain how prefix caching works?"},
288+
{Role: "system", Content: types.Content{Raw: "You are a helpful assistant"}},
289+
{Role: "user", Content: types.Content{Raw: "Hello, how are you?"}},
290+
{Role: "assistant", Content: types.Content{Raw: "I'm doing well, thank you! How can I help you today?"}},
291+
{Role: "user", Content: types.Content{Raw: "Can you explain how prefix caching works?"}},
292292
},
293293
},
294294
},
@@ -318,12 +318,12 @@ func TestPrefixPluginChatCompletionsGrowth(t *testing.T) {
318318
Body: &types.LLMRequestBody{
319319
ChatCompletions: &types.ChatCompletionsRequest{
320320
Messages: []types.Message{
321-
{Role: "system", Content: "You are a helpful assistant"},
322-
{Role: "user", Content: "Hello, how are you?"},
323-
{Role: "assistant", Content: "I'm doing well, thank you! How can I help you today?"},
324-
{Role: "user", Content: "Can you explain how prefix caching works?"},
325-
{Role: "assistant", Content: "Prefix caching is a technique where..."},
326-
{Role: "user", Content: "That's very helpful, thank you!"},
321+
{Role: "system", Content: types.Content{Raw: "You are a helpful assistant"}},
322+
{Role: "user", Content: types.Content{Raw: "Hello, how are you?"}},
323+
{Role: "assistant", Content: types.Content{Raw: "I'm doing well, thank you! How can I help you today?"}},
324+
{Role: "user", Content: types.Content{Raw: "Can you explain how prefix caching works?"}},
325+
{Role: "assistant", Content: types.Content{Raw: "Prefix caching is a technique where..."}},
326+
{Role: "user", Content: types.Content{Raw: "That's very helpful, thank you!"}},
327327
},
328328
},
329329
},
@@ -437,15 +437,15 @@ func BenchmarkPrefixPluginChatCompletionsStress(b *testing.B) {
437437
b.Run(fmt.Sprintf("messages_%d_length_%d", scenario.messageCount, scenario.messageLength), func(b *testing.B) {
438438
// Generate messages for this scenario
439439
messages := make([]types.Message, scenario.messageCount)
440-
messages[0] = types.Message{Role: "system", Content: "You are a helpful assistant."}
440+
messages[0] = types.Message{Role: "system", Content: types.Content{Raw: "You are a helpful assistant."}}
441441

442442
for i := 1; i < scenario.messageCount; i++ {
443443
role := "user"
444444
if i%2 == 0 {
445445
role = "assistant"
446446
}
447447
content := randomPrompt(scenario.messageLength)
448-
messages[i] = types.Message{Role: role, Content: content}
448+
messages[i] = types.Message{Role: role, Content: types.Content{Raw: content}}
449449
}
450450

451451
pod := &types.PodMetrics{

‎pkg/epp/scheduling/types/types.go‎

Lines changed: 66 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,10 @@ limitations under the License.
1717
package types
1818

1919
import (
20+
"encoding/json"
21+
"errors"
2022
"fmt"
23+
"strings"
2124

2225
"sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend"
2326
backendmetrics "sigs.k8s.io/gateway-api-inference-extension/pkg/epp/backend/metrics"
@@ -130,16 +133,75 @@ func (r *ChatCompletionsRequest) String() string {
130133

131134
messagesLen := 0
132135
for _, msg := range r.Messages {
133-
messagesLen += len(msg.Content)
136+
messagesLen += len(msg.Content.PlainText())
134137
}
135-
136138
return fmt.Sprintf("{MessagesLength: %d}", messagesLen)
137139
}
138140

139141
// Message represents a single message in a chat-completions request.
140142
type Message struct {
141-
Role string
142-
Content string // TODO: support multi-modal content
143+
// Role is the message Role, optional values are 'user', 'assistant', ...
144+
Role string `json:"role,omitempty"`
145+
// Content defines text of this message
146+
Content Content `json:"content,omitempty"`
147+
}
148+
149+
type Content struct {
150+
Raw string
151+
Structured []ContentBlock
152+
}
153+
154+
type ContentBlock struct {
155+
Type string `json:"type"`
156+
Text string `json:"text,omitempty"`
157+
ImageURL ImageBlock `json:"image_url,omitempty"`
158+
}
159+
160+
type ImageBlock struct {
161+
Url string `json:"url,omitempty"`
162+
}
163+
164+
// UnmarshalJSON allow use both format
165+
func (mc *Content) UnmarshalJSON(data []byte) error {
166+
// Raw format
167+
var str string
168+
if err := json.Unmarshal(data, &str); err == nil {
169+
mc.Raw = str
170+
return nil
171+
}
172+
173+
// Block format
174+
var blocks []ContentBlock
175+
if err := json.Unmarshal(data, &blocks); err == nil {
176+
mc.Structured = blocks
177+
return nil
178+
}
179+
180+
return errors.New("content format not supported")
181+
}
182+
183+
func (mc Content) MarshalJSON() ([]byte, error) {
184+
if mc.Raw != "" {
185+
return json.Marshal(mc.Raw)
186+
}
187+
if mc.Structured != nil {
188+
return json.Marshal(mc.Structured)
189+
}
190+
return json.Marshal("")
191+
}
192+
193+
func (mc Content) PlainText() string {
194+
if mc.Raw != "" {
195+
return mc.Raw
196+
}
197+
var sb strings.Builder
198+
for _, block := range mc.Structured {
199+
if block.Type == "text" {
200+
sb.WriteString(block.Text)
201+
sb.WriteString(" ")
202+
}
203+
}
204+
return sb.String()
143205
}
144206

145207
type Pod interface {

‎pkg/epp/util/request/body_test.go‎

Lines changed: 55 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,58 @@ func TestExtractRequestData(t *testing.T) {
5858
want: &types.LLMRequestBody{
5959
ChatCompletions: &types.ChatCompletionsRequest{
6060
Messages: []types.Message{
61-
{Role: "system", Content: "this is a system message"},
62-
{Role: "user", Content: "hello"},
61+
{Role: "system", Content: types.Content{Raw: "this is a system message"}},
62+
{Role: "user", Content: types.Content{Raw: "hello"}},
63+
},
64+
},
65+
},
66+
},
67+
{
68+
name: "chat completions request body with multi-modal content",
69+
body: map[string]any{
70+
"model": "test",
71+
"messages": []any{
72+
map[string]any{
73+
"role": "system",
74+
"content": []map[string]any{
75+
{
76+
"type": "text",
77+
"text": "Describe this image in one sentence.",
78+
},
79+
},
80+
},
81+
map[string]any{
82+
"role": "user",
83+
"content": []map[string]any{
84+
{
85+
"type": "image_url",
86+
"image_url": map[string]any{
87+
"url": "https://example.com/images/dui.jpg.",
88+
},
89+
},
90+
},
91+
},
92+
},
93+
},
94+
want: &types.LLMRequestBody{
95+
ChatCompletions: &types.ChatCompletionsRequest{
96+
Messages: []types.Message{
97+
{Role: "system", Content: types.Content{
98+
Structured: []types.ContentBlock{
99+
{
100+
Text: "Describe this image in one sentence.",
101+
Type: "text",
102+
},
103+
},
104+
}},
105+
{Role: "user", Content: types.Content{
106+
Structured: []types.ContentBlock{
107+
{
108+
Type: "image_url",
109+
ImageURL: types.ImageBlock{Url: "https://example.com/images/dui.jpg."},
110+
},
111+
},
112+
}},
63113
},
64114
},
65115
},
@@ -81,7 +131,7 @@ func TestExtractRequestData(t *testing.T) {
81131
},
82132
want: &types.LLMRequestBody{
83133
ChatCompletions: &types.ChatCompletionsRequest{
84-
Messages: []types.Message{{Role: "user", Content: "hello"}},
134+
Messages: []types.Message{{Role: "user", Content: types.Content{Raw: "hello"}}},
85135
Tools: []any{map[string]any{"type": "function"}},
86136
Documents: []any{map[string]any{"content": "doc"}},
87137
ChatTemplate: "custom template",
@@ -256,8 +306,8 @@ func TestExtractRequestData(t *testing.T) {
256306
want: &types.LLMRequestBody{
257307
ChatCompletions: &types.ChatCompletionsRequest{
258308
Messages: []types.Message{
259-
{Role: "system", Content: "this is a system message"},
260-
{Role: "user", Content: "hello"},
309+
{Role: "system", Content: types.Content{Raw: "this is a system message"}},
310+
{Role: "user", Content: types.Content{Raw: "hello"}},
261311
},
262312
CacheSalt: "Z3V2bmV3aGxza3ZubGFoZ3Zud3V3ZWZ2bmd0b3V2bnZmc2xpZ3RoZ2x2aQ==",
263313
},

0 commit comments

Comments
 (0)