Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 10 additions & 7 deletions pkg/chat/media.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,17 @@ type MediaDelta struct {
// runtime accumulator synthesizes one when needed.
Name string `json:"name,omitempty"`

// RequestedPath is the prompt-directed target path the model asked for
// (e.g. echoed from an "as sunshine.jpg" instruction), when one exists.
// It is untrusted model input: the runtime routes it through
// RequestedPath is the prompt-directed target path for this blob, when
// one exists. It is untrusted input: the runtime routes it through
// workspacemedia.ClassifyRequestedPath, and a path escaping the workspace
// requires an explicit user confirmation before it is honored. Response
// marker extraction (the "[media-file: ...]" protocol) will populate it;
// until that lands, providers leave it empty and materialization falls
// back to Name.
// requires an explicit user confirmation before it is honored. The
// runtime's response marker filter (the "[media-file: ...]" protocol,
// pkg/runtime/generated_media_markers.go) populates it by pairing marker
// paths with blobs in response order; a single otherwise-unnamed blob may
// instead get it from deterministic explicit-filename extraction on the
// triggering user message (pkg/runtime/generated_media_prompt_filename.go).
// Blobs neither source names keep it empty and materialization falls back
// to Name, then a generic name.
RequestedPath string `json:"requested_path,omitempty"`

// Size is the byte length of Data, cached because Data itself is
Expand Down
1 change: 1 addition & 0 deletions pkg/model/provider/gemini/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -761,6 +761,7 @@ func (c *Client) CreateChatCompletionStream(

if c.wantsImageResponseModalities(imageOutputEnabled) {
config.ResponseModalities = []string{string(genai.ModalityText), string(genai.ModalityImage)}
applyImageOutputMediaFileInstruction(config)
}

// Start with Google built-in tools (search, maps, code execution) from provider_opts
Expand Down
31 changes: 31 additions & 0 deletions pkg/model/provider/gemini/image_output_instruction.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
package gemini

import "google.golang.org/genai"

// imageOutputMediaFileInstruction is appended as a system instruction on
// declared image-output chat requests (see wantsImageResponseModalities) so
// generated images arrive with a machine-readable filename: the runtime
// strips these exact marker lines from the reply and uses the paths to name
// the materialized workspace files (pkg/runtime/generated_media_markers.go).
// The single-image steering keeps one request yielding one predictably named
// file; every blob the model actually returns is still persisted.
const imageOutputMediaFileInstruction = `When you generate images, name each one with a marker line, placed alone on its own line, in this exact format:
[media-file: relative/path.ext]
Rules:
- Emit exactly one marker line per generated image, in the same order as the images.
- If the user asked for a specific file name or path, echo it in the marker exactly as requested.
- Otherwise choose a short, meaningful, kebab-case file name.
- Generate a single image unless the user explicitly asks for multiple images or variations.
- Never emit a marker line for an image you did not generate.`

// applyImageOutputMediaFileInstruction appends the marker-protocol
// instruction to the request's system instruction, preserving any parts
// already present. Callers gate it on wantsImageResponseModalities so only
// declared image-output chat requests carry it.
func applyImageOutputMediaFileInstruction(config *genai.GenerateContentConfig) {
if config.SystemInstruction == nil {
config.SystemInstruction = &genai.Content{}
}
config.SystemInstruction.Parts = append(config.SystemInstruction.Parts,
genai.NewPartFromText(imageOutputMediaFileInstruction))
}
188 changes: 188 additions & 0 deletions pkg/model/provider/gemini/image_output_instruction_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,188 @@
package gemini

import (
"encoding/json"
"strings"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"

"github.com/docker/docker-agent/pkg/chat"
"github.com/docker/docker-agent/pkg/config/latest"
"github.com/docker/docker-agent/pkg/environment"
"github.com/docker/docker-agent/pkg/model/provider/options"
)

// systemInstructionTextsInBody decodes body's systemInstruction part texts
// (genai serializes GenerateContentConfig.SystemInstruction under the
// top-level "systemInstruction" key for the Gemini Developer API), returning
// nil when the key is absent.
func systemInstructionTextsInBody(t *testing.T, body []byte) []string {
t.Helper()

var req map[string]any
require.NoError(t, json.Unmarshal(body, &req))

si, ok := req["systemInstruction"].(map[string]any)
if !ok {
return nil
}
parts, ok := si["parts"].([]any)
if !ok {
return nil
}
var out []string
for _, p := range parts {
partMap, ok := p.(map[string]any)
if !ok {
continue
}
if text, ok := partMap["text"].(string); ok {
out = append(out, text)
}
}
return out
}

// TestCreateChatCompletionStream_MediaFileInstruction_PositiveRoutes pins
// that ordinary image-output chat requests carry the media-file marker
// instruction exactly once on every supported Google surface.
func TestCreateChatCompletionStream_MediaFileInstruction_PositiveRoutes(t *testing.T) {
t.Parallel()

tests := []struct {
name string
cfg func(serverURL string) *latest.ModelConfig
env map[string]string
gateway bool
}{
{
name: "gateway",
cfg: func(string) *latest.ModelConfig {
return &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image", OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}}
},
env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"},
gateway: true,
},
{
name: "direct Gemini API",
cfg: func(serverURL string) *latest.ModelConfig {
return &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image", BaseURL: serverURL, OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)}}
},
env: map[string]string{"GOOGLE_API_KEY": "test-key"},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()

server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse)
opts := []options.Opt(nil)
if tt.gateway {
opts = append(opts, options.WithGateway(server.URL))
}
client, err := NewClient(t.Context(), tt.cfg(server.URL), environment.NewMapEnvProvider(tt.env), opts...)
require.NoError(t, err)

stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{{Role: chat.MessageRoleUser, Content: "generate an image of a red panda"}}, nil)
require.NoError(t, err)
drainStream(t, stream)

bodies := captured.all()
require.Len(t, bodies, 1)
texts := systemInstructionTextsInBody(t, bodies[0])
require.Len(t, texts, 1, "the instruction must be sent exactly once")
assert.Equal(t, imageOutputMediaFileInstruction, texts[0])
assert.Equal(t, 1, strings.Count(texts[0], "[media-file: "), "the instruction must show the marker format exactly once")
})
}
}

// TestCreateChatCompletionStream_MediaFileInstruction_AbsentOnOtherRoutes
// pins that non-image-output and internal text-only requests send no marker
// instruction: gateway calls without image output enabled, plus gateway
// title-generation and compaction calls.
func TestCreateChatCompletionStream_MediaFileInstruction_AbsentOnOtherRoutes(t *testing.T) {
t.Parallel()

tests := []struct {
name string
cfg func(serverURL string) *latest.ModelConfig
env map[string]string
gateway bool
opts []options.Opt
}{
{
name: "gateway, declared false: absent",
cfg: func(string) *latest.ModelConfig {
return &latest.ModelConfig{
Provider: "google", Model: "gemini-2.5-flash-image",
OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(false)},
}
},
env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"},
gateway: true,
},
{
name: "gateway, declaration missing: absent",
cfg: func(string) *latest.ModelConfig {
return &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash-image"}
},
env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"},
gateway: true,
},
{
name: "gateway, declared true, generating title: absent",
cfg: func(string) *latest.ModelConfig {
return &latest.ModelConfig{
Provider: "google", Model: "gemini-2.5-flash-image",
OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)},
}
},
env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"},
gateway: true,
opts: []options.Opt{options.WithGeneratingTitle()},
},
{
name: "gateway, declared true, compacting: absent",
cfg: func(string) *latest.ModelConfig {
return &latest.ModelConfig{
Provider: "google", Model: "gemini-2.5-flash-image",
OutputCapabilities: &latest.OutputCapabilitiesConfig{Image: new(true)},
}
},
env: map[string]string{environment.DockerDesktopTokenEnv: "test-dd-token"},
gateway: true,
opts: []options.Opt{options.WithCompacting()},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
t.Parallel()

server, captured := newBodyCapturingGeminiServer(t, writeGeminiSSEResponse)

cfg := tt.cfg(server.URL)
env := environment.NewMapEnvProvider(tt.env)
opts := tt.opts
if tt.gateway {
opts = append(opts, options.WithGateway(server.URL))
}
client, err := NewClient(t.Context(), cfg, env, opts...)
require.NoError(t, err)

stream, err := client.CreateChatCompletionStream(t.Context(), []chat.Message{
{Role: chat.MessageRoleUser, Content: "hello"},
}, nil)
require.NoError(t, err)
drainStream(t, stream)

bodies := captured.all()
require.Len(t, bodies, 1)
assert.Nil(t, systemInstructionTextsInBody(t, bodies[0]), "the marker instruction must be absent on this route")
})
}
}
Loading
Loading