fix(growth): harden AI jobs and cost controls
This commit is contained in:
@@ -131,6 +131,17 @@ func (s *Server) growthReportEvents(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
merchantID := sessionFrom(r.Context()).MerchantID
|
merchantID := sessionFrom(r.Context()).MerchantID
|
||||||
|
slots := s.growthSSESlots
|
||||||
|
if slots == nil {
|
||||||
|
slots = defaultGrowthSSESlots
|
||||||
|
}
|
||||||
|
select {
|
||||||
|
case slots <- struct{}{}:
|
||||||
|
defer func() { <-slots }()
|
||||||
|
default:
|
||||||
|
s.problem(w, r, http.StatusTooManyRequests, "RATE_LIMITED", "Too many growth report streams")
|
||||||
|
return
|
||||||
|
}
|
||||||
record, err := s.growth.Get(r.Context(), merchantID, id)
|
record, err := s.growth.Get(r.Context(), merchantID, id)
|
||||||
if errors.Is(err, growthreport.ErrNotFound) {
|
if errors.Is(err, growthreport.ErrNotFound) {
|
||||||
s.problem(w, r, 404, "RESOURCE_NOT_FOUND", "Resource not found")
|
s.problem(w, r, 404, "RESOURCE_NOT_FOUND", "Resource not found")
|
||||||
@@ -148,7 +159,17 @@ func (s *Server) growthReportEvents(w http.ResponseWriter, r *http.Request) {
|
|||||||
w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
|
w.Header().Set("Content-Type", "text/event-stream; charset=utf-8")
|
||||||
w.Header().Set("Cache-Control", "no-cache, no-transform")
|
w.Header().Set("Cache-Control", "no-cache, no-transform")
|
||||||
w.Header().Set("X-Accel-Buffering", "no")
|
w.Header().Set("X-Accel-Buffering", "no")
|
||||||
ticker := time.NewTicker(2 * time.Second)
|
maxDuration := s.growthSSEMaxDuration
|
||||||
|
if maxDuration <= 0 {
|
||||||
|
maxDuration = defaultGrowthSSEMaxDuration
|
||||||
|
}
|
||||||
|
pollInterval := s.growthSSEPollInterval
|
||||||
|
if pollInterval <= 0 {
|
||||||
|
pollInterval = defaultGrowthSSEPollInterval
|
||||||
|
}
|
||||||
|
streamContext, cancel := context.WithTimeout(r.Context(), maxDuration)
|
||||||
|
defer cancel()
|
||||||
|
ticker := time.NewTicker(pollInterval)
|
||||||
defer ticker.Stop()
|
defer ticker.Stop()
|
||||||
heartbeat := time.NewTicker(15 * time.Second)
|
heartbeat := time.NewTicker(15 * time.Second)
|
||||||
defer heartbeat.Stop()
|
defer heartbeat.Stop()
|
||||||
@@ -165,10 +186,10 @@ func (s *Server) growthReportEvents(w http.ResponseWriter, r *http.Request) {
|
|||||||
return
|
return
|
||||||
}
|
}
|
||||||
select {
|
select {
|
||||||
case <-r.Context().Done():
|
case <-streamContext.Done():
|
||||||
return
|
return
|
||||||
case <-ticker.C:
|
case <-ticker.C:
|
||||||
record, err = s.growth.Get(r.Context(), merchantID, id)
|
record, err = s.growth.Get(streamContext, merchantID, id)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -127,3 +127,45 @@ func TestGrowthSSEAuthorizesBeforeSettingStreamHeaders(t *testing.T) {
|
|||||||
t.Fatalf("status=%d content-type=%q", response.Code, response.Header().Get("Content-Type"))
|
t.Fatalf("status=%d content-type=%q", response.Code, response.Header().Get("Content-Type"))
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestGrowthSSERejectsWhenGlobalCapacityIsFull(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
slots := make(chan struct{}, 1)
|
||||||
|
slots <- struct{}{}
|
||||||
|
reads := 0
|
||||||
|
service := fakeGrowthService{get: func(int64, string) (growthreport.Record, error) {
|
||||||
|
reads++
|
||||||
|
return growthreport.Record{Report: growthreport.Report{ReportID: "growth_0123456789abcdef01234567", Status: growthreport.StatusQueued}}, nil
|
||||||
|
}}
|
||||||
|
server := &Server{logger: slog.New(slog.NewTextHandler(io.Discard, nil)), growth: service, growthSSESlots: slots}
|
||||||
|
request := httptest.NewRequest("GET", "/", nil)
|
||||||
|
request.SetPathValue("reportId", "growth_0123456789abcdef01234567")
|
||||||
|
request = request.WithContext(context.WithValue(request.Context(), sessionKey{}, identity.Session{MerchantID: 7}))
|
||||||
|
response := httptest.NewRecorder()
|
||||||
|
server.growthReportEvents(response, request)
|
||||||
|
if response.Code != http.StatusTooManyRequests || reads != 0 {
|
||||||
|
t.Fatalf("status=%d database reads=%d", response.Code, reads)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestGrowthSSEEndsAtConfiguredDeadline(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
reads := 0
|
||||||
|
service := fakeGrowthService{get: func(int64, string) (growthreport.Record, error) {
|
||||||
|
reads++
|
||||||
|
return growthreport.Record{Report: growthreport.Report{ReportID: "growth_0123456789abcdef01234567", Status: growthreport.StatusQueued, UpdatedAt: time.Unix(int64(reads), 0)}}, nil
|
||||||
|
}}
|
||||||
|
server := &Server{logger: slog.New(slog.NewTextHandler(io.Discard, nil)), growth: service, growthSSESlots: make(chan struct{}, 1), growthSSEMaxDuration: 5 * time.Millisecond, growthSSEPollInterval: time.Millisecond}
|
||||||
|
request := httptest.NewRequest("GET", "/", nil)
|
||||||
|
request.SetPathValue("reportId", "growth_0123456789abcdef01234567")
|
||||||
|
request = request.WithContext(context.WithValue(request.Context(), sessionKey{}, identity.Session{MerchantID: 7}))
|
||||||
|
response := httptest.NewRecorder()
|
||||||
|
started := time.Now()
|
||||||
|
server.growthReportEvents(response, request)
|
||||||
|
if elapsed := time.Since(started); elapsed > 250*time.Millisecond {
|
||||||
|
t.Fatalf("stream remained open for %s", elapsed)
|
||||||
|
}
|
||||||
|
if reads < 1 || reads > 10 {
|
||||||
|
t.Fatalf("database reads=%d", reads)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -23,22 +23,33 @@ import (
|
|||||||
|
|
||||||
const maxExportBytes = 50 << 20
|
const maxExportBytes = 50 << 20
|
||||||
|
|
||||||
|
const (
|
||||||
|
defaultGrowthSSECapacity = 64
|
||||||
|
defaultGrowthSSEMaxDuration = 5 * time.Minute
|
||||||
|
defaultGrowthSSEPollInterval = 2 * time.Second
|
||||||
|
)
|
||||||
|
|
||||||
|
var defaultGrowthSSESlots = make(chan struct{}, defaultGrowthSSECapacity)
|
||||||
|
|
||||||
var errExportBytesExceeded = errors.New("export exceeds byte limit")
|
var errExportBytesExceeded = errors.New("export exceeds byte limit")
|
||||||
|
|
||||||
type Server struct {
|
type Server struct {
|
||||||
identity *identity.Service
|
identity *identity.Service
|
||||||
analytics *analytics.Service
|
analytics *analytics.Service
|
||||||
growth growthReportService
|
growth growthReportService
|
||||||
logger *slog.Logger
|
logger *slog.Logger
|
||||||
frontendOrigin string
|
frontendOrigin string
|
||||||
cookieSecure bool
|
cookieSecure bool
|
||||||
|
growthSSESlots chan struct{}
|
||||||
|
growthSSEMaxDuration time.Duration
|
||||||
|
growthSSEPollInterval time.Duration
|
||||||
}
|
}
|
||||||
|
|
||||||
type sessionKey struct{}
|
type sessionKey struct{}
|
||||||
type tokenKey struct{}
|
type tokenKey struct{}
|
||||||
|
|
||||||
func NewServer(identityService *identity.Service, analyticsService *analytics.Service, logger *slog.Logger, frontendOrigin string, cookieSecure bool, growthService ...growthReportService) *Server {
|
func NewServer(identityService *identity.Service, analyticsService *analytics.Service, logger *slog.Logger, frontendOrigin string, cookieSecure bool, growthService ...growthReportService) *Server {
|
||||||
server := &Server{identity: identityService, analytics: analyticsService, logger: logger, frontendOrigin: frontendOrigin, cookieSecure: cookieSecure}
|
server := &Server{identity: identityService, analytics: analyticsService, logger: logger, frontendOrigin: frontendOrigin, cookieSecure: cookieSecure, growthSSESlots: defaultGrowthSSESlots, growthSSEMaxDuration: defaultGrowthSSEMaxDuration, growthSSEPollInterval: defaultGrowthSSEPollInterval}
|
||||||
if len(growthService) > 0 {
|
if len(growthService) > 0 {
|
||||||
server.growth = growthService[0]
|
server.growth = growthService[0]
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ import (
|
|||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
"io"
|
||||||
|
"math/big"
|
||||||
"net/http"
|
"net/http"
|
||||||
"strings"
|
"strings"
|
||||||
"time"
|
"time"
|
||||||
@@ -60,17 +61,19 @@ func (c *AvalAIClient) GenerateWithTools(ctx context.Context, evidence json.RawM
|
|||||||
"text": map[string]any{"format": map[string]any{"type": "json_schema", "name": "growth_copy", "strict": true, "schema": aiCopySchema()}},
|
"text": map[string]any{"format": map[string]any{"type": "json_schema", "name": "growth_copy", "strict": true, "schema": aiCopySchema()}},
|
||||||
}
|
}
|
||||||
var last error
|
var last error
|
||||||
|
var total Usage
|
||||||
for attempt := 0; attempt < 2; attempt++ {
|
for attempt := 0; attempt < 2; attempt++ {
|
||||||
copy, usage, transient, err := c.workflow(ctx, base, evidence, known, executor)
|
copy, usage, transient, err := c.workflow(ctx, base, evidence, known, executor)
|
||||||
|
total = mergeUsage(total, usage)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
return copy, usage, nil
|
return copy, total, nil
|
||||||
}
|
}
|
||||||
last = err
|
last = err
|
||||||
if !transient || ctx.Err() != nil {
|
if !transient || ctx.Err() != nil {
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
return AICopy{}, Usage{}, last
|
return AICopy{}, total, last
|
||||||
}
|
}
|
||||||
|
|
||||||
type responseEnvelope struct {
|
type responseEnvelope struct {
|
||||||
@@ -99,9 +102,9 @@ type outputItem struct {
|
|||||||
func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, evidence json.RawMessage, known map[string]struct{}, executor GrowthAnalyticsTools) (AICopy, Usage, bool, error) {
|
func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, evidence json.RawMessage, known map[string]struct{}, executor GrowthAnalyticsTools) (AICopy, Usage, bool, error) {
|
||||||
firstBody := cloneRequest(base)
|
firstBody := cloneRequest(base)
|
||||||
firstBody["input"] = string(evidence)
|
firstBody["input"] = string(evidence)
|
||||||
first, latency, transient, err := c.do(ctx, firstBody)
|
first, usage, transient, err := c.do(ctx, firstBody)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return AICopy{}, Usage{}, transient, err
|
return AICopy{}, usage, transient, err
|
||||||
}
|
}
|
||||||
calls := make([]outputItem, 0, 4)
|
calls := make([]outputItem, 0, 4)
|
||||||
for _, raw := range first.Output {
|
for _, raw := range first.Output {
|
||||||
@@ -110,9 +113,6 @@ func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, eviden
|
|||||||
calls = append(calls, item)
|
calls = append(calls, item)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
usage := Usage{InputTokens: first.Usage.InputTokens, OutputTokens: first.Usage.OutputTokens, LatencyMS: latency}
|
|
||||||
usage.ProviderRequestID = first.ProviderRequestID
|
|
||||||
usage.EstimatedCost = costString(first.EstimatedCost)
|
|
||||||
final := first
|
final := first
|
||||||
if len(calls) > 0 {
|
if len(calls) > 0 {
|
||||||
if executor == nil {
|
if executor == nil {
|
||||||
@@ -141,19 +141,11 @@ func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, eviden
|
|||||||
}
|
}
|
||||||
secondBody := cloneRequest(base)
|
secondBody := cloneRequest(base)
|
||||||
secondBody["input"] = input
|
secondBody["input"] = input
|
||||||
second, secondLatency, secondTransient, err := c.do(ctx, secondBody)
|
second, secondUsage, secondTransient, err := c.do(ctx, secondBody)
|
||||||
|
usage = mergeUsage(usage, secondUsage)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return AICopy{}, usage, secondTransient, err
|
return AICopy{}, usage, secondTransient, err
|
||||||
}
|
}
|
||||||
usage.InputTokens += second.Usage.InputTokens
|
|
||||||
usage.OutputTokens += second.Usage.OutputTokens
|
|
||||||
usage.LatencyMS += secondLatency
|
|
||||||
if second.ProviderRequestID != "" {
|
|
||||||
usage.ProviderRequestID = second.ProviderRequestID
|
|
||||||
}
|
|
||||||
if cost := costString(second.EstimatedCost); cost != "" {
|
|
||||||
usage.EstimatedCost = cost
|
|
||||||
}
|
|
||||||
final = second
|
final = second
|
||||||
for _, raw := range second.Output {
|
for _, raw := range second.Output {
|
||||||
var item outputItem
|
var item outputItem
|
||||||
@@ -169,44 +161,76 @@ func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, eviden
|
|||||||
return copy, usage, false, nil
|
return copy, usage, false, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (c *AvalAIClient) do(parent context.Context, bodyValue map[string]any) (responseEnvelope, int64, bool, error) {
|
func (c *AvalAIClient) do(parent context.Context, bodyValue map[string]any) (responseEnvelope, Usage, bool, error) {
|
||||||
body, err := json.Marshal(bodyValue)
|
body, err := json.Marshal(bodyValue)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return responseEnvelope{}, 0, false, fmt.Errorf("encoding AvalAI request: %w", err)
|
return responseEnvelope{}, Usage{}, false, fmt.Errorf("encoding AvalAI request: %w", err)
|
||||||
}
|
}
|
||||||
ctx, cancel := context.WithTimeout(parent, c.config.Timeout)
|
ctx, cancel := context.WithTimeout(parent, c.config.Timeout)
|
||||||
defer cancel()
|
defer cancel()
|
||||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.config.BaseURL, "/")+"/v1/responses", bytes.NewReader(body))
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(c.config.BaseURL, "/")+"/v1/responses", bytes.NewReader(body))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return responseEnvelope{}, 0, false, fmt.Errorf("creating AvalAI request: %w", err)
|
return responseEnvelope{}, Usage{}, false, fmt.Errorf("creating AvalAI request: %w", err)
|
||||||
}
|
}
|
||||||
req.Header.Set("Authorization", "Bearer "+c.config.APIKey)
|
req.Header.Set("Authorization", "Bearer "+c.config.APIKey)
|
||||||
req.Header.Set("Content-Type", "application/json")
|
req.Header.Set("Content-Type", "application/json")
|
||||||
started := time.Now()
|
started := time.Now()
|
||||||
resp, err := c.client.Do(req)
|
resp, err := c.client.Do(req)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return responseEnvelope{}, 0, true, fmt.Errorf("calling AvalAI: %w", err)
|
return responseEnvelope{}, Usage{ProviderRequests: 1, LatencyMS: time.Since(started).Milliseconds()}, true, fmt.Errorf("calling AvalAI: %w", err)
|
||||||
}
|
}
|
||||||
defer resp.Body.Close()
|
defer resp.Body.Close()
|
||||||
responseBody, err := io.ReadAll(io.LimitReader(resp.Body, maxAvalAIResponseBytes+1))
|
responseBody, err := io.ReadAll(io.LimitReader(resp.Body, maxAvalAIResponseBytes+1))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return responseEnvelope{}, 0, true, fmt.Errorf("reading AvalAI response: %w", err)
|
return responseEnvelope{}, Usage{ProviderRequests: 1, LatencyMS: time.Since(started).Milliseconds(), ProviderRequestID: resp.Header.Get("X-Request-ID")}, true, fmt.Errorf("reading AvalAI response: %w", err)
|
||||||
}
|
}
|
||||||
if len(responseBody) > maxAvalAIResponseBytes {
|
if len(responseBody) > maxAvalAIResponseBytes {
|
||||||
return responseEnvelope{}, 0, false, errors.New("avalai: response too large")
|
return responseEnvelope{}, Usage{ProviderRequests: 1, LatencyMS: time.Since(started).Milliseconds(), ProviderRequestID: resp.Header.Get("X-Request-ID")}, false, errors.New("avalai: response too large")
|
||||||
}
|
|
||||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
||||||
return responseEnvelope{}, 0, resp.StatusCode == 429 || resp.StatusCode >= 500, fmt.Errorf("avalai: HTTP %d", resp.StatusCode)
|
|
||||||
}
|
}
|
||||||
var result responseEnvelope
|
var result responseEnvelope
|
||||||
if err := json.Unmarshal(responseBody, &result); err != nil {
|
decodeErr := json.Unmarshal(responseBody, &result)
|
||||||
return responseEnvelope{}, 0, false, fmt.Errorf("decoding AvalAI response: %w", err)
|
result.ProviderRequestID = resp.Header.Get("X-Request-ID")
|
||||||
|
usage := Usage{InputTokens: result.Usage.InputTokens, OutputTokens: result.Usage.OutputTokens, LatencyMS: time.Since(started).Milliseconds(), ProviderRequests: 1, ProviderRequestID: result.ProviderRequestID, EstimatedCost: costString(result.EstimatedCost)}
|
||||||
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||||
|
return result, usage, resp.StatusCode == 429 || resp.StatusCode >= 500, fmt.Errorf("avalai: HTTP %d", resp.StatusCode)
|
||||||
|
}
|
||||||
|
if decodeErr != nil {
|
||||||
|
return responseEnvelope{}, usage, false, fmt.Errorf("decoding AvalAI response: %w", decodeErr)
|
||||||
}
|
}
|
||||||
if result.Status != "completed" {
|
if result.Status != "completed" {
|
||||||
return responseEnvelope{}, 0, false, fmt.Errorf("avalai: incomplete response status %q", result.Status)
|
return result, usage, false, fmt.Errorf("avalai: incomplete response status %q", result.Status)
|
||||||
}
|
}
|
||||||
result.ProviderRequestID = resp.Header.Get("X-Request-ID")
|
return result, usage, false, nil
|
||||||
return result, time.Since(started).Milliseconds(), false, nil
|
}
|
||||||
|
|
||||||
|
func mergeUsage(total, next Usage) Usage {
|
||||||
|
total.InputTokens += next.InputTokens
|
||||||
|
total.OutputTokens += next.OutputTokens
|
||||||
|
total.LatencyMS += next.LatencyMS
|
||||||
|
total.ProviderRequests += next.ProviderRequests
|
||||||
|
total.ReservedTokens += next.ReservedTokens
|
||||||
|
if next.ProviderRequestID != "" {
|
||||||
|
total.ProviderRequestID = next.ProviderRequestID
|
||||||
|
}
|
||||||
|
total.EstimatedCost = addDecimalStrings(total.EstimatedCost, next.EstimatedCost)
|
||||||
|
return total
|
||||||
|
}
|
||||||
|
|
||||||
|
func addDecimalStrings(left, right string) string {
|
||||||
|
if left == "" {
|
||||||
|
return right
|
||||||
|
}
|
||||||
|
if right == "" {
|
||||||
|
return left
|
||||||
|
}
|
||||||
|
a, okA := new(big.Rat).SetString(left)
|
||||||
|
b, okB := new(big.Rat).SetString(right)
|
||||||
|
if !okA || !okB {
|
||||||
|
return left
|
||||||
|
}
|
||||||
|
text := new(big.Rat).Add(a, b).FloatString(12)
|
||||||
|
text = strings.TrimRight(strings.TrimRight(text, "0"), ".")
|
||||||
|
return text
|
||||||
}
|
}
|
||||||
|
|
||||||
func costString(raw json.RawMessage) string {
|
func costString(raw json.RawMessage) string {
|
||||||
@@ -287,6 +311,9 @@ func ValidateAICopy(copy AICopy, known map[string]struct{}) error {
|
|||||||
return errors.New("avalai: numeric claim rejected")
|
return errors.New("avalai: numeric claim rejected")
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if containsQuantityWord(field) {
|
||||||
|
return errors.New("avalai: written numeric claim rejected")
|
||||||
|
}
|
||||||
lower := strings.ToLower(field)
|
lower := strings.ToLower(field)
|
||||||
for _, forbidden := range []string{"payer", "card", "رمز", "شماره کارت", "sql", "merchant_id"} {
|
for _, forbidden := range []string{"payer", "card", "رمز", "شماره کارت", "sql", "merchant_id"} {
|
||||||
if strings.Contains(lower, forbidden) {
|
if strings.Contains(lower, forbidden) {
|
||||||
@@ -306,6 +333,20 @@ func ValidateAICopy(copy AICopy, known map[string]struct{}) error {
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func containsQuantityWord(value string) bool {
|
||||||
|
quantityWords := map[string]struct{}{
|
||||||
|
"zero": {}, "one": {}, "two": {}, "three": {}, "four": {}, "five": {}, "six": {}, "seven": {}, "eight": {}, "nine": {}, "ten": {}, "eleven": {}, "twelve": {}, "thirteen": {}, "fourteen": {}, "fifteen": {}, "sixteen": {}, "seventeen": {}, "eighteen": {}, "nineteen": {}, "twenty": {}, "thirty": {}, "forty": {}, "fifty": {}, "sixty": {}, "seventy": {}, "eighty": {}, "ninety": {}, "hundred": {}, "thousand": {}, "million": {}, "billion": {}, "percent": {}, "percentage": {}, "double": {}, "triple": {}, "half": {},
|
||||||
|
"صفر": {}, "یک": {}, "دو": {}, "سه": {}, "چهار": {}, "پنج": {}, "شش": {}, "هفت": {}, "هشت": {}, "نه": {}, "ده": {}, "یازده": {}, "دوازده": {}, "سیزده": {}, "چهارده": {}, "پانزده": {}, "شانزده": {}, "هفده": {}, "هجده": {}, "نوزده": {}, "بیست": {}, "سی": {}, "چهل": {}, "پنجاه": {}, "شصت": {}, "هفتاد": {}, "هشتاد": {}, "نود": {}, "صد": {}, "هزار": {}, "میلیون": {}, "میلیارد": {}, "درصد": {}, "برابر": {}, "نصف": {},
|
||||||
|
}
|
||||||
|
for _, word := range strings.Fields(strings.ToLower(value)) {
|
||||||
|
word = strings.Trim(word, ".,;:!?،؛؟٪%()[]{}«»\"'")
|
||||||
|
if _, found := quantityWords[word]; found {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
func growthTools() []map[string]any {
|
func growthTools() []map[string]any {
|
||||||
names := []string{"get_growth_context", "get_event_evidence", "get_peer_context", "get_campaign_playbook"}
|
names := []string{"get_growth_context", "get_event_evidence", "get_peer_context", "get_campaign_playbook"}
|
||||||
tools := make([]map[string]any, 0, len(names))
|
tools := make([]map[string]any, 0, len(names))
|
||||||
|
|||||||
@@ -3,6 +3,7 @@ package growthreport
|
|||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
|
"fmt"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
"strings"
|
"strings"
|
||||||
@@ -10,6 +11,31 @@ import (
|
|||||||
"time"
|
"time"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
func TestAvalAICountsFailedRetryRequestsAndUsage(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
requests := 0
|
||||||
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||||
|
requests++
|
||||||
|
w.Header().Set("X-Request-ID", fmt.Sprintf("request-%d", requests))
|
||||||
|
if requests == 1 {
|
||||||
|
w.WriteHeader(http.StatusServiceUnavailable)
|
||||||
|
_, _ = w.Write([]byte(`{"usage":{"input_tokens":7,"output_tokens":3},"estimated_cost":"0.01"}`))
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_, _ = w.Write([]byte(`{"status":"completed","output":[{"type":"message","content":[{"type":"output_text","text":"{\"briefs\":[{\"headline\":\"فرصت مناسب\",\"interpretation\":\"همراستا با شواهد است.\",\"offerSuggestion\":\"پیشنهاد آماده کنید.\",\"messageSuggestion\":\"پیام روشن ارسال کنید.\",\"preparationAction\":\"مسیر پرداخت را بررسی کنید.\",\"limitation\":\"نتیجه قطعی نیست.\",\"evidenceIds\":[\"event_teacher\"]}]}"}]}],"usage":{"input_tokens":11,"output_tokens":5},"estimated_cost":"0.02"}`))
|
||||||
|
}))
|
||||||
|
defer server.Close()
|
||||||
|
|
||||||
|
client := NewAvalAIClient(AvalAIConfig{BaseURL: server.URL, APIKey: "test", Model: "test", Timeout: time.Second}, server.Client())
|
||||||
|
_, usage, err := client.Generate(context.Background(), json.RawMessage(`{"evidence":"event_teacher"}`), map[string]struct{}{"event_teacher": {}})
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if requests != 2 || usage.ProviderRequests != 2 || usage.InputTokens != 18 || usage.OutputTokens != 8 || usage.EstimatedCost != "0.03" {
|
||||||
|
t.Fatalf("requests=%d usage=%+v", requests, usage)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func TestValidateAICopyRejectsNumbersAndUnknownEvidence(t *testing.T) {
|
func TestValidateAICopyRejectsNumbersAndUnknownEvidence(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
tests := []struct {
|
tests := []struct {
|
||||||
@@ -19,6 +45,8 @@ func TestValidateAICopyRejectsNumbersAndUnknownEvidence(t *testing.T) {
|
|||||||
}{
|
}{
|
||||||
{name: "grounded Persian copy", copy: AICopy{Briefs: []CampaignBrief{{Headline: "فرصت مناسبی پیش رو دارید", Interpretation: "این مناسبت با الگوی موفق گذشته همراستا است.", EvidenceIDs: []string{"event_teacher"}}}}, want: true},
|
{name: "grounded Persian copy", copy: AICopy{Briefs: []CampaignBrief{{Headline: "فرصت مناسبی پیش رو دارید", Interpretation: "این مناسبت با الگوی موفق گذشته همراستا است.", EvidenceIDs: []string{"event_teacher"}}}}, want: true},
|
||||||
{name: "numeric claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "رشد ۲۰ درصدی", EvidenceIDs: []string{"event_teacher"}}}}, want: false},
|
{name: "numeric claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "رشد ۲۰ درصدی", EvidenceIDs: []string{"event_teacher"}}}}, want: false},
|
||||||
|
{name: "Persian number-word claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "رشد بیست درصدی", EvidenceIDs: []string{"event_teacher"}}}}, want: false},
|
||||||
|
{name: "English number-word claim", copy: AICopy{Briefs: []CampaignBrief{{Headline: "growth by twenty percent", EvidenceIDs: []string{"event_teacher"}}}}, want: false},
|
||||||
{name: "unknown evidence", copy: AICopy{Briefs: []CampaignBrief{{Headline: "فرصت مناسب", EvidenceIDs: []string{"unknown"}}}}, want: false},
|
{name: "unknown evidence", copy: AICopy{Briefs: []CampaignBrief{{Headline: "فرصت مناسب", EvidenceIDs: []string{"unknown"}}}}, want: false},
|
||||||
}
|
}
|
||||||
known := map[string]struct{}{"event_teacher": {}}
|
known := map[string]struct{}{"event_teacher": {}}
|
||||||
|
|||||||
@@ -37,17 +37,17 @@ func (c *Coordinator) AllowCreate(ctx context.Context, userID int64, now time.Ti
|
|||||||
func (c *Coordinator) AllowUncached(ctx context.Context, merchantID int64, now time.Time) error {
|
func (c *Coordinator) AllowUncached(ctx context.Context, merchantID int64, now time.Time) error {
|
||||||
return c.increment(ctx, fmt.Sprintf("growth:scope:%d:%s", merchantID, tehranDate(now)), c.config.UncachedPerDay, 36*time.Hour, now)
|
return c.increment(ctx, fmt.Sprintf("growth:scope:%d:%s", merchantID, tehranDate(now)), c.config.UncachedPerDay, 36*time.Hour, now)
|
||||||
}
|
}
|
||||||
func (c *Coordinator) ReserveBudget(ctx context.Context, tokens int64, now time.Time) error {
|
func (c *Coordinator) ReserveBudget(ctx context.Context, requests, tokens int64, now time.Time) error {
|
||||||
if c.config.DailyRequests < 1 || c.config.DailyTokens < 1 {
|
if c.config.DailyRequests < 1 || c.config.DailyTokens < 1 {
|
||||||
return ErrBudgetExhausted
|
return ErrBudgetExhausted
|
||||||
}
|
}
|
||||||
if c.redis == nil {
|
if c.redis == nil || requests < 1 || tokens < 1 || requests > c.config.DailyRequests || tokens > c.config.DailyTokens {
|
||||||
return ErrBudgetExhausted
|
return ErrBudgetExhausted
|
||||||
}
|
}
|
||||||
if err := c.incrementRedis(ctx, "growth:budget:req:"+tehranDate(now), 1, c.config.DailyRequests, 36*time.Hour); err != nil {
|
script := redis.NewScript(`local req=tonumber(redis.call('GET',KEYS[1]) or '0');local tok=tonumber(redis.call('GET',KEYS[2]) or '0');local addReq=tonumber(ARGV[1]);local addTok=tonumber(ARGV[2]);if req+addReq>tonumber(ARGV[4]) or tok+addTok>tonumber(ARGV[5]) then return 0 end;redis.call('INCRBY',KEYS[1],addReq);redis.call('INCRBY',KEYS[2],addTok);redis.call('PEXPIRE',KEYS[1],ARGV[3]);redis.call('PEXPIRE',KEYS[2],ARGV[3]);return 1`)
|
||||||
return ErrBudgetExhausted
|
date := tehranDate(now)
|
||||||
}
|
allowed, err := script.Run(ctx, c.redis, []string{"growth:budget:req:" + date, "growth:budget:tok:" + date}, requests, tokens, (36 * time.Hour).Milliseconds(), c.config.DailyRequests, c.config.DailyTokens).Int()
|
||||||
if err := c.incrementRedis(ctx, "growth:budget:tok:"+tehranDate(now), tokens, c.config.DailyTokens, 36*time.Hour); err != nil {
|
if err != nil || allowed != 1 {
|
||||||
return ErrBudgetExhausted
|
return ErrBudgetExhausted
|
||||||
}
|
}
|
||||||
return nil
|
return nil
|
||||||
|
|||||||
@@ -20,7 +20,7 @@ func TestCoordinatorFallsBackToBoundedProcessCounters(t *testing.T) {
|
|||||||
if err := coordinator.AllowCreate(context.Background(), 7, now); !errors.Is(err, ErrRateLimited) {
|
if err := coordinator.AllowCreate(context.Background(), 7, now); !errors.Is(err, ErrRateLimited) {
|
||||||
t.Fatalf("third request error=%v", err)
|
t.Fatalf("third request error=%v", err)
|
||||||
}
|
}
|
||||||
if err := coordinator.ReserveBudget(context.Background(), 1600, now); !errors.Is(err, ErrBudgetExhausted) {
|
if err := coordinator.ReserveBudget(context.Background(), MaxProviderRequests, 1600, now); !errors.Is(err, ErrBudgetExhausted) {
|
||||||
t.Fatalf("budget without Redis error=%v", err)
|
t.Fatalf("budget without Redis error=%v", err)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -41,6 +41,8 @@ type Usage struct {
|
|||||||
InputTokens int64 `json:"inputTokens"`
|
InputTokens int64 `json:"inputTokens"`
|
||||||
OutputTokens int64 `json:"outputTokens"`
|
OutputTokens int64 `json:"outputTokens"`
|
||||||
LatencyMS int64 `json:"latencyMs"`
|
LatencyMS int64 `json:"latencyMs"`
|
||||||
|
ProviderRequests int64 `json:"providerRequests"`
|
||||||
|
ReservedTokens int64 `json:"reservedTokens"`
|
||||||
ProviderRequestID string `json:"providerRequestId,omitempty"`
|
ProviderRequestID string `json:"providerRequestId,omitempty"`
|
||||||
EstimatedCost string `json:"estimatedCost,omitempty"`
|
EstimatedCost string `json:"estimatedCost,omitempty"`
|
||||||
}
|
}
|
||||||
@@ -77,6 +79,7 @@ type Record struct {
|
|||||||
Usage Usage
|
Usage Usage
|
||||||
ManualRefresh bool
|
ManualRefresh bool
|
||||||
TehranRefreshDay *time.Time
|
TehranRefreshDay *time.Time
|
||||||
|
ClaimToken string
|
||||||
}
|
}
|
||||||
|
|
||||||
type CreateRequest struct {
|
type CreateRequest struct {
|
||||||
|
|||||||
@@ -15,6 +15,8 @@ import (
|
|||||||
|
|
||||||
var ErrDisabled = errors.New("growth report: disabled")
|
var ErrDisabled = errors.New("growth report: disabled")
|
||||||
|
|
||||||
|
const MaxProviderRequests int64 = 4
|
||||||
|
|
||||||
type ServiceConfig struct {
|
type ServiceConfig struct {
|
||||||
Enabled bool
|
Enabled bool
|
||||||
Model, PromptVersion, SchemaVersion string
|
Model, PromptVersion, SchemaVersion string
|
||||||
@@ -42,13 +44,15 @@ func (s *Service) Create(ctx context.Context, userID, merchantID int64, request
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return Record{}, false, err
|
return Record{}, false, err
|
||||||
}
|
}
|
||||||
fingerprint := growth.Fingerprint(growth.FingerprintInput{MerchantID: merchantID, TerminalID: request.TerminalID, DatasetWatermark: watermark, ForecastStart: start, HorizonDays: 90, CatalogVersion: growth.CatalogVersion, AnalyticsVersion: growth.AnalyticsVersion, PromptVersion: s.config.PromptVersion, SchemaVersion: s.config.SchemaVersion, Model: s.config.Model, ModelConfigVersion: growth.DefaultModelConfigVersion})
|
fingerprint := s.fingerprint(merchantID, request.TerminalID, watermark, start)
|
||||||
var refreshDay *time.Time
|
var refreshDay *time.Time
|
||||||
if request.Refresh {
|
if request.Refresh {
|
||||||
day := tehranMidnight(now)
|
day := tehranMidnight(now)
|
||||||
refreshDay = &day
|
refreshDay = &day
|
||||||
}
|
}
|
||||||
record, cached, err := s.repository.CreateOrReuse(ctx, CreateRecord{ID: newReportID(), MerchantID: merchantID, TerminalID: request.TerminalID, Fingerprint: fingerprint, DatasetWatermark: watermark, ForecastStartDate: &start, CatalogVersion: growth.CatalogVersion, AnalyticsVersion: growth.AnalyticsVersion, PromptVersion: s.config.PromptVersion, SchemaVersion: s.config.SchemaVersion, Model: s.config.Model, ManualRefresh: request.Refresh, RefreshDay: refreshDay, IdempotencyKey: idempotencyKey, RequestHash: RequestHash(merchantID, request.TerminalID, request.Refresh)})
|
record, cached, err := s.repository.CreateOrReuse(ctx, CreateRecord{ID: newReportID(), MerchantID: merchantID, TerminalID: request.TerminalID, Fingerprint: fingerprint, DatasetWatermark: watermark, ForecastStartDate: &start, CatalogVersion: growth.CatalogVersion, AnalyticsVersion: growth.AnalyticsVersion, PromptVersion: s.config.PromptVersion, SchemaVersion: s.config.SchemaVersion, Model: s.config.Model, ManualRefresh: request.Refresh, RefreshDay: refreshDay, IdempotencyKey: idempotencyKey, RequestHash: RequestHash(merchantID, request.TerminalID, request.Refresh)}, func() error {
|
||||||
|
return s.coordinator.AllowUncached(ctx, merchantID, now)
|
||||||
|
})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
if errors.Is(err, ErrRefreshLimited) {
|
if errors.Is(err, ErrRefreshLimited) {
|
||||||
next, quotaErr := s.repository.NextRefreshAt(ctx, merchantID, now)
|
next, quotaErr := s.repository.NextRefreshAt(ctx, merchantID, now)
|
||||||
@@ -59,12 +63,6 @@ func (s *Service) Create(ctx context.Context, userID, merchantID int64, request
|
|||||||
}
|
}
|
||||||
return Record{}, false, err
|
return Record{}, false, err
|
||||||
}
|
}
|
||||||
if !cached {
|
|
||||||
if err := s.coordinator.AllowUncached(ctx, merchantID, now); err != nil {
|
|
||||||
_ = s.repository.Complete(ctx, record.ReportID, StatusFailed, json.RawMessage(`{}`), nil, Usage{}, "UNCACHED_SCOPE_LIMIT")
|
|
||||||
return Record{}, false, err
|
|
||||||
}
|
|
||||||
}
|
|
||||||
record.Cached = cached
|
record.Cached = cached
|
||||||
record.NextRefreshAt, err = s.repository.NextRefreshAt(ctx, merchantID, now)
|
record.NextRefreshAt, err = s.repository.NextRefreshAt(ctx, merchantID, now)
|
||||||
return record, cached, err
|
return record, cached, err
|
||||||
@@ -80,11 +78,21 @@ func (s *Service) Latest(ctx context.Context, merchantID int64, terminalID strin
|
|||||||
record, err := s.repository.Latest(ctx, merchantID, terminalID)
|
record, err := s.repository.Latest(ctx, merchantID, terminalID)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
record.Cached = true
|
record.Cached = true
|
||||||
|
watermark, start, metadataErr := s.repository.DatasetMetadata(ctx)
|
||||||
|
if metadataErr != nil {
|
||||||
|
record.Stale = true
|
||||||
|
} else {
|
||||||
|
record.Stale = record.Fingerprint != s.fingerprint(merchantID, terminalID, watermark, start)
|
||||||
|
}
|
||||||
record.NextRefreshAt, _ = s.repository.NextRefreshAt(ctx, merchantID, s.now())
|
record.NextRefreshAt, _ = s.repository.NextRefreshAt(ctx, merchantID, s.now())
|
||||||
}
|
}
|
||||||
return record, err
|
return record, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (s *Service) fingerprint(merchantID int64, terminalID, watermark string, start time.Time) string {
|
||||||
|
return growth.Fingerprint(growth.FingerprintInput{MerchantID: merchantID, TerminalID: terminalID, DatasetWatermark: watermark, ForecastStart: start, HorizonDays: 90, CatalogVersion: growth.CatalogVersion, AnalyticsVersion: growth.AnalyticsVersion, PromptVersion: s.config.PromptVersion, SchemaVersion: s.config.SchemaVersion, Model: s.config.Model, ModelConfigVersion: growth.DefaultModelConfigVersion})
|
||||||
|
}
|
||||||
|
|
||||||
type Processor struct {
|
type Processor struct {
|
||||||
repository *Repository
|
repository *Repository
|
||||||
builder Builder
|
builder Builder
|
||||||
@@ -108,10 +116,10 @@ func (p *Processor) ProcessNext(ctx context.Context) error {
|
|||||||
}
|
}
|
||||||
analysis, err := p.builder.Build(ctx, record.MerchantID, terminal)
|
analysis, err := p.builder.Build(ctx, record.MerchantID, terminal)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
_ = p.repository.Complete(ctx, record.ReportID, StatusFailed, json.RawMessage(`{}`), nil, Usage{}, "ANALYTICS_FAILED")
|
_ = p.repository.CompleteClaimed(ctx, record.ReportID, record.ClaimToken, StatusFailed, json.RawMessage(`{}`), nil, Usage{}, "ANALYTICS_FAILED")
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
if err := p.repository.StoreEvidence(ctx, record.ReportID, analysis.Encoded); err != nil {
|
if err := p.repository.StoreEvidence(ctx, record.ReportID, record.ClaimToken, analysis.Encoded); err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
copy := analysis.Fallback
|
copy := analysis.Fallback
|
||||||
@@ -119,30 +127,41 @@ func (p *Processor) ProcessNext(ctx context.Context) error {
|
|||||||
failure := "AI_DISABLED_OR_UNAVAILABLE"
|
failure := "AI_DISABLED_OR_UNAVAILABLE"
|
||||||
if analysis.ShouldCallAI && p.aiEnabled && p.avalai != nil {
|
if analysis.ShouldCallAI && p.aiEnabled && p.avalai != nil {
|
||||||
reservedTokens := EstimatedTokenReservation(analysis.Encoded)
|
reservedTokens := EstimatedTokenReservation(analysis.Encoded)
|
||||||
if err := p.coordinator.ReserveBudget(ctx, reservedTokens, p.now()); err == nil {
|
release, enterErr := p.coordinator.Enter(ctx)
|
||||||
release, enterErr := p.coordinator.Enter(ctx)
|
if enterErr == nil {
|
||||||
if enterErr == nil {
|
defer release()
|
||||||
|
if err := p.coordinator.ReserveBudget(ctx, MaxProviderRequests, reservedTokens, p.now()); err == nil {
|
||||||
generated, usage, aiErr := p.avalai.GenerateWithTools(ctx, analysis.Encoded, analysis.KnownEvidence, NewEvidenceTools(analysis))
|
generated, usage, aiErr := p.avalai.GenerateWithTools(ctx, analysis.Encoded, analysis.KnownEvidence, NewEvidenceTools(analysis))
|
||||||
release()
|
usage.ReservedTokens = reservedTokens
|
||||||
if aiErr == nil {
|
if aiErr == nil {
|
||||||
slog.Info("growth AI generation complete", "report_id", record.ReportID, "input_tokens", usage.InputTokens, "output_tokens", usage.OutputTokens, "latency_ms", usage.LatencyMS)
|
slog.Info("growth AI generation complete", "report_id", record.ReportID, "provider_requests", usage.ProviderRequests, "reserved_tokens", usage.ReservedTokens, "input_tokens", usage.InputTokens, "output_tokens", usage.OutputTokens, "latency_ms", usage.LatencyMS)
|
||||||
return p.repository.Complete(ctx, record.ReportID, StatusReady, analysis.Encoded, &generated, usage, "")
|
return p.repository.CompleteClaimed(ctx, record.ReportID, record.ClaimToken, StatusReady, analysis.Encoded, &generated, usage, "")
|
||||||
}
|
}
|
||||||
slog.Warn("growth AI output rejected", "report_id", record.ReportID, "reason", "AI_REJECTED")
|
slog.Warn("growth AI output rejected", "report_id", record.ReportID, "reason", "AI_REJECTED", "provider_requests", usage.ProviderRequests, "reserved_tokens", usage.ReservedTokens, "input_tokens", usage.InputTokens, "output_tokens", usage.OutputTokens, "latency_ms", usage.LatencyMS)
|
||||||
failure = "AI_REJECTED"
|
failure = "AI_REJECTED"
|
||||||
|
return p.repository.CompleteClaimed(ctx, record.ReportID, record.ClaimToken, status, analysis.Encoded, ©, usage, failure)
|
||||||
|
} else {
|
||||||
|
failure = "AI_BUDGET_EXHAUSTED"
|
||||||
|
slog.Warn("growth AI skipped", "report_id", record.ReportID, "reason", failure)
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
failure = "AI_BUDGET_EXHAUSTED"
|
failure = "AI_CONCURRENCY_UNAVAILABLE"
|
||||||
slog.Warn("growth AI skipped", "report_id", record.ReportID, "reason", failure)
|
slog.Warn("growth AI skipped", "report_id", record.ReportID, "reason", failure)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if !analysis.ShouldCallAI || !p.aiEnabled {
|
if !analysis.ShouldCallAI || !p.aiEnabled {
|
||||||
slog.Info("growth deterministic fallback used", "report_id", record.ReportID, "reason", failure)
|
slog.Info("growth deterministic fallback used", "report_id", record.ReportID, "reason", failure)
|
||||||
}
|
}
|
||||||
return p.repository.Complete(ctx, record.ReportID, status, analysis.Encoded, ©, Usage{}, failure)
|
return p.repository.CompleteClaimed(ctx, record.ReportID, record.ClaimToken, status, analysis.Encoded, ©, Usage{}, failure)
|
||||||
}
|
}
|
||||||
|
|
||||||
func EstimatedTokenReservation(evidence json.RawMessage) int64 { return int64(2*len(evidence) + 3200) }
|
// EstimatedTokenReservation reserves for two complete attempts, each with two
|
||||||
|
// provider turns, before any network call. Bytes are deliberately counted as
|
||||||
|
// tokens and the fixed allowance covers tool schemas, instructions, tool
|
||||||
|
// outputs, and every turn's 1,600-token output ceiling.
|
||||||
|
func EstimatedTokenReservation(evidence json.RawMessage) int64 {
|
||||||
|
return int64(12*len(evidence) + 32000)
|
||||||
|
}
|
||||||
|
|
||||||
func newReportID() string {
|
func newReportID() string {
|
||||||
data := make([]byte, 12)
|
data := make([]byte, 12)
|
||||||
|
|||||||
@@ -20,6 +20,8 @@ var (
|
|||||||
|
|
||||||
type Repository struct{ pool *pgxpool.Pool }
|
type Repository struct{ pool *pgxpool.Pool }
|
||||||
|
|
||||||
|
const reportLeaseDuration = 3 * time.Minute
|
||||||
|
|
||||||
func NewRepository(pool *pgxpool.Pool) *Repository { return &Repository{pool: pool} }
|
func NewRepository(pool *pgxpool.Pool) *Repository { return &Repository{pool: pool} }
|
||||||
|
|
||||||
type CreateRecord struct {
|
type CreateRecord struct {
|
||||||
@@ -40,7 +42,7 @@ type CreateRecord struct {
|
|||||||
RequestHash string
|
RequestHash string
|
||||||
}
|
}
|
||||||
|
|
||||||
func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord) (Record, bool, error) {
|
func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord, beforeInsert func() error) (Record, bool, error) {
|
||||||
tx, err := r.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
|
tx, err := r.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return Record{}, false, fmt.Errorf("beginning report transaction: %w", err)
|
return Record{}, false, fmt.Errorf("beginning report transaction: %w", err)
|
||||||
@@ -96,6 +98,11 @@ func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord) (Record
|
|||||||
return Record{}, false, ErrRefreshLimited
|
return Record{}, false, ErrRefreshLimited
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
if beforeInsert != nil {
|
||||||
|
if err := beforeInsert(); err != nil {
|
||||||
|
return Record{}, false, err
|
||||||
|
}
|
||||||
|
}
|
||||||
var revision int
|
var revision int
|
||||||
if err := tx.QueryRow(ctx, `SELECT coalesce(max(revision),0)+1 FROM growth_reports WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text`, in.MerchantID, nilIfEmpty(in.TerminalID)).Scan(&revision); err != nil {
|
if err := tx.QueryRow(ctx, `SELECT coalesce(max(revision),0)+1 FROM growth_reports WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text`, in.MerchantID, nilIfEmpty(in.TerminalID)).Scan(&revision); err != nil {
|
||||||
return Record{}, false, fmt.Errorf("calculating growth report revision: %w", err)
|
return Record{}, false, fmt.Errorf("calculating growth report revision: %w", err)
|
||||||
@@ -139,6 +146,9 @@ func (r *Repository) Claim(ctx context.Context) (Record, error) {
|
|||||||
return Record{}, fmt.Errorf("beginning report claim: %w", err)
|
return Record{}, fmt.Errorf("beginning report claim: %w", err)
|
||||||
}
|
}
|
||||||
defer func() { _ = tx.Rollback(ctx) }()
|
defer func() { _ = tx.Rollback(ctx) }()
|
||||||
|
if _, err := tx.Exec(ctx, `UPDATE growth_reports SET status='QUEUED',lease_token=NULL,lease_expires_at=NULL,updated_at=now(),failure_code='WORKER_LEASE_EXPIRED' WHERE status IN ('ANALYZING','GENERATING') AND (lease_expires_at IS NULL OR lease_expires_at<now())`); err != nil {
|
||||||
|
return Record{}, fmt.Errorf("recovering expired report leases: %w", err)
|
||||||
|
}
|
||||||
row := tx.QueryRow(ctx, reportSelect+` WHERE id=(SELECT id FROM growth_reports WHERE status='QUEUED' ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1)`)
|
row := tx.QueryRow(ctx, reportSelect+` WHERE id=(SELECT id FROM growth_reports WHERE status='QUEUED' ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1)`)
|
||||||
record, err := scanRecord(row)
|
record, err := scanRecord(row)
|
||||||
if errors.Is(err, ErrNotFound) {
|
if errors.Is(err, ErrNotFound) {
|
||||||
@@ -147,10 +157,13 @@ func (r *Repository) Claim(ctx context.Context) (Record, error) {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return Record{}, err
|
return Record{}, err
|
||||||
}
|
}
|
||||||
if _, err := tx.Exec(ctx, `UPDATE growth_reports SET status='ANALYZING',updated_at=now() WHERE id=$1`, record.ReportID); err != nil {
|
claimToken := newReportID()
|
||||||
|
leaseExpiresAt := time.Now().UTC().Add(reportLeaseDuration)
|
||||||
|
if _, err := tx.Exec(ctx, `UPDATE growth_reports SET status='ANALYZING',lease_token=$2,lease_expires_at=$3,updated_at=now() WHERE id=$1`, record.ReportID, claimToken, leaseExpiresAt); err != nil {
|
||||||
return Record{}, fmt.Errorf("claiming growth report: %w", err)
|
return Record{}, fmt.Errorf("claiming growth report: %w", err)
|
||||||
}
|
}
|
||||||
record.Status = StatusAnalyzing
|
record.Status = StatusAnalyzing
|
||||||
|
record.ClaimToken = claimToken
|
||||||
record.UpdatedAt = time.Now().UTC()
|
record.UpdatedAt = time.Now().UTC()
|
||||||
if err := tx.Commit(ctx); err != nil {
|
if err := tx.Commit(ctx); err != nil {
|
||||||
return Record{}, fmt.Errorf("committing report claim: %w", err)
|
return Record{}, fmt.Errorf("committing report claim: %w", err)
|
||||||
@@ -158,8 +171,8 @@ func (r *Repository) Claim(ctx context.Context) (Record, error) {
|
|||||||
return record, nil
|
return record, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (r *Repository) StoreEvidence(ctx context.Context, reportID string, evidence json.RawMessage) error {
|
func (r *Repository) StoreEvidence(ctx context.Context, reportID, claimToken string, evidence json.RawMessage) error {
|
||||||
command, err := r.pool.Exec(ctx, `UPDATE growth_reports SET deterministic_evidence=$2,status='GENERATING',updated_at=now() WHERE id=$1 AND status='ANALYZING'`, reportID, evidence)
|
command, err := r.pool.Exec(ctx, `UPDATE growth_reports SET deterministic_evidence=$3,status='GENERATING',lease_expires_at=now()+make_interval(secs => $4),updated_at=now() WHERE id=$1 AND lease_token=$2 AND status='ANALYZING'`, reportID, claimToken, evidence, int(reportLeaseDuration/time.Second))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("storing deterministic evidence: %w", err)
|
return fmt.Errorf("storing deterministic evidence: %w", err)
|
||||||
}
|
}
|
||||||
@@ -169,7 +182,7 @@ func (r *Repository) StoreEvidence(ctx context.Context, reportID string, evidenc
|
|||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
func (r *Repository) Complete(ctx context.Context, reportID string, status Status, evidence json.RawMessage, copy *AICopy, usage Usage, failureCode string) error {
|
func (r *Repository) CompleteClaimed(ctx context.Context, reportID, claimToken string, status Status, evidence json.RawMessage, copy *AICopy, usage Usage, failureCode string) error {
|
||||||
if status != StatusReady && status != StatusPartial && status != StatusFailed {
|
if status != StatusReady && status != StatusPartial && status != StatusFailed {
|
||||||
return errors.New("growth report: invalid terminal status")
|
return errors.New("growth report: invalid terminal status")
|
||||||
}
|
}
|
||||||
@@ -177,14 +190,17 @@ func (r *Repository) Complete(ctx context.Context, reportID string, status Statu
|
|||||||
if copy != nil {
|
if copy != nil {
|
||||||
ai = copy
|
ai = copy
|
||||||
}
|
}
|
||||||
_, err := r.pool.Exec(ctx, `UPDATE growth_reports SET status=$2,deterministic_evidence=$3,ai_response=$4,failure_code=$5,input_tokens=$6,output_tokens=$7,latency_ms=$8,provider_request_id=$9,estimated_cost=$10,updated_at=now(),completed_at=now() WHERE id=$1`, reportID, status, evidence, ai, nilIfEmpty(failureCode), usage.InputTokens, usage.OutputTokens, usage.LatencyMS, nilIfEmpty(usage.ProviderRequestID), nilIfEmpty(usage.EstimatedCost))
|
command, err := r.pool.Exec(ctx, `UPDATE growth_reports SET status=$3,deterministic_evidence=$4,ai_response=$5,failure_code=$6,input_tokens=$7,output_tokens=$8,latency_ms=$9,provider_requests=$10,reserved_tokens=$11,provider_request_id=$12,estimated_cost=$13,lease_token=NULL,lease_expires_at=NULL,updated_at=now(),completed_at=now() WHERE id=$1 AND lease_token=$2`, reportID, claimToken, status, evidence, ai, nilIfEmpty(failureCode), usage.InputTokens, usage.OutputTokens, usage.LatencyMS, usage.ProviderRequests, usage.ReservedTokens, nilIfEmpty(usage.ProviderRequestID), nilIfEmpty(usage.EstimatedCost))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("completing growth report: %w", err)
|
return fmt.Errorf("completing growth report: %w", err)
|
||||||
}
|
}
|
||||||
|
if command.RowsAffected() != 1 {
|
||||||
|
return ErrNotFound
|
||||||
|
}
|
||||||
return nil
|
return nil
|
||||||
}
|
}
|
||||||
|
|
||||||
const reportSelect = `SELECT id,revision,merchant_id,terminal_id,terminal_source_key,fingerprint,status,horizon_days,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,deterministic_evidence,ai_response,failure_code,input_tokens,output_tokens,latency_ms,provider_request_id,estimated_cost,manual_refresh,tehran_refresh_day,created_at,updated_at FROM growth_reports`
|
const reportSelect = `SELECT id,revision,merchant_id,terminal_id,terminal_source_key,fingerprint,status,horizon_days,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,deterministic_evidence,ai_response,failure_code,input_tokens,output_tokens,latency_ms,provider_requests,reserved_tokens,provider_request_id,estimated_cost,manual_refresh,tehran_refresh_day,created_at,updated_at FROM growth_reports`
|
||||||
|
|
||||||
type rowScanner interface{ Scan(...any) error }
|
type rowScanner interface{ Scan(...any) error }
|
||||||
type queryer interface {
|
type queryer interface {
|
||||||
@@ -205,7 +221,7 @@ func scanRecord(row rowScanner) (Record, error) {
|
|||||||
var ai []byte
|
var ai []byte
|
||||||
var terminalSource *string
|
var terminalSource *string
|
||||||
var providerRequestID, estimatedCost *string
|
var providerRequestID, estimatedCost *string
|
||||||
err := row.Scan(&r.ReportID, &r.Revision, &r.MerchantID, &r.TerminalDBID, &terminalSource, &r.Fingerprint, &status, &r.HorizonDays, &r.ForecastStartDate, &r.DatasetWatermark, &r.CatalogVersion, &r.AnalyticsVersion, &r.PromptVersion, &r.SchemaVersion, &r.Model, &evidence, &ai, &r.FailureCode, &r.Usage.InputTokens, &r.Usage.OutputTokens, &r.Usage.LatencyMS, &providerRequestID, &estimatedCost, &r.ManualRefresh, &r.TehranRefreshDay, &r.CreatedAt, &r.UpdatedAt)
|
err := row.Scan(&r.ReportID, &r.Revision, &r.MerchantID, &r.TerminalDBID, &terminalSource, &r.Fingerprint, &status, &r.HorizonDays, &r.ForecastStartDate, &r.DatasetWatermark, &r.CatalogVersion, &r.AnalyticsVersion, &r.PromptVersion, &r.SchemaVersion, &r.Model, &evidence, &ai, &r.FailureCode, &r.Usage.InputTokens, &r.Usage.OutputTokens, &r.Usage.LatencyMS, &r.Usage.ProviderRequests, &r.Usage.ReservedTokens, &providerRequestID, &estimatedCost, &r.ManualRefresh, &r.TehranRefreshDay, &r.CreatedAt, &r.UpdatedAt)
|
||||||
if errors.Is(err, pgx.ErrNoRows) {
|
if errors.Is(err, pgx.ErrNoRows) {
|
||||||
return Record{}, ErrNotFound
|
return Record{}, ErrNotFound
|
||||||
}
|
}
|
||||||
|
|||||||
81
internal/growthreport/repository_integration_test.go
Normal file
81
internal/growthreport/repository_integration_test.go
Normal file
@@ -0,0 +1,81 @@
|
|||||||
|
package growthreport
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"errors"
|
||||||
|
"fmt"
|
||||||
|
"os"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/jackc/pgx/v5/pgxpool"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestRepositoryRecoversExpiredActiveLeases(t *testing.T) {
|
||||||
|
pool, merchantID := growthRepositoryTestPool(t)
|
||||||
|
ctx := context.Background()
|
||||||
|
now := time.Now().UTC()
|
||||||
|
for index, status := range []Status{StatusAnalyzing, StatusGenerating} {
|
||||||
|
id := fmt.Sprintf("growth_%024x", index+1)
|
||||||
|
_, err := pool.Exec(ctx, `INSERT INTO growth_reports(id,merchant_id,fingerprint,status,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,lease_token,lease_expires_at,created_at,updated_at) VALUES($1,$2,$3,$4,current_date,'test','test','test','test','test','test',$5,$6,$7,$7)`, id, merchantID, id, status, "expired-token", now.Add(-time.Minute), now.Add(time.Duration(index)*time.Second))
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
repository := NewRepository(pool)
|
||||||
|
first, err := repository.Claim(ctx)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
if first.Status != StatusAnalyzing || first.ClaimToken == "" || first.ClaimToken == "expired-token" {
|
||||||
|
t.Fatalf("recovered claim=%+v", first)
|
||||||
|
}
|
||||||
|
if err := repository.CompleteClaimed(ctx, first.ReportID, "expired-token", StatusPartial, json.RawMessage(`{}`), nil, Usage{}, "stale"); !errors.Is(err, ErrNotFound) {
|
||||||
|
t.Fatalf("stale completion error=%v", err)
|
||||||
|
}
|
||||||
|
second, err := repository.Claim(ctx)
|
||||||
|
if err != nil || second.ReportID == first.ReportID || second.ClaimToken == "" {
|
||||||
|
t.Fatalf("second claim=%+v error=%v", second, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestRepositoryChecksUncachedAllowanceBeforeInsert(t *testing.T) {
|
||||||
|
pool, merchantID := growthRepositoryTestPool(t)
|
||||||
|
repository := NewRepository(pool)
|
||||||
|
start := time.Now().UTC()
|
||||||
|
input := CreateRecord{ID: "growth_ffffffffffffffffffffffff", MerchantID: merchantID, Fingerprint: "blocked", DatasetWatermark: "test", ForecastStartDate: &start, CatalogVersion: "test", AnalyticsVersion: "test", PromptVersion: "test", SchemaVersion: "test", Model: "test"}
|
||||||
|
_, _, err := repository.CreateOrReuse(context.Background(), input, func() error { return ErrRateLimited })
|
||||||
|
if !errors.Is(err, ErrRateLimited) {
|
||||||
|
t.Fatalf("CreateOrReuse error=%v", err)
|
||||||
|
}
|
||||||
|
var count int
|
||||||
|
if err := pool.QueryRow(context.Background(), `SELECT count(*) FROM growth_reports WHERE id=$1`, input.ID).Scan(&count); err != nil || count != 0 {
|
||||||
|
t.Fatalf("queued reports=%d error=%v", count, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func growthRepositoryTestPool(t *testing.T) (*pgxpool.Pool, int64) {
|
||||||
|
t.Helper()
|
||||||
|
databaseURL := os.Getenv("TEST_DATABASE_URL")
|
||||||
|
if databaseURL == "" {
|
||||||
|
t.Skip("TEST_DATABASE_URL is not configured")
|
||||||
|
}
|
||||||
|
pool, err := pgxpool.New(context.Background(), databaseURL)
|
||||||
|
if err != nil {
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
sourceKey := fmt.Sprintf("growth-test-%d", time.Now().UnixNano())
|
||||||
|
var merchantID int64
|
||||||
|
if err := pool.QueryRow(context.Background(), `INSERT INTO merchants(source_key,category_id,category_title,display_name) VALUES($1,1,'test','test') RETURNING id`, sourceKey).Scan(&merchantID); err != nil {
|
||||||
|
pool.Close()
|
||||||
|
t.Fatal(err)
|
||||||
|
}
|
||||||
|
t.Cleanup(func() {
|
||||||
|
_, _ = pool.Exec(context.Background(), `DELETE FROM growth_reports WHERE merchant_id=$1`, merchantID)
|
||||||
|
_, _ = pool.Exec(context.Background(), `DELETE FROM merchants WHERE id=$1`, merchantID)
|
||||||
|
pool.Close()
|
||||||
|
})
|
||||||
|
return pool, merchantID
|
||||||
|
}
|
||||||
@@ -27,11 +27,25 @@ func TestNextTehranRefreshUsesCalendarBoundary(t *testing.T) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func TestEstimatedTokenReservationCoversTwoTurnsAndGrowsWithEvidence(t *testing.T) {
|
func TestEstimatedTokenReservationCoversEveryRetryAndTurn(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
small := EstimatedTokenReservation([]byte("{}"))
|
small := EstimatedTokenReservation([]byte("{}"))
|
||||||
large := EstimatedTokenReservation(make([]byte, 1000))
|
large := EstimatedTokenReservation(make([]byte, 1000))
|
||||||
if small < 3200 || large <= small || large != 5200 {
|
if small < 4*1600 || large <= small || large != 44000 {
|
||||||
t.Fatalf("reservations small=%d large=%d", small, large)
|
t.Fatalf("reservations small=%d large=%d", small, large)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestServiceFingerprintChangesWithDatasetAndConfiguration(t *testing.T) {
|
||||||
|
t.Parallel()
|
||||||
|
start := time.Date(2026, 8, 21, 0, 0, 0, 0, time.UTC)
|
||||||
|
service := &Service{config: ServiceConfig{Model: "model-a", PromptVersion: "prompt-a", SchemaVersion: "schema-a"}}
|
||||||
|
current := service.fingerprint(7, "T-1", "watermark-a", start)
|
||||||
|
if current == service.fingerprint(7, "T-1", "watermark-b", start) {
|
||||||
|
t.Fatal("dataset watermark did not invalidate fingerprint")
|
||||||
|
}
|
||||||
|
service.config.PromptVersion = "prompt-b"
|
||||||
|
if current == service.fingerprint(7, "T-1", "watermark-a", start) {
|
||||||
|
t.Fatal("prompt version did not invalidate fingerprint")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -20,8 +20,12 @@ CREATE TABLE IF NOT EXISTS growth_reports (
|
|||||||
input_tokens bigint NOT NULL DEFAULT 0 CHECK (input_tokens >= 0),
|
input_tokens bigint NOT NULL DEFAULT 0 CHECK (input_tokens >= 0),
|
||||||
output_tokens bigint NOT NULL DEFAULT 0 CHECK (output_tokens >= 0),
|
output_tokens bigint NOT NULL DEFAULT 0 CHECK (output_tokens >= 0),
|
||||||
latency_ms bigint NOT NULL DEFAULT 0 CHECK (latency_ms >= 0),
|
latency_ms bigint NOT NULL DEFAULT 0 CHECK (latency_ms >= 0),
|
||||||
|
provider_requests bigint NOT NULL DEFAULT 0 CHECK (provider_requests >= 0),
|
||||||
|
reserved_tokens bigint NOT NULL DEFAULT 0 CHECK (reserved_tokens >= 0),
|
||||||
provider_request_id text,
|
provider_request_id text,
|
||||||
estimated_cost text,
|
estimated_cost text,
|
||||||
|
lease_token text,
|
||||||
|
lease_expires_at timestamptz,
|
||||||
manual_refresh boolean NOT NULL DEFAULT false,
|
manual_refresh boolean NOT NULL DEFAULT false,
|
||||||
tehran_refresh_day date,
|
tehran_refresh_day date,
|
||||||
created_at timestamptz NOT NULL DEFAULT now(),
|
created_at timestamptz NOT NULL DEFAULT now(),
|
||||||
@@ -33,6 +37,14 @@ ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS provider_request_id text;
|
|||||||
-- statement-break
|
-- statement-break
|
||||||
ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS estimated_cost text;
|
ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS estimated_cost text;
|
||||||
-- statement-break
|
-- statement-break
|
||||||
|
ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS provider_requests bigint NOT NULL DEFAULT 0 CHECK (provider_requests >= 0);
|
||||||
|
-- statement-break
|
||||||
|
ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS reserved_tokens bigint NOT NULL DEFAULT 0 CHECK (reserved_tokens >= 0);
|
||||||
|
-- statement-break
|
||||||
|
ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS lease_token text;
|
||||||
|
-- statement-break
|
||||||
|
ALTER TABLE growth_reports ADD COLUMN IF NOT EXISTS lease_expires_at timestamptz;
|
||||||
|
-- statement-break
|
||||||
ALTER TABLE growth_reports ALTER COLUMN deterministic_evidence DROP NOT NULL;
|
ALTER TABLE growth_reports ALTER COLUMN deterministic_evidence DROP NOT NULL;
|
||||||
-- statement-break
|
-- statement-break
|
||||||
CREATE UNIQUE INDEX IF NOT EXISTS growth_reports_current_fingerprint_idx
|
CREATE UNIQUE INDEX IF NOT EXISTS growth_reports_current_fingerprint_idx
|
||||||
|
|||||||
Reference in New Issue
Block a user