Skip to content

Commit

Permalink
internal/gaby: allow policy checks via flag for overviews in Gaby
Browse files Browse the repository at this point in the history
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>
  • Loading branch information
tatianab committed Dec 20, 2024
1 parent a315118 commit 5d12b58
Show file tree
Hide file tree
Showing 2 changed files with 16 additions and 1 deletion.
15 changes: 14 additions & 1 deletion internal/gaby/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -57,6 +58,7 @@ type gabyFlags struct {
level string
overlay string
requireApproval string
enforcePolicy bool
}

var flags gabyFlags
Expand All @@ -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.
Expand All @@ -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
Expand Down Expand Up @@ -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/")
Expand Down Expand Up @@ -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{
Expand Down
2 changes: 2 additions & 0 deletions internal/llmapp/check.go
Original file line number Diff line number Diff line change
Expand Up @@ -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}
}
Expand Down

0 comments on commit 5d12b58

Please sign in to comment.