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

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

@@ -53,6 +53,10 @@ type ProviderConfig struct {
 	cacheKeyPrefix string
 }
 
+func (c *ProviderConfig) GetProviderType() string {
+	return c.typ
+}
+
 func (c *ProviderConfig) FromJson(json gjson.Result) {
 	c.typ = json.Get("type").String()
 	c.serviceName = json.Get("serviceName").String()

plugins/wasm-go/extensions/ai-cache/config/config.go

@@ -31,8 +31,8 @@ type PluginConfig struct {
 }
 
 func (c *PluginConfig) FromJson(json gjson.Result) {
-	c.embeddingProviderConfig.FromJson(json.Get("embedding"))
 	c.vectorProviderConfig.FromJson(json.Get("vector"))
+	c.embeddingProviderConfig.FromJson(json.Get("embedding"))
 	c.cacheProviderConfig.FromJson(json.Get("cache"))
 
 	c.CacheKeyFrom = json.Get("cacheKeyFrom").String()
@@ -54,20 +54,25 @@ func (c *PluginConfig) FromJson(json gjson.Result) {
 
 	c.StreamResponseTemplate = json.Get("streamResponseTemplate").String()
 	if c.StreamResponseTemplate == "" {
-		c.StreamResponseTemplate = `data:{"id":"ai-cache.hit","choices":[{"index":0,"delta":{"role":"assistant","content":%s},"finish_reason":"stop"}],"model":"gpt-4o","object":"chat.completion","usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}}` + "\n\ndata:[DONE]\n\n"
+		c.StreamResponseTemplate = `data:{"id":"from-cache","choices":[{"index":0,"delta":{"role":"assistant","content":"%s"},"finish_reason":"stop"}],"model":"gpt-4o","object":"chat.completion","usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}}` + "\n\ndata:[DONE]\n\n"
 	}
 	c.ResponseTemplate = json.Get("responseTemplate").String()
 	if c.ResponseTemplate == "" {
-		c.ResponseTemplate = `{"id":"ai-cache.hit","choices":[{"index":0,"message":{"role":"assistant","content":%s},"finish_reason":"stop"}],"model":"gpt-4o","object":"chat.completion","usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}}`
+		c.ResponseTemplate = `{"id":"from-cache","choices":[{"index":0,"message":{"role":"assistant","content":"%s"},"finish_reason":"stop"}],"model":"gpt-4o","object":"chat.completion","usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0}}`
 	}
 }
 
 func (c *PluginConfig) Validate() error {
-	if err := c.cacheProviderConfig.Validate(); err != nil {
-		return err
+	// if cache provider is configured, validate it
+	if c.cacheProviderConfig.GetProviderType() != "" {
+		if err := c.cacheProviderConfig.Validate(); err != nil {
+			return err
+		}
 	}
-	if err := c.embeddingProviderConfig.Validate(); err != nil {
-		return err
+	if c.embeddingProviderConfig.GetProviderType() != "" {
+		if err := c.embeddingProviderConfig.Validate(); err != nil {
+			return err
+		}
 	}
 	if err := c.vectorProviderConfig.Validate(); err != nil {
 		return err
@@ -77,15 +82,25 @@ func (c *PluginConfig) Validate() error {
 
 func (c *PluginConfig) Complete(log wrapper.Log) error {
 	var err error
-	c.embeddingProvider, err = embedding.CreateProvider(c.embeddingProviderConfig)
-	if err != nil {
-		return err
+	if c.embeddingProviderConfig.GetProviderType() != "" {
+		c.embeddingProvider, err = embedding.CreateProvider(c.embeddingProviderConfig)
+		if err != nil {
+			return err
+		}
+	} else {
+		log.Info("embedding provider is not configured")
+		c.embeddingProvider = nil
 	}
-	c.vectorProvider, err = vector.CreateProvider(c.vectorProviderConfig)
-	if err != nil {
-		return err
+	if c.cacheProviderConfig.GetProviderType() != "" {
+		c.cacheProvider, err = cache.CreateProvider(c.cacheProviderConfig)
+		if err != nil {
+			return err
+		}
+	} else {
+		log.Info("cache provider is not configured")
+		c.cacheProvider = nil
 	}
-	c.cacheProvider, err = cache.CreateProvider(c.cacheProviderConfig)
+	c.vectorProvider, err = vector.CreateProvider(c.vectorProviderConfig)
 	if err != nil {
 		return err
 	}

plugins/wasm-go/extensions/ai-cache/core.go

@@ -1,13 +1,12 @@
 package main
 
 import (
-	"encoding/json"
+	"errors"
 	"fmt"
 
 	"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-cache/config"
 	"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-cache/vector"
 	"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
-	"github.com/go-errors/errors"
 	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
 	"github.com/tidwall/resp"
 )
@@ -50,29 +49,32 @@ func handleCacheResponse(key string, response resp.Value, ctx wrapper.HttpContex
 	if useSimilaritySearch {
 		if err := performSimilaritySearch(key, ctx, config, log, key, stream); err != nil {
 			log.Errorf("Failed to perform similarity search for key: %s, error: %v", key, err)
+			proxywasm.ResumeHttpRequest()
 		}
+	} else {
+		proxywasm.ResumeHttpRequest()
 	}
-	proxywasm.ResumeHttpRequest()
 }
 
 // processCacheHit handles a successful cache hit.
 func processCacheHit(key string, response string, stream bool, ctx wrapper.HttpContext, config config.PluginConfig, log wrapper.Log) {
-	escapedResponse, err := json.Marshal(response)
-	log.Debugf("Cached response for key %s: %s", key, escapedResponse)
-
-	if err != nil {
-		handleInternalError(err, "Failed to marshal cached response", log)
-		return
+	if stream {
+		log.Debug("streaming response is not supported for cache hit yet")
+		stream = false
 	}
+	// escapedResponse, err := json.Marshal(response)
+	// log.Debugf("Cached response for key %s: %s", key, escapedResponse)
 
-	ctx.SetContext(CACHE_KEY_CONTEXT_KEY, nil)
+	// if err != nil {
+	// 	handleInternalError(err, "Failed to marshal cached response", log)
+	// 	return
+	// }
+	log.Debugf("Cached response for key %s: %s", key, response)
 
-	contentType := "application/json; charset=utf-8"
-	if stream {
-		contentType = "text/event-stream; charset=utf-8"
-	}
+	ctx.SetContext(CACHE_KEY_CONTEXT_KEY, nil)
 
-	proxywasm.SendHttpResponse(200, [][2]string{{"content-type", contentType}}, []byte(fmt.Sprintf(config.ResponseTemplate, escapedResponse)), -1)
+	// proxywasm.SendHttpResponseWithDetail(200, "ai-cache.hit", [][2]string{{"content-type", "application/json; charset=utf-8"}}, []byte(fmt.Sprintf(config.ResponseTemplate, escapedResponse)), -1)
+	proxywasm.SendHttpResponseWithDetail(200, "ai-cache.hit", [][2]string{{"content-type", "application/json; charset=utf-8"}}, []byte(fmt.Sprintf(config.ResponseTemplate, response)), -1)
 }
 
 // performSimilaritySearch determines the appropriate similarity search method to use.
@@ -149,7 +151,7 @@ func handleQueryResults(key string, results []vector.QueryResult, ctx wrapper.Ht
 	}
 
 	mostSimilarData := results[0]
-	log.Debugf("Most similar key found: %s with score: %f", mostSimilarData.Text, mostSimilarData.Score)
+	log.Debugf("For key: %s, the most similar key found: %s with score: %f", key, mostSimilarData.Text, mostSimilarData.Score)
 
 	simThresholdProvider, ok := config.GetVectorProvider().(vector.SimilarityThresholdProvider)
 	if !ok {
@@ -165,7 +167,11 @@ func handleQueryResults(key string, results []vector.QueryResult, ctx wrapper.Ht
 			processCacheHit(key, mostSimilarData.Answer, stream, ctx, config, log)
 		} else {
 			// otherwise, continue to check cache for the most similar key
-			CheckCacheForKey(mostSimilarData.Text, ctx, config, log, stream, false)
+			err = CheckCacheForKey(mostSimilarData.Text, ctx, config, log, stream, false)
+			if err != nil {
+				log.Errorf("check cache for key: %s failed, error: %v", mostSimilarData.Text, err)
+				proxywasm.ResumeHttpRequest()
+			}
 		}
 	} else {
 		log.Infof("Score too high for key: %s with score: %f above threshold", mostSimilarData.Text, mostSimilarData.Score)
@@ -220,9 +226,9 @@ func uploadEmbeddingAndAnswer(ctx wrapper.HttpContext, config config.PluginConfi
 	}
 
 	// Attempt to upload answer embedding first
-	if ansEmbUploader, ok := activeVectorProvider.(vector.AnswerEmbeddingUploader); ok {
+	if ansEmbUploader, ok := activeVectorProvider.(vector.AnswerAndEmbeddingUploader); ok {
 		log.Infof("[onHttpResponseBody] uploading answer embedding for key: %s", key)
-		err := ansEmbUploader.UploadAnswerEmbedding(key, emb, value, ctx, log, nil)
+		err := ansEmbUploader.UploadAnswerAndEmbedding(key, emb, value, ctx, log, nil)
 		if err != nil {
 			log.Warnf("[onHttpResponseBody] failed to upload answer embedding for key: %s, error: %v", key, err)
 		} else {

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

@@ -3,17 +3,18 @@ package embedding
 import (
 	"encoding/json"
 	"errors"
+	"fmt"
 	"net/http"
 	"strconv"
 
 	"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
 )
 
 const (
-	DOMAIN             = "dashscope.aliyuncs.com"
-	PORT               = 443
-	DEFAULT_MODEL_NAME = "text-embedding-v1"
-	ENDPOINT           = "/api/v1/services/embeddings/text-embedding/text-embedding"
+	DASHSCOPE_DOMAIN             = "dashscope.aliyuncs.com"
+	DASHSCOPE_PORT               = 443
+	DASHSCOPE_DEFAULT_MODEL_NAME = "text-embedding-v1"
+	DASHSCOPE_ENDPOINT           = "/api/v1/services/embeddings/text-embedding/text-embedding"
 )
 
 type dashScopeProviderInitializer struct {
@@ -28,10 +29,10 @@ func (d *dashScopeProviderInitializer) ValidateConfig(config ProviderConfig) err
 
 func (d *dashScopeProviderInitializer) CreateProvider(c ProviderConfig) (Provider, error) {
 	if c.servicePort == 0 {
-		c.servicePort = PORT
+		c.servicePort = DASHSCOPE_PORT
 	}
 	if c.serviceDomain == "" {
-		c.serviceDomain = DOMAIN
+		c.serviceDomain = DASHSCOPE_DOMAIN
 	}
 	return &DSProvider{
 		config: c,
@@ -95,7 +96,7 @@ func (d *DSProvider) constructParameters(texts []string, log wrapper.Log) (strin
 	model := d.config.model
 
 	if model == "" {
-		model = DEFAULT_MODEL_NAME
+		model = DASHSCOPE_DEFAULT_MODEL_NAME
 	}
 	data := EmbeddingRequest{
 		Model: model,
@@ -124,7 +125,7 @@ func (d *DSProvider) constructParameters(texts []string, log wrapper.Log) (strin
 		{"Content-Type", "application/json"},
 	}
 
-	return ENDPOINT, headers, requestBody, err
+	return DASHSCOPE_ENDPOINT, headers, requestBody, err
 }
 
 type Result struct {
@@ -168,7 +169,7 @@ func (d *DSProvider) GetEmbedding(
 
 			resp, err = d.parseTextEmbedding(responseBody)
 			if err != nil {
-				err = errors.New("failed to parse response: " + err.Error())
+				err = fmt.Errorf("failed to parse response: %v", err)
 				callback(nil, err)
 				return
 			}

plugins/wasm-go/extensions/ai-cache/main.go

@@ -118,16 +118,21 @@ func onHttpResponseHeaders(ctx wrapper.HttpContext, config config.PluginConfig,
 }
 
 func onHttpResponseBody(ctx wrapper.HttpContext, config config.PluginConfig, chunk []byte, isLastChunk bool, log wrapper.Log) []byte {
-	log.Debugf("[onHttpResponseBody] chunk: %s", string(chunk))
+	log.Debugf("[onHttpResponseBody] escaped chunk: %q", string(chunk))
 	log.Debugf("[onHttpResponseBody] isLastChunk: %v", isLastChunk)
 
+	// if strings.HasSuffix(string(chunk), "[DONE] \n\n") {
+	// 	isLastChunk = true
+	// }
+
 	// If the context contains TOOL_CALLS_CONTEXT_KEY, bypass caching
 	if ctx.GetContext(TOOL_CALLS_CONTEXT_KEY) != nil {
 		return chunk
 	}
 
-	keyI := ctx.GetContext(CACHE_KEY_CONTEXT_KEY)
-	if keyI == nil {
+	key := ctx.GetContext(CACHE_KEY_CONTEXT_KEY)
+	if key == nil {
+		log.Debug("[onHttpResponseBody] key is nil, bypass caching")
 		return chunk
 	}
 
@@ -137,17 +142,24 @@ func onHttpResponseBody(ctx wrapper.HttpContext, config config.PluginConfig, chu
 	}
 
 	// Handle last chunk
-	value, err := processFinalChunk(ctx, config, chunk, log)
+	var value string
+	var err error
+
+	if len(chunk) > 0 {
+		value, err = processNonEmptyChunk(ctx, config, chunk, log)
+	} else {
+		value, err = processEmptyChunk(ctx, config, chunk, log)
+	}
+
 	if err != nil {
-		log.Warnf("[onHttpResponseBody] failed to process final chunk: %v", err)
+		log.Warnf("[onHttpResponseBody] failed to process chunk: %v", err)
 		return chunk
 	}
-
 	// Cache the final value
-	cacheResponse(ctx, config, keyI.(string), value, log)
+	cacheResponse(ctx, config, key.(string), value, log)
 
 	// Handle embedding upload if available
-	uploadEmbeddingAndAnswer(ctx, config, keyI.(string), value, log)
+	uploadEmbeddingAndAnswer(ctx, config, key.(string), value, log)
 
 	return chunk
 }

plugins/wasm-go/extensions/ai-cache/util.go

@@ -99,7 +99,7 @@ func savePartialMessage(ctx wrapper.HttpContext, partialMessage []byte, messages
 }
 
 // Processes the final chunk and returns the parsed value or an error
-func processFinalChunk(ctx wrapper.HttpContext, config config.PluginConfig, chunk []byte, log wrapper.Log) (string, error) {
+func processNonEmptyChunk(ctx wrapper.HttpContext, config config.PluginConfig, chunk []byte, log wrapper.Log) (string, error) {
 	stream := ctx.GetContext(STREAM_CONTEXT_KEY)
 	var value string
 
@@ -122,6 +122,15 @@ func processFinalChunk(ctx wrapper.HttpContext, config config.PluginConfig, chun
 	return value, nil
 }
 
+func processEmptyChunk(ctx wrapper.HttpContext, config config.PluginConfig, chunk []byte, log wrapper.Log) (string, error) {
+	tempContentI := ctx.GetContext(CACHE_CONTENT_CONTEXT_KEY)
+	if tempContentI == nil {
+		return string(chunk), nil
+	}
+	value := tempContentI.(string)
+	return value, nil
+}
+
 // Appends the final body chunk to the existing body content
 func appendFinalBody(ctx wrapper.HttpContext, chunk []byte) []byte {
 	tempContentI := ctx.GetContext(CACHE_CONTENT_CONTEXT_KEY)

plugins/wasm-go/extensions/ai-cache/vector/dashvector.go

@@ -157,6 +157,13 @@ func (d *DvProvider) QueryEmbedding(
 	return err
 }
 
+func checkField(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 {
@@ -171,9 +178,10 @@ func (d *DvProvider) ParseQueryResponse(responseBody []byte, ctx wrapper.HttpCon
 
 	for _, output := range resp.Output {
 		result := QueryResult{
-			Text:      output.Fields["query"].(string),
+			Text:      checkField(output.Fields, "query"),
 			Embedding: output.Vector,
 			Score:     output.Score,
+			Answer:    checkField(output.Fields, "answer"),
 		}
 		results = append(results, result)
 	}
@@ -234,7 +242,11 @@ func (d *DvProvider) UploadEmbedding(queryString string, queryEmb []float64, ctx
 	return err
 }
 
-func (d *DvProvider) UploadAnswerEmbedding(queryString string, queryEmb []float64, queryAnswer string, ctx wrapper.HttpContext, log wrapper.Log, callback func(ctx wrapper.HttpContext, log wrapper.Log, err error)) error {
+func (d *DvProvider) GetSimilarityThreshold() float64 {
+	return threshold
+}
+
+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

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

@@ -53,8 +53,8 @@ type EmbeddingUploader interface {
 		callback func(ctx wrapper.HttpContext, log wrapper.Log, err error)) error
 }
 
-type AnswerEmbeddingUploader interface {
-	UploadAnswerEmbedding(
+type AnswerAndEmbeddingUploader interface {
+	UploadAnswerAndEmbedding(
 		queryString string,
 		queryEmb []float64,
 		answer string,
@@ -118,6 +118,10 @@ type ProviderConfig struct {
 	// ChromaTimeout uint32 `require:"false" yaml:"ChromaTimeout" json:"ChromaTimeout"`
 }
 
+func (c *ProviderConfig) GetProviderType() string {
+	return c.typ
+}
+
 func (c *ProviderConfig) FromJson(json gjson.Result) {
 	c.typ = json.Get("type").String()
 	// DashVector