diff --git a/internal/api/handler/generation_ai_variable.go b/internal/api/handler/generation_ai_variable.go new file mode 100644 index 00000000..d194f19a --- /dev/null +++ b/internal/api/handler/generation_ai_variable.go @@ -0,0 +1,282 @@ +// AI-variable preview endpoint: generate the recipient-specific snippet a +// per-recipient AI block would produce, so the campaign editor's "Preview" +// button shows real output. Same credit flow as /generation/write and +// /generation/edit (gate, consume up front, refund on provider failure, usage +// settle), and the exact snippet framing the send path uses (tasks.BuildAIVariablePrompt) +// so preview and send never drift. The prompt renders against a supplied contact +// (org-scoped; 404 if not in the org) or a sample contact. + +package handler + +import ( + "context" + "errors" + "net/http" + "strings" + "time" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" + + "github.com/warmbly/warmbly/internal/api/middleware" + "github.com/warmbly/warmbly/internal/app/credits" + "github.com/warmbly/warmbly/internal/errx" + "github.com/warmbly/warmbly/internal/models" + "github.com/warmbly/warmbly/internal/pkg/generation" + "github.com/warmbly/warmbly/internal/tasks" +) + +// aiVarPreviewMaxPromptLen bounds the block prompt; aiVarPreviewMaxTokens caps +// the completion (matches the send-path resolver). +const ( + aiVarPreviewMaxPromptLen = 8000 + aiVarPreviewMaxTokens = 400 + aiVarPreviewTimeout = 20 * time.Second +) + +type generationAIVariableRequest struct { + Mode string `json:"mode"` // "instant" | "research" + Prompt string `json:"prompt"` + Tone string `json:"tone"` + WebSearch bool `json:"web_search"` + ContactID string `json:"contact_id"` + // The email text on either side of the block, so the fragment fits the + // sentence it lands in (matches the send path). Optional. + ContextBefore string `json:"context_before"` + ContextAfter string `json:"context_after"` +} + +// GenerateAIVariable — POST /generation/ai-variable +func (h *Handler) GenerateAIVariable(c *gin.Context) { + orgID := middleware.GetOrganizationID(c) + if orgID == nil { + errx.JSON(c, errx.New(errx.BadRequest, "no organization selected")) + return + } + + var req generationAIVariableRequest + if err := c.ShouldBindJSON(&req); err != nil { + errx.JSON(c, errx.ErrInvalid) + return + } + req.Prompt = strings.TrimSpace(req.Prompt) + if req.Prompt == "" { + errx.JSON(c, errx.New(errx.BadRequest, "prompt is required")) + return + } + if len(req.Prompt) > aiVarPreviewMaxPromptLen { + errx.JSON(c, errx.New(errx.BadRequest, "prompt is too long")) + return + } + + // Feature gate: paid orgs and free-trial orgs may use AI generation. + allowed, xerr := h.FeatureGateService.CanUseWritingAssistant(c.Request.Context(), *orgID) + if xerr != nil { + errx.JSON(c, xerr) + return + } + if !allowed { + errx.JSON(c, errx.New(errx.Forbidden, "AI variables require an active plan or trial.")) + return + } + if h.AIProvider == nil { + errx.JSON(c, errx.New(errx.ServiceUnavailable, "AI generation is not configured.")) + return + } + + // Resolve the contact to render the prompt against: a supplied id (must be in + // the org) or a sample contact. + contact := sampleContact() + if strings.TrimSpace(req.ContactID) != "" { + cid, perr := uuid.Parse(strings.TrimSpace(req.ContactID)) + if perr != nil { + errx.JSON(c, errx.New(errx.BadRequest, "invalid contact_id")) + return + } + if h.ContactRepo == nil { + errx.JSON(c, errx.ErrNotFound) + return + } + found, cxerr := h.ContactRepo.GetByIDsAndOrganization(c.Request.Context(), *orgID, []uuid.UUID{cid}) + if cxerr != nil { + errx.JSON(c, cxerr) + return + } + if len(found) == 0 { + errx.JSON(c, errx.ErrNotFound) + return + } + contact = found[0] + } + + rendered := strings.TrimSpace(tasks.RenderTemplate(req.Prompt, contact)) + if rendered == "" { + // A prompt that renders empty against the contact charges nothing. + remaining := 0 + if bal, berr := h.CreditService.GetBalance(c.Request.Context(), *orgID); berr == nil { + remaining = bal + } + c.JSON(http.StatusOK, gin.H{ + "text": "", "credits_remaining": remaining, "credits_charged": 0, + "tokens_used": 0, "model": "", + }) + return + } + + // Cost mirrors the resolver: instant -> CostWritingAssistant, research -> + // CostResearchRun. Research degrades to the same completion path here (no agent + // tool loop on the preview) but still charges the research price. + research := strings.EqualFold(strings.TrimSpace(req.Mode), "research") + cost := credits.CostWritingAssistant + reason := "campaign_ai_var_preview" + if research { + cost = credits.CostResearchRun + } + + paid, xerr := h.FeatureGateService.IsPaidOrganization(c.Request.Context(), *orgID) + if xerr != nil { + errx.JSON(c, xerr) + return + } + model := h.AIProvider.ModelForTier(paid) + + idemKey := strings.TrimSpace(c.GetHeader("Idempotency-Key")) + local := h.AIProvider.IsLocal() + reqCtx := c.Request.Context() + if actor, aerr := middleware.GetUserUUID(c); aerr == nil { + reqCtx = models.WithCreditMeta(reqCtx, models.CreditMeta{ActorID: actor}) + } + + var remaining int + if local { + if bal, berr := h.CreditService.GetBalance(reqCtx, *orgID); berr == nil { + remaining = bal + } + } else { + var err error + remaining, err = h.CreditService.Consume(reqCtx, *orgID, cost, reason, model, 0, idemKey) + if err != nil { + switch { + case errors.Is(err, credits.ErrInsufficientCredits): + paymentRequiredJSON(c, "You're out of AI credits. Upgrade or purchase more to keep using AI variables.") + case errors.Is(err, credits.ErrCapExceeded): + errx.JSON(c, errx.New(errx.TooManyRequests, "AI usage limit reached, please try again later.")) + default: + errx.JSON(c, errx.InternalError()) + } + return + } + } + + // Optional web enrichment (one bounded lookup, charged only when it returns + // results), matching the send-path resolver. Research implies web on. + web := "" + if (req.WebSearch || research) && h.AISearch != nil { + if q := previewSearchQuery(contact); q != "" { + sctx, scancel := context.WithTimeout(reqCtx, 15*time.Second) + results, serr := h.AISearch.Search(sctx, q, 3) + scancel() + if serr == nil && len(results) > 0 { + web = generation.FormatSearchResults(results) + if !local { + if _, cerr := h.CreditService.Consume(reqCtx, *orgID, credits.CostWebSearch, reason+"_search", "", 0, searchIdem(idemKey)); cerr == nil { + remaining -= credits.CostWebSearch + } + } + } + } + } + + // Ground the humanizer with the org's voice and the standard merge variables, + // so preview matches the send path's shared voice rules. The standard five + // tokens are enough here; per-contact custom keys are the send path's concern. + vc := h.orgVoice(reqCtx, *orgID, req.Tone) + vc.AvailableVars = generation.StandardMergeVars + system, prompt := tasks.BuildAIVariablePrompt(vc, contact, rendered, web, + clampPreviewContext(req.ContextBefore), clampPreviewContext(req.ContextAfter)) + + cctx, cancel := context.WithTimeout(reqCtx, aiVarPreviewTimeout) + defer cancel() + result, gerr := h.AIProvider.Complete(cctx, generation.CompletionRequest{ + System: system, + Prompt: prompt, + Model: model, + MaxTokens: aiVarPreviewMaxTokens, + Temperature: generation.Deterministic(), + }) + if gerr != nil || result == nil { + if !local { + if bal, rerr := h.CreditService.Grant(reqCtx, *orgID, cost, reason+"_refund"); rerr == nil { + remaining = bal + } + } + errx.JSON(c, errx.New(errx.ServiceUnavailable, "AI generation is temporarily unavailable. Your credit was not charged.")) + return + } + + charged := 0 + if !local { + charged = cost + if extra, serr := h.CreditService.SettleUsage(reqCtx, *orgID, cost, result.Model, result.TokensUsed, reason, settleKey(idemKey)); serr == nil && extra > 0 { + remaining -= extra + charged += extra + } + } + + c.JSON(http.StatusOK, gin.H{ + "text": strings.TrimSpace(result.Text), + "credits_remaining": remaining, + "credits_charged": charged, + "tokens_used": result.TokensUsed, + "model": result.Model, + }) +} + +// aiVarPreviewContextMax bounds each side of the surrounding-email context sent +// with a preview, so a caller cannot inflate the prompt. +const aiVarPreviewContextMax = 1200 + +func clampPreviewContext(s string) string { + s = strings.TrimSpace(s) + r := []rune(s) + if len(r) > aiVarPreviewContextMax { + return string(r[:aiVarPreviewContextMax]) + } + return s +} + +// searchIdem derives the web-search settle/consume key from the call's key. An +// empty key stays empty (non-idempotent call, non-idempotent search charge). +func searchIdem(idemKey string) string { + if idemKey == "" { + return "" + } + return idemKey + ":search" +} + +// previewSearchQuery derives a web-search query from the contact's own fields +// (company + name, then a corporate email domain), never from any external text, +// so a hostile field cannot steer the search. Free-mail domains are useless as a +// company signal and are skipped. +func previewSearchQuery(contact models.Contact) string { + company := strings.TrimSpace(contact.Company) + name := strings.TrimSpace(strings.TrimSpace(contact.FirstName) + " " + strings.TrimSpace(contact.LastName)) + if company != "" { + return strings.TrimSpace(company + " " + name) + } + if at := strings.LastIndex(contact.Email, "@"); at >= 0 { + domain := strings.ToLower(strings.TrimSpace(contact.Email[at+1:])) + if domain != "" && !previewFreeMailDomains[domain] { + return domain + } + } + return "" +} + +// previewFreeMailDomains never identify a company, so they are useless as a +// search fallback (mirrors the send-path resolver's list). +var previewFreeMailDomains = map[string]bool{ + "gmail.com": true, "googlemail.com": true, "outlook.com": true, "hotmail.com": true, + "live.com": true, "yahoo.com": true, "icloud.com": true, "me.com": true, "aol.com": true, + "proton.me": true, "protonmail.com": true, "gmx.com": true, "mail.com": true, +}