mirror of
https://github.com/alibaba/higress.git
synced 2026-03-17 00:40:48 +08:00
[ai-cache] Implement a WASM plugin for LLM result retrieval based on vector similarity (#1290)
This commit is contained in:
Yang Beining
GitHub
256
plugins/wasm-go/extensions/ai-cache/vector/dashvector.go
Normal file
256
plugins/wasm-go/extensions/ai-cache/vector/dashvector.go
Normal file
@@ -0,0 +1,256 @@
|
||||
package vector
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||
)
|
||||
|
||||
type dashVectorProviderInitializer struct {
|
||||
}
|
||||
|
||||
func (d *dashVectorProviderInitializer) ValidateConfig(config ProviderConfig) error {
|
||||
if len(config.apiKey) == 0 {
|
||||
return errors.New("[DashVector] apiKey is required")
|
||||
}
|
||||
if len(config.collectionID) == 0 {
|
||||
return errors.New("[DashVector] collectionID is required")
|
||||
}
|
||||
if len(config.serviceName) == 0 {
|
||||
return errors.New("[DashVector] serviceName is required")
|
||||
}
|
||||
if len(config.serviceHost) == 0 {
|
||||
return errors.New("[DashVector] serviceHost is required")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (d *dashVectorProviderInitializer) CreateProvider(config ProviderConfig) (Provider, error) {
|
||||
return &DvProvider{
|
||||
config: config,
|
||||
client: wrapper.NewClusterClient(wrapper.FQDNCluster{
|
||||
FQDN: config.serviceName,
|
||||
Host: config.serviceHost,
|
||||
Port: int64(config.servicePort),
|
||||
}),
|
||||
}, nil
|
||||
}
|
||||
|
||||
type DvProvider struct {
|
||||
config ProviderConfig
|
||||
client wrapper.HttpClient
|
||||
}
|
||||
|
||||
func (d *DvProvider) GetProviderType() string {
|
||||
return PROVIDER_TYPE_DASH_VECTOR
|
||||
}
|
||||
|
||||
// type embeddingRequest struct {
|
||||
// Model string `json:"model"`
|
||||
// Input input `json:"input"`
|
||||
// Parameters params `json:"parameters"`
|
||||
// }
|
||||
|
||||
// type params struct {
|
||||
// TextType string `json:"text_type"`
|
||||
// }
|
||||
|
||||
// type input struct {
|
||||
// Texts []string `json:"texts"`
|
||||
// }
|
||||
|
||||
// queryResponse 定义查询响应的结构
|
||||
type queryResponse struct {
|
||||
Code int `json:"code"`
|
||||
RequestID string `json:"request_id"`
|
||||
Message string `json:"message"`
|
||||
Output []result `json:"output"`
|
||||
}
|
||||
|
||||
// queryRequest 定义查询请求的结构
|
||||
type queryRequest struct {
|
||||
Vector []float64 `json:"vector"`
|
||||
TopK int `json:"topk"`
|
||||
IncludeVector bool `json:"include_vector"`
|
||||
}
|
||||
|
||||
// result 定义查询结果的结构
|
||||
type result struct {
|
||||
ID string `json:"id"`
|
||||
Vector []float64 `json:"vector,omitempty"` // omitempty 使得如果 vector 是空，它将不会被序列化
|
||||
Fields map[string]interface{} `json:"fields"`
|
||||
Score float64 `json:"score"`
|
||||
}
|
||||
|
||||
func (d *DvProvider) constructEmbeddingQueryParameters(vector []float64) (string, []byte, [][2]string, error) {
|
||||
url := fmt.Sprintf("/v1/collections/%s/query", d.config.collectionID)
|
||||
|
||||
requestData := queryRequest{
|
||||
Vector: vector,
|
||||
TopK: d.config.topK,
|
||||
IncludeVector: false,
|
||||
}
|
||||
|
||||
requestBody, err := json.Marshal(requestData)
|
||||
if err != nil {
|
||||
return "", nil, nil, err
|
||||
}
|
||||
|
||||
header := [][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{"dashvector-auth-token", d.config.apiKey},
|
||||
}
|
||||
|
||||
return url, requestBody, header, nil
|
||||
}
|
||||
|
||||
func (d *DvProvider) parseQueryResponse(responseBody []byte) (queryResponse, error) {
|
||||
var queryResp queryResponse
|
||||
err := json.Unmarshal(responseBody, &queryResp)
|
||||
if err != nil {
|
||||
return queryResponse{}, err
|
||||
}
|
||||
return queryResp, nil
|
||||
}
|
||||
|
||||
func (d *DvProvider) QueryEmbedding(
|
||||
emb []float64,
|
||||
ctx wrapper.HttpContext,
|
||||
log wrapper.Log,
|
||||
callback func(results []QueryResult, ctx wrapper.HttpContext, log wrapper.Log, err error)) error {
|
||||
url, body, headers, err := d.constructEmbeddingQueryParameters(emb)
|
||||
log.Debugf("url:%s, body:%s, headers:%v", url, string(body), headers)
|
||||
if err != nil {
|
||||
err = fmt.Errorf("failed to construct embedding query parameters: %v", err)
|
||||
return err
|
||||
}
|
||||
|
||||
err = d.client.Post(url, headers, body,
|
||||
func(statusCode int, responseHeaders http.Header, responseBody []byte) {
|
||||
err = nil
|
||||
if statusCode != http.StatusOK {
|
||||
err = fmt.Errorf("failed to query embedding: %d", statusCode)
|
||||
callback(nil, ctx, log, err)
|
||||
return
|
||||
}
|
||||
log.Debugf("query embedding response: %d, %s", statusCode, responseBody)
|
||||
results, err := d.ParseQueryResponse(responseBody, ctx, log)
|
||||
if err != nil {
|
||||
err = fmt.Errorf("failed to parse query response: %v", err)
|
||||
}
|
||||
callback(results, ctx, log, err)
|
||||
},
|
||||
d.config.timeout)
|
||||
if err != nil {
|
||||
err = fmt.Errorf("failed to query embedding: %v", err)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func getStringValue(fields map[string]interface{}, key string) string {
|
||||
if val, ok := fields[key]; ok {
|
||||
return val.(string)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (d *DvProvider) ParseQueryResponse(responseBody []byte, ctx wrapper.HttpContext, log wrapper.Log) ([]QueryResult, error) {
|
||||
resp, err := d.parseQueryResponse(responseBody)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if len(resp.Output) == 0 {
|
||||
return nil, errors.New("no query results found in response")
|
||||
}
|
||||
|
||||
results := make([]QueryResult, 0, len(resp.Output))
|
||||
|
||||
for _, output := range resp.Output {
|
||||
result := QueryResult{
|
||||
Text: getStringValue(output.Fields, "query"),
|
||||
Embedding: output.Vector,
|
||||
Score: output.Score,
|
||||
Answer: getStringValue(output.Fields, "answer"),
|
||||
}
|
||||
results = append(results, result)
|
||||
}
|
||||
|
||||
return results, nil
|
||||
}
|
||||
|
||||
type document struct {
|
||||
Vector []float64 `json:"vector"`
|
||||
Fields map[string]string `json:"fields"`
|
||||
}
|
||||
|
||||
type insertRequest struct {
|
||||
Docs []document `json:"docs"`
|
||||
}
|
||||
|
||||
func (d *DvProvider) constructUploadParameters(emb []float64, queryString string, answer string) (string, []byte, [][2]string, error) {
|
||||
url := "/v1/collections/" + d.config.collectionID + "/docs"
|
||||
|
||||
doc := document{
|
||||
Vector: emb,
|
||||
Fields: map[string]string{
|
||||
"query": queryString,
|
||||
"answer": answer,
|
||||
},
|
||||
}
|
||||
|
||||
requestBody, err := json.Marshal(insertRequest{Docs: []document{doc}})
|
||||
if err != nil {
|
||||
return "", nil, nil, err
|
||||
}
|
||||
|
||||
header := [][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{"dashvector-auth-token", d.config.apiKey},
|
||||
}
|
||||
|
||||
return url, requestBody, header, err
|
||||
}
|
||||
|
||||
func (d *DvProvider) UploadEmbedding(queryString string, queryEmb []float64, ctx wrapper.HttpContext, log wrapper.Log, callback func(ctx wrapper.HttpContext, log wrapper.Log, err error)) error {
|
||||
url, body, headers, err := d.constructUploadParameters(queryEmb, queryString, "")
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = d.client.Post(
|
||||
url,
|
||||
headers,
|
||||
body,
|
||||
func(statusCode int, responseHeaders http.Header, responseBody []byte) {
|
||||
log.Debugf("statusCode:%d, responseBody:%s", statusCode, string(responseBody))
|
||||
if statusCode != http.StatusOK {
|
||||
err = fmt.Errorf("failed to upload embedding: %d", statusCode)
|
||||
}
|
||||
callback(ctx, log, err)
|
||||
},
|
||||
d.config.timeout)
|
||||
return err
|
||||
}
|
||||
|
||||
func (d *DvProvider) UploadAnswerAndEmbedding(queryString string, queryEmb []float64, queryAnswer string, ctx wrapper.HttpContext, log wrapper.Log, callback func(ctx wrapper.HttpContext, log wrapper.Log, err error)) error {
|
||||
url, body, headers, err := d.constructUploadParameters(queryEmb, queryString, queryAnswer)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = d.client.Post(
|
||||
url,
|
||||
headers,
|
||||
body,
|
||||
func(statusCode int, responseHeaders http.Header, responseBody []byte) {
|
||||
log.Debugf("statusCode:%d, responseBody:%s", statusCode, string(responseBody))
|
||||
if statusCode != http.StatusOK {
|
||||
err = fmt.Errorf("failed to upload embedding: %d", statusCode)
|
||||
}
|
||||
callback(ctx, log, err)
|
||||
},
|
||||
d.config.timeout)
|
||||
return err
|
||||
}
|
||||
Reference in New Issue
Block a user