Merge pull request #555 from utopeadia/main

Continue fixing Ollama embedding return issue
This commit is contained in:
Calcium-Ion 2024-11-06 21:13:06 +08:00 committed by GitHub
commit be652fa3c2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 2 deletions

View File

@ -37,5 +37,5 @@ type OllamaEmbeddingRequest struct {
type OllamaEmbeddingResponse struct { type OllamaEmbeddingResponse struct {
Error string `json:"error,omitempty"` Error string `json:"error,omitempty"`
Model string `json:"model"` Model string `json:"model"`
Embedding []float64 `json:"embeddings,omitempty"` Embedding [][]float64 `json:"embeddings,omitempty"`
} }

View File

@ -73,9 +73,10 @@ func ollamaEmbeddingHandler(c *gin.Context, resp *http.Response, promptTokens in
if ollamaEmbeddingResponse.Error != "" { if ollamaEmbeddingResponse.Error != "" {
return service.OpenAIErrorWrapper(err, "ollama_error", resp.StatusCode), nil return service.OpenAIErrorWrapper(err, "ollama_error", resp.StatusCode), nil
} }
flattenedEmbeddings := flattenEmbeddings(ollamaEmbeddingResponse.Embedding)
data := make([]dto.OpenAIEmbeddingResponseItem, 0, 1) data := make([]dto.OpenAIEmbeddingResponseItem, 0, 1)
data = append(data, dto.OpenAIEmbeddingResponseItem{ data = append(data, dto.OpenAIEmbeddingResponseItem{
Embedding: ollamaEmbeddingResponse.Embedding, Embedding: flattenedEmbeddings,
Object: "embedding", Object: "embedding",
}) })
usage := &dto.Usage{ usage := &dto.Usage{
@ -120,3 +121,11 @@ func ollamaEmbeddingHandler(c *gin.Context, resp *http.Response, promptTokens in
} }
return nil, usage return nil, usage
} }
func flattenEmbeddings(embeddings [][]float64) []float64 {
flattened := []float64{}
for _, row := range embeddings {
flattened = append(flattened, row...)
}
return flattened
}