cogcore/longcontext/kv_cache.py

"""Low-budget token cache for sliding long-context inference."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Sequence


@dataclass
class TokenKVCache:
    """A tiny token-history cache with bounded capacity.

    This is not a Transformer KV cache yet; it is a stable interface and smoke
    implementation used by the long-context bridge before attention KV tensors
    are added.
    """

    capacity: int = 2048
    tokens: list[int] = field(default_factory=list)

    def append(self, new_tokens: Sequence[int]) -> None:
        self.tokens.extend(int(t) for t in new_tokens)
        if len(self.tokens) > self.capacity:
            self.tokens = self.tokens[-self.capacity:]

    def window(self, size: int | None = None) -> list[int]:
        n = self.capacity if size is None else int(size)
        return list(self.tokens[-n:])

    def clear(self) -> None:
        self.tokens.clear()