starcoach-hackathon/zarinpal-analytics-backend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: finalize growth cost controls

Browse Source
This commit is contained in:
Amirhossein Khalili
2026-08-21 00:56:59 +03:30
parent bb4a85cd65
commit 3e2d1f32c8
7 changed files with 86 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
internal/growthreport/avalai.go
View File

@@ -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 {

46
internal/growthreport/coordinator.go
View File

@@ -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)
}

8
internal/growthreport/model.go
View File

@@ -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 {

5
internal/growthreport/orchestrator.go
View File

@@ -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, &copy, 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 {

8
internal/growthreport/repository.go
View File

@@ -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
}

9
internal/growthreport/service_test.go
View File

@@ -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)
}
}

2
migrations/000002_growth_analyst.sql
View File

@@ -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(),