From 15d9f76ff90d01aeffbf84d252ad975ff9372777 Mon Sep 17 00:00:00 2001 From: rinfx <893383980@qq.com> Date: Mon, 17 Jun 2024 15:31:34 +0800 Subject: [PATCH] add plugin: ai-prompt-decorator (#1021) --- .../extensions/ai-prompt-decorator/.gitignore | 3 + .../extensions/ai-prompt-decorator/README.md | 82 ++++++++++++++++ .../extensions/ai-prompt-decorator/go.mod | 19 ++++ .../extensions/ai-prompt-decorator/go.sum | 25 +++++ .../extensions/ai-prompt-decorator/main.go | 94 +++++++++++++++++++ 5 files changed, 223 insertions(+) create mode 100644 plugins/wasm-go/extensions/ai-prompt-decorator/.gitignore create mode 100644 plugins/wasm-go/extensions/ai-prompt-decorator/README.md create mode 100644 plugins/wasm-go/extensions/ai-prompt-decorator/go.mod create mode 100644 plugins/wasm-go/extensions/ai-prompt-decorator/go.sum create mode 100644 plugins/wasm-go/extensions/ai-prompt-decorator/main.go diff --git a/plugins/wasm-go/extensions/ai-prompt-decorator/.gitignore b/plugins/wasm-go/extensions/ai-prompt-decorator/.gitignore new file mode 100644 index 000000000..c9f9dc52b --- /dev/null +++ b/plugins/wasm-go/extensions/ai-prompt-decorator/.gitignore @@ -0,0 +1,3 @@ +config.yaml +main.wasm +tmp/ \ No newline at end of file diff --git a/plugins/wasm-go/extensions/ai-prompt-decorator/README.md b/plugins/wasm-go/extensions/ai-prompt-decorator/README.md new file mode 100644 index 000000000..9522de1c5 --- /dev/null +++ b/plugins/wasm-go/extensions/ai-prompt-decorator/README.md @@ -0,0 +1,82 @@ +# 简介 +AI提示词修饰插件，通过在与大模型发起的请求前后插入指定信息来调整大模型的输出。 + +# 配置说明 +| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 | +|----------------|-----------------|------|-----|----------------------------------| +| `decorators` | array of object | 必填 | - | 修饰设置 | + +template object 配置说明： + +| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 | +|----------------|-----------------|------|-----|----------------------------------| +| `name` | string | 必填 | - | 修饰名称 | +| `decorator.prepend` | array of message object | 必填 | - | 在初始输入之前插入的语句 | +| `decorator.append` | array of message object | 必填 | - | 在初始输入之后插入的语句 | + +message object 配置说明： + +| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 | +|----------------|-----------------|------|-----|----------------------------------| +| `role` | string | 必填 | - | 角色 | +| `content` | string | 必填 | - | 消息 | + +# 示例 + +配置示例如下： + +```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." +``` + +使用以上配置发起请求： + +```bash +{ + "model": "gpt-3.5-turbo", + "messages": [ + { + "role": "user", + "content": "Please introduce your home." + } + ] +} +``` + +响应如下： + +``` +{ + "id": "chatcmpl-9UYwQlEg6GwAswEZBDYXl41RU4gab", + "object": "chat.completion", + "created": 1717071182, + "model": "gpt-3.5-turbo-0125", + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "杭州是一个美丽的城市，有着悠久的历史和富有特色的文化。这里风景优美，有西湖、雷峰塔等著名景点，吸引着许多游客前来观光。杭州人民热情好客，城市宁静安逸，是一个适合居住和旅游的地方。" + }, + "logprobs": null, + "finish_reason": "stop" + } + ], + "usage": { + "prompt_tokens": 49, + "completion_tokens": 117, + "total_tokens": 166 + }, + "system_fingerprint": null +} +``` \ No newline at end of file diff --git a/plugins/wasm-go/extensions/ai-prompt-decorator/go.mod b/plugins/wasm-go/extensions/ai-prompt-decorator/go.mod new file mode 100644 index 000000000..8aea59ffb --- /dev/null +++ b/plugins/wasm-go/extensions/ai-prompt-decorator/go.mod @@ -0,0 +1,19 @@ +module ai-prompt-decorator + +go 1.18 + +require ( + github.com/alibaba/higress/plugins/wasm-go v1.3.5 + github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240226064518-b3dc4646a35a + github.com/tidwall/gjson v1.14.3 +) + +require ( + github.com/google/uuid v1.3.0 // indirect + github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 // indirect + github.com/magefile/mage v1.14.0 // indirect + github.com/tidwall/match v1.1.1 // indirect + github.com/tidwall/pretty v1.2.0 // indirect + github.com/tidwall/resp v0.1.1 // indirect + github.com/tidwall/sjson v1.2.5 +) diff --git a/plugins/wasm-go/extensions/ai-prompt-decorator/go.sum b/plugins/wasm-go/extensions/ai-prompt-decorator/go.sum new file mode 100644 index 000000000..94fcae90c --- /dev/null +++ b/plugins/wasm-go/extensions/ai-prompt-decorator/go.sum @@ -0,0 +1,25 @@ +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/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo= +github.com/magefile/mage v1.14.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk= +github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/gjson v1.14.3 h1:9jvXn7olKEHU1S9vwoMGliaT8jq1vJ7IH/n9zD9Dnlw= +github.com/tidwall/gjson v1.14.3/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= +github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= +github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM= +github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs= +github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= +github.com/tidwall/resp v0.1.1 h1:Ly20wkhqKTmDUPlyM1S7pWo5kk0tDu8OoC/vFArXmwE= +github.com/tidwall/resp v0.1.1/go.mod h1:3/FrruOBAxPTPtundW0VXgmsQ4ZBA0Aw714lVYgwFa0= +github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= +github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= diff --git a/plugins/wasm-go/extensions/ai-prompt-decorator/main.go b/plugins/wasm-go/extensions/ai-prompt-decorator/main.go new file mode 100644 index 000000000..6fad232d0 --- /dev/null +++ b/plugins/wasm-go/extensions/ai-prompt-decorator/main.go @@ -0,0 +1,94 @@ +package main + +import ( + "errors" + "strings" + + "github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper" + "github.com/higress-group/proxy-wasm-go-sdk/proxywasm" + "github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types" + "github.com/tidwall/gjson" + "github.com/tidwall/sjson" +) + +func main() { + wrapper.SetCtx( + "ai-prompt-decorator", + wrapper.ParseConfigBy(parseConfig), + wrapper.ProcessRequestHeadersBy(onHttpRequestHeaders), + wrapper.ProcessRequestBodyBy(onHttpRequestBody), + ) +} + +type AIPromptDecoratorConfig struct { + decorators map[string]string +} + +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 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) + } + } + + rawMessage := gjson.GetBytes(body, "messages") + if rawMessage.Exists() { + 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) + } + } + + newbody, err := sjson.SetRaw(string(body), "messages", gjson.Get(messageJson, "messages").Raw) + if err != nil { + log.Error("modify body failed") + } + if err = proxywasm.ReplaceHttpRequestBody([]byte(newbody)); err != nil { + log.Error("rewrite body failed") + } + + return types.ActionContinue +}