feat: support allowScript

promptc · May 17, 2023 · d6c3719 · d6c3719
1 parent 584a89c
commit d6c3719
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/prompt/file.go b/prompt/file.go
@@ -104,6 +104,10 @@ func ParseBasicPrompt(content string) *PromptC {
 }
 
 func (f *PromptC) Compile(vars map[string]string) *CompiledPromptC {
+	return f.CompileWithOption(vars, true)
+}
+
+func (f *PromptC) CompileWithOption(vars map[string]string, allowScript bool) *CompiledPromptC {
 	//varMap := make(map[string]string)
 	fileFatal := false
 	compiledVars := make(map[string]string)
@@ -139,7 +143,7 @@ func (f *PromptC) Compile(vars map[string]string) *CompiledPromptC {
 			result = append(result, compiled...)
 			continue
 		}
-		compiled, exp, fatal := p.Compile(compiledVars)
+		compiled, exp, fatal := p.CompileWithOption(compiledVars, allowScript)
 		if len(exp) > 0 {
 			errs = append(errs, exp...)
 		}

diff --git a/prompt/parsed_block.go b/prompt/parsed_block.go
@@ -52,8 +52,10 @@ func (p *ParsedBlock) ToMap() map[string]any {
 	m["extra"] = p.Extra
 	return m
 }
-
 func (p *ParsedBlock) Compile(varMap map[string]string) (compiled string, exceptions []error, fatal bool) {
+	return p.CompileWithOption(varMap, true)
+}
+func (p *ParsedBlock) CompileWithOption(varMap map[string]string, allowScript bool) (compiled string, exceptions []error, fatal bool) {
 	fatal = false
 	sb := strings.Builder{}
 	vm := otto.New()
@@ -91,6 +93,11 @@ func (p *ParsedBlock) Compile(varMap map[string]string) (compiled string, except
 		case BlockTokenKindReservedQuota:
 			sb.WriteString("'''")
 		case BlockTokenKindScript:
+			if !allowScript {
+				exceptions = append(exceptions, errors.New("script is not allowed"))
+				fatal = true
+				return
+			}
 			script := token.Text
 			easyMod := false
 			if strings.HasPrefix(script, "E") {