ADD json escape checker for template.

illacloud · Jan 10, 2024 · 2b3a670 · 2b3a670
1 parent 5573933
commit 2b3a670
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/src/utils/parser/template/parser.go b/src/utils/parser/template/parser.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"strconv"
+	"strings"
 )
 
 type JSONNumberConvertor struct {
@@ -128,6 +129,15 @@ func ExtractVariableNameConst(template string) []string {
 }
 
 func AssembleTemplateWithVariable(template string, variableLT map[string]interface{}) (string, error) {
+	// check if template is json
+	templateIsJSON := false
+	var templateInJSONObject interface{}
+	errInUnmarshalTemplate := json.Unmarshal([]byte(template), &templateInJSONObject)
+	if errInUnmarshalTemplate == nil {
+		templateIsJSON = true
+	}
+
+	// process start
 	processesPrompt := ""
 	variable := ""
 	escapedBracketWithVariable := ""
@@ -171,7 +181,12 @@ func AssembleTemplateWithVariable(template string, variableLT map[string]interfa
 			dataInFloat64 := data.(float64)
 			return ExportFloat64ToNumberInString(dataInFloat64), nil
 		case string:
-			return data.(string), nil
+			finalStr := data.(string)
+			if templateIsJSON {
+				finalStr = strings.Replace(finalStr, "\"", "\\\"", -1)
+				finalStr = strings.Replace(finalStr, "\n", "\\n", -1)
+			}
+			return finalStr, nil
 		case bool:
 			dataInBool := data.(bool)
 			if dataInBool {

diff --git a/src/utils/parser/template/parser_test.go b/src/utils/parser/template/parser_test.go
@@ -63,3 +63,14 @@ func TestAssembleTemplateWithVariable_case3_WarppedString(t *testing.T) {
 	assert.Equal(t, `{"mode": "sql", "query": "select * \nfrom users\njoin orders\non users.id = orders.id\nwhere \"BIG APPLE\" or lower(users.name) like '%[A\nAA]%'"}`, finalTemplate, "it should be equal ")
 
 }
+
+func TestAssembleTemplateWithVariable_case4_EscapeCase(t *testing.T) {
+	actionTemplate := `{"msg_type":"post","content":{"post":{"zh_cn":{"title":"☀️ 早上好","content":[[{"tag":"text","text":"{{restapi3.data[0].choices[0].message.content}}"}]]}}}}`
+	dataLT := map[string]interface{}{
+		"restapi3.data[0].choices[0].message.content": "\"Success is not final, failure is not fatal: It is the courage to continue that counts.\" - Winston Churchill\n（成功不是终点，失败也不是致命的：真正重要的是勇于继续前进。）",
+	}
+	finalTemplate, errInAssemble := AssembleTemplateWithVariable(actionTemplate, dataLT)
+	assert.Nil(t, errInAssemble)
+	assert.Equal(t, `{"mode": "sql", "query": "select * \nfrom users\njoin orders\non users.id = orders.id\nwhere \"BIG APPLE\" or lower(users.name) like '%[A\nAA]%'"}`, finalTemplate, "it should be equal ")
+
+}