Skip to content

Commit e0ead1a

Browse files
npardalnpardal
authored
embeddings: base64 encoding fix (ollama#12715)
1 parent d515aed commit e0ead1a

File tree

4 files changed

+386
-12
lines changed

4 files changed

+386
-12
lines changed

‎middleware/openai.go‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import (
77
"io"
88
"math/rand"
99
"net/http"
10+
"strings"
1011

1112
"github.com/gin-gonic/gin"
1213

@@ -44,7 +45,8 @@ type RetrieveWriter struct {
4445

4546
type EmbedWriter struct {
4647
BaseWriter
47-
model string
48+
model string
49+
encodingFormat string
4850
}
4951

5052
func (w *BaseWriter) writeError(data []byte) (int, error) {
@@ -254,7 +256,7 @@ func (w *EmbedWriter) writeResponse(data []byte) (int, error) {
254256
}
255257

256258
w.ResponseWriter.Header().Set("Content-Type", "application/json")
257-
err = json.NewEncoder(w.ResponseWriter).Encode(openai.ToEmbeddingList(w.model, embedResponse))
259+
err = json.NewEncoder(w.ResponseWriter).Encode(openai.ToEmbeddingList(w.model, embedResponse, w.encodingFormat))
258260
if err != nil {
259261
return 0, err
260262
}
@@ -348,6 +350,14 @@ func EmbeddingsMiddleware() gin.HandlerFunc {
348350
return
349351
}
350352

353+
// Validate encoding_format parameter
354+
if req.EncodingFormat != "" {
355+
if !strings.EqualFold(req.EncodingFormat, "float") && !strings.EqualFold(req.EncodingFormat, "base64") {
356+
c.AbortWithStatusJSON(http.StatusBadRequest, openai.NewError(http.StatusBadRequest, fmt.Sprintf("Invalid value for 'encoding_format' = %s. Supported values: ['float', 'base64'].", req.EncodingFormat)))
357+
return
358+
}
359+
}
360+
351361
if req.Input == "" {
352362
req.Input = []string{""}
353363
}
@@ -371,8 +381,9 @@ func EmbeddingsMiddleware() gin.HandlerFunc {
371381
c.Request.Body = io.NopCloser(&b)
372382

373383
w := &EmbedWriter{
374-
BaseWriter: BaseWriter{ResponseWriter: c.Writer},
375-
model: req.Model,
384+
BaseWriter: BaseWriter{ResponseWriter: c.Writer},
385+
model: req.Model,
386+
encodingFormat: req.EncodingFormat,
376387
}
377388

378389
c.Writer = w

‎middleware/openai_encoding_format_test.go‎

Lines changed: 220 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,220 @@
1+
package middleware
2+
3+
import (
4+
"encoding/base64"
5+
"encoding/json"
6+
"net/http"
7+
"net/http/httptest"
8+
"strings"
9+
"testing"
10+
11+
"github.com/gin-gonic/gin"
12+
"github.com/ollama/ollama/api"
13+
"github.com/ollama/ollama/openai"
14+
)
15+
16+
func TestEmbeddingsMiddleware_EncodingFormats(t *testing.T) {
17+
testCases := []struct {
18+
name string
19+
encodingFormat string
20+
expectType string // "array" or "string"
21+
verifyBase64 bool
22+
}{
23+
{"float format", "float", "array", false},
24+
{"base64 format", "base64", "string", true},
25+
{"default format", "", "array", false},
26+
}
27+
28+
gin.SetMode(gin.TestMode)
29+
30+
endpoint := func(c *gin.Context) {
31+
resp := api.EmbedResponse{
32+
Embeddings: [][]float32{{0.1, -0.2, 0.3}},
33+
PromptEvalCount: 5,
34+
}
35+
c.JSON(http.StatusOK, resp)
36+
}
37+
38+
router := gin.New()
39+
router.Use(EmbeddingsMiddleware())
40+
router.Handle(http.MethodPost, "/api/embed", endpoint)
41+
42+
for _, tc := range testCases {
43+
t.Run(tc.name, func(t *testing.T) {
44+
body := `{"input": "test", "model": "test-model"`
45+
if tc.encodingFormat != "" {
46+
body += `, "encoding_format": "` + tc.encodingFormat + `"`
47+
}
48+
body += `}`
49+
50+
req, _ := http.NewRequest(http.MethodPost, "/api/embed", strings.NewReader(body))
51+
req.Header.Set("Content-Type", "application/json")
52+
53+
resp := httptest.NewRecorder()
54+
router.ServeHTTP(resp, req)
55+
56+
if resp.Code != http.StatusOK {
57+
t.Fatalf("expected status 200, got %d", resp.Code)
58+
}
59+
60+
var result openai.EmbeddingList
61+
if err := json.Unmarshal(resp.Body.Bytes(), &result); err != nil {
62+
t.Fatalf("failed to unmarshal response: %v", err)
63+
}
64+
65+
if len(result.Data) != 1 {
66+
t.Fatalf("expected 1 embedding, got %d", len(result.Data))
67+
}
68+
69+
switch tc.expectType {
70+
case "array":
71+
if _, ok := result.Data[0].Embedding.([]interface{}); !ok {
72+
t.Errorf("expected array, got %T", result.Data[0].Embedding)
73+
}
74+
case "string":
75+
embStr, ok := result.Data[0].Embedding.(string)
76+
if !ok {
77+
t.Errorf("expected string, got %T", result.Data[0].Embedding)
78+
} else if tc.verifyBase64 {
79+
decoded, err := base64.StdEncoding.DecodeString(embStr)
80+
if err != nil {
81+
t.Errorf("invalid base64: %v", err)
82+
} else if len(decoded) != 12 {
83+
t.Errorf("expected 12 bytes, got %d", len(decoded))
84+
}
85+
}
86+
}
87+
})
88+
}
89+
}
90+
91+
func TestEmbeddingsMiddleware_BatchWithBase64(t *testing.T) {
92+
gin.SetMode(gin.TestMode)
93+
94+
endpoint := func(c *gin.Context) {
95+
resp := api.EmbedResponse{
96+
Embeddings: [][]float32{
97+
{0.1, 0.2},
98+
{0.3, 0.4},
99+
{0.5, 0.6},
100+
},
101+
PromptEvalCount: 10,
102+
}
103+
c.JSON(http.StatusOK, resp)
104+
}
105+
106+
router := gin.New()
107+
router.Use(EmbeddingsMiddleware())
108+
router.Handle(http.MethodPost, "/api/embed", endpoint)
109+
110+
body := `{
111+
"input": ["hello", "world", "test"],
112+
"model": "test-model",
113+
"encoding_format": "base64"
114+
}`
115+
116+
req, _ := http.NewRequest(http.MethodPost, "/api/embed", strings.NewReader(body))
117+
req.Header.Set("Content-Type", "application/json")
118+
119+
resp := httptest.NewRecorder()
120+
router.ServeHTTP(resp, req)
121+
122+
if resp.Code != http.StatusOK {
123+
t.Fatalf("expected status 200, got %d", resp.Code)
124+
}
125+
126+
var result openai.EmbeddingList
127+
if err := json.Unmarshal(resp.Body.Bytes(), &result); err != nil {
128+
t.Fatalf("failed to unmarshal response: %v", err)
129+
}
130+
131+
if len(result.Data) != 3 {
132+
t.Fatalf("expected 3 embeddings, got %d", len(result.Data))
133+
}
134+
135+
// All should be base64 strings
136+
for i := range 3 {
137+
embeddingStr, ok := result.Data[i].Embedding.(string)
138+
if !ok {
139+
t.Errorf("embedding %d: expected string, got %T", i, result.Data[i].Embedding)
140+
continue
141+
}
142+
143+
// Verify it's valid base64
144+
if _, err := base64.StdEncoding.DecodeString(embeddingStr); err != nil {
145+
t.Errorf("embedding %d: invalid base64: %v", i, err)
146+
}
147+
148+
// Check index
149+
if result.Data[i].Index != i {
150+
t.Errorf("embedding %d: expected index %d, got %d", i, i, result.Data[i].Index)
151+
}
152+
}
153+
}
154+
155+
func TestEmbeddingsMiddleware_InvalidEncodingFormat(t *testing.T) {
156+
gin.SetMode(gin.TestMode)
157+
158+
endpoint := func(c *gin.Context) {
159+
c.Status(http.StatusOK)
160+
}
161+
162+
router := gin.New()
163+
router.Use(EmbeddingsMiddleware())
164+
router.Handle(http.MethodPost, "/api/embed", endpoint)
165+
166+
testCases := []struct {
167+
name string
168+
encodingFormat string
169+
shouldFail bool
170+
}{
171+
{"valid: float", "float", false},
172+
{"valid: base64", "base64", false},
173+
{"valid: FLOAT (uppercase)", "FLOAT", false},
174+
{"valid: BASE64 (uppercase)", "BASE64", false},
175+
{"valid: Float (mixed)", "Float", false},
176+
{"valid: Base64 (mixed)", "Base64", false},
177+
{"invalid: json", "json", true},
178+
{"invalid: hex", "hex", true},
179+
{"invalid: invalid_format", "invalid_format", true},
180+
}
181+
182+
for _, tc := range testCases {
183+
t.Run(tc.name, func(t *testing.T) {
184+
body := `{
185+
"input": "test",
186+
"model": "test-model",
187+
"encoding_format": "` + tc.encodingFormat + `"
188+
}`
189+
190+
req, _ := http.NewRequest(http.MethodPost, "/api/embed", strings.NewReader(body))
191+
req.Header.Set("Content-Type", "application/json")
192+
193+
resp := httptest.NewRecorder()
194+
router.ServeHTTP(resp, req)
195+
196+
if tc.shouldFail {
197+
if resp.Code != http.StatusBadRequest {
198+
t.Errorf("expected status 400, got %d", resp.Code)
199+
}
200+
201+
var errResp openai.ErrorResponse
202+
if err := json.Unmarshal(resp.Body.Bytes(), &errResp); err != nil {
203+
t.Fatalf("failed to unmarshal error response: %v", err)
204+
}
205+
206+
if errResp.Error.Type != "invalid_request_error" {
207+
t.Errorf("expected error type 'invalid_request_error', got %q", errResp.Error.Type)
208+
}
209+
210+
if !strings.Contains(errResp.Error.Message, "encoding_format") {
211+
t.Errorf("expected error message to mention encoding_format, got %q", errResp.Error.Message)
212+
}
213+
} else {
214+
if resp.Code != http.StatusOK {
215+
t.Errorf("expected status 200, got %d: %s", resp.Code, resp.Body.String())
216+
}
217+
}
218+
})
219+
}
220+
}

‎openai/openai.go‎

Lines changed: 26 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,9 @@
22
package openai
33

44
import (
5+
"bytes"
56
"encoding/base64"
7+
"encoding/binary"
68
"encoding/json"
79
"errors"
810
"fmt"
@@ -73,9 +75,10 @@ type JsonSchema struct {
7375
}
7476

7577
type EmbedRequest struct {
76-
Input any `json:"input"`
77-
Model string `json:"model"`
78-
Dimensions int `json:"dimensions,omitempty"`
78+
Input any `json:"input"`
79+
Model string `json:"model"`
80+
Dimensions int `json:"dimensions,omitempty"`
81+
EncodingFormat string `json:"encoding_format,omitempty"` // "float" or "base64"
7982
}
8083

8184
type StreamOptions struct {
@@ -181,9 +184,9 @@ type Model struct {
181184
}
182185

183186
type Embedding struct {
184-
Object string `json:"object"`
185-
Embedding []float32 `json:"embedding"`
186-
Index int `json:"index"`
187+
Object string `json:"object"`
188+
Embedding any `json:"embedding"` // Can be []float32 (float format) or string (base64 format)
189+
Index int `json:"index"`
187190
}
188191

189192
type ListCompletion struct {
@@ -377,13 +380,21 @@ func ToListCompletion(r api.ListResponse) ListCompletion {
377380
}
378381

379382
// ToEmbeddingList converts an api.EmbedResponse to EmbeddingList
380-
func ToEmbeddingList(model string, r api.EmbedResponse) EmbeddingList {
383+
// encodingFormat can be "float", "base64", or empty (defaults to "float")
384+
func ToEmbeddingList(model string, r api.EmbedResponse, encodingFormat string) EmbeddingList {
381385
if r.Embeddings != nil {
382386
var data []Embedding
383387
for i, e := range r.Embeddings {
388+
var embedding any
389+
if strings.EqualFold(encodingFormat, "base64") {
390+
embedding = floatsToBase64(e)
391+
} else {
392+
embedding = e
393+
}
394+
384395
data = append(data, Embedding{
385396
Object: "embedding",
386-
Embedding: e,
397+
Embedding: embedding,
387398
Index: i,
388399
})
389400
}
@@ -402,6 +413,13 @@ func ToEmbeddingList(model string, r api.EmbedResponse) EmbeddingList {
402413
return EmbeddingList{}
403414
}
404415

416+
// floatsToBase64 encodes a []float32 to a base64 string
417+
func floatsToBase64(floats []float32) string {
418+
var buf bytes.Buffer
419+
binary.Write(&buf, binary.LittleEndian, floats)
420+
return base64.StdEncoding.EncodeToString(buf.Bytes())
421+
}
422+
405423
// ToModel converts an api.ShowResponse to Model
406424
func ToModel(r api.ShowResponse, m string) Model {
407425
return Model{

0 commit comments

Comments
 (0)