fix(go/ai): enable use of media URIs in prompts (#3630)

Author: hugoaguirre

go/ai/prompt.go

@@ -468,11 +468,11 @@ func convertToPartPointers(parts []dotprompt.Part) ([]*Part, error) {
 				result[i] = NewTextPart(p.Text)
 			}
 		case *dotprompt.MediaPart:
-			ct, err := contentType(p.Media.URL)
+			ct, data, err := contentType(p.Media.ContentType, p.Media.URL)
 			if err != nil {
 				return nil, err
 			}
-			result[i] = NewMediaPart(ct, p.Media.URL)
+			result[i] = NewMediaPart(ct, string(data))
 		}
 	}
 	return result, nil
@@ -644,24 +644,30 @@ func variantKey(variant string) string {
 }
 
 // contentType determines the MIME content type of the given data URI
-func contentType(uri string) (string, error) {
+func contentType(ct, uri string) (string, []byte, error) {
 	if uri == "" {
-		return "", errors.New("found empty URI in part")
+		return "", nil, errors.New("found empty URI in part")
 	}
 
 	if strings.HasPrefix(uri, "gs://") || strings.HasPrefix(uri, "http") {
-		return "", errors.New("data URI is the only media type supported")
+		if ct == "" {
+			return "", nil, errors.New("must supply contentType when using media from gs:// or http(s):// URLs")
+		}
+		return ct, []byte(uri), nil
 	}
 	if contents, isData := strings.CutPrefix(uri, "data:"); isData {
 		prefix, _, found := strings.Cut(contents, ",")
 		if !found {
-			return "", errors.New("failed to parse data URI: missing comma")
+			return "", nil, errors.New("failed to parse data URI: missing comma")
 		}
 
 		if p, isBase64 := strings.CutSuffix(prefix, ";base64"); isBase64 {
-			return p, nil
+			if ct == "" {
+				ct = p
+			}
+			return ct, []byte(uri), nil
 		}
 	}
 
-	return "", errors.New("uri content type not found")
+	return "", nil, errors.New("uri content type not found")
 }

go/samples/prompts/main.go

@@ -206,7 +206,10 @@ func PromptWithMultiMessage(ctx context.Context, g *genkit.Genkit) {
 	if prompt == nil {
 		log.Fatal("empty prompt")
 	}
-	resp, err := prompt.Execute(ctx)
+	resp, err := prompt.Execute(ctx,
+		ai.WithModelName("googleai/gemini-2.5-pro"),
+		ai.WithInput(map[string]any{"videoUrl": "https://www.youtube.com/watch?v=K-hY0E6cGfo video/mp4"}),
+	)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -252,6 +255,7 @@ func PromptWithMessageHistory(ctx context.Context, g *genkit.Genkit) {
 	helloPrompt := genkit.DefinePrompt(
 		g, "PromptWithMessageHistory",
 		ai.WithSystem("You are a helpful AI assistant named Walt"),
+		ai.WithModelName("googleai/gemini-2.5-flash-lite"),
 		ai.WithMessages(
 			ai.NewUserTextMessage("Hi, my name is Bob"),
 			ai.NewModelTextMessage("Hi, my name is Walt, what can I help you with?"),
@@ -277,7 +281,7 @@ func PromptWithExecuteOverrides(ctx context.Context, g *genkit.Genkit) {
 
 	// Call the model and add additional messages from the user.
 	resp, err := helloPrompt.Execute(ctx,
-		ai.WithModel(googlegenai.GoogleAIModel(g, "gemini-2.5-pro")),
+		ai.WithModel(googlegenai.GoogleAIModel(g, "gemini-2.5-flash-lite")),
 		ai.WithMessages(ai.NewUserTextMessage("And I like turtles.")),
 	)
 	if err != nil {
@@ -324,8 +328,8 @@ func PromptWithMediaType(ctx context.Context, g *genkit.Genkit) {
 		log.Fatal("empty prompt")
 	}
 	resp, err := prompt.Execute(ctx,
-		ai.WithModelName("googleai/gemini-2.0-flash"),
-		ai.WithInput(map[string]any{"imageUrl": "data:image/jpg;base64," + img}),
+		ai.WithModelName("googleai/gemini-2.5-flash"),
+		ai.WithInput(map[string]any{"imageUrl": "data:image/jpeg;base64," + img}),
 	)
 	if err != nil {
 		log.Fatal(err)

go/samples/prompts/prompts/media.prompt

@@ -1,5 +1,5 @@
 ---
-model: googleai/gemini-2.5-flash-preview-04-17
+model: googleai/gemini-2.5-flash
 config:
   temperature: 0.1
 input:

go/samples/prompts/prompts/multi-msg.prompt

@@ -1,10 +1,16 @@
 ---
-model: googleai/gemini-2.0-flash
+model: googleai/gemini-2.5-flash
+input:
+  schema:
+    videoUrl: string
+output:
+  summary: string
 ---
 {{ role "system" }}
 
-You are a great pirate and an AI assistant.
+You are a great AI assistant that summarizes videos talking as a pirate
 
 {{ role "user" }}
 
-Say hi
+Give me a summary of this video
+{{media url=videoUrl}}