mirror of
https://github.com/ollama/ollama.git
synced 2025-11-11 16:37:03 +01:00
embeddings: base64 encoding fix (#12715)
This commit is contained in:
nicole pardal
GitHub
220
middleware/openai_encoding_format_test.go
Normal file
220
middleware/openai_encoding_format_test.go
Normal file
@@ -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())
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user