headroomlabs-ai--headroom/tests/test_cortex_code_compression.py

#!/usr/bin/env python3
"""End-to-end token-savings test for Cortex Code (CoCo) + Headroom.

Simulates a real Cortex Code session using JSON-format tool results —
the format Snowflake's Python connector and most tool wrappers actually
emit.  Headroom's SmartCrusher compresses JSON natively without any ML
model, so this test works with the base install (no [ml] extra needed).

No API key required. Compression runs fully local.

Usage:
    # Benchmark (pretty-printed report):
    cd headroom && uv run python tests/test_cortex_code_compression.py

    # Pytest (CI-friendly assertions):
    cd headroom && uv run --with pytest pytest tests/test_cortex_code_compression.py -v -s
"""

from __future__ import annotations

import json
import time

MODEL = "claude-sonnet-4-5-20250929"

# ── Realistic CoCo JSON payload builders ─────────────────────────────────────


def snowflake_tables_json() -> str:
    """JSON array returned by INFORMATION_SCHEMA.TABLES — SmartCrusher target."""
    rows = [
        {
            "TABLE_CATALOG": "PROD_DB",
            "TABLE_SCHEMA": "ANALYTICS",
            "TABLE_NAME": f"FACT_ORDERS_{i:03d}",
            "TABLE_TYPE": "BASE TABLE",
            "ROW_COUNT": i * 1_423_001,
            "BYTES": i * 8_192_000,
            "CREATED": "2024-01-15T08:00:00Z",
            "LAST_ALTERED": "2025-06-10T14:22:00Z",
            "COMMENT": f"Daily order fact partition {i:03d}",
        }
        for i in range(1, 80)
    ]
    return json.dumps(rows, indent=2)


def snowflake_schema_json() -> str:
    """JSON array from DESCRIBE TABLE — repeated structure SmartCrusher loves."""
    base = [
        {
            "COLUMN_NAME": "order_id",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 36,
            "NULLABLE": False,
            "PRIMARY_KEY": True,
            "COMMENT": "UUID primary key",
        },
        {
            "COLUMN_NAME": "order_date",
            "DATA_TYPE": "DATE",
            "LENGTH": None,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Order placement date",
        },
        {
            "COLUMN_NAME": "customer_id",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 36,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "FK to dim_customers",
        },
        {
            "COLUMN_NAME": "region",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 50,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Sales region code",
        },
        {
            "COLUMN_NAME": "product_category",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 100,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Top-level product category",
        },
        {
            "COLUMN_NAME": "product_sku",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 50,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "FK to dim_products",
        },
        {
            "COLUMN_NAME": "quantity",
            "DATA_TYPE": "NUMBER",
            "LENGTH": None,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Units ordered",
        },
        {
            "COLUMN_NAME": "unit_price",
            "DATA_TYPE": "NUMBER",
            "LENGTH": None,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Price per unit USD",
        },
        {
            "COLUMN_NAME": "discount_pct",
            "DATA_TYPE": "NUMBER",
            "LENGTH": None,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Discount percentage 0-100",
        },
        {
            "COLUMN_NAME": "status",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 20,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Order lifecycle status",
        },
        {
            "COLUMN_NAME": "net_revenue",
            "DATA_TYPE": "NUMBER",
            "LENGTH": None,
            "NULLABLE": True,
            "PRIMARY_KEY": False,
            "COMMENT": "qty * price * (1-disc)",
        },
        {
            "COLUMN_NAME": "gross_profit",
            "DATA_TYPE": "NUMBER",
            "LENGTH": None,
            "NULLABLE": True,
            "PRIMARY_KEY": False,
            "COMMENT": "net_revenue - COGS",
        },
        {
            "COLUMN_NAME": "customer_tier",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 20,
            "NULLABLE": True,
            "PRIMARY_KEY": False,
            "COMMENT": "Gold/Silver/Bronze",
        },
        {
            "COLUMN_NAME": "acquisition_channel",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 50,
            "NULLABLE": True,
            "PRIMARY_KEY": False,
            "COMMENT": "How customer was acquired",
        },
        {
            "COLUMN_NAME": "created_at",
            "DATA_TYPE": "TIMESTAMP_NTZ",
            "LENGTH": None,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Row creation timestamp",
        },
        {
            "COLUMN_NAME": "updated_at",
            "DATA_TYPE": "TIMESTAMP_NTZ",
            "LENGTH": None,
            "NULLABLE": False,
            "PRIMARY_KEY": False,
            "COMMENT": "Last modified timestamp",
        },
        {
            "COLUMN_NAME": "_dbt_scd_id",
            "DATA_TYPE": "VARCHAR",
            "LENGTH": 36,
            "NULLABLE": True,
            "PRIMARY_KEY": False,
            "COMMENT": "dbt SCD type-2 surrogate key",
        },
        {
            "COLUMN_NAME": "_dbt_updated_at",
            "DATA_TYPE": "TIMESTAMP_NTZ",
            "LENGTH": None,
            "NULLABLE": True,
            "PRIMARY_KEY": False,
            "COMMENT": "dbt update marker",
        },
        {
            "COLUMN_NAME": "_dbt_valid_from",
            "DATA_TYPE": "TIMESTAMP_NTZ",
            "LENGTH": None,
            "NULLABLE": True,
            "PRIMARY_KEY": False,
            "COMMENT": "SCD validity start",
        },
        {
            "COLUMN_NAME": "_dbt_valid_to",
            "DATA_TYPE": "TIMESTAMP_NTZ",
            "LENGTH": None,
            "NULLABLE": True,
            "PRIMARY_KEY": False,
            "COMMENT": "SCD validity end",
        },
    ]
    # Three tables introspected in sequence — same schema, different table names
    result = []
    for table in ["stg_orders", "int_orders_enriched", "fct_revenue"]:
        for col in base:
            result.append({**col, "TABLE_NAME": table})
    return json.dumps(result, indent=2)


def dbt_run_results_json() -> str:
    """JSON run-results.json from a dbt invocation — realistic CoCo tool output."""
    nodes = [
        {
            "unique_id": f"model.analytics.{'stg_' if i < 10 else 'fct_'}model_{i:03d}",
            "status": "success" if i % 7 != 0 else "error",
            "execution_time": round(0.8 + i * 0.12, 3),
            "rows_affected": i * 12_500,
            "compiled_code": f"SELECT * FROM raw.orders_{i:03d} WHERE status = 'active'",
            "failures": None
            if i % 7 != 0
            else [{"message": f"Invalid identifier 'col_{i}' in select list", "line": i % 40 + 1}],
            "adapter_response": {
                "query_id": f"01b{i:06x}-0000-0001-0000-000300000001",
                "rows_produced": i * 12_500,
                "bytes_scanned": i * 8_192,
                "compilation_time": 0.05,
                "execution_time": round(0.8 + i * 0.12, 3),
            },
        }
        for i in range(40)
    ]
    return json.dumps(
        {"metadata": {"dbt_version": "1.8.0", "invocation_id": "abc123"}, "results": nodes},
        indent=2,
    )


def rag_cortex_search_json() -> str:
    """JSON results from a Cortex Search query — common in CoCo sessions."""
    docs = [
        {
            "rank": i + 1,
            "score": round(0.98 - i * 0.02, 4),
            "document_id": f"doc_{i:04d}",
            "source_table": "PROD_DB.DOCS.ENGINEERING_WIKI",
            "chunk_index": i % 5,
            "content": (
                "The revenue pipeline processes approximately 2.3 million orders per day "
                "across 14 regional data centers. Each order record contains pricing "
                "information, customer segmentation data, and fulfillment status. "
                "The dbt transformation layer applies discount calculations and joins "
                "to the customer dimension table to derive net revenue and gross profit "
                "metrics. Incremental models refresh every 4 hours using Snowflake "
                "dynamic tables as the upstream source. Known issue: the product_family "
                "column was renamed to product_group in Q3 2024; models referencing "
                "the old column name will fail with SQL compilation error 001003. "
                "Migration guide: update all references from product_family to product_group "
                "in models/marts/revenue/ and run dbt run --full-refresh."
            ),
            "metadata": {
                "author": f"engineer_{i % 8}@company.com",
                "last_updated": "2025-05-20",
                "tags": ["dbt", "revenue", "snowflake", "migration"],
            },
        }
        for i in range(15)
    ]
    return json.dumps(docs, indent=2)


def build_coco_session_messages() -> list[dict]:
    """Multi-turn CoCo session: diagnose a failing dbt model via Snowflake tools.

    Turn structure mirrors what CoCo actually does:
      1. User asks to fix fct_revenue
      2. CoCo queries table catalog  (→ large JSON tool result)
      3. CoCo introspects schema     (→ large JSON tool result)
      4. CoCo runs dbt, reads results (→ large JSON tool result)
      5. CoCo searches the wiki      (→ large JSON tool result)
      6. User asks follow-up
    """
    return [
        {
            "role": "user",
            "content": (
                "My dbt model fct_revenue is failing in prod with SQL compilation error 001003. "
                "Check the table catalog, inspect the schema, run dbt, and search the wiki for any "
                "known migration guides. Then tell me exactly what to fix."
            ),
        },
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_tables",
                    "type": "function",
                    "function": {
                        "name": "snowflake_query",
                        "arguments": json.dumps(
                            {
                                "sql": "SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = 'ANALYTICS'"
                            }
                        ),
                    },
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_tables",
            "content": snowflake_tables_json(),
        },
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_schema",
                    "type": "function",
                    "function": {
                        "name": "snowflake_query",
                        "arguments": json.dumps(
                            {"sql": "DESCRIBE TABLE PROD_DB.ANALYTICS.FCT_REVENUE"}
                        ),
                    },
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_schema",
            "content": snowflake_schema_json(),
        },
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_dbt",
                    "type": "function",
                    "function": {
                        "name": "bash",
                        "arguments": json.dumps(
                            {"command": "dbt run --select fct_revenue --target prod 2>&1"}
                        ),
                    },
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_dbt",
            "content": dbt_run_results_json(),
        },
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "call_search",
                    "type": "function",
                    "function": {
                        "name": "cortex_search",
                        "arguments": json.dumps(
                            {"query": "product_family column rename migration fct_revenue"}
                        ),
                    },
                }
            ],
        },
        {
            "role": "tool",
            "tool_call_id": "call_search",
            "content": rag_cortex_search_json(),
        },
        {
            "role": "assistant",
            "content": (
                "Found it. The column `product_family` was renamed to `product_group` in Q3 2024. "
                "The fix is to update line 47 of `models/marts/revenue/fct_revenue.sql` and run "
                "`dbt run --select fct_revenue --full-refresh`."
            ),
        },
        {
            "role": "user",
            "content": "Perfect. Are there any other models in models/marts/revenue/ that reference product_family?",
        },
    ]


# ── Helpers ───────────────────────────────────────────────────────────────────


def _count_tokens_approx(messages: list[dict]) -> int:
    """Approximate token count from serialised JSON (~4 chars/token)."""
    return len(json.dumps(messages)) // 4


def _table_row(label: str, before: int, after: int) -> str:
    saved = before - after
    pct = saved / max(before, 1) * 100
    bar = "█" * int(pct / 5)
    return f"  {label:<35} {before:>7,} → {after:>7,}   {pct:>5.1f}%  {bar}"


# ── Pytest tests ──────────────────────────────────────────────────────────────


def test_cortex_code_headroom_compression_saves_tokens() -> None:
    """Headroom must compress a realistic multi-turn CoCo session."""
    from headroom import compress

    messages = build_coco_session_messages()

    t0 = time.perf_counter()
    result = compress(messages, model=MODEL)
    latency_ms = (time.perf_counter() - t0) * 1000

    _ = result.tokens_saved / max(result.tokens_before, 1) * 100
    print(f"\n{_table_row('Full CoCo session', result.tokens_before, result.tokens_after)}")
    print(f"  Latency: {latency_ms:.0f} ms   Transforms: {', '.join(result.transforms_applied)}")

    assert result.tokens_saved > 0, (
        f"Expected compression on the multi-turn CoCo session. "
        f"before={result.tokens_before}, after={result.tokens_after}. "
        f"Transforms: {result.transforms_applied}"
    )
    assert len(result.messages) == len(messages), "Message count must not change"
    assert result.messages[0]["content"] == messages[0]["content"], "User prompt must be verbatim"


def test_cortex_code_tool_results_are_compressed_not_user_turns() -> None:
    """User turn content must be identical before and after compression."""
    from headroom import compress

    messages = build_coco_session_messages()
    result = compress(messages, model=MODEL)

    user_orig = [m for m in messages if m.get("role") == "user"]
    user_comp = [m for m in result.messages if m.get("role") == "user"]

    assert len(user_orig) == len(user_comp)
    for orig, comp in zip(user_orig, user_comp):
        assert orig["content"] == comp["content"], (
            f"User turn was mutated:\n  before: {orig['content'][:80]!r}"
        )


def test_cortex_code_tables_json_compresses() -> None:
    """Large Snowflake INFORMATION_SCHEMA result (JSON) must compress."""
    from headroom import compress

    messages = [
        {"role": "user", "content": "List all tables in ANALYTICS schema."},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c1",
                    "type": "function",
                    "function": {
                        "name": "snowflake_query",
                        "arguments": json.dumps({"sql": "SELECT * FROM INFORMATION_SCHEMA.TABLES"}),
                    },
                }
            ],
        },
        {"role": "tool", "tool_call_id": "c1", "content": snowflake_tables_json()},
    ]

    result = compress(messages, model=MODEL)
    _ = result.tokens_saved / max(result.tokens_before, 1) * 100
    print(f"\n{_table_row('Tables JSON (79 rows)', result.tokens_before, result.tokens_after)}")

    assert result.tokens_saved > 0, (
        f"INFORMATION_SCHEMA tables JSON was not compressed. "
        f"before={result.tokens_before}, after={result.tokens_after}. "
        f"Payload size: {len(snowflake_tables_json())} chars."
    )


def test_cortex_code_rag_search_json_compresses() -> None:
    """Cortex Search JSON results (repeated structure) must compress."""
    from headroom import compress

    messages = [
        {"role": "user", "content": "Search for product_family migration guide."},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c2",
                    "type": "function",
                    "function": {
                        "name": "cortex_search",
                        "arguments": json.dumps({"query": "product_family rename"}),
                    },
                }
            ],
        },
        {"role": "tool", "tool_call_id": "c2", "content": rag_cortex_search_json()},
    ]

    result = compress(messages, model=MODEL)
    _ = result.tokens_saved / max(result.tokens_before, 1) * 100
    print(
        f"\n{_table_row('Cortex Search JSON (15 docs)', result.tokens_before, result.tokens_after)}"
    )

    assert result.tokens_saved > 0, (
        f"Cortex Search JSON was not compressed. "
        f"before={result.tokens_before}, after={result.tokens_after}."
    )


def test_cortex_code_compression_is_lossless_on_key_content() -> None:
    """Key answer tokens must survive compression (the model can still answer)."""
    from headroom import compress

    messages = [
        {"role": "user", "content": "Search wiki for product_family rename."},
        {
            "role": "assistant",
            "content": None,
            "tool_calls": [
                {
                    "id": "c3",
                    "type": "function",
                    "function": {
                        "name": "cortex_search",
                        "arguments": json.dumps({"query": "product_family"}),
                    },
                }
            ],
        },
        {"role": "tool", "tool_call_id": "c3", "content": rag_cortex_search_json()},
    ]

    result = compress(messages, model=MODEL)
    compressed_tool = next(
        (m.get("content", "") for m in result.messages if m.get("role") == "tool"), ""
    )

    # The critical answer ("product_group") must survive
    key_terms = ["product_group", "migration", "dbt", "fct_revenue"]
    found = [t for t in key_terms if t in str(compressed_tool)]
    assert len(found) >= 2, (
        f"Too many key terms lost in compression. "
        f"Found: {found}, missing: {[t for t in key_terms if t not in found]}. "
        f"Compressed output (first 500 chars): {str(compressed_tool)[:500]}"
    )


# ── Standalone benchmark ──────────────────────────────────────────────────────


if __name__ == "__main__":
    from headroom import compress

    print()
    print("=" * 65)
    print("  Cortex Code × Headroom  —  token savings benchmark")
    print("  (No API key needed — compression is fully local)")
    print("=" * 65)

    payloads = [
        ("Full CoCo session (10 turns)", build_coco_session_messages),
        (
            "INFORMATION_SCHEMA tables (79 rows)",
            lambda: [
                {"role": "user", "content": "List tables."},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "c1",
                            "type": "function",
                            "function": {"name": "q", "arguments": "{}"},
                        }
                    ],
                },
                {"role": "tool", "tool_call_id": "c1", "content": snowflake_tables_json()},
            ],
        ),
        (
            "Schema JSON (3 tables × 20 cols)",
            lambda: [
                {"role": "user", "content": "Describe schema."},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "c1",
                            "type": "function",
                            "function": {"name": "q", "arguments": "{}"},
                        }
                    ],
                },
                {"role": "tool", "tool_call_id": "c1", "content": snowflake_schema_json()},
            ],
        ),
        (
            "dbt run-results JSON (40 models)",
            lambda: [
                {"role": "user", "content": "Run dbt."},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "c1",
                            "type": "function",
                            "function": {"name": "q", "arguments": "{}"},
                        }
                    ],
                },
                {"role": "tool", "tool_call_id": "c1", "content": dbt_run_results_json()},
            ],
        ),
        (
            "Cortex Search JSON (15 docs)",
            lambda: [
                {"role": "user", "content": "Search wiki."},
                {
                    "role": "assistant",
                    "content": None,
                    "tool_calls": [
                        {
                            "id": "c1",
                            "type": "function",
                            "function": {"name": "q", "arguments": "{}"},
                        }
                    ],
                },
                {"role": "tool", "tool_call_id": "c1", "content": rag_cortex_search_json()},
            ],
        ),
    ]

    print(f"\n  {'Payload':<35} {'Before':>7}   {'After':>7}   {'Saved%':>6}  Bar")
    print(f"  {'─' * 35} {'─' * 7}   {'─' * 7}   {'─' * 6}  {'─' * 20}")

    total_before = total_after = 0
    for label, builder in payloads:
        msgs = builder()
        t0 = time.perf_counter()
        r = compress(msgs, model=MODEL)
        ms = (time.perf_counter() - t0) * 1000
        total_before += r.tokens_before
        total_after += r.tokens_after
        print(f"{_table_row(label, r.tokens_before, r.tokens_after)}  ({ms:.0f}ms)")

    total_saved = total_before - total_after
    total_pct = total_saved / max(total_before, 1) * 100
    print(f"\n  {'─' * 65}")
    print(f"{_table_row('TOTAL', total_before, total_after)}")
    print()
    if total_saved > 0:
        print(
            f"  PASS  headroom saved {total_saved:,} tokens ({total_pct:.0f}%) across all CoCo payload types"
        )
    else:
        print("  FAIL  no compression — run: pip install 'headroom-ai[all]'")
    print()