embeddings: base64 encoding fix (#12715)

middleware/openai.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"math/rand"
 	"net/http"
+	"strings"
 
 	"github.com/gin-gonic/gin"
 
@@ -44,7 +45,8 @@ type RetrieveWriter struct {
 
 type EmbedWriter struct {
 	BaseWriter
-	model string
+	model          string
+	encodingFormat string
 }
 
 func (w *BaseWriter) writeError(data []byte) (int, error) {
@@ -254,7 +256,7 @@ func (w *EmbedWriter) writeResponse(data []byte) (int, error) {
 	}
 
 	w.ResponseWriter.Header().Set("Content-Type", "application/json")
-	err = json.NewEncoder(w.ResponseWriter).Encode(openai.ToEmbeddingList(w.model, embedResponse))
+	err = json.NewEncoder(w.ResponseWriter).Encode(openai.ToEmbeddingList(w.model, embedResponse, w.encodingFormat))
 	if err != nil {
 		return 0, err
 	}
@@ -348,6 +350,14 @@ func EmbeddingsMiddleware() gin.HandlerFunc {
 			return
 		}
 
+		// Validate encoding_format parameter
+		if req.EncodingFormat != "" {
+			if !strings.EqualFold(req.EncodingFormat, "float") && !strings.EqualFold(req.EncodingFormat, "base64") {
+				c.AbortWithStatusJSON(http.StatusBadRequest, openai.NewError(http.StatusBadRequest, fmt.Sprintf("Invalid value for 'encoding_format' = %s. Supported values: ['float', 'base64'].", req.EncodingFormat)))
+				return
+			}
+		}
+
 		if req.Input == "" {
 			req.Input = []string{""}
 		}
@@ -371,8 +381,9 @@ func EmbeddingsMiddleware() gin.HandlerFunc {
 		c.Request.Body = io.NopCloser(&b)
 
 		w := &EmbedWriter{
-			BaseWriter: BaseWriter{ResponseWriter: c.Writer},
-			model:      req.Model,
+			BaseWriter:     BaseWriter{ResponseWriter: c.Writer},
+			model:          req.Model,
+			encodingFormat: req.EncodingFormat,
 		}
 
 		c.Writer = w

middleware/openai_encoding_format_test.go

@@ -0,0 +1,220 @@
+package middleware
+
+import (
+	"encoding/base64"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/openai"
+)
+
+func TestEmbeddingsMiddleware_EncodingFormats(t *testing.T) {
+	testCases := []struct {
+		name           string
+		encodingFormat string
+		expectType     string // "array" or "string"
+		verifyBase64   bool
+	}{
+		{"float format", "float", "array", false},
+		{"base64 format", "base64", "string", true},
+		{"default format", "", "array", false},
+	}
+
+	gin.SetMode(gin.TestMode)
+
+	endpoint := func(c *gin.Context) {
+		resp := api.EmbedResponse{
+			Embeddings:      [][]float32{{0.1, -0.2, 0.3}},
+			PromptEvalCount: 5,
+		}
+		c.JSON(http.StatusOK, resp)
+	}
+
+	router := gin.New()
+	router.Use(EmbeddingsMiddleware())
+	router.Handle(http.MethodPost, "/api/embed", endpoint)
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			body := `{"input": "test", "model": "test-model"`
+			if tc.encodingFormat != "" {
+				body += `, "encoding_format": "` + tc.encodingFormat + `"`
+			}
+			body += `}`
+
+			req, _ := http.NewRequest(http.MethodPost, "/api/embed", strings.NewReader(body))
+			req.Header.Set("Content-Type", "application/json")
+
+			resp := httptest.NewRecorder()
+			router.ServeHTTP(resp, req)
+
+			if resp.Code != http.StatusOK {
+				t.Fatalf("expected status 200, got %d", resp.Code)
+			}
+
+			var result openai.EmbeddingList
+			if err := json.Unmarshal(resp.Body.Bytes(), &result); err != nil {
+				t.Fatalf("failed to unmarshal response: %v", err)
+			}
+
+			if len(result.Data) != 1 {
+				t.Fatalf("expected 1 embedding, got %d", len(result.Data))
+			}
+
+			switch tc.expectType {
+			case "array":
+				if _, ok := result.Data[0].Embedding.([]interface{}); !ok {
+					t.Errorf("expected array, got %T", result.Data[0].Embedding)
+				}
+			case "string":
+				embStr, ok := result.Data[0].Embedding.(string)
+				if !ok {
+					t.Errorf("expected string, got %T", result.Data[0].Embedding)
+				} else if tc.verifyBase64 {
+					decoded, err := base64.StdEncoding.DecodeString(embStr)
+					if err != nil {
+						t.Errorf("invalid base64: %v", err)
+					} else if len(decoded) != 12 {
+						t.Errorf("expected 12 bytes, got %d", len(decoded))
+					}
+				}
+			}
+		})
+	}
+}
+
+func TestEmbeddingsMiddleware_BatchWithBase64(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	endpoint := func(c *gin.Context) {
+		resp := api.EmbedResponse{
+			Embeddings: [][]float32{
+				{0.1, 0.2},
+				{0.3, 0.4},
+				{0.5, 0.6},
+			},
+			PromptEvalCount: 10,
+		}
+		c.JSON(http.StatusOK, resp)
+	}
+
+	router := gin.New()
+	router.Use(EmbeddingsMiddleware())
+	router.Handle(http.MethodPost, "/api/embed", endpoint)
+
+	body := `{
+		"input": ["hello", "world", "test"],
+		"model": "test-model",
+		"encoding_format": "base64"
+	}`
+
+	req, _ := http.NewRequest(http.MethodPost, "/api/embed", strings.NewReader(body))
+	req.Header.Set("Content-Type", "application/json")
+
+	resp := httptest.NewRecorder()
+	router.ServeHTTP(resp, req)
+
+	if resp.Code != http.StatusOK {
+		t.Fatalf("expected status 200, got %d", resp.Code)
+	}
+
+	var result openai.EmbeddingList
+	if err := json.Unmarshal(resp.Body.Bytes(), &result); err != nil {
+		t.Fatalf("failed to unmarshal response: %v", err)
+	}
+
+	if len(result.Data) != 3 {
+		t.Fatalf("expected 3 embeddings, got %d", len(result.Data))
+	}
+
+	// All should be base64 strings
+	for i := range 3 {
+		embeddingStr, ok := result.Data[i].Embedding.(string)
+		if !ok {
+			t.Errorf("embedding %d: expected string, got %T", i, result.Data[i].Embedding)
+			continue
+		}
+
+		// Verify it's valid base64
+		if _, err := base64.StdEncoding.DecodeString(embeddingStr); err != nil {
+			t.Errorf("embedding %d: invalid base64: %v", i, err)
+		}
+
+		// Check index
+		if result.Data[i].Index != i {
+			t.Errorf("embedding %d: expected index %d, got %d", i, i, result.Data[i].Index)
+		}
+	}
+}
+
+func TestEmbeddingsMiddleware_InvalidEncodingFormat(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	endpoint := func(c *gin.Context) {
+		c.Status(http.StatusOK)
+	}
+
+	router := gin.New()
+	router.Use(EmbeddingsMiddleware())
+	router.Handle(http.MethodPost, "/api/embed", endpoint)
+
+	testCases := []struct {
+		name           string
+		encodingFormat string
+		shouldFail     bool
+	}{
+		{"valid: float", "float", false},
+		{"valid: base64", "base64", false},
+		{"valid: FLOAT (uppercase)", "FLOAT", false},
+		{"valid: BASE64 (uppercase)", "BASE64", false},
+		{"valid: Float (mixed)", "Float", false},
+		{"valid: Base64 (mixed)", "Base64", false},
+		{"invalid: json", "json", true},
+		{"invalid: hex", "hex", true},
+		{"invalid: invalid_format", "invalid_format", true},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			body := `{
+				"input": "test",
+				"model": "test-model",
+				"encoding_format": "` + tc.encodingFormat + `"
+			}`
+
+			req, _ := http.NewRequest(http.MethodPost, "/api/embed", strings.NewReader(body))
+			req.Header.Set("Content-Type", "application/json")
+
+			resp := httptest.NewRecorder()
+			router.ServeHTTP(resp, req)
+
+			if tc.shouldFail {
+				if resp.Code != http.StatusBadRequest {
+					t.Errorf("expected status 400, got %d", resp.Code)
+				}
+
+				var errResp openai.ErrorResponse
+				if err := json.Unmarshal(resp.Body.Bytes(), &errResp); err != nil {
+					t.Fatalf("failed to unmarshal error response: %v", err)
+				}
+
+				if errResp.Error.Type != "invalid_request_error" {
+					t.Errorf("expected error type 'invalid_request_error', got %q", errResp.Error.Type)
+				}
+
+				if !strings.Contains(errResp.Error.Message, "encoding_format") {
+					t.Errorf("expected error message to mention encoding_format, got %q", errResp.Error.Message)
+				}
+			} else {
+				if resp.Code != http.StatusOK {
+					t.Errorf("expected status 200, got %d: %s", resp.Code, resp.Body.String())
+				}
+			}
+		})
+	}
+}