[ai-cache] Implement a WASM plugin for LLM result retrieval based on vector similarity (#1290)

plugins/wasm-go/extensions/ai-cache/embedding/dashscope.go

@@ -0,0 +1,187 @@
+package embedding
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"strconv"
+
+	"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
+)
+
+const (
+	DASHSCOPE_DOMAIN             = "dashscope.aliyuncs.com"
+	DASHSCOPE_PORT               = 443
+	DASHSCOPE_DEFAULT_MODEL_NAME = "text-embedding-v2"
+	DASHSCOPE_ENDPOINT           = "/api/v1/services/embeddings/text-embedding/text-embedding"
+)
+
+type dashScopeProviderInitializer struct {
+}
+
+func (d *dashScopeProviderInitializer) ValidateConfig(config ProviderConfig) error {
+	if config.apiKey == "" {
+		return errors.New("[DashScope] apiKey is required")
+	}
+	return nil
+}
+
+func (d *dashScopeProviderInitializer) CreateProvider(c ProviderConfig) (Provider, error) {
+	if c.servicePort == 0 {
+		c.servicePort = DASHSCOPE_PORT
+	}
+	if c.serviceHost == "" {
+		c.serviceHost = DASHSCOPE_DOMAIN
+	}
+	return &DSProvider{
+		config: c,
+		client: wrapper.NewClusterClient(wrapper.FQDNCluster{
+			FQDN: c.serviceName,
+			Host: c.serviceHost,
+			Port: int64(c.servicePort),
+		}),
+	}, nil
+}
+
+func (d *DSProvider) GetProviderType() string {
+	return PROVIDER_TYPE_DASHSCOPE
+}
+
+type Embedding struct {
+	Embedding []float64 `json:"embedding"`
+	TextIndex int       `json:"text_index"`
+}
+
+type Input struct {
+	Texts []string `json:"texts"`
+}
+
+type Params struct {
+	TextType string `json:"text_type"`
+}
+
+type Response struct {
+	RequestID string `json:"request_id"`
+	Output    Output `json:"output"`
+	Usage     Usage  `json:"usage"`
+}
+
+type Output struct {
+	Embeddings []Embedding `json:"embeddings"`
+}
+
+type Usage struct {
+	TotalTokens int `json:"total_tokens"`
+}
+
+type EmbeddingRequest struct {
+	Model      string `json:"model"`
+	Input      Input  `json:"input"`
+	Parameters Params `json:"parameters"`
+}
+
+type Document struct {
+	Vector []float64         `json:"vector"`
+	Fields map[string]string `json:"fields"`
+}
+
+type DSProvider struct {
+	config ProviderConfig
+	client wrapper.HttpClient
+}
+
+func (d *DSProvider) constructParameters(texts []string, log wrapper.Log) (string, [][2]string, []byte, error) {
+
+	model := d.config.model
+
+	if model == "" {
+		model = DASHSCOPE_DEFAULT_MODEL_NAME
+	}
+	data := EmbeddingRequest{
+		Model: model,
+		Input: Input{
+			Texts: texts,
+		},
+		Parameters: Params{
+			TextType: "query",
+		},
+	}
+
+	requestBody, err := json.Marshal(data)
+	if err != nil {
+		log.Errorf("failed to marshal request data: %v", err)
+		return "", nil, nil, err
+	}
+
+	if d.config.apiKey == "" {
+		err := errors.New("dashScopeKey is empty")
+		log.Errorf("failed to construct headers: %v", err)
+		return "", nil, nil, err
+	}
+
+	headers := [][2]string{
+		{"Authorization", "Bearer " + d.config.apiKey},
+		{"Content-Type", "application/json"},
+	}
+
+	return DASHSCOPE_ENDPOINT, headers, requestBody, err
+}
+
+type Result struct {
+	ID     string                 `json:"id"`
+	Vector []float64              `json:"vector,omitempty"`
+	Fields map[string]interface{} `json:"fields"`
+	Score  float64                `json:"score"`
+}
+
+func (d *DSProvider) parseTextEmbedding(responseBody []byte) (*Response, error) {
+	var resp Response
+	err := json.Unmarshal(responseBody, &resp)
+	if err != nil {
+		return nil, err
+	}
+	return &resp, nil
+}
+
+func (d *DSProvider) GetEmbedding(
+	queryString string,
+	ctx wrapper.HttpContext,
+	log wrapper.Log,
+	callback func(emb []float64, err error)) error {
+	embUrl, embHeaders, embRequestBody, err := d.constructParameters([]string{queryString}, log)
+	if err != nil {
+		log.Errorf("failed to construct parameters: %v", err)
+		return err
+	}
+
+	var resp *Response
+	err = d.client.Post(embUrl, embHeaders, embRequestBody,
+		func(statusCode int, responseHeaders http.Header, responseBody []byte) {
+
+			if statusCode != http.StatusOK {
+				err = errors.New("failed to get embedding due to status code: " + strconv.Itoa(statusCode))
+				callback(nil, err)
+				return
+			}
+
+			log.Debugf("get embedding response: %d, %s", statusCode, responseBody)
+
+			resp, err = d.parseTextEmbedding(responseBody)
+			if err != nil {
+				err = fmt.Errorf("failed to parse response: %v", err)
+				callback(nil, err)
+				return
+			}
+
+			if len(resp.Output.Embeddings) == 0 {
+				err = errors.New("no embedding found in response")
+				callback(nil, err)
+				return
+			}
+
+			callback(resp.Output.Embeddings[0].Embedding, nil)
+
+		}, d.config.timeout)
+	return err
+}

plugins/wasm-go/extensions/ai-cache/embedding/provider.go

@@ -0,0 +1,101 @@
+package embedding
+
+import (
+	"errors"
+
+	"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
+	"github.com/tidwall/gjson"
+)
+
+const (
+	PROVIDER_TYPE_DASHSCOPE = "dashscope"
+)
+
+type providerInitializer interface {
+	ValidateConfig(ProviderConfig) error
+	CreateProvider(ProviderConfig) (Provider, error)
+}
+
+var (
+	providerInitializers = map[string]providerInitializer{
+		PROVIDER_TYPE_DASHSCOPE: &dashScopeProviderInitializer{},
+	}
+)
+
+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 文本特征提取服务 API Key
+	// @Description zh-CN 文本特征提取服务 API Key
+	apiKey string
+	// @Title zh-CN 文本特征提取服务超时时间
+	// @Description zh-CN 文本特征提取服务超时时间
+	timeout uint32
+	// @Title zh-CN 文本特征提取服务使用的模型
+	// @Description zh-CN 用于文本特征提取的模型名称, 在 DashScope 中默认为 "text-embedding-v1"
+	model string
+}
+
+func (c *ProviderConfig) FromJson(json gjson.Result) {
+	c.typ = json.Get("type").String()
+	c.serviceName = json.Get("serviceName").String()
+	c.serviceHost = json.Get("serviceHost").String()
+	c.servicePort = json.Get("servicePort").Int()
+	c.apiKey = json.Get("apiKey").String()
+	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.apiKey == "" {
+		return errors.New("embedding service API key is required")
+	}
+	if c.typ == "" {
+		return errors.New("embedding service type is required")
+	}
+	initializer, has := providerInitializers[c.typ]
+	if !has {
+		return errors.New("unknown embedding service provider type: " + c.typ)
+	}
+	if err := initializer.ValidateConfig(*c); 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,
+		log wrapper.Log,
+		callback func(emb []float64, err error)) error
+}

plugins/wasm-go/extensions/ai-cache/embedding/weaviate.go

@@ -0,0 +1,27 @@
+package embedding
+
+// import (
+// 	"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
+// )
+
+// const (
+// 	weaviateURL = "172.17.0.1:8081"
+// )
+
+// type weaviateProviderInitializer struct {
+// }
+
+// func (d *weaviateProviderInitializer) ValidateConfig(config ProviderConfig) error {
+// 	return nil
+// }
+
+// func (d *weaviateProviderInitializer) CreateProvider(config ProviderConfig) (Provider, error) {
+// 	return &DSProvider{
+// 		config: config,
+// 		client: wrapper.NewClusterClient(wrapper.DnsCluster{
+// 			ServiceName: config.ServiceName,
+// 			Port:        dashScopePort,
+// 			Domain:      dashScopeDomain,
+// 		}),
+// 	}, nil
+// }