internal/gaby: allow policy checks via flag for overviews in Gaby

If the -enforcepolicy flag is set (default false), check all safety
categories when generating overviews. Since overviews are not yet
published anywhere, this would only affect users of the web UI.

For #70

Change-Id: I100070b5726ca0ff21cea4dec9f7f68e74018f08
Reviewed-on: https://go-review.googlesource.com/c/oscar/+/637978
Reviewed-by: Hyang-Ah Hana Kim <hyangah@gmail.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>

Author: tatianab

internal/gaby/main.go

@@ -30,6 +30,7 @@ import (
 	"golang.org/x/oscar/internal/discussion"
 	"golang.org/x/oscar/internal/docs"
 	"golang.org/x/oscar/internal/embeddocs"
+	"golang.org/x/oscar/internal/gcp/checks"
 	"golang.org/x/oscar/internal/gcp/firestore"
 	"golang.org/x/oscar/internal/gcp/gcphandler"
 	"golang.org/x/oscar/internal/gcp/gcpmetrics"
@@ -57,6 +58,7 @@ type gabyFlags struct {
 	level           string
 	overlay         string
 	requireApproval string
+	enforcePolicy   bool
 }
 
 var flags gabyFlags
@@ -70,6 +72,7 @@ func init() {
 	flag.StringVar(&flags.level, "level", "info", "initial log level")
 	flag.StringVar(&flags.overlay, "overlay", "", "spec for overlay to DB; see internal/dbspec for syntax")
 	flag.StringVar(&flags.requireApproval, "requireapproval", "", "comma-separated list of packages whose actions require approval")
+	flag.BoolVar(&flags.enforcePolicy, "enforcepolicy", false, "whether to enforce safety policies on LLM inputs and outputs")
 }
 
 // Gaby holds the state for gaby's execution.
@@ -91,6 +94,7 @@ type Gaby struct {
 	docs      *docs.Corpus           // document corpus to use
 	embed     llm.Embedder           // LLM embedder to use
 	llm       llm.ContentGenerator   // LLM content generator to use
+	policy    llm.PolicyChecker      // LLM checker to use
 	llmapp    *llmapp.Client         // LLM client to use
 	github    *github.Client         // github client to use
 	disc      *discussion.Client     // github discussion client to use
@@ -167,7 +171,7 @@ func main() {
 	}
 	g.embed = ai
 	g.llm = ai
-	g.llmapp = llmapp.New(g.slog, ai, g.db)
+	g.llmapp = llmapp.NewWithChecker(g.slog, ai, g.policy, g.db)
 
 	cr := crawl.New(g.slog, g.db, g.http)
 	cr.Add("https://go.dev/")
@@ -339,6 +343,15 @@ func (g *Gaby) initGCP() (shutdown func()) {
 	}
 	g.secret = sdb
 
+	if flags.enforcePolicy {
+		llmchecker, err := checks.New(g.ctx, g.slog, flags.project)
+		if err != nil {
+			log.Fatal(err)
+		}
+		llmchecker.SetPolicies(llm.AllPolicyTypes())
+		g.policy = llmchecker
+	}
+
 	// Initialize error reporting if we are running on Cloud Run.
 	if g.cloud {
 		rep, err := errorreporting.NewClient(g.ctx, flags.project, errorreporting.Config{

internal/llmapp/check.go

@@ -18,6 +18,8 @@ import (
 //
 // When any of the Overview functions are called, the prompts and outputs of the LLM
 // will be checked for safety violations.
+//
+// If the checker is nil, [NewWithChecker] is identical to [New].
 func NewWithChecker(lg *slog.Logger, g llm.ContentGenerator, checker llm.PolicyChecker, db storage.DB) *Client {
 	return &Client{slog: lg, g: g, checker: checker, db: db}
 }