mirror of
https://github.com/alibaba/higress.git
synced 2026-02-25 21:21:01 +08:00
115 lines
3.3 KiB
Go
115 lines
3.3 KiB
Go
package embedding
|
|
|
|
import (
|
|
"errors"
|
|
|
|
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
|
"github.com/tidwall/gjson"
|
|
)
|
|
|
|
const (
|
|
PROVIDER_TYPE_DASHSCOPE = "dashscope"
|
|
PROVIDER_TYPE_TEXTIN = "textin"
|
|
PROVIDER_TYPE_COHERE = "cohere"
|
|
PROVIDER_TYPE_OPENAI = "openai"
|
|
PROVIDER_TYPE_OLLAMA = "ollama"
|
|
PROVIDER_TYPE_HUGGINGFACE = "huggingface"
|
|
PROVIDER_TYPE_XFYUN = "xfyun"
|
|
PROVIDER_TYPE_AZURE = "azure"
|
|
)
|
|
|
|
type providerInitializer interface {
|
|
InitConfig(json gjson.Result)
|
|
ValidateConfig() error
|
|
CreateProvider(ProviderConfig) (Provider, error)
|
|
}
|
|
|
|
var (
|
|
providerInitializers = map[string]providerInitializer{
|
|
PROVIDER_TYPE_DASHSCOPE: &dashScopeProviderInitializer{},
|
|
PROVIDER_TYPE_TEXTIN: &textInProviderInitializer{},
|
|
PROVIDER_TYPE_COHERE: &cohereProviderInitializer{},
|
|
PROVIDER_TYPE_OPENAI: &openAIProviderInitializer{},
|
|
PROVIDER_TYPE_OLLAMA: &ollamaProviderInitializer{},
|
|
PROVIDER_TYPE_HUGGINGFACE: &huggingfaceProviderInitializer{},
|
|
PROVIDER_TYPE_XFYUN: &xfyunProviderInitializer{},
|
|
PROVIDER_TYPE_AZURE: &azureProviderInitializer{},
|
|
}
|
|
)
|
|
|
|
type ProviderConfig struct {
|
|
// @Title zh-CN 文本特征提取服务提供者类型
|
|
// @Description zh-CN 文本特征提取服务提供者类型，例如 DashScope
|
|
typ string
|
|
// @Title zh-CN DashScope 文本特征提取服务名称
|
|
// @Description zh-CN 文本特征提取服务名称
|
|
serviceName string
|
|
// @Title zh-CN 文本特征提取服务域名
|
|
// @Description zh-CN 文本特征提取服务域名
|
|
serviceHost string
|
|
// @Title zh-CN 文本特征提取服务端口
|
|
// @Description zh-CN 文本特征提取服务端口
|
|
servicePort int64
|
|
// @Title zh-CN 文本特征提取服务超时时间
|
|
// @Description zh-CN 文本特征提取服务超时时间
|
|
timeout uint32
|
|
// @Title zh-CN 文本特征提取服务使用的模型
|
|
// @Description zh-CN 用于文本特征提取的模型名称, 在 DashScope 中默认为 "text-embedding-v1"
|
|
model string
|
|
|
|
initializer providerInitializer
|
|
}
|
|
|
|
func (c *ProviderConfig) FromJson(json gjson.Result) {
|
|
c.typ = json.Get("type").String()
|
|
i, has := providerInitializers[c.typ]
|
|
if has {
|
|
i.InitConfig(json)
|
|
c.initializer = i
|
|
}
|
|
c.serviceName = json.Get("serviceName").String()
|
|
c.serviceHost = json.Get("serviceHost").String()
|
|
c.servicePort = json.Get("servicePort").Int()
|
|
c.timeout = uint32(json.Get("timeout").Int())
|
|
c.model = json.Get("model").String()
|
|
if c.timeout == 0 {
|
|
c.timeout = 10000
|
|
}
|
|
}
|
|
|
|
func (c *ProviderConfig) Validate() error {
|
|
if c.serviceName == "" {
|
|
return errors.New("embedding service name is required")
|
|
}
|
|
if c.typ == "" {
|
|
return errors.New("embedding service type is required")
|
|
}
|
|
if c.initializer == nil {
|
|
return errors.New("unknown embedding service provider type: " + c.typ)
|
|
}
|
|
if err := c.initializer.ValidateConfig(); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *ProviderConfig) GetProviderType() string {
|
|
return c.typ
|
|
}
|
|
|
|
func CreateProvider(pc ProviderConfig) (Provider, error) {
|
|
initializer, has := providerInitializers[pc.typ]
|
|
if !has {
|
|
return nil, errors.New("unknown provider type: " + pc.typ)
|
|
}
|
|
return initializer.CreateProvider(pc)
|
|
}
|
|
|
|
type Provider interface {
|
|
GetProviderType() string
|
|
GetEmbedding(
|
|
queryString string,
|
|
ctx wrapper.HttpContext,
|
|
callback func(emb []float64, err error)) error
|
|
}
|