Function bodies 22 total

def _load_jsonl(path: str) -> list[dict[str, object]]:
    """Load a JSONL file into a list of dicts."""
    records: list[dict[str, object]] = []
    with open(path, encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if line:
                records.append(json.loads(line))
    return records

def _load_context_json(path: str) -> list[ContextBlock]:
    """Load a JSON file containing a list of ContextBlock objects."""
    raw = Path(path).read_text(encoding="utf-8")
    data = json.loads(raw)
    if isinstance(data, list):
        return [ContextBlock.model_validate(item) for item in data]
    # Support {"blocks": [...]} format too.
    blocks_raw = data.get("blocks", [])
    return [ContextBlock.model_validate(item) for item in blocks_raw]

def manage_cmd(max_tokens: int, input_path: str, output: str) -> None:
    """Load conversation turns and return the highest-priority blocks that fit the budget."""
    records = _load_jsonl(input_path)
    manager = ContextManager(max_tokens=max_tokens)

    for record in records:
        content = str(record.get("content", record.get("text", "")))
        if not content:
            continue
        priority_raw = str(record.get("priority", "medium")).lower()
        # Map string name to ContextPriority enum.
        priority_map: dict[str, ContextPriority] = {
            "critical": ContextPriority.critical,
            "high": ContextPriority.high,
            "medium": ContextPriority.medium,
            "low": ContextPriority.low,
            "ephemeral": ContextPriority.ephemeral,
        }
        priority = priority_map.get(priority_raw, ContextPriority.medium)
        metadata_raw = record.get("metadata", {})
        metadata: dict[str, object] = (
            metadata_raw i

def compress_cmd(
    strategy: str,
    input_path: str,
    output: str,
    max_tokens: int,
) -> None:
    """Compress a context window using the chosen strategy."""
    blocks = _load_context_json(input_path)
    manager = ContextManager(max_tokens=max_tokens)

    for block in blocks:
        manager.window.blocks.append(block)
    manager.window.used_tokens = sum(b.token_count for b in blocks)

    before_tokens = manager.window.used_tokens
    freed = manager.compress(CompressionStrategy(strategy))
    after_tokens = manager.window.used_tokens

    click.echo(
        f"Compression strategy '{strategy}': "
        f"{before_tokens} -> {after_tokens} tokens (freed {freed}).",
        err=True,
    )

    compressed_blocks = manager.window.blocks
    output_data = [b.model_dump(mode="json") for b in compressed_blocks]
    out_text = json.dumps(output_data, indent=2, default=str)

    if output == "-":
        click.echo(out_text)
    else:
        Path(output).write_text(out_text

def stats_cmd(input_path: str, max_tokens: int) -> None:
    """Print statistics about a context window."""
    blocks = _load_context_json(input_path)
    counter = TokenCounter()

    total_tokens = sum(b.token_count for b in blocks)
    priority_counts: dict[str, int] = {}
    priority_tokens: dict[str, int] = {}

    for block in blocks:
        pri_name = block.priority.name
        priority_counts[pri_name] = priority_counts.get(pri_name, 0) + 1
        priority_tokens[pri_name] = priority_tokens.get(pri_name, 0) + block.token_count

    utilisation = (total_tokens / max_tokens * 100) if max_tokens > 0 else 0.0

    stats: dict[str, object] = {
        "total_blocks": len(blocks),
        "total_tokens": total_tokens,
        "max_tokens": max_tokens,
        "utilisation_pct": round(utilisation, 2),
        "by_priority_count": priority_counts,
        "by_priority_tokens": priority_tokens,
        "avg_tokens_per_block": round(total_tokens / len(blocks), 2) if blocks else 0,
  

    def __init__(self) -> None:
        self._encoding: object = None
        self._use_tiktoken = False
        try:
            import tiktoken  # type: ignore[import-untyped]

            self._encoding = tiktoken.get_encoding("cl100k_base")
            self._use_tiktoken = True
        except ImportError:
            pass

    def count(self, text: str) -> int:
        """Return the estimated token count for *text*."""
        if self._use_tiktoken and self._encoding is not None:
            # tiktoken Encoding has an .encode() method returning a list.
            import tiktoken  # type: ignore[import-untyped]

            enc: tiktoken.Encoding = self._encoding  # type: ignore[assignment]
            return len(enc.encode(text))
        # Word-based heuristic: average English token is ~4 chars.
        return max(1, int(len(text.split()) * 1.3))

    def add_block(
        self,
        content: str,
        priority: ContextPriority = ContextPriority.medium,
        metadata: dict[str, object] | None = None,
    ) -> ContextBlock:
        """Create and register a new ContextBlock.

        Raises:
            ValueError: If the block alone exceeds the maximum window size.
        """
        token_count = self._counter.count(content)
        now = datetime.now(tz=timezone.utc)
        block = ContextBlock(
            block_id=str(uuid.uuid4()),
            content=content,
            token_count=token_count,
            priority=priority,
            created_at=now,
            last_accessed=now,
            metadata=metadata or {},
        )
        if token_count > self._window.max_tokens:
            raise ValueError(
                f"Block token count {token_count} exceeds max_tokens {self._window.max_tokens}."
            )
        self._window.blocks.append(block)
        self._window.used_tokens += token_count
      

    def get_context(self, max_tokens: int | None = None) -> list[ContextBlock]:
        """Return the highest-priority blocks that fit within *max_tokens*.

        Blocks are selected using a greedy knapsack: sorted by priority
        (descending) then by recency, filling the budget.

        Args:
            max_tokens: Token budget. Defaults to the window's max_tokens.

        Returns:
            List of selected ContextBlocks, preserving original insertion order.
        """
        budget = max_tokens if max_tokens is not None else self._window.max_tokens
        # Sort candidates: primary = priority descending, secondary = last_accessed descending.
        candidates = sorted(
            self._window.blocks,
            key=lambda b: (b.priority.value, b.last_accessed.timestamp()),
            reverse=True,
        )
        selected_ids: set[str] = set()
        remaining = budget
        for block in candidates:
            if block.token_count <= remaining:
              

    def compress(self, strategy: CompressionStrategy) -> int:
        """Free tokens using *strategy*.

        Returns:
            Number of tokens freed.
        """
        before = self._window.used_tokens

        if strategy == CompressionStrategy.drop_low_priority:
            self._compress_drop_low_priority()
        elif strategy == CompressionStrategy.truncate:
            self._compress_truncate()
        elif strategy == CompressionStrategy.sliding_window:
            self._compress_sliding_window()
        elif strategy == CompressionStrategy.summarize:
            self._compress_summarize()

        self._recalculate_used()
        freed = before - self._window.used_tokens
        return max(0, freed)

    def evict_expired(self, max_age_seconds: int) -> int:
        """Remove blocks older than *max_age_seconds*.

        Returns:
            Number of tokens freed.
        """
        cutoff = time.time() - max_age_seconds
        before = self._window.used_tokens
        surviving: list[ContextBlock] = []
        for block in self._window.blocks:
            if block.created_at.timestamp() >= cutoff:
                surviving.append(block)
        self._window.blocks = surviving
        self._recalculate_used()
        return max(0, before - self._window.used_tokens)

    def _compress_drop_low_priority(self) -> None:
        """Remove all ephemeral and low-priority blocks."""
        surviving = [
            b
            for b in self._window.blocks
            if b.priority not in (ContextPriority.ephemeral, ContextPriority.low)
        ]
        self._window.blocks = surviving

    def _compress_truncate(self) -> None:
        """Truncate the content of each block to 50% and update token counts."""
        for block in self._window.blocks:
            half = len(block.content) // 2
            block.content = block.content[:half] + "…"
            block.token_count = self._counter.count(block.content)

    def _compress_sliding_window(self) -> None:
        """Keep only the most recently accessed 50% of blocks."""
        if not self._window.blocks:
            return
        blocks_sorted = sorted(
            self._window.blocks,
            key=lambda b: b.last_accessed.timestamp(),
            reverse=True,
        )
        keep_count = max(1, len(blocks_sorted) // 2)
        kept_ids = {b.block_id for b in blocks_sorted[:keep_count]}
        self._window.blocks = [b for b in self._window.blocks if b.block_id in kept_ids]

    def _compress_summarize(self) -> None:
        """Simulate summarisation by truncating low-priority blocks to 30%."""
        for block in self._window.blocks:
            if block.priority in (ContextPriority.low, ContextPriority.ephemeral):
                keep = max(1, len(block.content) * 30 // 100)
                block.content = block.content[:keep] + " [summarised]"
                block.token_count = self._counter.count(block.content)

    def __init__(
        self,
        max_tokens: int = 4096,
        overlap_blocks: int = 2,
    ) -> None:
        self._manager = ContextManager(max_tokens=max_tokens)
        self._overlap_blocks = overlap_blocks

    def add(
        self,
        content: str,
        priority: ContextPriority = ContextPriority.medium,
        metadata: dict[str, object] | None = None,
    ) -> ContextBlock:
        """Add a block, evicting old ones if the window is full."""
        token_count = self._manager.token_counter.count(content)
        window = self._manager.window

        # If adding this block would overflow, slide: remove oldest until room exists.
        while (
            window.used_tokens + token_count > window.max_tokens
            and len(window.blocks) > self._overlap_blocks
        ):
            oldest = window.blocks.pop(0)
            window.used_tokens -= oldest.token_count

        return self._manager.add_block(content, priority, metadata)

    def weave(
        self,
        blocks: list[ContextBlock],
        max_tokens: int,
    ) -> list[ContextBlock]:
        """Return the optimal subset of *blocks* fitting in *max_tokens*, interleaved.

        Uses a DP knapsack (capacity = max_tokens, weight = token_count,
        value = priority value).  Because max_tokens can be large, we cap
        the DP at 8192 and fall back to greedy for larger budgets.

        Args:
            blocks:     Candidate blocks.
            max_tokens: Token budget.

        Returns:
            Selected blocks in priority-interleaved order.
        """
        if not blocks:
            return []

        dp_limit = 8192
        _knapsack_cell_limit = 10_000_000
        use_dp = (
            max_tokens <= dp_limit
            and len(blocks) * max_tokens <= _knapsack_cell_limit
        )
        if use_dp:
            selected = self._knapsack(blocks, max_tokens)
        else:
            selected = self._greedy(blocks, max_tokens)

      

    def _knapsack(self, blocks: list[ContextBlock], capacity: int) -> list[ContextBlock]:
        """0/1 knapsack DP to maximise priority-weighted value within *capacity*.

        Uses a full 2D table so backtracking is exact.
        """
        n = len(blocks)
        # dp[i][w] = max value using first i items with weight budget w.
        dp: list[list[float]] = [[0.0] * (capacity + 1) for _ in range(n + 1)]

        for i in range(1, n + 1):
            block = blocks[i - 1]
            weight = block.token_count
            value = float(block.priority.value)
            for w in range(capacity + 1):
                # Block is only eligible when its true weight fits in the budget.
                if weight <= w:
                    dp[i][w] = max(dp[i - 1][w], dp[i - 1][w - weight] + value)
                else:
                    dp[i][w] = dp[i - 1][w]

        # Backtrack to determine which items were selected.
        selected: list[ContextBlock] = []
        w = capacity
  

    def _greedy(self, blocks: list[ContextBlock], capacity: int) -> list[ContextBlock]:
        """Greedy selection by priority/weight ratio."""
        candidates = sorted(
            blocks,
            key=lambda b: b.priority.value / max(b.token_count, 1),
            reverse=True,
        )
        selected: list[ContextBlock] = []
        remaining = capacity
        for block in candidates:
            if block.token_count <= remaining:
                selected.append(block)
                remaining -= block.token_count
        return selected

    def _interleave(self, blocks: list[ContextBlock]) -> list[ContextBlock]:
        """Interleave blocks by priority tier for natural reading order."""
        tiers: dict[int, list[ContextBlock]] = {}
        for block in blocks:
            tiers.setdefault(block.priority.value, []).append(block)

        result: list[ContextBlock] = []
        # Descending priority: critical first, then high, medium, low, ephemeral.
        for priority_value in sorted(tiers, reverse=True):
            result.extend(tiers[priority_value])
        return result

    def _coerce_priority(cls, value: Any) -> Any:
        """Accept priority by name (e.g. 'medium') or integer value (e.g. 3)."""
        if isinstance(value, str):
            # Try by name first; fall back to int parse.
            name_map = {member.name: member for member in ContextPriority}
            if value in name_map:
                return name_map[value].value
            try:
                return int(value)
            except ValueError:
                pass
        return value