diff --git a/backend/pkg/controller/assistant.go b/backend/pkg/controller/assistant.go index 3f41906a..4627b991 100644 --- a/backend/pkg/controller/assistant.go +++ b/backend/pkg/controller/assistant.go @@ -16,7 +16,6 @@ import ( "pentagi/pkg/providers" "pentagi/pkg/providers/pconfig" "pentagi/pkg/providers/provider" - "pentagi/pkg/templates" "pentagi/pkg/tools" "github.com/sirupsen/logrus" @@ -154,7 +153,10 @@ func NewAssistantWorker(ctx context.Context, awc newAssistantWorkerCtx) (Assista return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow assistant log worker", err) } - prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB + prompter, err := newUserPrompter(ctx, awc.db, awc.userID) + if err != nil { + return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to build user prompter", err) + } executor, err := tools.NewFlowToolsExecutor(awc.db, awc.cfg, awc.docker, awc.functions, awc.flowID) if err != nil { return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow tools executor", err) @@ -319,7 +321,10 @@ func LoadAssistantWorker( return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow assistant log worker", err) } - prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB + prompter, err := newUserPrompter(ctx, awc.db, awc.userID) + if err != nil { + return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to build user prompter", err) + } executor, err := tools.NewFlowToolsExecutor(awc.db, awc.cfg, awc.docker, functions, awc.flowID) if err != nil { return nil, wrapErrorEndSpan(ctx, assistantSpan, "failed to create flow tools executor", err) diff --git a/backend/pkg/controller/flow.go b/backend/pkg/controller/flow.go index e8909ffe..ce167089 100644 --- a/backend/pkg/controller/flow.go +++ b/backend/pkg/controller/flow.go @@ -19,7 +19,6 @@ import ( "pentagi/pkg/providers" "pentagi/pkg/providers/pconfig" "pentagi/pkg/providers/provider" - "pentagi/pkg/templates" "pentagi/pkg/tools" "github.com/sirupsen/logrus" @@ -166,7 +165,10 @@ func NewFlowWorker( flowSpan := observation.Span(langfuse.WithSpanName("prepare flow worker")) ctx, _ = flowSpan.Observation(ctx) - prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB + prompter, err := newUserPrompter(ctx, fwc.db, fwc.userID) + if err != nil { + return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to build user prompter", err) + } executor, err := tools.NewFlowToolsExecutor(fwc.db, fwc.cfg, fwc.docker, fwc.functions, flow.ID) if err != nil { return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to create flow tools executor", err) @@ -332,7 +334,10 @@ func LoadFlowWorker(ctx context.Context, flow database.Flow, fwc flowWorkerCtx) return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to unmarshal functions", err) } - prompter := templates.NewDefaultPrompter() // TODO: change to flow prompter by userID from DB + prompter, err := newUserPrompter(ctx, fwc.db, flow.UserID) + if err != nil { + return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to build user prompter", err) + } executor, err := tools.NewFlowToolsExecutor(fwc.db, fwc.cfg, fwc.docker, functions, flow.ID) if err != nil { return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to create flow tools executor", err) diff --git a/backend/pkg/controller/prompter.go b/backend/pkg/controller/prompter.go new file mode 100644 index 00000000..76376fda --- /dev/null +++ b/backend/pkg/controller/prompter.go @@ -0,0 +1,53 @@ +package controller + +import ( + "context" + "encoding/json" + "fmt" + + "pentagi/pkg/database" + "pentagi/pkg/templates" +) + +// newUserPrompter loads the user's custom prompts from the database and +// overlays them onto the compiled default templates. Prompt types that +// the user has not customized continue to use the defaults. A database +// error is returned to the caller so that session creation fails +// explicitly instead of silently falling back to defaults. +func newUserPrompter(ctx context.Context, db database.Querier, userID int64) (templates.Prompter, error) { + userPrompts, err := db.GetUserPrompts(ctx, userID) + if err != nil { + return nil, fmt.Errorf("failed to load user prompts: %w", err) + } + return buildUserPrompter(templates.NewDefaultPrompter(), userPrompts) +} + +// buildUserPrompter is the pure merge step extracted from newUserPrompter +// so it can be unit-tested without a database fake. It seeds the result +// with every default template and then overrides each prompt type the +// user has saved a non-empty body for. +func buildUserPrompter(defaults templates.Prompter, userPrompts []database.Prompt) (templates.Prompter, error) { + blob, err := defaults.DumpTemplates() + if err != nil { + return nil, fmt.Errorf("failed to dump default templates: %w", err) + } + + merged := templates.PromptsMap{} + if err := json.Unmarshal(blob, &merged); err != nil { + return nil, fmt.Errorf("failed to parse default templates: %w", err) + } + + for _, p := range userPrompts { + if p.Prompt == "" { + // The Prompts UI uses delete (or reset, which writes the + // default body back) to remove a customization, so an empty + // body is unexpected. Skip it instead of clobbering the + // default with an empty string that would later surface as + // ErrTemplateNotFound deep inside agent rendering. + continue + } + merged[templates.PromptType(p.Type)] = p.Prompt + } + + return templates.NewFlowPrompter(merged), nil +} diff --git a/backend/pkg/controller/prompter_test.go b/backend/pkg/controller/prompter_test.go new file mode 100644 index 00000000..e744ab4e --- /dev/null +++ b/backend/pkg/controller/prompter_test.go @@ -0,0 +1,211 @@ +package controller + +import ( + "context" + "errors" + "testing" + + "pentagi/pkg/database" + "pentagi/pkg/templates" +) + +// fakeQuerier satisfies database.Querier by embedding the interface +// (so the unused methods stay nil) and overrides only GetUserPrompts. +// Calling any other method would panic, which is exactly what we want +// for a unit test that should not touch unrelated DB code. +type fakeQuerier struct { + database.Querier + prompts []database.Prompt + err error +} + +func (f *fakeQuerier) GetUserPrompts(ctx context.Context, userID int64) ([]database.Prompt, error) { + if f.err != nil { + return nil, f.err + } + return f.prompts, nil +} + +// getDefaultTemplate returns the compiled default body for a prompt +// type, used to compare overridden vs. preserved prompts in tests. +func getDefaultTemplate(t *testing.T, pt templates.PromptType) string { + t.Helper() + body, err := templates.NewDefaultPrompter().GetTemplate(pt) + if err != nil { + t.Fatalf("default prompter has no template for %q: %v", pt, err) + } + return body +} + +func TestBuildUserPrompter_NoUserPrompts(t *testing.T) { + defaults := templates.NewDefaultPrompter() + prompter, err := buildUserPrompter(defaults, nil) + if err != nil { + t.Fatalf("buildUserPrompter returned error: %v", err) + } + + for _, pt := range []templates.PromptType{ + templates.PromptTypePrimaryAgent, + templates.PromptTypeAssistant, + } { + got, err := prompter.GetTemplate(pt) + if err != nil { + t.Fatalf("GetTemplate(%q) error: %v", pt, err) + } + want := getDefaultTemplate(t, pt) + if got != want { + t.Errorf("prompt %q: expected default body when user has no overrides, got divergent body", pt) + } + } +} + +func TestBuildUserPrompter_SingleOverride(t *testing.T) { + const customBody = "custom primary agent prompt body" + userPrompts := []database.Prompt{ + {Type: database.PromptTypePrimaryAgent, Prompt: customBody}, + } + + prompter, err := buildUserPrompter(templates.NewDefaultPrompter(), userPrompts) + if err != nil { + t.Fatalf("buildUserPrompter error: %v", err) + } + + got, err := prompter.GetTemplate(templates.PromptTypePrimaryAgent) + if err != nil { + t.Fatalf("GetTemplate(primary_agent) error: %v", err) + } + if got != customBody { + t.Errorf("primary_agent: expected custom body %q, got %q", customBody, got) + } + + // Spot-check that an unrelated type still resolves to the default. + for _, pt := range []templates.PromptType{ + templates.PromptTypeAssistant, + templates.PromptTypePentester, + } { + got, err := prompter.GetTemplate(pt) + if err != nil { + t.Fatalf("GetTemplate(%q) error: %v", pt, err) + } + if got != getDefaultTemplate(t, pt) { + t.Errorf("prompt %q should still match default after a single unrelated override", pt) + } + } +} + +func TestBuildUserPrompter_PartialOverridesPreserveDefaults(t *testing.T) { + overrides := map[database.PromptType]string{ + database.PromptTypePrimaryAgent: "custom primary", + database.PromptTypeCoder: "custom coder", + database.PromptTypeReporter: "custom reporter", + } + + userPrompts := make([]database.Prompt, 0, len(overrides)) + for pt, body := range overrides { + userPrompts = append(userPrompts, database.Prompt{Type: pt, Prompt: body}) + } + + prompter, err := buildUserPrompter(templates.NewDefaultPrompter(), userPrompts) + if err != nil { + t.Fatalf("buildUserPrompter error: %v", err) + } + + for pt, want := range overrides { + got, err := prompter.GetTemplate(templates.PromptType(pt)) + if err != nil { + t.Fatalf("GetTemplate(%q) error: %v", pt, err) + } + if got != want { + t.Errorf("prompt %q: expected custom body %q, got %q", pt, want, got) + } + } + + // Types that were not overridden must still match the defaults. + for _, pt := range []templates.PromptType{ + templates.PromptTypeAssistant, + templates.PromptTypePentester, + templates.PromptTypeSearcher, + } { + got, err := prompter.GetTemplate(pt) + if err != nil { + t.Fatalf("GetTemplate(%q) error: %v", pt, err) + } + if got != getDefaultTemplate(t, pt) { + t.Errorf("prompt %q should still match default after partial overrides", pt) + } + } +} + +func TestBuildUserPrompter_EmptyBodyIsIgnored(t *testing.T) { + userPrompts := []database.Prompt{ + {Type: database.PromptTypePrimaryAgent, Prompt: ""}, + } + + prompter, err := buildUserPrompter(templates.NewDefaultPrompter(), userPrompts) + if err != nil { + t.Fatalf("buildUserPrompter error: %v", err) + } + + got, err := prompter.GetTemplate(templates.PromptTypePrimaryAgent) + if err != nil { + t.Fatalf("GetTemplate(primary_agent) error: %v", err) + } + if got != getDefaultTemplate(t, templates.PromptTypePrimaryAgent) { + t.Errorf("primary_agent: empty user body must not clobber default") + } +} + +func TestNewUserPrompter_HappyPath(t *testing.T) { + const customAssistant = "custom assistant prompt" + const customCoder = "custom coder prompt" + + db := &fakeQuerier{ + prompts: []database.Prompt{ + {Type: database.PromptTypeAssistant, Prompt: customAssistant}, + {Type: database.PromptTypeCoder, Prompt: customCoder}, + }, + } + + prompter, err := newUserPrompter(context.Background(), db, 42) + if err != nil { + t.Fatalf("newUserPrompter error: %v", err) + } + + for pt, want := range map[templates.PromptType]string{ + templates.PromptTypeAssistant: customAssistant, + templates.PromptTypeCoder: customCoder, + } { + got, err := prompter.GetTemplate(pt) + if err != nil { + t.Fatalf("GetTemplate(%q) error: %v", pt, err) + } + if got != want { + t.Errorf("prompt %q: expected custom body %q, got %q", pt, want, got) + } + } + + // Sanity check that an un-customized type still falls back. + got, err := prompter.GetTemplate(templates.PromptTypePentester) + if err != nil { + t.Fatalf("GetTemplate(pentester) error: %v", err) + } + if got != getDefaultTemplate(t, templates.PromptTypePentester) { + t.Errorf("pentester: expected default body when user has no override") + } +} + +func TestNewUserPrompter_DBErrorPropagates(t *testing.T) { + sentinel := errors.New("db connection lost") + db := &fakeQuerier{err: sentinel} + + prompter, err := newUserPrompter(context.Background(), db, 42) + if err == nil { + t.Fatalf("expected error, got nil prompter=%v", prompter) + } + if prompter != nil { + t.Errorf("expected nil prompter on DB error, got %v", prompter) + } + if !errors.Is(err, sentinel) { + t.Errorf("expected wrapped sentinel error, got %v", err) + } +}