feat(go): Add ollama vision support (#2795)

Author: huangjeff5

go/plugins/ollama/ollama.go

@@ -39,7 +39,7 @@ import (
 const provider = "ollama"
 
 var (
-	mediaSupportedModels = []string{"llava"}
+	mediaSupportedModels = []string{"llava", "bakllava", "llava-llama3", "llava:13b", "llava:7b", "llava:latest"}
 	roleMapping          = map[ai.Role]string{
 		ai.RoleUser:   "user",
 		ai.RoleModel:  "assistant",
@@ -54,6 +54,7 @@ func (o *Ollama) DefineModel(g *genkit.Genkit, model ModelDefinition, info *ai.M
 		panic("ollama.Init not called")
 	}
 	var mi ai.ModelInfo
+
 	if info != nil {
 		mi = *info
 	} else {
@@ -120,8 +121,10 @@ keep_alive: controls how long the model will stay loaded into memory following t
 */
 type ollamaChatRequest struct {
 	Messages []*ollamaMessage `json:"messages"`
+	Images   []string         `json:"images,omitempty"`
 	Model    string           `json:"model"`
 	Stream   bool             `json:"stream"`
+	Format   string           `json:"format,omitempty"`
 }
 
 type ollamaModelRequest struct {
@@ -130,6 +133,7 @@ type ollamaModelRequest struct {
 	Model  string   `json:"model"`
 	Prompt string   `json:"prompt"`
 	Stream bool     `json:"stream"`
+	Format string   `json:"format,omitempty"`
 }
 
 // TODO: Add optional parameters (images, format, options, etc.) based on your use case
@@ -181,11 +185,21 @@ func (g *generator) generate(ctx context.Context, input *ai.ModelRequest, cb fun
 	stream := cb != nil
 	var payload any
 	isChatModel := g.model.Type == "chat"
-	if !isChatModel {
-		images, err := concatImages(input, []ai.Role{ai.RoleUser, ai.RoleModel})
+
+	// Check if this is an image model
+	hasMediaSupport := slices.Contains(mediaSupportedModels, g.model.Name)
+
+	// Extract images if the model supports them
+	var images []string
+	var err error
+	if hasMediaSupport {
+		images, err = concatImages(input, []ai.Role{ai.RoleUser, ai.RoleModel})
 		if err != nil {
 			return nil, fmt.Errorf("failed to grab image parts: %v", err)
 		}
+	}
+
+	if !isChatModel {
 		payload = ollamaModelRequest{
 			Model:  g.model.Name,
 			Prompt: concatMessages(input, []ai.Role{ai.RoleUser, ai.RoleModel, ai.RoleTool}),
@@ -207,6 +221,7 @@ func (g *generator) generate(ctx context.Context, input *ai.ModelRequest, cb fun
 			Messages: messages,
 			Model:    g.model.Name,
 			Stream:   stream,
+			Images:   images,
 		}
 	}
 	client := &http.Client{Timeout: 30 * time.Second}
@@ -293,21 +308,27 @@ func convertParts(role ai.Role, parts []*ai.Part) (*ollamaMessage, error) {
 		Role: roleMapping[role],
 	}
 	var contentBuilder strings.Builder
+	var images []string
+
 	for _, part := range parts {
 		if part.IsText() {
 			contentBuilder.WriteString(part.Text)
 		} else if part.IsMedia() {
 			_, data, err := uri.Data(part)
 			if err != nil {
-				return nil, err
+				return nil, fmt.Errorf("failed to extract media data: %v", err)
 			}
 			base64Encoded := base64.StdEncoding.EncodeToString(data)
-			message.Images = append(message.Images, base64Encoded)
+			images = append(images, base64Encoded)
 		} else {
-			return nil, errors.New("unknown content type")
+			return nil, errors.New("unsupported content type")
 		}
 	}
+
 	message.Content = contentBuilder.String()
+	if len(images) > 0 {
+		message.Images = images
+	}
 	return message, nil
 }
 
@@ -414,10 +435,18 @@ func concatImages(input *ai.ModelRequest, roleFilter []ai.Role) ([]string, error
 				if !part.IsMedia() {
 					continue
 				}
-				_, data, err := uri.Data(part)
+
+				// Get the media type and data
+				mediaType, data, err := uri.Data(part)
 				if err != nil {
-					return nil, err
+					return nil, fmt.Errorf("failed to extract image data: %v", err)
 				}
+
+				// Only include image media types
+				if !strings.HasPrefix(mediaType, "image/") {
+					continue
+				}
+
 				base64Encoded := base64.StdEncoding.EncodeToString(data)
 				images = append(images, base64Encoded)
 			}

go/samples/ollama-vision/main.go

@@ -0,0 +1,137 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"context"
+	"encoding/base64"
+	"fmt"
+	"log"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/firebase/genkit/go/ai"
+	"github.com/firebase/genkit/go/genkit"
+	"github.com/firebase/genkit/go/plugins/ollama"
+)
+
+// getContentTypeFromExtension returns a MIME type based on file extension
+func getContentTypeFromExtension(filename string) string {
+	ext := strings.ToLower(filepath.Ext(filename))
+	switch ext {
+	case ".jpg", ".jpeg":
+		return "image/jpeg"
+	case ".png":
+		return "image/png"
+	case ".gif":
+		return "image/gif"
+	case ".webp":
+		return "image/webp"
+	case ".bmp":
+		return "image/bmp"
+	case ".svg":
+		return "image/svg+xml"
+	default:
+		return "image/png" // Default fallback
+	}
+}
+
+func main() {
+	// Get the image path from command line argument or use a default path
+	imagePath := "test.png"
+	if len(os.Args) > 1 {
+		imagePath = os.Args[1]
+	}
+
+	// Check if image exists
+	if _, err := os.Stat(imagePath); os.IsNotExist(err) {
+		log.Fatalf("Image file not found: %s", imagePath)
+	}
+
+	// Read the image file
+	imageData, err := os.ReadFile(imagePath)
+	if err != nil {
+		log.Fatalf("Failed to read image file: %v", err)
+	}
+
+	// Detect content type (MIME type) from the file's binary signature
+	contentType := http.DetectContentType(imageData)
+
+	// If content type is generic/unknown, try to infer from file extension
+	if contentType == "application/octet-stream" {
+		contentType = getContentTypeFromExtension(imagePath)
+	}
+
+	// Encode image to base64
+	base64Image := base64.StdEncoding.EncodeToString(imageData)
+	dataURI := fmt.Sprintf("data:%s;base64,%s", contentType, base64Image)
+
+	// Create a new Genkit instance
+	g, err := genkit.Init(context.Background())
+	if err != nil {
+		log.Fatalf("Failed to initialize Genkit: %v", err)
+	}
+
+	// Initialize the Ollama plugin
+	ollamaPlugin := &ollama.Ollama{
+		ServerAddress: "http://localhost:11434", // Default Ollama server address
+	}
+
+	// Initialize the plugin
+	err = ollamaPlugin.Init(context.Background(), g)
+	if err != nil {
+		log.Fatalf("Failed to initialize Ollama plugin: %v", err)
+	}
+
+	// Define a model that supports images (llava is one of the supported models)
+	modelName := "llava"
+	model := ollamaPlugin.DefineModel(g, ollama.ModelDefinition{
+		Name: modelName,
+		Type: "chat",
+	}, nil)
+
+	// Create a context
+	ctx := context.Background()
+
+	// Create a request with text and image
+	request := &ai.ModelRequest{
+		Messages: []*ai.Message{
+			{
+				Role: ai.RoleUser,
+				Content: []*ai.Part{
+					ai.NewTextPart("Describe what you see in this image:"),
+					ai.NewMediaPart(contentType, dataURI),
+				},
+			},
+		},
+	}
+
+	// Call the model
+	fmt.Printf("Sending request to %s model...\n", modelName)
+	response, err := model.Generate(ctx, request, nil)
+	if err != nil {
+		log.Fatalf("Error generating response: %v", err)
+	}
+
+	// Print the response
+	fmt.Println("\nModel Response:")
+	for _, part := range response.Message.Content {
+		if part.IsText() {
+			fmt.Println(part.Text)
+		}
+	}
+}