feat: support jina rerank

This commit is contained in:
CalciumIon 2024-07-06 18:42:48 +08:00
parent 3ed4f2f0a9
commit 8a730cfe12
8 changed files with 120 additions and 3 deletions

View File

@ -34,7 +34,7 @@
3. 选择你的bot然后输入http(s)://你的网站地址/login 3. 选择你的bot然后输入http(s)://你的网站地址/login
4. Telegram Bot 名称是bot username 去掉@后的字符串 4. Telegram Bot 名称是bot username 去掉@后的字符串
13. 添加 [Suno API](https://github.com/Suno-API/Suno-API)接口的支持,[对接文档](Suno.md) 13. 添加 [Suno API](https://github.com/Suno-API/Suno-API)接口的支持,[对接文档](Suno.md)
14. 支持Rerank模型目前仅兼容Cohere的rerank可接入Dify[对接文档](Rerank.md) 14. 支持Rerank模型目前仅兼容Cohere和Jina可接入Dify[对接文档](Rerank.md)
## 模型支持 ## 模型支持
此版本额外支持以下模型: 此版本额外支持以下模型:
@ -46,7 +46,7 @@
6. [零一万物](https://platform.lingyiwanwu.com/) 6. [零一万物](https://platform.lingyiwanwu.com/)
7. 自定义渠道,支持填入完整调用地址 7. 自定义渠道,支持填入完整调用地址
8. [Suno API](https://github.com/Suno-API/Suno-API) 接口,[对接文档](Suno.md) 8. [Suno API](https://github.com/Suno-API/Suno-API) 接口,[对接文档](Suno.md)
9. Rerank模型目前支持[Cohere](https://cohere.ai/) 9. Rerank模型目前支持[Cohere](https://cohere.ai/)和[Jina](https://jina.ai/)[对接文档](Rerank.md)
10. Dify 10. Dify
您可以在渠道中添加自定义模型gpt-4-gizmo-*此模型并非OpenAI官方模型而是第三方模型使用官方key无法调用。 您可以在渠道中添加自定义模型gpt-4-gizmo-*此模型并非OpenAI官方模型而是第三方模型使用官方key无法调用。

View File

@ -211,6 +211,7 @@ const (
ChannelTypeMiniMax = 35 ChannelTypeMiniMax = 35
ChannelTypeSunoAPI = 36 ChannelTypeSunoAPI = 36
ChannelTypeDify = 37 ChannelTypeDify = 37
ChannelTypeJina = 38
ChannelTypeDummy // this one is only for count, do not add any channel after this ChannelTypeDummy // this one is only for count, do not add any channel after this
@ -254,4 +255,6 @@ var ChannelBaseURLs = []string{
"https://api.cohere.ai", //34 "https://api.cohere.ai", //34
"https://api.minimax.chat", //35 "https://api.minimax.chat", //35
"", //36 "", //36
"", //37
"https://api.jina.ai", //38
} }

View File

@ -0,0 +1,64 @@
package jina
import (
"errors"
"fmt"
"github.com/gin-gonic/gin"
"io"
"net/http"
"one-api/dto"
"one-api/relay/channel"
relaycommon "one-api/relay/common"
"one-api/relay/constant"
)
type Adaptor struct {
}
func (a *Adaptor) InitRerank(info *relaycommon.RelayInfo, request dto.RerankRequest) {
}
func (a *Adaptor) Init(info *relaycommon.RelayInfo, request dto.GeneralOpenAIRequest) {
}
func (a *Adaptor) GetRequestURL(info *relaycommon.RelayInfo) (string, error) {
if info.RelayMode == constant.RelayModeRerank {
return fmt.Sprintf("%s/v1/rerank", info.BaseUrl), nil
} else if info.RelayMode == constant.RelayModeEmbeddings {
return fmt.Sprintf("%s/v1/embeddings ", info.BaseUrl), nil
}
return "", errors.New("invalid relay mode")
}
func (a *Adaptor) SetupRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error {
channel.SetupApiRequestHeader(info, c, req)
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", info.ApiKey))
return nil
}
func (a *Adaptor) ConvertRequest(c *gin.Context, relayMode int, request *dto.GeneralOpenAIRequest) (any, error) {
return request, nil
}
func (a *Adaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) {
return channel.DoApiRequest(a, c, info, requestBody)
}
func (a *Adaptor) ConvertRerankRequest(c *gin.Context, relayMode int, request dto.RerankRequest) (any, error) {
return request, nil
}
func (a *Adaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (usage *dto.Usage, err *dto.OpenAIErrorWithStatusCode) {
if info.RelayMode == constant.RelayModeRerank {
err, usage = jinaRerankHandler(c, resp)
}
return
}
func (a *Adaptor) GetModelList() []string {
return ModelList
}
func (a *Adaptor) GetChannelName() string {
return ChannelName
}

View File

@ -0,0 +1,8 @@
package jina
var ModelList = []string{
"jina-clip-v1",
"jina-reranker-v2-base-multilingual",
}
var ChannelName = "jina"

View File

@ -0,0 +1,35 @@
package jina
import (
"encoding/json"
"github.com/gin-gonic/gin"
"io"
"net/http"
"one-api/dto"
"one-api/service"
)
func jinaRerankHandler(c *gin.Context, resp *http.Response) (*dto.OpenAIErrorWithStatusCode, *dto.Usage) {
responseBody, err := io.ReadAll(resp.Body)
if err != nil {
return service.OpenAIErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError), nil
}
err = resp.Body.Close()
if err != nil {
return service.OpenAIErrorWrapper(err, "close_response_body_failed", http.StatusInternalServerError), nil
}
var jinaResp dto.RerankResponse
err = json.Unmarshal(responseBody, &jinaResp)
if err != nil {
return service.OpenAIErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError), nil
}
jsonResponse, err := json.Marshal(jinaResp)
if err != nil {
return service.OpenAIErrorWrapper(err, "marshal_response_body_failed", http.StatusInternalServerError), nil
}
c.Writer.Header().Set("Content-Type", "application/json")
c.Writer.WriteHeader(resp.StatusCode)
_, err = c.Writer.Write(jsonResponse)
return nil, &jinaResp.Usage
}

View File

@ -21,6 +21,7 @@ const (
APITypeAws APITypeAws
APITypeCohere APITypeCohere
APITypeDify APITypeDify
APITypeJina
APITypeDummy // this one is only for count, do not add any channel after this APITypeDummy // this one is only for count, do not add any channel after this
) )
@ -60,6 +61,8 @@ func ChannelType2APIType(channelType int) (int, bool) {
apiType = APITypeCohere apiType = APITypeCohere
case common.ChannelTypeDify: case common.ChannelTypeDify:
apiType = APITypeDify apiType = APITypeDify
case common.ChannelTypeJina:
apiType = APITypeJina
} }
if apiType == -1 { if apiType == -1 {
return APITypeOpenAI, false return APITypeOpenAI, false

View File

@ -10,6 +10,7 @@ import (
"one-api/relay/channel/cohere" "one-api/relay/channel/cohere"
"one-api/relay/channel/dify" "one-api/relay/channel/dify"
"one-api/relay/channel/gemini" "one-api/relay/channel/gemini"
"one-api/relay/channel/jina"
"one-api/relay/channel/ollama" "one-api/relay/channel/ollama"
"one-api/relay/channel/openai" "one-api/relay/channel/openai"
"one-api/relay/channel/palm" "one-api/relay/channel/palm"
@ -56,6 +57,8 @@ func GetAdaptor(apiType int) channel.Adaptor {
return &cohere.Adaptor{} return &cohere.Adaptor{}
case constant.APITypeDify: case constant.APITypeDify:
return &dify.Adaptor{} return &dify.Adaptor{}
case constant.APITypeJina:
return &jina.Adaptor{}
} }
return nil return nil
} }

View File

@ -104,7 +104,8 @@ export const CHANNEL_OPTIONS = [
{ key: 23, text: '腾讯混元', value: 23, color: 'teal', label: '腾讯混元' }, { key: 23, text: '腾讯混元', value: 23, color: 'teal', label: '腾讯混元' },
{ key: 31, text: '零一万物', value: 31, color: 'green', label: '零一万物' }, { key: 31, text: '零一万物', value: 31, color: 'green', label: '零一万物' },
{ key: 35, text: 'MiniMax', value: 35, color: 'green', label: 'MiniMax' }, { key: 35, text: 'MiniMax', value: 35, color: 'green', label: 'MiniMax' },
{ key: 37, text: 'Dify', value: 37, color: 'green', label: 'Dify' }, { key: 37, text: 'Dify', value: 37, color: 'teal', label: 'Dify' },
{ key: 38, text: 'Jina', value: 38, color: 'blue', label: 'Jina' },
{ key: 8, text: '自定义渠道', value: 8, color: 'pink', label: '自定义渠道' }, { key: 8, text: '自定义渠道', value: 8, color: 'pink', label: '自定义渠道' },
{ {
key: 22, key: 22,