Prompt generation for AI-assisted authoring based on a CEL environment (#1160)

* AI authoring prompt generation

Author: TristonianJones

cel/BUILD.bazel

@@ -18,8 +18,10 @@ go_library(
         "optimizer.go",
         "options.go",
         "program.go",
+        "prompt.go",
         "validator.go",
     ],
+    embedsrcs = ["//cel/templates"],
     importpath = "github.com/google/cel-go/cel",
     visibility = ["//visibility:public"],
     deps = [
@@ -65,6 +67,7 @@ go_test(
         "io_test.go",
         "inlining_test.go",
         "optimizer_test.go",
+        "prompt_test.go",
         "validator_test.go",
     ],
     data = [
@@ -73,6 +76,9 @@ go_test(
     embed = [
         ":go_default_library",
     ],
+    embedsrcs = [
+        "//cel/testdata:prompts",
+    ],
     deps = [
         "//common/operators:go_default_library",
         "//common/overloads:go_default_library",

cel/decls.go

@@ -148,6 +148,14 @@ func Variable(name string, t *Type) EnvOption {
 	}
 }
 
+// VariableWithDoc creates an instance of a variable declaration with a variable name, type, and doc string.
+func VariableWithDoc(name string, t *Type, doc string) EnvOption {
+	return func(e *Env) (*Env, error) {
+		e.variables = append(e.variables, decls.NewVariableWithDoc(name, t, doc))
+		return e, nil
+	}
+}
+
 // VariableDecls configures a set of fully defined cel.VariableDecl instances in the environment.
 func VariableDecls(vars ...*decls.VariableDecl) EnvOption {
 	return func(e *Env) (*Env, error) {
@@ -239,6 +247,13 @@ func FunctionDecls(funcs ...*decls.FunctionDecl) EnvOption {
 // FunctionOpt defines a functional  option for configuring a function declaration.
 type FunctionOpt = decls.FunctionOpt
 
+// FunctionDocs provides a general usage documentation for the function.
+//
+// Use OverloadExamples to provide example usage instructions for specific overloads.
+func FunctionDocs(docs ...string) FunctionOpt {
+	return decls.FunctionDocs(docs...)
+}
+
 // SingletonUnaryBinding creates a singleton function definition to be used for all function overloads.
 //
 // Note, this approach works well if operand is expected to have a specific trait which it implements,
@@ -312,6 +327,11 @@ func MemberOverload(overloadID string, args []*Type, resultType *Type, opts ...O
 // OverloadOpt is a functional option for configuring a function overload.
 type OverloadOpt = decls.OverloadOpt
 
+// OverloadExamples configures an example of how to invoke the overload.
+func OverloadExamples(docs ...string) OverloadOpt {
+	return decls.OverloadExamples(docs...)
+}
+
 // UnaryBinding provides the implementation of a unary overload. The provided function is protected by a runtime
 // type-guard which ensures runtime type agreement between the overload signature and runtime argument types.
 func UnaryBinding(binding functions.UnaryOp) OverloadOpt {

cel/env.go

@@ -553,14 +553,27 @@ func (e *Env) HasFunction(functionName string) bool {
 	return ok
 }
 
-// Functions returns map of Functions, keyed by function name, that have been configured in the environment.
+// Functions returns a shallow copy of the Functions, keyed by function name, that have been configured in the environment.
 func (e *Env) Functions() map[string]*decls.FunctionDecl {
-	return e.functions
+	shallowCopy := make(map[string]*decls.FunctionDecl, len(e.functions))
+	for nm, fn := range e.functions {
+		shallowCopy[nm] = fn
+	}
+	return shallowCopy
 }
 
-// Variables returns the set of variables associated with the environment.
+// Variables returns a shallow copy of the variables associated with the environment.
 func (e *Env) Variables() []*decls.VariableDecl {
-	return e.variables[:]
+	shallowCopy := make([]*decls.VariableDecl, len(e.variables))
+	copy(shallowCopy, e.variables)
+	return shallowCopy
+}
+
+// Macros returns a shallow copy of macros associated with the environment.
+func (e *Env) Macros() []Macro {
+	shallowCopy := make([]Macro, len(e.macros))
+	copy(shallowCopy, e.macros)
+	return shallowCopy
 }
 
 // HasValidator returns whether a specific ASTValidator has been configured in the environment.

cel/macro.go

@@ -142,24 +142,38 @@ type MacroExprHelper interface {
 	NewError(exprID int64, message string) *Error
 }
 
+// MacroOpt defines a functional option for configuring macro behavior.
+type MacroOpt = parser.MacroOpt
+
+// MacroDocs configures a list of strings into a multiline description for the macro.
+func MacroDocs(docs ...string) MacroOpt {
+	return parser.MacroDocs(docs...)
+}
+
+// MacroExamples configures a list of examples, either as a string or common.MultilineString,
+// into an example set to be provided with the macro Documentation() call.
+func MacroExamples(examples ...string) MacroOpt {
+	return parser.MacroExamples(examples...)
+}
+
 // GlobalMacro creates a Macro for a global function with the specified arg count.
-func GlobalMacro(function string, argCount int, factory MacroFactory) Macro {
-	return parser.NewGlobalMacro(function, argCount, factory)
+func GlobalMacro(function string, argCount int, factory MacroFactory, opts ...MacroOpt) Macro {
+	return parser.NewGlobalMacro(function, argCount, factory, opts...)
 }
 
 // ReceiverMacro creates a Macro for a receiver function matching the specified arg count.
-func ReceiverMacro(function string, argCount int, factory MacroFactory) Macro {
-	return parser.NewReceiverMacro(function, argCount, factory)
+func ReceiverMacro(function string, argCount int, factory MacroFactory, opts ...MacroOpt) Macro {
+	return parser.NewReceiverMacro(function, argCount, factory, opts...)
 }
 
 // GlobalVarArgMacro creates a Macro for a global function with a variable arg count.
-func GlobalVarArgMacro(function string, factory MacroFactory) Macro {
-	return parser.NewGlobalVarArgMacro(function, factory)
+func GlobalVarArgMacro(function string, factory MacroFactory, opts ...MacroOpt) Macro {
+	return parser.NewGlobalVarArgMacro(function, factory, opts...)
 }
 
 // ReceiverVarArgMacro creates a Macro for a receiver function matching a variable arg count.
-func ReceiverVarArgMacro(function string, factory MacroFactory) Macro {
-	return parser.NewReceiverVarArgMacro(function, factory)
+func ReceiverVarArgMacro(function string, factory MacroFactory, opts ...MacroOpt) Macro {
+	return parser.NewReceiverVarArgMacro(function, factory, opts...)
 }
 
 // NewGlobalMacro creates a Macro for a global function with the specified arg count.

cel/prompt.go

@@ -0,0 +1,155 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cel
+
+import (
+	_ "embed"
+	"sort"
+	"strings"
+	"text/template"
+
+	"github.com/google/cel-go/common"
+	"github.com/google/cel-go/common/operators"
+	"github.com/google/cel-go/common/overloads"
+)
+
+//go:embed templates/authoring.tmpl
+var authoringPrompt string
+
+// AuthoringPrompt creates a prompt template from a CEL environment for the purpose of AI-assisted authoring.
+func AuthoringPrompt(env *Env) (*Prompt, error) {
+	funcMap := template.FuncMap{
+		"split": func(str string) []string { return strings.Split(str, "\n") },
+	}
+	tmpl := template.New("cel").Funcs(funcMap)
+	tmpl, err := tmpl.Parse(authoringPrompt)
+	if err != nil {
+		return nil, err
+	}
+	return &Prompt{
+		Persona:      defaultPersona,
+		FormatRules:  defaultFormatRules,
+		GeneralUsage: defaultGeneralUsage,
+		tmpl:         tmpl,
+		env:          env,
+	}, nil
+}
+
+// Prompt represents the core components of an LLM prompt based on a CEL environment.
+//
+// All fields of the prompt may be overwritten / modified with support for rendering the
+// prompt to a human-readable string.
+type Prompt struct {
+	// Persona indicates something about the kind of user making the request
+	Persona string
+
+	// FormatRules indicate how the LLM should generate its output
+	FormatRules string
+
+	// GeneralUsage specifies additional context on how CEL should be used.
+	GeneralUsage string
+
+	// tmpl is the text template base-configuration for rendering text.
+	tmpl *template.Template
+
+	// env reference used to collect variables, functions, and macros available to the prompt.
+	env *Env
+}
+
+type promptInst struct {
+	*Prompt
+
+	Variables  []*common.Doc
+	Macros     []*common.Doc
+	Functions  []*common.Doc
+	UserPrompt string
+}
+
+// Render renders the user prompt with the associated context from the prompt template
+// for use with LLM generators.
+func (p *Prompt) Render(userPrompt string) string {
+	var buffer strings.Builder
+	vars := make([]*common.Doc, len(p.env.Variables()))
+	for i, v := range p.env.Variables() {
+		vars[i] = v.Documentation()
+	}
+	sort.SliceStable(vars, func(i, j int) bool {
+		return vars[i].Name < vars[j].Name
+	})
+	macs := make([]*common.Doc, len(p.env.Macros()))
+	for i, m := range p.env.Macros() {
+		macs[i] = m.(common.Documentor).Documentation()
+	}
+	funcs := make([]*common.Doc, 0, len(p.env.Functions()))
+	for _, f := range p.env.Functions() {
+		if _, hidden := hiddenFunctions[f.Name()]; hidden {
+			continue
+		}
+		funcs = append(funcs, f.Documentation())
+	}
+	sort.SliceStable(funcs, func(i, j int) bool {
+		return funcs[i].Name < funcs[j].Name
+	})
+	inst := &promptInst{
+		Prompt:     p,
+		Variables:  vars,
+		Macros:     macs,
+		Functions:  funcs,
+		UserPrompt: userPrompt}
+	p.tmpl.Execute(&buffer, inst)
+	return buffer.String()
+}
+
+const (
+	defaultPersona = `You are a software engineer with expertise in networking and application security
+authoring boolean Common Expression Language (CEL) expressions to ensure firewall,
+networking, authentication, and data access is only permitted when all conditions
+are satisified.`
+
+	defaultFormatRules = `Output your response as a CEL expression.
+
+Write the expression with the comment on the first line and the expression on the
+subsequent lines. Format the expression using 80-character line limits commonly
+found in C++ or Java code.`
+
+	defaultGeneralUsage = `CEL supports Protocol Buffer and JSON types, as well as simple types and aggregate types.
+
+Simple types include bool, bytes, double, int, string, and uint:
+
+* double literals must always include a decimal point: 1.0, 3.5, -2.2
+* uint literals must be positive values suffixed with a 'u': 42u
+* byte literals are strings prefixed with a 'b': b'1235'
+* string literals can use either single quotes or double quotes: 'hello', "world"
+* string literals can also be treated as raw strings that do not require any
+  escaping within the string by using the 'R' prefix: R"""quote: "hi" """
+
+Aggregate types include list and map:
+
+* list literals consist of zero or more values between brackets: "['a', 'b', 'c']"
+* map literal consist of colon-separated key-value pairs within braces: "{'key1': 1, 'key2': 2}"
+* Only int, uint, string, and bool types are valid map keys.
+* Maps containing HTTP headers must always use lower-cased string keys.
+
+Comments start with two-forward slashes followed by text and a newline.`
+)
+
+var (
+	hiddenFunctions = map[string]bool{
+		overloads.DeprecatedIn:        true,
+		operators.OldIn:               true,
+		operators.OldNotStrictlyFalse: true,
+		operators.NotStrictlyFalse:    true,
+	}
+)

cel/prompt_test.go

@@ -0,0 +1,73 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package cel
+
+import (
+	_ "embed"
+	"testing"
+
+	"github.com/google/cel-go/common/env"
+	"github.com/google/cel-go/test"
+)
+
+//go:embed testdata/basic.prompt.md
+var wantBasicPrompt string
+
+//go:embed testdata/macros.prompt.md
+var wantMacrosPrompt string
+
+//go:embed testdata/standard_env.prompt.md
+var wantStandardEnvPrompt string
+
+func TestPromptTemplate(t *testing.T) {
+	tests := []struct {
+		name    string
+		envOpts []EnvOption
+		out     string
+	}{
+		{
+			name: "basic",
+			out:  wantBasicPrompt,
+		},
+		{
+			name:    "macros",
+			envOpts: []EnvOption{Macros(StandardMacros...)},
+			out:     wantMacrosPrompt,
+		},
+		{
+			name:    "standard_env",
+			envOpts: []EnvOption{StdLib(StdLibSubset(env.NewLibrarySubset().SetDisableMacros(true)))},
+			out:     wantStandardEnvPrompt,
+		},
+	}
+
+	for _, tst := range tests {
+		tc := tst
+		t.Run(tc.name, func(t *testing.T) {
+			env, err := NewCustomEnv(tc.envOpts...)
+			if err != nil {
+				t.Fatalf("cel.NewCustomEnv() failed: %v", err)
+			}
+			prompt, err := AuthoringPrompt(env)
+			if err != nil {
+				t.Fatalf("cel.AuthoringPrompt() failed: %v", err)
+			}
+			out := prompt.Render("<USER_PROMPT>")
+			if !test.Compare(out, tc.out) {
+				t.Errorf("got %s, wanted %s", out, tc.out)
+			}
+		})
+	}
+}

cel/templates/BUILD.bazel

@@ -0,0 +1,7 @@
+licenses(["notice"])  # Apache 2.0
+
+filegroup(
+    name = "templates",
+    srcs = glob(["*.tmpl"]),
+    visibility = ["//visibility:public"],
+)