kyorlin2001/code-analysis-tool

1from __future__ import annotations2 3from dataclasses import dataclass, field4 5from models.retrieval_chunk import RetrievalChunk6 7 8@dataclass(frozen=True)9class ContextBudgetResult:10    """11    Result of selecting chunks within a context budget.12    """13 14    selected_chunks: list[RetrievalChunk] = field(default_factory=list)15    total_characters: int = 016    truncated: bool = False17 18 19class ContextBudgetManager:20    """21    Keeps retrieved context within a maximum character budget.22    """23 24    def __init__(self, max_context_chars: int = 24_000) -> None:25        self.max_context_chars = max_context_chars26 27    def apply(self, chunks: list[RetrievalChunk]) -> ContextBudgetResult:28        selected: list[RetrievalChunk] = []29        total_chars = 030        truncated = False31 32        for chunk in chunks:33            chunk_length = len(chunk.text)34 35            if selected and total_chars + chunk_length > self.max_context_chars:36                truncated = True37                break38 39            if not selected and chunk_length > self.max_context_chars:40                # Always keep at least one chunk, even if it exceeds the budget.41                selected.append(chunk)42                total_chars += chunk_length43                truncated = True44                break45 46            selected.append(chunk)47            total_chars += chunk_length48 49        if len(selected) < len(chunks):50            truncated = True51 52        return ContextBudgetResult(53            selected_chunks=selected,54            total_characters=total_chars,55            truncated=truncated,56        )