feat: update Ollama embedding API to latest version with multi-text embedding support (#1715)

This commit is contained in:
SLKun 2024-08-06 23:43:20 +08:00 committed by GitHub
parent 402fcf7f79
commit e7e99e558a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 26 additions and 13 deletions

View File

@ -24,7 +24,7 @@ func (a *Adaptor) GetRequestURL(meta *meta.Meta) (string, error) {
// https://github.com/ollama/ollama/blob/main/docs/api.md // https://github.com/ollama/ollama/blob/main/docs/api.md
fullRequestURL := fmt.Sprintf("%s/api/chat", meta.BaseURL) fullRequestURL := fmt.Sprintf("%s/api/chat", meta.BaseURL)
if meta.Mode == relaymode.Embeddings { if meta.Mode == relaymode.Embeddings {
fullRequestURL = fmt.Sprintf("%s/api/embeddings", meta.BaseURL) fullRequestURL = fmt.Sprintf("%s/api/embed", meta.BaseURL)
} }
return fullRequestURL, nil return fullRequestURL, nil
} }

View File

@ -157,8 +157,15 @@ func StreamHandler(c *gin.Context, resp *http.Response) (*model.ErrorWithStatusC
func ConvertEmbeddingRequest(request model.GeneralOpenAIRequest) *EmbeddingRequest { func ConvertEmbeddingRequest(request model.GeneralOpenAIRequest) *EmbeddingRequest {
return &EmbeddingRequest{ return &EmbeddingRequest{
Model: request.Model, Model: request.Model,
Prompt: strings.Join(request.ParseInput(), " "), Input: request.ParseInput(),
Options: &Options{
Seed: int(request.Seed),
Temperature: request.Temperature,
TopP: request.TopP,
FrequencyPenalty: request.FrequencyPenalty,
PresencePenalty: request.PresencePenalty,
},
} }
} }
@ -201,15 +208,17 @@ func embeddingResponseOllama2OpenAI(response *EmbeddingResponse) *openai.Embeddi
openAIEmbeddingResponse := openai.EmbeddingResponse{ openAIEmbeddingResponse := openai.EmbeddingResponse{
Object: "list", Object: "list",
Data: make([]openai.EmbeddingResponseItem, 0, 1), Data: make([]openai.EmbeddingResponseItem, 0, 1),
Model: "text-embedding-v1", Model: response.Model,
Usage: model.Usage{TotalTokens: 0}, Usage: model.Usage{TotalTokens: 0},
} }
openAIEmbeddingResponse.Data = append(openAIEmbeddingResponse.Data, openai.EmbeddingResponseItem{ for i, embedding := range response.Embeddings {
Object: `embedding`, openAIEmbeddingResponse.Data = append(openAIEmbeddingResponse.Data, openai.EmbeddingResponseItem{
Index: 0, Object: `embedding`,
Embedding: response.Embedding, Index: i,
}) Embedding: embedding,
})
}
return &openAIEmbeddingResponse return &openAIEmbeddingResponse
} }

View File

@ -37,11 +37,15 @@ type ChatResponse struct {
} }
type EmbeddingRequest struct { type EmbeddingRequest struct {
Model string `json:"model"` Model string `json:"model"`
Prompt string `json:"prompt"` Input []string `json:"input"`
// Truncate bool `json:"truncate,omitempty"`
Options *Options `json:"options,omitempty"`
// KeepAlive string `json:"keep_alive,omitempty"`
} }
type EmbeddingResponse struct { type EmbeddingResponse struct {
Error string `json:"error,omitempty"` Error string `json:"error,omitempty"`
Embedding []float64 `json:"embedding,omitempty"` Model string `json:"model"`
Embeddings [][]float64 `json:"embeddings"`
} }