update template decorator (#1142)

2026-03-10 11:40:49 +08:00 · 2024-07-22 17:04:55 +08:00
parent ef31e09310
commit 8c48fcb423
3 changed files with 60 additions and 90 deletions
--- a/plugins/wasm-go/extensions/ai-prompt-decorator/README.md
+++ b/plugins/wasm-go/extensions/ai-prompt-decorator/README.md
@@ -1,18 +1,12 @@
 # 简介
-AI提示词修饰插件，通过在与大模型发起的请求前后插入指定信息来调整大模型的输出。
+AI提示词装饰器插件，支持在LLM的请求前后插入prompt。

 # 配置说明
-| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
-|----------------|-----------------|------|-----|----------------------------------|
-| `decorators` | array of object | 必填 | - | 修饰设置 |
-
-template object 配置说明：

 | 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
 |----------------|-----------------|------|-----|----------------------------------|
-| `name` | string | 必填 | - | 修饰名称 |
-| `decorator.prepend` | array of message object | 必填 | - | 在初始输入之前插入的语句 |
-| `decorator.append` | array of message object | 必填 | - | 在初始输入之后插入的语句 |
+| `prepend` | array of message object | optional | - | 在初始输入之前插入的语句 |
+| `append` | array of message object | optional | - | 在初始输入之后插入的语句 |

 message object 配置说明：

@@ -26,57 +20,50 @@ message object 配置说明：
 配置示例如下：

 ```yaml
-decorators:
- name: "hangzhou-guide"
-  decorator:
-    prepend:
-    - role: system
-      content: "You will always respond in the Chinese language."
-    - role: user
-      content: "Assume you are from Hangzhou."
-    append:
-    - role: user
-      content: "Don't introduce Hangzhou's food."
+prepend:
+- role: system
+  content: "请使用英语回答问题"
+append:
+- role: user
+  content: "每次回答完问题，尝试进行反问"
 ```

 使用以上配置发起请求：

 ```bash
-{
+curl http://localhost/test \
+-H "content-type: application/json" \
+-d '{
  "model": "gpt-3.5-turbo",
  "messages": [
    {
      "role": "user",
-      "content": "Please introduce your home."
+      "content": "你是谁？"
    }
  ]
 }
 ```

-响应如下：
+经过插件处理后，实际请求为：

-```
-{
-  "id": "chatcmpl-9UYwQlEg6GwAswEZBDYXl41RU4gab",
-  "object": "chat.completion",
-  "created": 1717071182,
-  "model": "gpt-3.5-turbo-0125",
-  "choices": [
+```bash
+curl http://localhost/test \
+-H "content-type: application/json" \
+-d '{
+  "model": "gpt-3.5-turbo",
+  "messages": [
    {
-      "index": 0,
-      "message": {
-        "role": "assistant",
-        "content": "杭州是一个美丽的城市，有着悠久的历史和富有特色的文化。这里风景优美，有西湖、雷峰塔等著名景点，吸引着许多游客前来观光。杭州人民热情好客，城市宁静安逸，是一个适合居住和旅游的地方。"
-      },
-      "logprobs": null,
-      "finish_reason": "stop"
+      "role": "system",
+      "content": "请使用英语回答问题"
+    },
+    {
+      "role": "user",
+      "content": "你是谁？"
+    },
+    {
+      "role": "user",
+      "content": "每次回答完问题，尝试进行反问"
    }
-  ],
-  "usage": {
-    "prompt_tokens": 49,
-    "completion_tokens": 117,
-    "total_tokens": 166
-  },
-  "system_fingerprint": null
+  ]
 }
 ```
--- a/plugins/wasm-go/extensions/ai-prompt-decorator/go.sum
+++ b/plugins/wasm-go/extensions/ai-prompt-decorator/go.sum
@@ -1,12 +1,9 @@
-github.com/alibaba/higress/plugins/wasm-go v1.3.5 h1:VOLL3m442IHCSu8mR5AZ4sc6LVT9X0w1hdqDI7oB9jY=
-github.com/alibaba/higress/plugins/wasm-go v1.3.5/go.mod h1:kr3V9Ntbspj1eSrX8rgjBsdMXkGupYEf+LM72caGPQc=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 h1:IHDghbGQ2DTIXHBHxWfqCYQW1fKjyJ/I7W1pMyUDeEA=
 github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520/go.mod h1:Nz8ORLaFiLWotg6GeKlJMhv8cci8mM43uEnLA5t8iew=
-github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240226064518-b3dc4646a35a h1:luYRvxLTE1xYxrXYj7nmjd1U0HHh8pUPiKfdZ0MhCGE=
-github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240226064518-b3dc4646a35a/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
+github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f h1:ZIiIBRvIw62gA5MJhuwp1+2wWbqL9IGElQ499rUsYYg=
 github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
 github.com/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo=
 github.com/magefile/mage v1.14.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
--- a/plugins/wasm-go/extensions/ai-prompt-decorator/main.go
+++ b/plugins/wasm-go/extensions/ai-prompt-decorator/main.go
@@ -1,8 +1,7 @@
 package main

 import (
-	"errors"
-	"strings"
+	"encoding/json"

 	"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
 	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
@@ -20,66 +19,53 @@ func main() {
 	)
 }

+type Message struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
 type AIPromptDecoratorConfig struct {
-	decorators map[string]string
+	Prepend []Message `json:"prepend"`
+	Append  []Message `json:"append"`
 }

-func removeBrackets(raw string) (string, error) {
-	startIndex := strings.Index(raw, "{")
-	endIndex := strings.LastIndex(raw, "}")
-	if startIndex == -1 || endIndex == -1 {
-		return raw, errors.New("message format is wrong!")
-	} else {
-		return raw[startIndex : endIndex+1], nil
-	}
-}
-
-func parseConfig(json gjson.Result, config *AIPromptDecoratorConfig, log wrapper.Log) error {
-	config.decorators = make(map[string]string)
-	for _, v := range json.Get("decorators").Array() {
-		config.decorators[v.Get("name").String()] = v.Get("decorator").Raw
-		// log.Info(v.Get("decorator").Raw)
-	}
-	return nil
+func parseConfig(jsonConfig gjson.Result, config *AIPromptDecoratorConfig, log wrapper.Log) error {
+	return json.Unmarshal([]byte(jsonConfig.Raw), config)
 }

 func onHttpRequestHeaders(ctx wrapper.HttpContext, config AIPromptDecoratorConfig, log wrapper.Log) types.Action {
-	decorator, _ := proxywasm.GetHttpRequestHeader("decorator")
-	if decorator == "" {
-		ctx.DontReadRequestBody()
-		return types.ActionContinue
-	}
-	ctx.SetContext("decorator", decorator)
-	proxywasm.RemoveHttpRequestHeader("decorator")
 	proxywasm.RemoveHttpRequestHeader("content-length")
 	return types.ActionContinue
 }

 func onHttpRequestBody(ctx wrapper.HttpContext, config AIPromptDecoratorConfig, body []byte, log wrapper.Log) types.Action {
-	decoratorName := ctx.GetContext("decorator").(string)
-	decorator := config.decorators[decoratorName]
-
 	messageJson := `{"messages":[]}`

-	prependMessage := gjson.Get(decorator, "prepend")
-	if prependMessage.Exists() {
-		for _, entry := range prependMessage.Array() {
-			messageJson, _ = sjson.SetRaw(messageJson, "messages.-1", entry.Raw)
+	for _, entry := range config.Prepend {
+		msg, err := json.Marshal(entry)
+		if err != nil {
+			log.Errorf("Failed to add prepend message, error: %v", err)
+			return types.ActionContinue
 		}
+		messageJson, _ = sjson.SetRaw(messageJson, "messages.-1", string(msg))
 	}

 	rawMessage := gjson.GetBytes(body, "messages")
-	if rawMessage.Exists() {
-		for _, entry := range rawMessage.Array() {
-			messageJson, _ = sjson.SetRaw(messageJson, "messages.-1", entry.Raw)
-		}
+	if !rawMessage.Exists() {
+		log.Errorf("Cannot find messages field in request body")
+		return types.ActionContinue
+	}
+	for _, entry := range rawMessage.Array() {
+		messageJson, _ = sjson.SetRaw(messageJson, "messages.-1", entry.Raw)
 	}

-	appendMessage := gjson.Get(decorator, "append")
-	if appendMessage.Exists() {
-		for _, entry := range appendMessage.Array() {
-			messageJson, _ = sjson.SetRaw(messageJson, "messages.-1", entry.Raw)
+	for _, entry := range config.Append {
+		msg, err := json.Marshal(entry)
+		if err != nil {
+			log.Errorf("Failed to add prepend message, error: %v", err)
+			return types.ActionContinue
 		}
+		messageJson, _ = sjson.SetRaw(messageJson, "messages.-1", string(msg))
 	}

 	newbody, err := sjson.SetRaw(string(body), "messages", gjson.Get(messageJson, "messages").Raw)