modules/cypher_context_compression.py

"""CYPHER V12 M31 — Context Compression (LLMLingua-inspired).

Compress long contexts (>2048 tokens) to fit within encoder max_enc=2048
while preserving important information. Uses importance scoring per token
or per chunk based on:
  - TF-IDF style rarity
  - Domain keyword density
  - Question entity overlap
  - Sentence position (lead+tail bias)

Replaces low-importance tokens with [...] markers or drops them.
Allows CYPHER to effectively use 4K+ contexts despite arch limit.
"""
from __future__ import annotations

import logging
import re
from collections import Counter
from typing import Any

logger = logging.getLogger(__name__)


# Stop words (light, FR+EN)
_STOPWORDS = {
    "the", "a", "an", "is", "are", "was", "were", "be", "been", "being",
    "of", "in", "on", "at", "for", "with", "and", "or", "but", "to",
    "from", "by", "this", "that", "these", "those", "it", "its", "i",
    "you", "we", "they", "he", "she", "him", "her", "his", "their",
    "le", "la", "les", "un", "une", "des", "du", "de", "et", "ou",
    "qui", "que", "où", "ce", "cette", "ces", "il", "elle", "ils",
    "elles", "je", "tu", "nous", "vous", "leur", "leurs", "mon", "ton",
    "son", "ma", "ta", "sa", "mes", "tes", "ses", "pour", "avec",
    "dans", "sur", "sans", "par", "est", "sont",
}


# High-value domain terms (auto-boost)
_VALUE_TERMS = {
    "cve", "cvss", "mitre", "att&ck", "yara", "sigma", "ioc", "rce",
    "sqli", "xss", "lateral", "persistence", "privilege", "exfiltration",
    "log4shell", "ransomware", "phishing", "malware", "apt",
    "btc", "eth", "smc", "fvg", "choch", "bos", "order block",
    "premium", "discount", "liquidity", "killzone",
}


def split_into_chunks(text: str, chunk_size: int = 100) -> list[str]:
    """Split by sentences first, then fallback to fixed-size words chunks."""
    sentences = re.split(r"(?<=[.!?])\s+", text)
    chunks: list[str] = []
    current = []
    cur_len = 0
    for s in sentences:
        s_tokens = s.split()
        if cur_len + len(s_tokens) > chunk_size:
            if current:
                chunks.append(" ".join(current))
                current = []
                cur_len = 0
        current.append(s)
        cur_len += len(s_tokens)
    if current:
        chunks.append(" ".join(current))
    return chunks


def importance_score(
    chunk: str,
    query: str,
    chunk_idx: int,
    total_chunks: int,
    global_freq: Counter,
) -> float:
    """Compute importance 0..1 for a chunk."""
    score = 0.0
    c_lower = chunk.lower()
    q_lower = query.lower() if query else ""

    # 1. Query entity overlap (most important)
    q_tokens = {t for t in re.findall(r"[a-zA-Z0-9-]+", q_lower)
                if len(t) > 3 and t not in _STOPWORDS}
    c_tokens = {t for t in re.findall(r"[a-zA-Z0-9-]+", c_lower)
                if len(t) > 3 and t not in _STOPWORDS}
    if q_tokens:
        overlap = len(q_tokens & c_tokens) / len(q_tokens)
        score += 0.40 * overlap
    # 2. CVE/T-ID/MITRE entity presence
    has_cve = bool(re.search(r"cve-\d{4}-\d+", c_lower))
    has_tid = bool(re.search(r"\bt\d{4}", c_lower))
    if has_cve:
        score += 0.20
    if has_tid:
        score += 0.15
    # 3. Value terms
    n_value = sum(1 for t in _VALUE_TERMS if t in c_lower)
    score += min(0.15, 0.03 * n_value)
    # 4. Rare terms (low global freq among non-stopwords)
    rare_score = 0.0
    for t in c_tokens:
        if t in global_freq and global_freq[t] <= 2:
            rare_score += 0.05
    score += min(0.10, rare_score)
    # 5. Position bias (lead+tail)
    if chunk_idx == 0 or chunk_idx == total_chunks - 1:
        score += 0.05
    return min(1.0, score)


class ContextCompressor:
    """Compress a long context to fit a target token budget."""

    def __init__(
        self,
        target_tokens: int = 1500,
        chunk_size: int = 80,
        min_keep_ratio: float = 0.3,
        ellipsis: str = "[...]",
    ):
        self.target_tokens = target_tokens
        self.chunk_size = chunk_size
        self.min_keep_ratio = min_keep_ratio
        self.ellipsis = ellipsis

    @staticmethod
    def estimate_tokens(text: str) -> int:
        # Rough estimate: 1 token ≈ 4 chars for English/French mix
        return len(text) // 4

    def compress(self, text: str, query: str = "") -> dict:
        if not text:
            return {"compressed": "", "ratio": 1.0, "tokens_before": 0,
                    "tokens_after": 0, "n_chunks_kept": 0}
        tokens_before = self.estimate_tokens(text)
        if tokens_before <= self.target_tokens:
            return {"compressed": text, "ratio": 1.0,
                    "tokens_before": tokens_before, "tokens_after": tokens_before,
                    "n_chunks_kept": 1}
        chunks = split_into_chunks(text, chunk_size=self.chunk_size)
        if not chunks:
            return {"compressed": text[:self.target_tokens * 4], "ratio": 0.0,
                    "tokens_before": tokens_before, "tokens_after": self.target_tokens,
                    "n_chunks_kept": 1}
        # Build global freq
        global_freq: Counter = Counter()
        for c in chunks:
            for t in re.findall(r"[a-zA-Z0-9-]+", c.lower()):
                if len(t) > 3 and t not in _STOPWORDS:
                    global_freq[t] += 1
        # Score chunks
        scored = []
        for i, c in enumerate(chunks):
            s = importance_score(c, query, i, len(chunks), global_freq)
            scored.append((i, s, c))
        # Keep top until budget
        scored_sorted = sorted(scored, key=lambda x: x[1], reverse=True)
        kept: list[tuple[int, str]] = []
        cum_tokens = 0
        for idx, _score, chunk_text in scored_sorted:
            chunk_t = self.estimate_tokens(chunk_text)
            if cum_tokens + chunk_t > self.target_tokens:
                # Try truncating chunk
                budget_left = self.target_tokens - cum_tokens
                if budget_left > 20:
                    chunk_text = chunk_text[: budget_left * 4]
                    kept.append((idx, chunk_text))
                break
            kept.append((idx, chunk_text))
            cum_tokens += chunk_t
            if cum_tokens >= self.target_tokens * (1 - self.min_keep_ratio):
                pass  # continue greedy
        # Reorder by original index
        kept.sort(key=lambda x: x[0])
        # Glue with ellipsis between non-contiguous
        parts: list[str] = []
        prev_idx = -2
        for idx, ctext in kept:
            if idx != prev_idx + 1 and prev_idx >= 0:
                parts.append(self.ellipsis)
            parts.append(ctext)
            prev_idx = idx
        if parts and prev_idx < len(chunks) - 1:
            parts.append(self.ellipsis)
        compressed = " ".join(parts)
        tokens_after = self.estimate_tokens(compressed)
        return {
            "compressed": compressed,
            "ratio": tokens_after / max(1, tokens_before),
            "tokens_before": tokens_before,
            "tokens_after": tokens_after,
            "n_chunks_kept": len(kept),
            "n_chunks_total": len(chunks),
        }


__all__ = ["ContextCompressor", "split_into_chunks", "importance_score"]


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    print("=== M31 cypher_context_compression SMOKE ===")
    long_text = (
        "CVE-2021-44228 (Log4Shell) is a critical RCE in Apache Log4j2 with CVSS 10.0. "
        "It abuses JNDI lookups. Many companies were affected. "
    ) * 50 + (
        "MITRE T1059 is Command and Scripting Interpreter. Many techniques exist. "
        "Background talk about software development and history. Linux Torvalds invented Linux. "
    ) * 30 + (
        "T1190 Exploit Public-Facing Application. CVE-2024-3400 Palo Alto PAN-OS critical. "
    ) * 20
    print(f"Original tokens (estimated): {ContextCompressor.estimate_tokens(long_text)}")
    compressor = ContextCompressor(target_tokens=500)
    result = compressor.compress(long_text, query="Tell me about Log4Shell and T1190")
    print(f"\nCompressed:")
    print(f"  Before: {result['tokens_before']} tokens")
    print(f"  After: {result['tokens_after']} tokens (ratio {result['ratio']:.2f})")
    print(f"  Chunks kept: {result['n_chunks_kept']}/{result['n_chunks_total']}")
    print(f"  Preview: {result['compressed'][:400]}...")
    print("\n=== SMOKE PASS ===")