modules/cypher_long_context.py

"""CYPHER V12 M43 — Long Context Handling.

Combines M31 compression + just-in-time RAG + chunked streaming to handle
contexts >4K tokens effectively despite arch max_enc=2048.

Strategy:
  1. If prompt has long context attached → split into chunks
  2. For each chunk: embed + index in M4 memory temporarily
  3. Retrieve most relevant chunks (top-K) via similarity to question
  4. Compress retrieved chunks via M31 if still too long
  5. Pass to encoder
  6. Cleanup ephemeral chunks after query
"""
from __future__ import annotations

import logging
import re
import uuid
from typing import Any

logger = logging.getLogger(__name__)


def split_long_context(text: str, chunk_chars: int = 1500, overlap_chars: int = 200) -> list[str]:
    """Sliding window chunk split for long documents."""
    if len(text) <= chunk_chars:
        return [text]
    chunks: list[str] = []
    i = 0
    while i < len(text):
        chunk = text[i:i + chunk_chars]
        chunks.append(chunk)
        i += chunk_chars - overlap_chars
    return chunks


def estimate_token_count(text: str) -> int:
    return len(text) // 4


class LongContextHandler:
    """Just-in-time RAG over ephemeral large documents."""

    def __init__(
        self,
        compressor=None,           # M31 ContextCompressor
        hier_memory=None,          # M41 HierarchicalMemory
        max_chunk_chars: int = 1500,
        target_final_tokens: int = 1500,
        top_k_chunks: int = 4,
    ):
        self.compressor = compressor
        self.hier_memory = hier_memory
        self.max_chunk_chars = max_chunk_chars
        self.target_final_tokens = target_final_tokens
        self.top_k_chunks = top_k_chunks
        self._session_id = str(uuid.uuid4())[:8]

    def handle(self, long_context: str, question: str) -> dict:
        if not long_context:
            return {"text": question, "method": "no_context"}

        total_tokens = estimate_token_count(long_context)
        if total_tokens <= self.target_final_tokens:
            return {
                "text": f"{long_context}\n\nQuestion: {question}",
                "method": "no_compression_needed",
                "tokens": total_tokens,
            }

        # 1. Chunk
        chunks = split_long_context(long_context, self.max_chunk_chars, overlap_chars=200)
        method_steps: list[str] = ["chunked"]
        # 2. Just-in-time index via hier_memory (working tier)
        chunk_ids: list[str] = []
        if self.hier_memory:
            for i, chunk in enumerate(chunks):
                mid = self.hier_memory.store(
                    content=chunk,
                    tier="working",
                    metadata={"jit_session": self._session_id, "chunk_idx": i},
                    importance=0.3,
                )
                if mid:
                    chunk_ids.append(mid)
            method_steps.append(f"indexed_{len(chunk_ids)}")
            # 3. Retrieve top-k via question
            relevant = self.hier_memory.recall(
                question, k_per_tier=self.top_k_chunks, tiers=["working"]
            )
            # Filter to this session only
            relevant = [r for r in relevant if (r.get("metadata") or {}).get("jit_session") == self._session_id]
            top_chunks = [r["content"] for r in relevant[:self.top_k_chunks]]
            method_steps.append(f"retrieved_{len(top_chunks)}")
        else:
            top_chunks = chunks[:self.top_k_chunks]
            method_steps.append("no_retrieval_fallback_first_k")

        # 4. Compress if still too long
        merged = " | ".join(top_chunks)
        if self.compressor and estimate_token_count(merged) > self.target_final_tokens:
            comp_result = self.compressor.compress(merged, query=question)
            merged = comp_result["compressed"]
            method_steps.append(f"compressed_ratio_{comp_result.get('ratio', 0):.2f}")

        # 5. Cleanup ephemeral chunks (best-effort)
        cleanup_count = 0
        if self.hier_memory and chunk_ids:
            try:
                self.hier_memory.collections["working"].delete(ids=chunk_ids)
                cleanup_count = len(chunk_ids)
            except Exception:
                pass

        return {
            "text": f"[LONG_CTX_DIGEST: {merged}]\n\nQuestion: {question}",
            "method": "->".join(method_steps),
            "tokens_before": total_tokens,
            "tokens_after": estimate_token_count(merged),
            "n_chunks": len(chunks),
            "n_retrieved": len(top_chunks) if 'top_chunks' in dir() else 0,
            "cleanup_count": cleanup_count,
        }


__all__ = ["LongContextHandler", "split_long_context", "estimate_token_count"]


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    print("=== M43 cypher_long_context SMOKE ===")
    long_doc = (
        "CVE-2021-44228 (Log4Shell) is a critical RCE in Apache Log4j2. It affects "
        "Log4j 2.x versions before 2.17.1. Exploits use JNDI lookups via crafted "
        "log messages containing ${jndi:ldap://...} patterns. Many companies were affected. "
    ) * 30 + (
        "Unrelated: history of Linux kernel development by Linus Torvalds since 1991. "
    ) * 30
    print(f"Long doc tokens estimate: {estimate_token_count(long_doc)}")
    handler = LongContextHandler(target_final_tokens=500)
    result = handler.handle(long_doc, "What is Log4Shell and how to mitigate?")
    print(f"Method: {result['method']}")
    print(f"Tokens before: {result['tokens_before']} → after: {result['tokens_after']}")
    print(f"Final text (preview): {result['text'][:300]}")
    print("=== SMOKE PASS ===")