neuron/packages/stats/core/src/domain/inference.test.ts

import { describe, expect, test } from "bun:test"
import { buildStatsQueries, toGeoAggregate, toModelAggregate, toProviderAggregate } from "./inference"
import { modelAuthor, normalizeInferenceModel, statModel, statProvider } from "./model-normalization"

describe("inference stat normalization", () => {
  test("normalizes model suffixes used by router/provider variants", () => {
    expect(normalizeInferenceModel("deepseek-v4-flash-free")).toBe("deepseek-v4-flash")
    expect(normalizeInferenceModel("deepseek-v4-flash:global")).toBe("deepseek-v4-flash")
    expect(normalizeInferenceModel("mimo-v2.5-free")).toBe("mimo-v2.5")
    expect(normalizeInferenceModel("nemotron-3-super-free")).toBe("nemotron-3-super")
    expect(normalizeInferenceModel("mimo-v2.5-free:global")).toBe("mimo-v2.5")
    expect(normalizeInferenceModel("hy3-preview:free")).toBe("hy3-preview")
  })

  test("maps normalized model ids to public authors", () => {
    expect(modelAuthor("big-pickle")).toBe("unknown")
    expect(modelAuthor("claude-sonnet-4-5")).toBe("anthropic")
    expect(modelAuthor("deepseek-v4-pro")).toBe("deepseek")
    expect(modelAuthor("gemini-3.5-flash")).toBe("google")
    expect(modelAuthor("glm-5.1")).toBe("zhipu")
    expect(modelAuthor("gpt-5.5-pro")).toBe("openai")
    expect(modelAuthor("grok-build-0.1")).toBe("xai")
    expect(modelAuthor("hy3-preview")).toBe("tencent")
    expect(modelAuthor("kimi-k2.6")).toBe("moonshot")
    expect(modelAuthor("mimo-v2-omni")).toBe("xiaomi")
    expect(modelAuthor("minimax-m2.7")).toBe("minimax")
    expect(modelAuthor("muse-spark-1.2-contributor")).toBe("meta")
    expect(modelAuthor("nemotron-3-super-free")).toBe("nvidia")
    expect(modelAuthor("qwen3.7-max")).toBe("qwen")
    expect(modelAuthor("alpha-gpt-next")).toBeUndefined()
  })

  test("uses provider.model to resolve opencode route providers", () => {
    expect(statModel("big-pickle", "claude-sonnet-4-5")).toBe("claude-sonnet-4-5")
    expect(statModel("big-pickle", "gpt-5-free")).toBe("gpt-5")
    expect(statModel("big-pickle", "xiaomi/mimo-v2.5")).toBe("mimo-v2.5")
    expect(statModel("big-pickle", "")).toBe("unknown")
    expect(statProvider("big-pickle", "claude-sonnet-4-5", "opencode")).toBe("anthropic")
    expect(statProvider("big-pickle", "gpt-5", "opencode")).toBe("openai")
    expect(statProvider("big-pickle", "", "opencode")).toBe("unknown")
    expect(statProvider("unknown", "", "custom-provider")).toBe("custom-provider")
  })

  test("model aggregates prefer provider.model and use normalized model", () => {
    expect(toModelAggregate(aggregate("alpha-gpt-next", "openai"))).toEqual([])

    expect(toModelAggregate(aggregate("deepseek-v4-flash-free", "not-public-provider"))).toMatchObject([
      {
        period_key: "2026-05-20",
        provider: "deepseek",
        model: "deepseek-v4-flash",
      },
    ])

    expect(
      toModelAggregate({ ...aggregate("big-pickle", "opencode"), provider_model: "claude-sonnet-4-5" }),
    ).toMatchObject([
      {
        provider: "anthropic",
        model: "claude-sonnet-4-5",
        provider_model: "claude-sonnet-4-5",
      },
    ])
  })

  test("provider aggregates never keep opencode as the provider", () => {
    expect(toProviderAggregate({ ...aggregate("big-pickle", "opencode"), provider_model: "gpt-5" })).toMatchObject([
      { provider: "openai" },
    ])
    expect(toProviderAggregate(aggregate("big-pickle", "opencode"))).toMatchObject([{ provider: "unknown" }])
    expect(toProviderAggregate(aggregate("muse-spark-1.2-contributor", "unknown"))).toMatchObject([
      { provider: "meta" },
    ])
  })

  test("geo aggregates never keep opencode or big-pickle dimensions", () => {
    expect(toGeoAggregate({ ...aggregate("big-pickle", "opencode"), country: "US" })).toMatchObject([
      { provider: "unknown", model: "unknown", country: "US" },
    ])
  })

  test("model aggregates use ISO week period keys", () => {
    expect(
      toModelAggregate({
        ...aggregate("gpt-5.5-pro", "openai"),
        grain: "week",
        period_key: "2026-W20",
      }),
    ).toMatchObject([{ period_key: "2026-W20" }])
  })

  test("builds bounded R2 SQL queries for each day and week", () => {
    const queries = buildStatsQueries(new Date("2026-08-10T00:00:00.000Z"), new Date("2026-08-12T12:00:00.000Z"), {
      namespace: "inference",
      table: "generation",
      dataset: "zen",
    })

    expect(queries).toHaveLength(8)
    expect(queries[0]).toContain("'week' AS grain")
    expect(queries[0]).toContain("'2026-W33' AS period_key")
    expect(queries[2]).toContain("'2026-08-10' AS period_key")
    expect(queries[6]).toContain("'2026-08-12' AS period_key")
    expect(queries[0]).toContain('FROM "inference"."generation"')
    expect(queries[0]).toContain("event_type = 'generation.completed'")
    expect(queries[0]).toContain("AND (product = 'go' OR (lower(COALESCE(model_tier, '')) = 'free'")
    expect(queries[0]).toContain("COALESCE(NULLIF(lower(model_tier), ''), '') AS raw_tier")
    expect(queries[0]).toContain("WHEN lower(COALESCE(raw_tier, '')) = 'free'")
    expect(queries[0]).toContain("regexp_replace(NULLIF(route_model, ''), '^.*/', '')")
    expect(queries[0]).toContain("OR lower(raw_model) IN ('gpt-5-nano', 'grok-code', 'big-pickle')")
    expect(queries[0]).toContain("OR lower(raw_model) LIKE '%-free'")
    expect(queries[0]).toContain("THEN 'Free'")
    expect(queries[0]).toContain("LIMIT 10000")
    expect(queries[0]).toContain("approx_distinct(session) AS sessions")
    expect(queries[1]).toContain("'geo_model' ELSE 'geo'")
    expect(queries[1]).toContain("0 AS sessions")
  })

  test("aligns periods to UTC calendar boundaries", () => {
    const queries = buildStatsQueries(new Date("2026-06-17T15:56:00.000Z"), new Date("2026-06-19T15:56:00.000Z"), {
      namespace: "inference",
      table: "generation",
      dataset: "zen",
    })

    expect(queries).toHaveLength(8)
    expect(queries[0]).toContain("'2026-W25' AS period_key")
    expect(queries[0]).toContain("started_at >= '2026-06-15T00:00:00.000Z'")
    expect(queries[2]).toContain("'2026-06-17' AS period_key")
    expect(queries[2]).toContain("started_at >= '2026-06-17T00:00:00.000Z'")
    expect(queries[2]).toContain("started_at < '2026-06-18T00:00:00.000Z'")
    expect(queries[6]).toContain("'2026-06-19' AS period_key")
    expect(queries[6]).toContain("started_at < '2026-06-19T15:56:00.000Z'")
  })

  test("uses an exclusive live and legacy source handoff", () => {
    const [query] = buildStatsQueries(new Date("2026-08-11T00:00:00.000Z"), new Date("2026-08-12T00:00:00.000Z"), {
      namespace: "inference",
      table: "generation",
      dataset: "zen",
    })

    expect(query).toContain("(source = 'inference-legacy' AND started_at < '2026-08-11T10:57:48.186Z')")
    expect(query).toContain("(source = 'inference' AND started_at >= '2026-08-11T10:57:48.186Z')")
  })
})

function aggregate(model: string, provider: string) {
  return {
    grain: "day",
    period_key: "2026-05-20",
    dataset: "zen",
    tier: "Paid",
    provider,
    model,
    sessions: "1",
    requests: "1",
    sample_count: "1",
  }
}