chore: import upstream snapshot with attribution

tests/test_cortex_code_compression.py

@@ -0,0 +1,679 @@
+#!/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()