lkn
/
ollie-core
Archived
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

codewhisperer: image content block support in user messages and tool results

This commit is contained in:
Levi Neely 2026-05-25 23:14:40 +02:00
parent 4b4dab3f02
commit 2b914ce150
4 changed files with 124 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
pkg/agent/core_test.go
View File

@ -677,7 +677,11 @@ func TestCommand_I_SetsInject(t *testing.T) {
return textStream("ok"), nil
}
c := newCore(t, be, nil)
go c.Submit(context.Background(), "hello")
submitDone := make(chan struct{})
go func() {
c.Submit(context.Background(), "hello")
close(submitDone)
}()
waitState(t, c, "thinking")
c.handleCommand(context.Background(), "/i my inject")
p := c.pendingInject.Load()
@ -685,7 +689,7 @@ func TestCommand_I_SetsInject(t *testing.T) {
t.Errorf("pendingInject = %v; want %q", p, "my inject")
}
close(done)
waitState(t, c, "idle")
<-submitDone
}
func TestCommand_I_SubmitsWhenIdle(t *testing.T) {

103
pkg/backend/codewhisperer.go
View File

@ -282,8 +282,8 @@ func buildKiroRequest(model string, messages []Message, tools []Tool, profileARN
if err != nil {
return nil, err
}
toolResults := encodeKiroToolResults(messages[cutoff+1:])
current = kiroChatMessage{UserInputMessage: emptyKiroUserMessage(model, encodedTools, toolResults)}
toolResults, images := encodeKiroToolResults(messages[cutoff+1:])
current = kiroChatMessage{UserInputMessage: emptyKiroUserMessage(model, encodedTools, toolResults, images)}
default:
return nil, fmt.Errorf("last message role %q is not supported (expected user, system, or tool)", lastRole)
@ -326,8 +326,8 @@ func encodeKiroHistory(messages []Message) ([]kiroChatMessage, error) {
for i+1 < len(messages) && strings.EqualFold(messages[i+1].Role, "tool") {
i++
}
results := encodeKiroToolResults(messages[start : i+1])
out = append(out, kiroChatMessage{UserInputMessage: emptyKiroUserMessage("", nil, results)})
results, imgs := encodeKiroToolResults(messages[start : i+1])
out = append(out, kiroChatMessage{UserInputMessage: emptyKiroUserMessage("", nil, results, imgs)})
}
}
return out, nil
@ -335,7 +335,7 @@ func encodeKiroHistory(messages []Message) ([]kiroChatMessage, error) {
func encodeKiroUserMessage(m Message, modelID string, tools []kiroTool, toolResults []kiroToolResult) *kiroUserInputMessage {
msg := &kiroUserInputMessage{
Content: m.Content,
Content: kiroEncodeContent(m),
Origin: kiroDefaultOrigin,
ModelID: modelID,
}
@ -352,11 +352,12 @@ func encodeKiroUserMessage(m Message, modelID string, tools []kiroTool, toolResu
return msg
}
func emptyKiroUserMessage(modelID string, tools []kiroTool, toolResults []kiroToolResult) *kiroUserInputMessage {
func emptyKiroUserMessage(modelID string, tools []kiroTool, toolResults []kiroToolResult, images []kiroImageBlock) *kiroUserInputMessage {
msg := &kiroUserInputMessage{
Content: "",
Content: kiroEncodeStringContent(""),
Origin: kiroDefaultOrigin,
ModelID: modelID,
Images: images,
}
ctx := &kiroUserInputContext{}
if len(tools) > 0 {
@ -412,19 +413,49 @@ func encodeKiroToolDefinitions(tools []Tool) []kiroTool {
return out
}
func encodeKiroToolResults(messages []Message) []kiroToolResult {
func encodeKiroToolResults(messages []Message) ([]kiroToolResult, []kiroImageBlock) {
results := make([]kiroToolResult, 0, len(messages))
var images []kiroImageBlock
for _, m := range messages {
if !strings.EqualFold(m.Role, "tool") {
continue
}
toolUseID := strings.TrimSpace(m.ToolCallID)
content, imgs := encodeKiroToolResultBlocks(m)
results = append(results, kiroToolResult{
ToolUseID: toolUseID,
Content: []kiroToolResultContent{encodeKiroToolResultContent(m.Content)},
Content: content,
})
images = append(images, imgs...)
}
return results
return results, images
}
func encodeKiroToolResultBlocks(m Message) ([]kiroToolResultContent, []kiroImageBlock) {
if len(m.ContentBlocks) == 0 {
return []kiroToolResultContent{encodeKiroToolResultContent(m.Content)}, nil
}
out := make([]kiroToolResultContent, 0, len(m.ContentBlocks))
var images []kiroImageBlock
for _, cb := range m.ContentBlocks {
switch cb.Type {
case "text":
out = append(out, encodeKiroToolResultContent(cb.Text))
case "image":
if cb.ImageSource == nil {
continue
}
format := strings.TrimPrefix(cb.ImageSource.MediaType, "image/")
images = append(images, kiroImageBlock{
Format: format,
Source: kiroImageSource{Bytes: cb.ImageSource.Data},
})
}
}
if len(out) == 0 {
out = []kiroToolResultContent{{Text: ""}}
}
return out, images
}
func encodeKiroToolResultContent(content string) kiroToolResultContent {
@ -435,6 +466,58 @@ func encodeKiroToolResultContent(content string) kiroToolResultContent {
return kiroToolResultContent{Text: content}
}
// kiroEncodeContent returns a json.RawMessage that is either a JSON string
// (plain text) or a JSON array of Bedrock Converse content blocks (when images
// are present).
func kiroEncodeContent(m Message) json.RawMessage {
if len(m.ContentBlocks) == 0 {
return kiroEncodeStringContent(m.Content)
}
hasImage := false
for _, cb := range m.ContentBlocks {
if cb.Type == "image" && cb.ImageSource != nil {
hasImage = true
break
}
}
if !hasImage {
// Text-only blocks: concatenate and send as string.
var sb strings.Builder
for _, cb := range m.ContentBlocks {
if cb.Type == "text" {
sb.WriteString(cb.Text)
}
}
return kiroEncodeStringContent(sb.String())
}
// Mixed content: encode as Bedrock Converse content block array.
blocks := make([]kiroContentBlock, 0, len(m.ContentBlocks))
for _, cb := range m.ContentBlocks {
switch cb.Type {
case "text":
blocks = append(blocks, kiroContentBlock{Text: cb.Text})
case "image":
if cb.ImageSource == nil {
continue
}
format := strings.TrimPrefix(cb.ImageSource.MediaType, "image/")
blocks = append(blocks, kiroContentBlock{
Image: &kiroImageBlock{
Format: format,
Source: kiroImageSource{Bytes: cb.ImageSource.Data},
},
})
}
}
data, _ := json.Marshal(blocks)
return json.RawMessage(data)
}
func kiroEncodeStringContent(s string) json.RawMessage {
data, _ := json.Marshal(s)
return json.RawMessage(data)
}
// ── Auth sources ──────────────────────────────────────────────────────────────
type kiroAuthSource interface {

21
pkg/backend/codewhisperer_internal.go
View File

@ -73,12 +73,27 @@ type kiroChatMessage struct {
}
type kiroUserInputMessage struct {
Content string `json:"content"`
Content json.RawMessage `json:"content"`
UserInputMessageContext *kiroUserInputContext `json:"userInputMessageContext,omitempty"`
Origin string `json:"origin,omitempty"`
Images []kiroImageBlock `json:"images,omitempty"`
ModelID string `json:"modelId,omitempty"`
}
type kiroContentBlock struct {
Text string `json:"text,omitempty"`
Image *kiroImageBlock `json:"image,omitempty"`
}
type kiroImageBlock struct {
Format string `json:"format"`
Source kiroImageSource `json:"source"`
}
type kiroImageSource struct {
Bytes string `json:"bytes"`
}
type kiroUserInputContext struct {
ToolResults []kiroToolResult `json:"toolResults,omitempty"`
Tools []kiroTool `json:"tools,omitempty"`
@ -117,8 +132,8 @@ type kiroToolResult struct {
}
type kiroToolResultContent struct {
Text string `json:"text,omitempty"`
JSON json.RawMessage `json:"json,omitempty"`
Text string `json:"text,omitempty"`
JSON json.RawMessage `json:"json,omitempty"`
}
// stream event types

7
pkg/tools/execute/server.go
View File

@ -315,6 +315,13 @@ func (e *Server) CallTool(ctx context.Context, tool string, args json.RawMessage
"content": []map[string]string{{"type": "text", "text": err.Error()}},
})
}
// If the tool output is already structured content-block JSON, pass it through.
var structured map[string]json.RawMessage
if json.Unmarshal([]byte(result), &structured) == nil {
if _, ok := structured["content"]; ok {
return []byte(result), nil
}
}
return json.Marshal(map[string]any{
"content": []map[string]string{{"type": "text", "text": result}},
})