ai-search support quark (#1811)

This commit is contained in:
rinfx
2025-02-26 18:42:22 +08:00
committed by GitHub
parent 1787553294
commit 3eda7def89
6 changed files with 250 additions and 7 deletions

View File

@@ -34,6 +34,7 @@ import (
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine/bing"
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine/elasticsearch"
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine/google"
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine/quark"
)
type SearchRewrite struct {
@@ -166,6 +167,13 @@ func parseConfig(json gjson.Result, config *Config, log wrapper.Log) error {
}
config.engine = append(config.engine, searchEngine)
privateExists = true
case "quark":
searchEngine, err := quark.NewQuarkSearch(&e)
if err != nil {
return fmt.Errorf("elasticsearch search engine init failed:%s", err)
}
config.engine = append(config.engine, searchEngine)
internetExists = true
default:
return fmt.Errorf("unkown search engine:%s", e.Get("type").String())
}
@@ -541,7 +549,8 @@ func setReferencesToFirstMessage(ctx wrapper.HttpContext, chunk []byte, referenc
if len(messages) > 1 {
firstMessage := messages[0]
log.Debugf("first message: %s", firstMessage)
firstMessage = strings.TrimPrefix(firstMessage, "data: ")
firstMessage = strings.TrimPrefix(firstMessage, "data:")
firstMessage = strings.TrimPrefix(firstMessage, " ")
firstMessage = strings.TrimSuffix(firstMessage, "\n")
deltaContent := gjson.Get(firstMessage, "choices.0.delta.content")
modifiedMessage, err := sjson.Set(firstMessage, "choices.0.delta.content", fmt.Sprintf("%s\n\n%s", references, deltaContent))