Skip to content

Commit df46c46

Browse files
hugoaguirrehugoaguirre
authored
feat(go/dotprompt): add support for media in dotprompt (#3006)
1 parent 7df3fa9 commit df46c46

File tree

6 files changed

+110
-22
lines changed

6 files changed

+110
-22
lines changed

‎go/ai/prompt.go

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -443,12 +443,17 @@ func renderDotpromptToParts(ctx context.Context, promptFn dotprompt.PromptFuncti
443443
func convertToPartPointers(parts []dotprompt.Part) ([]*Part, error) {
444444
result := make([]*Part, len(parts))
445445
for i, part := range parts {
446-
if p, ok := part.(*dotprompt.TextPart); ok {
446+
switch p := part.(type) {
447+
case *dotprompt.TextPart:
447448
if p.Text != "" {
448449
result[i] = NewTextPart(p.Text)
449450
}
450-
} else {
451-
return nil, fmt.Errorf("unsupported prompt format: %T", part)
451+
case *dotprompt.MediaPart:
452+
ct, err := contentType(p.Media.URL)
453+
if err != nil {
454+
return nil, err
455+
}
456+
result[i] = NewMediaPart(ct, p.Media.URL)
452457
}
453458
}
454459
return result, nil
@@ -625,3 +630,26 @@ func variantKey(variant string) string {
625630
}
626631
return ""
627632
}
633+
634+
// contentType determines the MIME content type of the given data URI
635+
func contentType(uri string) (string, error) {
636+
if uri == "" {
637+
return "", errors.New("found empty URI in part")
638+
}
639+
640+
if strings.HasPrefix(uri, "gs://") || strings.HasPrefix(uri, "http") {
641+
return "", errors.New("data URI is the only media type supported")
642+
}
643+
if contents, isData := strings.CutPrefix(uri, "data:"); isData {
644+
prefix, _, found := strings.Cut(contents, ",")
645+
if !found {
646+
return "", errors.New("failed to parse data URI: missing comma")
647+
}
648+
649+
if p, isBase64 := strings.CutSuffix(prefix, ";base64"); isBase64 {
650+
return p, nil
651+
}
652+
}
653+
654+
return "", errors.New("uri content type not found")
655+
}

‎go/plugins/googlegenai/googleai_live_test.go

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -293,19 +293,19 @@ func TestGoogleAILive(t *testing.T) {
293293
t.Fatal(err)
294294
}
295295
resp, err := genkit.Generate(ctx, g,
296-
ai.WithSystem("You are a pirate expert in TV Shows, your response should include the name of the character in the image provided"),
296+
ai.WithSystem("You are a pirate expert in animals, your response should include the name of the animal in the provided image"),
297297
ai.WithMessages(
298298
ai.NewUserMessage(
299-
ai.NewTextPart("do you know who's in the image?"),
300-
ai.NewMediaPart("image/png", "data:image/png;base64,"+i),
299+
ai.NewTextPart("do you what animal is in the image?"),
300+
ai.NewMediaPart("image/jpg", "data:image/jpg;base64,"+i),
301301
),
302302
),
303303
)
304304
if err != nil {
305305
t.Fatal(err)
306306
}
307-
if !strings.Contains(resp.Text(), "Bluey") {
308-
t.Fatalf("image detection failed, want: Bluey, got: %s", resp.Text())
307+
if !strings.Contains(strings.ToLower(resp.Text()), "donkey") {
308+
t.Fatalf("image detection failed, want: donkey, got: %s", resp.Text())
309309
}
310310
})
311311
t.Run("media content", func(t *testing.T) {
@@ -341,8 +341,8 @@ func TestGoogleAILive(t *testing.T) {
341341
if err != nil {
342342
t.Fatal(err)
343343
}
344-
if !strings.Contains(resp.Text(), "Bluey") {
345-
t.Fatalf("image detection failed, want: Bluey, got: %s", resp.Text())
344+
if !strings.Contains(resp.Text(), "donkey") {
345+
t.Fatalf("image detection failed, want: donkey, got: %s", resp.Text())
346346
}
347347
})
348348
t.Run("image generation", func(t *testing.T) {
@@ -539,7 +539,8 @@ func TestCacheHelper(t *testing.T) {
539539
}
540540

541541
func fetchImgAsBase64() (string, error) {
542-
imgUrl := "https://www.bluey.tv/wp-content/uploads/2023/07/Bluey.png"
542+
// CC0 license image
543+
imgUrl := "https://pd.w.org/2025/05/64268380a8c42af85.63713105-2048x1152.jpg"
543544
resp, err := http.Get(imgUrl)
544545
if err != nil {
545546
return "", err

‎go/plugins/googlegenai/vertexai_live_test.go

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,7 @@ func TestVertexAILive(t *testing.T) {
218218
t.Fatal(err)
219219
}
220220
resp, err := genkit.Generate(ctx, g,
221-
ai.WithSystem("You are a pirate expert in TV Shows, your response should include the name of the character in the image provided"),
221+
ai.WithSystem("You are a pirate expert in animals, your response should include the name of the animal in the provided image"),
222222
ai.WithMessages(
223223
ai.NewUserMessage(
224224
ai.NewTextPart("do you know who's in the image?"),
@@ -229,8 +229,8 @@ func TestVertexAILive(t *testing.T) {
229229
if err != nil {
230230
t.Fatal(err)
231231
}
232-
if !strings.Contains(resp.Text(), "Bluey") {
233-
t.Fatalf("image detection failed, want: Bluey, got: %s", resp.Text())
232+
if !strings.Contains(resp.Text(), "donkey") {
233+
t.Fatalf("image detection failed, want: donkey, got: %s", resp.Text())
234234
}
235235
})
236236
t.Run("media content", func(t *testing.T) {
@@ -266,8 +266,8 @@ func TestVertexAILive(t *testing.T) {
266266
if err != nil {
267267
t.Fatal(err)
268268
}
269-
if !strings.Contains(resp.Text(), "Bluey") {
270-
t.Fatalf("image detection failed, want: Bluey, got: %s", resp.Text())
269+
if !strings.Contains(resp.Text(), "donkey") {
270+
t.Fatalf("image detection failed, want: donkey, got: %s", resp.Text())
271271
}
272272
})
273273
t.Run("image generation", func(t *testing.T) {

‎go/plugins/vertexai/modelgarden/anthropic_live_test.go

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,7 @@ func TestAnthropicLive(t *testing.T) {
9696
}
9797
m := modelgarden.AnthropicModel(g, "claude-3-7-sonnet")
9898
resp, err := genkit.Generate(ctx, g,
99-
ai.WithSystem("You are a professional image detective that talks like an evil pirate that does not like tv shows, your task is to tell the name of the character in the image but be very short"),
99+
ai.WithSystem("You are a professional image detective that talks like an evil pirate that loves animals, your task is to tell the name of the animal in the image but be very short"),
100100
ai.WithModel(m),
101101
ai.WithMessages(
102102
ai.NewUserMessage(
@@ -106,8 +106,8 @@ func TestAnthropicLive(t *testing.T) {
106106
t.Fatal(err)
107107
}
108108

109-
if !strings.Contains(resp.Text(), "Bluey") {
110-
t.Fatalf("it should've said Bluey but got: %s", resp.Text())
109+
if !strings.Contains(resp.Text(), "donkey") {
110+
t.Fatalf("it should've said donkey but got: %s", resp.Text())
111111
}
112112
})
113113

@@ -207,9 +207,9 @@ func TestAnthropicLive(t *testing.T) {
207207
})
208208
}
209209

210-
// Bluey rocks
211210
func fetchImgAsBase64() (string, error) {
212-
imgUrl := "https://www.bluey.tv/wp-content/uploads/2023/07/Bluey.png"
211+
// CC0 license image
212+
imgUrl := "https://pd.w.org/2025/05/64268380a8c42af85.63713105-2048x1152.jpg"
213213
resp, err := http.Get(imgUrl)
214214
if err != nil {
215215
return "", err
@@ -220,7 +220,6 @@ func fetchImgAsBase64() (string, error) {
220220
return "", err
221221
}
222222

223-
// keep the img in memory
224223
imageBytes, err := io.ReadAll(resp.Body)
225224
if err != nil {
226225
return "", err

‎go/samples/prompts/main.go

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,10 @@ package main
1616

1717
import (
1818
"context"
19+
"encoding/base64"
1920
"encoding/json"
2021
"fmt"
22+
"io"
2123
"log"
2224
"math"
2325
"net/http"
@@ -48,6 +50,7 @@ func main() {
4850
PromptWithExecuteOverrides(ctx, g)
4951
PromptWithFunctions(ctx, g)
5052
PromptWithOutputTypeDotprompt(ctx, g)
53+
PromptWithMediaType(ctx, g)
5154

5255
mux := http.NewServeMux()
5356
for _, a := range genkit.ListFlows(g) {
@@ -319,3 +322,47 @@ func PromptWithFunctions(ctx context.Context, g *genkit.Genkit) {
319322

320323
fmt.Println(resp.Text())
321324
}
325+
326+
func PromptWithMediaType(ctx context.Context, g *genkit.Genkit) {
327+
img, err := fetchImgAsBase64()
328+
if err != nil {
329+
log.Fatal(err)
330+
}
331+
332+
prompt, err := genkit.LoadPrompt(g, "./prompts/media.prompt", "mediaspace")
333+
if err != nil {
334+
log.Fatal(err)
335+
}
336+
if prompt == nil {
337+
log.Fatal("empty prompt")
338+
}
339+
resp, err := prompt.Execute(ctx,
340+
341+
ai.WithModelName("vertexai/gemini-2.0-flash"),
342+
ai.WithInput(map[string]any{"imageUrl": "data:image/png;base64," + img}),
343+
)
344+
if err != nil {
345+
log.Fatal(err)
346+
}
347+
fmt.Println(resp.Text())
348+
}
349+
350+
func fetchImgAsBase64() (string, error) {
351+
imgUrl := "https://pd.w.org/2025/05/64268380a8c42af85.63713105-2048x1152.jpg"
352+
resp, err := http.Get(imgUrl)
353+
if err != nil {
354+
return "", err
355+
}
356+
defer resp.Body.Close()
357+
if resp.StatusCode != http.StatusOK {
358+
return "", err
359+
}
360+
361+
imageBytes, err := io.ReadAll(resp.Body)
362+
if err != nil {
363+
return "", err
364+
}
365+
366+
base64string := base64.StdEncoding.EncodeToString(imageBytes)
367+
return base64string, nil
368+
}

‎go/samples/prompts/prompts/media.prompt

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
---
2+
model: vertexai/gemini-2.5-flash-preview-04-17
3+
config:
4+
temperature: 0.1
5+
input:
6+
schema:
7+
imageUrl: string
8+
output:
9+
animal: string
10+
---
11+
Analyze the provided image of an animal, tell which animal is in the picture
12+
13+
{{media url=imageUrl}}

0 commit comments

Comments
 (0)