package aiprompts import ( "context" "errors" "fmt" "strings" "time" "github.com/descrybe/descrybe-v2/apps/api/internal/company" "github.com/google/uuid" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" ) // Service loads and stores per-company AI prompt templates. type Service struct { Pool *pgxpool.Pool } func NewService(pool *pgxpool.Pool) *Service { return &Service{Pool: pool} } type stored struct { key string language string systemTemplate string userTemplate string isEnabled bool updatedAt time.Time } // GetBundle returns all prompt keys with effective templates for language + variable catalog. func (s *Service) GetBundle(ctx context.Context, companyID uuid.UUID, language string) (Bundle, error) { lang, err := company.ParseLanguage(language, true) if err != nil { lang = company.LoadLanguage(ctx, s.Pool, companyID) } contentLangs := company.LoadContentLanguages(ctx, s.Pool, companyID) storedRows, err := s.loadAll(ctx, companyID) if err != nil { return Bundle{}, err } byKeyLang := map[string]stored{} customLangs := map[string][]string{} for _, st := range storedRows { byKeyLang[st.key+"\x00"+st.language] = st customLangs[st.key] = appendUnique(customLangs[st.key], st.language) } out := make([]Template, 0, len(BuiltInDefaults)) for _, def := range BuiltInDefaults { t := Template{ Key: def.Key, Language: lang, Label: def.Label, Description: def.Description, IsDefault: true, IsCustom: false, IsEnabled: true, } if st, ok := byKeyLang[def.Key+"\x00"+lang]; ok { t.IsCustom = true t.IsDefault = false t.IsEnabled = st.isEnabled t.UpdatedAt = st.updatedAt if st.isEnabled { t.SystemTemplate = st.systemTemplate t.UserTemplate = st.userTemplate } else { t.SystemTemplate = def.SystemTemplate t.UserTemplate = def.UserTemplate t.IsDefault = true } } else { t.SystemTemplate = def.SystemTemplate t.UserTemplate = def.UserTemplate } out = append(out, t) } return Bundle{ Language: lang, Prompts: out, Variables: VariableCatalog, CustomLanguages: customLangs, ContentLanguages: contentLangs, }, nil } // Resolve returns the effective templates for one key + language // (custom if enabled for lang, else built-in). No cross-language company fallback. func (s *Service) Resolve(ctx context.Context, companyID uuid.UUID, key, language string) (Resolved, error) { if !ValidPromptKey(key) { return Resolved{}, ErrInvalidKey } def, ok := DefaultFor(key) if !ok { return Resolved{}, ErrInvalidKey } lang, err := company.ParseLanguage(language, true) if err != nil { lang = company.DefaultLanguage } st, err := s.loadOne(ctx, companyID, key, lang) if err != nil && !errors.Is(err, pgx.ErrNoRows) { return Resolved{}, err } if err == nil && st.isEnabled { sys := strings.TrimSpace(st.systemTemplate) user := strings.TrimSpace(st.userTemplate) if sys == "" { sys = def.SystemTemplate } if user == "" { user = def.UserTemplate } return Resolved{ Key: key, Language: lang, SystemTemplate: sys, UserTemplate: user, IsCustom: true, }, nil } return Resolved{ Key: key, Language: lang, SystemTemplate: def.SystemTemplate, UserTemplate: def.UserTemplate, IsCustom: false, }, nil } // Update applies prompt updates (upsert or reset). Empty prompts list is a no-op. func (s *Service) Update(ctx context.Context, companyID uuid.UUID, in UpdateInput) (Bundle, error) { defaultLang := strings.TrimSpace(in.Language) if defaultLang == "" { defaultLang = company.LoadLanguage(ctx, s.Pool, companyID) } if len(in.Prompts) == 0 { return s.GetBundle(ctx, companyID, defaultLang) } tx, err := s.Pool.Begin(ctx) if err != nil { return Bundle{}, err } defer tx.Rollback(ctx) lastLang := defaultLang for _, item := range in.Prompts { key := strings.TrimSpace(strings.ToLower(item.Key)) if !ValidPromptKey(key) { return Bundle{}, fmt.Errorf("%w: %s", ErrInvalidKey, item.Key) } langRaw := strings.TrimSpace(item.Language) if langRaw == "" { langRaw = defaultLang } lang, err := company.ParseLanguage(langRaw, false) if err != nil { return Bundle{}, fmt.Errorf("%w: language %q", ErrInvalidInput, langRaw) } lastLang = lang if item.Reset { _, err := tx.Exec(ctx, ` DELETE FROM ai_prompt_templates WHERE company_id = $1 AND prompt_key = $2 AND language = $3`, companyID, key, lang) if err != nil { return Bundle{}, err } continue } sys := SanitizeTemplate(item.SystemTemplate, MaxSystemRunes) user := SanitizeTemplate(item.UserTemplate, MaxUserRunes) if sys == "" && user == "" { return Bundle{}, fmt.Errorf("%w: empty templates for %s", ErrInvalidInput, key) } enabled := true if item.IsEnabled != nil { enabled = *item.IsEnabled } _, err = tx.Exec(ctx, ` INSERT INTO ai_prompt_templates ( company_id, prompt_key, language, system_template, user_template, is_enabled, updated_at ) VALUES ($1,$2,$3,$4,$5,$6, now()) ON CONFLICT (company_id, prompt_key, language) DO UPDATE SET system_template = EXCLUDED.system_template, user_template = EXCLUDED.user_template, is_enabled = EXCLUDED.is_enabled, updated_at = now()`, companyID, key, lang, sys, user, enabled) if err != nil { return Bundle{}, err } } if err := tx.Commit(ctx); err != nil { return Bundle{}, err } return s.GetBundle(ctx, companyID, lastLang) } func (s *Service) loadAll(ctx context.Context, companyID uuid.UUID) ([]stored, error) { rows, err := s.Pool.Query(ctx, ` SELECT prompt_key, language, system_template, user_template, is_enabled, updated_at FROM ai_prompt_templates WHERE company_id = $1`, companyID) if err != nil { return nil, err } defer rows.Close() var out []stored for rows.Next() { var st stored if err := rows.Scan(&st.key, &st.language, &st.systemTemplate, &st.userTemplate, &st.isEnabled, &st.updatedAt); err != nil { return nil, err } out = append(out, st) } return out, rows.Err() } func (s *Service) loadOne(ctx context.Context, companyID uuid.UUID, key, language string) (stored, error) { var st stored err := s.Pool.QueryRow(ctx, ` SELECT prompt_key, language, system_template, user_template, is_enabled, updated_at FROM ai_prompt_templates WHERE company_id = $1 AND prompt_key = $2 AND language = $3`, companyID, key, language).Scan(&st.key, &st.language, &st.systemTemplate, &st.userTemplate, &st.isEnabled, &st.updatedAt) return st, err } func appendUnique(list []string, v string) []string { for _, x := range list { if x == v { return list } } return append(list, v) }