diff --git a/internal/growthreport/avalai.go b/internal/growthreport/avalai.go index e1e629a..45fb490 100644 --- a/internal/growthreport/avalai.go +++ b/internal/growthreport/avalai.go @@ -80,6 +80,8 @@ type responseEnvelope struct { InputTokens int64 `json:"input_tokens"` OutputTokens int64 `json:"output_tokens"` } `json:"usage"` + EstimatedCost json.RawMessage `json:"estimated_cost"` + ProviderRequestID string `json:"-"` } type outputItem struct { @@ -109,6 +111,8 @@ func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, eviden } } usage := Usage{InputTokens: first.Usage.InputTokens, OutputTokens: first.Usage.OutputTokens, LatencyMS: latency} + usage.ProviderRequestID = first.ProviderRequestID + usage.EstimatedCost = costString(first.EstimatedCost) final := first if len(calls) > 0 { if executor == nil { @@ -144,6 +148,12 @@ func (c *AvalAIClient) workflow(ctx context.Context, base map[string]any, eviden 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 for _, raw := range second.Output { var item outputItem @@ -195,9 +205,33 @@ func (c *AvalAIClient) do(parent context.Context, bodyValue map[string]any) (res if result.Status != "completed" { return responseEnvelope{}, 0, false, fmt.Errorf("avalai: incomplete response status %q", result.Status) } + result.ProviderRequestID = resp.Header.Get("X-Request-ID") return result, time.Since(started).Milliseconds(), false, nil } +func costString(raw json.RawMessage) string { + if len(raw) == 0 { + return "" + } + var text string + if json.Unmarshal(raw, &text) == nil { + if len(text) <= 64 { + return text + } + return "" + } + text = strings.TrimSpace(string(raw)) + if len(text) > 64 { + return "" + } + for _, r := range text { + if !(unicode.IsDigit(r) || r == '.' || r == '-' || r == 'e' || r == 'E') { + return "" + } + } + return text +} + func parseAICopy(result responseEnvelope, known map[string]struct{}) (AICopy, error) { var text string for _, raw := range result.Output { diff --git a/internal/growthreport/coordinator.go b/internal/growthreport/coordinator.go index 18cdcd8..6461afa 100644 --- a/internal/growthreport/coordinator.go +++ b/internal/growthreport/coordinator.go @@ -18,19 +18,18 @@ type CoordinatorConfig struct { MaxConcurrent int } type Coordinator struct { - redis *redis.Client - config CoordinatorConfig - mu sync.Mutex - counters map[string]int64 - expiries map[string]time.Time - semaphore chan struct{} + redis *redis.Client + config CoordinatorConfig + mu sync.Mutex + counters map[string]int64 + expiries map[string]time.Time } func NewCoordinator(client *redis.Client, config CoordinatorConfig) *Coordinator { if config.MaxConcurrent < 1 { config.MaxConcurrent = 4 } - return &Coordinator{redis: client, config: config, counters: map[string]int64{}, expiries: map[string]time.Time{}, semaphore: make(chan struct{}, config.MaxConcurrent)} + return &Coordinator{redis: client, config: config, counters: map[string]int64{}, expiries: map[string]time.Time{}} } func (c *Coordinator) AllowCreate(ctx context.Context, userID int64, now time.Time) error { return c.increment(ctx, fmt.Sprintf("growth:create:%d:%s", userID, now.UTC().Format("200601021504")), c.config.CreatePerMinute, time.Minute, now) @@ -66,19 +65,30 @@ func (c *Coordinator) incrementRedis(ctx context.Context, key string, amount, li return nil } func (c *Coordinator) Enter(ctx context.Context) (func(), error) { - select { - case c.semaphore <- struct{}{}: - return func() { <-c.semaphore }, nil - case <-ctx.Done(): - return nil, ctx.Err() + if c.redis == nil { + return nil, ErrBudgetExhausted + } + acquire := redis.NewScript(`local v=redis.call('INCR',KEYS[1]);if v==1 then redis.call('PEXPIRE',KEYS[1],ARGV[2]) end;if v>tonumber(ARGV[1]) then redis.call('DECR',KEYS[1]);return 0 end;return 1`) + for { + ok, err := acquire.Run(ctx, c.redis, []string{"growth:ai:concurrency"}, c.config.MaxConcurrent, (3 * time.Minute).Milliseconds()).Int() + if err != nil { + return nil, ErrBudgetExhausted + } + if ok == 1 { + return func() { + releaseCtx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + script := redis.NewScript(`local v=tonumber(redis.call('GET',KEYS[1]) or '0');if v>0 then return redis.call('DECR',KEYS[1]) end return 0`) + _ = script.Run(releaseCtx, c.redis, []string{"growth:ai:concurrency"}).Err() + }, nil + } + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-time.After(100 * time.Millisecond): + } } } -func (c *Coordinator) Publish(ctx context.Context, reportID string) { - if c.redis != nil { - _ = c.redis.LPush(ctx, "growth:jobs", reportID).Err() - } -} - func (c *Coordinator) increment(ctx context.Context, key string, limit int64, ttl time.Duration, now time.Time) error { return c.incrementBy(ctx, key, 1, limit, ttl, now) } diff --git a/internal/growthreport/model.go b/internal/growthreport/model.go index 16a36aa..8802772 100644 --- a/internal/growthreport/model.go +++ b/internal/growthreport/model.go @@ -38,9 +38,11 @@ type AICopy struct { } type Usage struct { - InputTokens int64 `json:"inputTokens"` - OutputTokens int64 `json:"outputTokens"` - LatencyMS int64 `json:"latencyMs"` + InputTokens int64 `json:"inputTokens"` + OutputTokens int64 `json:"outputTokens"` + LatencyMS int64 `json:"latencyMs"` + ProviderRequestID string `json:"providerRequestId,omitempty"` + EstimatedCost string `json:"estimatedCost,omitempty"` } type Report struct { diff --git a/internal/growthreport/orchestrator.go b/internal/growthreport/orchestrator.go index 993f938..dac1490 100644 --- a/internal/growthreport/orchestrator.go +++ b/internal/growthreport/orchestrator.go @@ -64,7 +64,6 @@ func (s *Service) Create(ctx context.Context, userID, merchantID int64, request _ = s.repository.Complete(ctx, record.ReportID, StatusFailed, json.RawMessage(`{}`), nil, Usage{}, "UNCACHED_SCOPE_LIMIT") return Record{}, false, err } - s.coordinator.Publish(ctx, record.ReportID) } record.Cached = cached record.NextRefreshAt, err = s.repository.NextRefreshAt(ctx, merchantID, now) @@ -119,7 +118,7 @@ func (p *Processor) ProcessNext(ctx context.Context) error { status := StatusPartial failure := "AI_DISABLED_OR_UNAVAILABLE" if analysis.ShouldCallAI && p.aiEnabled && p.avalai != nil { - reservedTokens := int64((len(analysis.Encoded)+3)/4*2 + 1600) + reservedTokens := EstimatedTokenReservation(analysis.Encoded) if err := p.coordinator.ReserveBudget(ctx, reservedTokens, p.now()); err == nil { release, enterErr := p.coordinator.Enter(ctx) if enterErr == nil { @@ -143,6 +142,8 @@ func (p *Processor) ProcessNext(ctx context.Context) error { return p.repository.Complete(ctx, record.ReportID, status, analysis.Encoded, ©, Usage{}, failure) } +func EstimatedTokenReservation(evidence json.RawMessage) int64 { return int64(2*len(evidence) + 3200) } + func newReportID() string { data := make([]byte, 12) if _, err := rand.Read(data); err != nil { diff --git a/internal/growthreport/repository.go b/internal/growthreport/repository.go index 7136cbd..fd15140 100644 --- a/internal/growthreport/repository.go +++ b/internal/growthreport/repository.go @@ -132,6 +132,8 @@ func (r *Repository) Latest(ctx context.Context, merchantID int64, terminalID st } func (r *Repository) Claim(ctx context.Context) (Record, error) { + // PostgreSQL is the durable queue and coordination authority. SKIP LOCKED plus + // the state transition prevents two worker processes from generating one row. tx, err := r.pool.Begin(ctx) if err != nil { return Record{}, fmt.Errorf("beginning report claim: %w", err) @@ -175,14 +177,14 @@ func (r *Repository) Complete(ctx context.Context, reportID string, status Statu if copy != nil { 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,updated_at=now(),completed_at=now() WHERE id=$1`, reportID, status, evidence, ai, nilIfEmpty(failureCode), usage.InputTokens, usage.OutputTokens, usage.LatencyMS) + _, 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)) if err != nil { return fmt.Errorf("completing growth report: %w", err) } 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,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_request_id,estimated_cost,manual_refresh,tehran_refresh_day,created_at,updated_at FROM growth_reports` type rowScanner interface{ Scan(...any) error } type queryer interface { @@ -202,7 +204,7 @@ func scanRecord(row rowScanner) (Record, error) { var evidence []byte var ai []byte var terminalSource *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, &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.ProviderRequestID, &r.Usage.EstimatedCost, &r.ManualRefresh, &r.TehranRefreshDay, &r.CreatedAt, &r.UpdatedAt) if errors.Is(err, pgx.ErrNoRows) { return Record{}, ErrNotFound } diff --git a/internal/growthreport/service_test.go b/internal/growthreport/service_test.go index 2ee7eff..60ddd46 100644 --- a/internal/growthreport/service_test.go +++ b/internal/growthreport/service_test.go @@ -26,3 +26,12 @@ func TestNextTehranRefreshUsesCalendarBoundary(t *testing.T) { t.Fatalf("NextTehranRefresh() = %s", next) } } + +func TestEstimatedTokenReservationCoversTwoTurnsAndGrowsWithEvidence(t *testing.T) { + t.Parallel() + small := EstimatedTokenReservation([]byte("{}")) + large := EstimatedTokenReservation(make([]byte, 1000)) + if small < 3200 || large <= small || large != 5200 { + t.Fatalf("reservations small=%d large=%d", small, large) + } +} diff --git a/migrations/000002_growth_analyst.sql b/migrations/000002_growth_analyst.sql index 7bc8335..813691d 100644 --- a/migrations/000002_growth_analyst.sql +++ b/migrations/000002_growth_analyst.sql @@ -20,6 +20,8 @@ CREATE TABLE IF NOT EXISTS growth_reports ( input_tokens bigint NOT NULL DEFAULT 0 CHECK (input_tokens >= 0), output_tokens bigint NOT NULL DEFAULT 0 CHECK (output_tokens >= 0), latency_ms bigint NOT NULL DEFAULT 0 CHECK (latency_ms >= 0), + provider_request_id text, + estimated_cost text, manual_refresh boolean NOT NULL DEFAULT false, tehran_refresh_day date, created_at timestamptz NOT NULL DEFAULT now(),