llm-intelligence/scripts/live_pricing_smoke_runner.go

//go:build llm_script

package main

import (
	"bytes"
	"flag"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"time"
)

type pricingSmokeCheck struct {
	Name string
	URL  string
	Run  func() (string, error)
}

type pricingSmokeSummary struct {
	Source     string
	ModelCount int
	Operator   string
}

type pricingSmokeResult struct {
	Name       string
	URL        string
	Source     string
	Operator   string
	ModelCount int
	DurationMS int64
	Success    bool
	Error      string
}

func main() {
	loadSubscriptionImportEnv()

	var timeoutSeconds int
	var vertexURL string
	var cloudflareURL string
	var perplexityURL string
	var vertexFixture string
	var cloudflareFixture string
	var perplexityFixture string

	flag.IntVar(&timeoutSeconds, "timeout", 20, "请求超时（秒）")
	flag.StringVar(&vertexURL, "vertex-url", defaultVertexPricingURL, "Vertex AI 官方价格页")
	flag.StringVar(&cloudflareURL, "cloudflare-url", defaultCloudflarePricingFetchURL, "Cloudflare Workers AI 官方价格页")
	flag.StringVar(&perplexityURL, "perplexity-url", defaultPerplexityPricingFetchURL, "Perplexity API 官方模型页")
	flag.StringVar(&vertexFixture, "vertex-fixture", "", "Vertex AI fixture 文件")
	flag.StringVar(&cloudflareFixture, "cloudflare-fixture", "", "Cloudflare fixture 文件")
	flag.StringVar(&perplexityFixture, "perplexity-fixture", "", "Perplexity fixture 文件")
	flag.Parse()

	timeout := time.Duration(timeoutSeconds) * time.Second
	checks := []pricingSmokeCheck{
		buildVertexSmokeCheck(vertexURL, vertexFixture, timeout),
		buildCloudflareSmokeCheck(cloudflareURL, cloudflareFixture, timeout),
		buildPerplexitySmokeCheck(perplexityURL, perplexityFixture, timeout),
	}
	results := runPricingSmokeChecks(checks, time.Now)
	renderPricingSmokeTextReport(os.Stdout, results, time.Now())
	if hasFailedPricingSmoke(results) {
		os.Exit(1)
	}
}

func buildVertexSmokeCheck(url, fixture string, timeout time.Duration) pricingSmokeCheck {
	return pricingSmokeCheck{
		Name: "Vertex",
		URL:  url,
		Run: func() (string, error) {
			var out bytes.Buffer
			err := runVertexPricingImport(vertexPricingImportConfig{
				URL:     url,
				Fixture: fixture,
				DryRun:  true,
				Timeout: timeout,
			}, nil, &out)
			return strings.TrimSpace(out.String()), err
		},
	}
}

func buildCloudflareSmokeCheck(url, fixture string, timeout time.Duration) pricingSmokeCheck {
	return pricingSmokeCheck{
		Name: "Cloudflare",
		URL:  url,
		Run: func() (string, error) {
			var out bytes.Buffer
			err := runCloudflarePricingImport(cloudflarePricingImportConfig{
				URL:     url,
				Fixture: fixture,
				DryRun:  true,
				Timeout: timeout,
			}, nil, &out)
			return strings.TrimSpace(out.String()), err
		},
	}
}

func buildPerplexitySmokeCheck(url, fixture string, timeout time.Duration) pricingSmokeCheck {
	return pricingSmokeCheck{
		Name: "Perplexity",
		URL:  url,
		Run: func() (string, error) {
			var out bytes.Buffer
			err := runPerplexityPricingImport(perplexityPricingImportConfig{
				URL:     url,
				Fixture: fixture,
				DryRun:  true,
				Timeout: timeout,
			}, nil, &out)
			return strings.TrimSpace(out.String()), err
		},
	}
}

func runPricingSmokeChecks(checks []pricingSmokeCheck, now func() time.Time) []pricingSmokeResult {
	results := make([]pricingSmokeResult, 0, len(checks))
	for _, check := range checks {
		start := now()
		result := pricingSmokeResult{
			Name: check.Name,
			URL:  check.URL,
		}
		summaryLine, err := check.Run()
		result.DurationMS = now().Sub(start).Milliseconds()
		if err != nil {
			result.Error = err.Error()
			results = append(results, result)
			continue
		}
		summary, ok := parsePricingSmokeSummaryLine(summaryLine)
		if !ok {
			result.Error = fmt.Sprintf("invalid dry-run summary: %q", summaryLine)
			results = append(results, result)
			continue
		}
		result.Success = true
		result.Source = summary.Source
		result.Operator = summary.Operator
		result.ModelCount = summary.ModelCount
		results = append(results, result)
	}
	return results
}

func parsePricingSmokeSummaryLine(line string) (pricingSmokeSummary, bool) {
	fields := strings.Fields(strings.TrimSpace(line))
	if len(fields) == 0 {
		return pricingSmokeSummary{}, false
	}

	source := ""
	modelCount := -1
	operatorParts := make([]string, 0)
	capturingOperator := false
	for _, field := range fields {
		switch {
		case strings.HasPrefix(field, "source="):
			source = strings.TrimPrefix(field, "source=")
			capturingOperator = false
		case strings.HasPrefix(field, "models="):
			value := strings.TrimPrefix(field, "models=")
			parsed, err := strconv.Atoi(value)
			if err != nil {
				return pricingSmokeSummary{}, false
			}
			modelCount = parsed
			capturingOperator = false
		case strings.HasPrefix(field, "operator="):
			capturingOperator = true
			operatorParts = append(operatorParts, strings.TrimPrefix(field, "operator="))
		case strings.HasPrefix(field, "dry_run="), strings.HasPrefix(field, "table_rows="):
			capturingOperator = false
		default:
			if capturingOperator {
				operatorParts = append(operatorParts, field)
			}
		}
	}
	if source == "" || modelCount < 0 || len(operatorParts) == 0 {
		return pricingSmokeSummary{}, false
	}
	return pricingSmokeSummary{
		Source:     source,
		ModelCount: modelCount,
		Operator:   strings.Join(operatorParts, " "),
	}, true
}

func renderPricingSmokeTextReport(out io.Writer, results []pricingSmokeResult, now time.Time) {
	passed := 0
	failed := 0
	_, _ = fmt.Fprintf(out, "=== Live Pricing Smoke Report (%s) ===\n", now.Format("2006-01-02 15:04"))
	for _, result := range results {
		if result.Success {
			passed++
			_, _ = fmt.Fprintf(out, "PASS %s source=%s models=%d operator=%s duration_ms=%d url=%s\n",
				result.Name, result.Source, result.ModelCount, result.Operator, result.DurationMS, result.URL)
			continue
		}
		failed++
		_, _ = fmt.Fprintf(out, "FAIL %s duration_ms=%d error=%s url=%s\n",
			result.Name, result.DurationMS, result.Error, result.URL)
	}
	_, _ = fmt.Fprintf(out, "Summary: %d passed, %d failed\n", passed, failed)
}

func hasFailedPricingSmoke(results []pricingSmokeResult) bool {
	for _, result := range results {
		if !result.Success {
			return true
		}
	}
	return false
}