gdmcode 0.1.0__py3-none-any.whl

src/agent/context_budget.py

@@ -0,0 +1,391 @@
+"""Context budget — tracks tokens and triggers compression.
+
+The golden rule: never send more than 80% of the model's context window
+in a single request. This module enforces that constraint and provides
+chunked file reading so large files never blow the budget in one shot.
+
+Token counting uses tiktoken cl100k_base (close enough for Grok/Gemini).
+"""
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import tiktoken
+
+from src._internal.constants import _COMPRESSION_TRIGGER_RATIO
+from src.agent.transcript import TranscriptStore
+
+__all__ = ["ContextBudget", "count_tokens", "chunk_file", "WholeCodebaseMode"]
+
+log = logging.getLogger(__name__)
+
+# cl100k_base is used by GPT-4/Claude/Grok — good-enough approximation
+_ENCODER = tiktoken.get_encoding("cl100k_base")
+
+# Chars-per-token approximation (used as fast path before tiktoken)
+_CHARS_PER_TOKEN = 3.5
+
+
+def count_tokens(text: str) -> int:
+    """Count tokens in *text* using tiktoken cl100k_base.
+
+    Falls back to char-based estimate on encode failure.
+    """
+    try:
+        return len(_ENCODER.encode(text, disallowed_special=()))
+    except Exception as exc:
+        log.debug("tiktoken encode failed, using estimate: %s", exc)
+        return max(1, int(len(text) / _CHARS_PER_TOKEN))
+
+
+def chunk_file(path: Path, chunk_tokens: int = 4_000) -> list[str]:
+    """Split a file into token-sized chunks for incremental injection.
+
+    Args:
+        path: file to read
+        chunk_tokens: max tokens per chunk
+
+    Returns:
+        list of text chunks, each ≤ chunk_tokens tokens
+    """
+    try:
+        text = path.read_text(encoding="utf-8", errors="replace")
+    except OSError as exc:
+        log.warning("Could not read %s for chunking: %s", path, exc)
+        return []
+
+    lines = text.splitlines(keepends=True)
+    chunks: list[str] = []
+    current_lines: list[str] = []
+    current_tokens = 0
+
+    for line in lines:
+        line_tokens = count_tokens(line)
+        if current_tokens + line_tokens > chunk_tokens and current_lines:
+            chunks.append("".join(current_lines))
+            current_lines = [line]
+            current_tokens = line_tokens
+        else:
+            current_lines.append(line)
+            current_tokens += line_tokens
+
+    if current_lines:
+        chunks.append("".join(current_lines))
+
+    return chunks
+
+
+@dataclass
+class BudgetSnapshot:
+    """A point-in-time snapshot of token usage."""
+
+    used_tokens: int
+    max_tokens: int
+    budget_ratio: float
+    needs_compression: bool
+
+
+class ContextBudget:
+    """Tracks token usage and enforces the 80% compression threshold.
+
+    Usage::
+
+        budget = ContextBudget(model_context_window=131_072)
+        tokens = count_tokens(user_message)
+        budget.add(tokens)
+
+        if budget.needs_compression:
+            n = transcript.maybe_evict()
+            budget.report_eviction(n_tokens_freed)
+    """
+
+    def __init__(
+        self,
+        model_context_window: int = 131_072,
+        compression_ratio: float = _COMPRESSION_TRIGGER_RATIO,
+    ) -> None:
+        self._max = model_context_window
+        self._threshold = compression_ratio
+        self._used: int = 0
+        self._evicted_total: int = 0
+
+    # ------------------------------------------------------------------
+    # Mutation
+    # ------------------------------------------------------------------
+
+    def add(self, tokens: int) -> None:
+        """Record *tokens* added to the context."""
+        self._used += max(0, tokens)
+
+    def subtract(self, tokens: int) -> None:
+        """Record *tokens* removed from the context (eviction/compression)."""
+        self._used = max(0, self._used - tokens)
+        self._evicted_total += max(0, tokens)
+
+    def reset(self) -> None:
+        """Reset used counter (e.g., new turn)."""
+        self._used = 0
+
+    def sync_from_transcript(self, store: TranscriptStore) -> None:
+        """Re-sync token count from the authoritative TranscriptStore."""
+        self._used = store.token_count
+
+    # ------------------------------------------------------------------
+    # Query
+    # ------------------------------------------------------------------
+
+    @property
+    def used_tokens(self) -> int:
+        return self._used
+
+    @property
+    def max_tokens(self) -> int:
+        return self._max
+
+    @property
+    def remaining_tokens(self) -> int:
+        return max(0, self._max - self._used)
+
+    @property
+    def budget_ratio(self) -> float:
+        return self._used / self._max if self._max else 0.0
+
+    @property
+    def needs_compression(self) -> bool:
+        """True when usage exceeds the compression trigger ratio."""
+        return self.budget_ratio >= self._threshold
+
+    def snapshot(self) -> BudgetSnapshot:
+        return BudgetSnapshot(
+            used_tokens=self._used,
+            max_tokens=self._max,
+            budget_ratio=self.budget_ratio,
+            needs_compression=self.needs_compression,
+        )
+
+    def is_near_limit(self, threshold: float = 0.8) -> bool:
+        """True when usage exceeds *threshold* fraction of max tokens.
+
+        Looser than needs_compression so callers can act earlier.  Separate
+        method so the default threshold can differ from the compression trigger.
+        """
+        return self._used / self._max >= threshold
+
+    def can_fit(self, tokens: int) -> bool:
+        """Check whether *tokens* more can be added without hitting the threshold."""
+        return (self._used + tokens) / self._max < self._threshold
+
+
+class WholeCodebaseMode:
+    """Whole-codebase loading — loads all project source files into context upfront.
+
+    When a project's estimated token count fits within 75 % of the model's context
+    window, all source files are loaded in one shot.  Deterministic inclusion rules
+    prevent binary/generated/vendor files from polluting the context.
+
+    Usage::
+
+        budget = ContextBudget(model_context_window=2_000_000)
+        wcm = WholeCodebaseMode(budget)
+        files = wcm.load_whole_codebase(root, tech_stack, mode="auto")
+    """
+
+    _EXCLUDE_DIRS: frozenset[str] = frozenset({
+        "node_modules", ".git", "dist", "build", "__pycache__",
+        "vendor", "third_party", "site-packages",
+    })
+
+    _STACK_EXTENSIONS: dict[str, frozenset[str]] = {
+        "python":     frozenset({".py", ".pyi", ".toml", ".cfg", ".md", ".txt"}),
+        "nodejs":     frozenset({".ts", ".tsx", ".js", ".jsx", ".json", ".md"}),
+        "typescript": frozenset({".ts", ".tsx", ".js", ".jsx", ".json", ".md"}),
+        "rust":       frozenset({".rs", ".toml"}),
+        "golang":     frozenset({".go", ".mod"}),
+    }
+
+    _LOCK_FILES: frozenset[str] = frozenset({
+        "package-lock.json", "yarn.lock", "pnpm-lock.yaml",
+        "poetry.lock", "Cargo.lock",
+    })
+
+    _MAX_FILE_BYTES: int = 200 * 1024  # 200 KB
+
+    def __init__(self, budget: ContextBudget) -> None:
+        self._budget = budget
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def estimate_token_count(self, root: Path, tech_stack: list[str]) -> int:
+        """Estimate total token count for *root* without per-file content reads.
+
+        Uses ``file_size_bytes / 3.5`` as a fast approximation.  Must complete
+        in <3 s for projects with up to 10 K files.
+        """
+        allowed_ext = self._get_allowed_extensions(tech_stack)
+        ignore_spec = self._load_ignore_spec(root)
+        total_bytes = 0
+        for fp in self._iter_files(root, allowed_ext, ignore_spec):
+            try:
+                total_bytes += fp.stat().st_size
+            except OSError:
+                continue
+        return max(1, int(total_bytes / _CHARS_PER_TOKEN))
+
+    def load_whole_codebase(
+        self,
+        root: Path,
+        tech_stack: list[str],
+        mode: str = "auto",
+    ) -> list[tuple[Path, str]]:
+        """Load all source files for the project.
+
+        Args:
+            root: project root directory.
+            tech_stack: detected tech stack (e.g. ``["python", "docker"]``).
+            mode: ``"auto"`` checks budget threshold; ``"always"`` forces load;
+                ``"never"`` returns an empty list immediately.
+
+        Returns:
+            List of ``(path, content)`` tuples in filesystem order.
+        """
+        if mode == "never":
+            return []
+
+        allowed_ext = self._get_allowed_extensions(tech_stack)
+        ignore_spec = self._load_ignore_spec(root)
+        threshold = self._get_threshold()
+
+        if mode == "auto":
+            estimated = self.estimate_token_count(root, tech_stack)
+            if estimated > threshold:
+                log.debug(
+                    "Project too large for whole-codebase mode (%dK / %dK threshold).",
+                    estimated // 1000,
+                    threshold // 1000,
+                )
+                return []
+            log.info(
+                "Project fits in context (%dK / %dK) — whole-codebase mode active.",
+                estimated // 1000,
+                threshold // 1000,
+            )
+
+        results: list[tuple[Path, str]] = []
+        for fp in self._iter_files(root, allowed_ext, ignore_spec):
+            if self._is_binary(fp):
+                continue
+            try:
+                text = fp.read_text(encoding="utf-8", errors="replace")
+                results.append((fp, text))
+            except OSError:
+                continue
+
+        if mode == "always":
+            log.info(
+                "Whole-codebase mode active (forced) — loaded %d files.", len(results)
+            )
+        return results
+
+    # ------------------------------------------------------------------
+    # Private helpers
+    # ------------------------------------------------------------------
+
+    def _get_threshold(self) -> int:
+        """75 % of model context window — leaves room for history + system prompt."""
+        return int(self._budget._max * 0.75)
+
+    def _get_allowed_extensions(self, tech_stack: list[str]) -> frozenset[str] | None:
+        """Return the set of allowed extensions for *tech_stack*, or None for all."""
+        if not tech_stack:
+            return None
+        extensions: set[str] = set()
+        found = False
+        for stack in tech_stack:
+            if stack in self._STACK_EXTENSIONS:
+                extensions.update(self._STACK_EXTENSIONS[stack])
+                found = True
+        return frozenset(extensions) if found else None
+
+    def _load_ignore_spec(self, root: Path):
+        """Load .gdmignore, falling back to .gitignore, via pathspec."""
+        import pathspec  # local import — optional dependency
+
+        for name in (".gdmignore", ".gitignore"):
+            ignore_file = root / name
+            if ignore_file.exists():
+                text = ignore_file.read_text(encoding="utf-8", errors="replace")
+                return pathspec.PathSpec.from_lines("gitignore", text.splitlines())
+        return None
+
+    def _is_binary(self, path: Path) -> bool:
+        """Return True when the first 8 192 bytes of *path* contain a null byte."""
+        try:
+            with path.open("rb") as fh:
+                return b"\x00" in fh.read(8192)
+        except OSError:
+            return True
+
+    def _iter_files(self, root: Path, allowed_ext: frozenset[str] | None, ignore_spec):
+        """Yield all eligible source files under *root*.
+
+        Applies:
+        - Directory exclusions (node_modules, .git, dist, …)
+        - File-size limit (>200 KB skipped)
+        - Extension filter derived from tech_stack
+        - Name-based exclusions (.pyc, .min.js, lock files)
+        - Ignore-spec patterns (.gdmignore / .gitignore)
+        - Resolved-path tracking to prevent symlink loops
+        """
+        visited: set[Path] = set()
+
+        def _walk(directory: Path):
+            try:
+                real = directory.resolve()
+            except OSError:
+                return
+            if real in visited:
+                return
+            visited.add(real)
+
+            try:
+                entries = sorted(directory.iterdir())
+            except OSError:
+                return
+
+            for entry in entries:
+                if entry.is_dir():
+                    if entry.name in self._EXCLUDE_DIRS:
+                        continue
+                    yield from _walk(entry)
+                elif entry.is_file():
+                    try:
+                        if entry.stat().st_size > self._MAX_FILE_BYTES:
+                            continue
+                    except OSError:
+                        continue
+
+                    name = entry.name
+                    if name.endswith((".min.js", ".min.css", ".pyc")):
+                        continue
+                    if name in self._LOCK_FILES:
+                        continue
+
+                    if allowed_ext is not None:
+                        if entry.suffix.lower() not in allowed_ext:
+                            continue
+
+                    if ignore_spec is not None:
+                        try:
+                            rel = str(entry.relative_to(root)).replace("\\", "/")
+                            if ignore_spec.match_file(rel):
+                                continue
+                        except ValueError:
+                            pass
+
+                    yield entry
+
+        yield from _walk(root)