velune-cli 0.9.0__py3-none-any.whl

velune/context/assembler.py

@@ -0,0 +1,240 @@
+"""Context assembly with canonical section ordering and budget enforcement.
+
+The ContextAssembler composes chunks into a final context string that respects
+budget constraints, enforces canonical ordering, and prioritizes trimming of
+low-trust, low-priority content.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections import defaultdict
+from collections.abc import Sequence
+
+from velune.context.budget import ContextBudget
+from velune.context.sections import (
+    ContextAssemblyReport,
+    ContextChunk,
+    ContextSection,
+)
+from velune.context.token_counter import TokenCounter
+from velune.core.types.model import ModelDescriptor
+
+logger = logging.getLogger("velune.context.assembler")
+
+
+class ContextAssembler:
+    """Assembles context chunks into final context respecting budget limits.
+
+    Algorithm:
+    1. Group chunks by section
+    2. Sort groups by ContextSection (canonical order 1-7)
+    3. For RETRIEVED_CONTEXT: trim to fit retrieval_allocation (drop low trust_score)
+    4. For WORKING_MEMORY: trim oldest turns to fit working_memory_allocation
+    5. For REPOSITORY_SNAPSHOT: reduce to architecture summary if > 2000 tokens
+    6. NEVER trim: SYSTEM_PROMPT, ARCHITECTURAL_DRIFT, CURRENT_PROMPT
+    7. Join with clear section separators
+    """
+
+    SECTION_SEPARATOR = "--- SECTION: {name} ({number} of 7) ---"
+    SECTION_END = "--- END SECTION ---"
+    UNTRIMMED_SECTIONS = {
+        ContextSection.SYSTEM_PROMPT,
+        ContextSection.ARCHITECTURAL_DRIFT,
+        ContextSection.CURRENT_PROMPT,
+    }
+
+    def assemble(
+        self,
+        chunks: Sequence[ContextChunk],
+        budget: ContextBudget,
+        model: ModelDescriptor | None = None,
+    ) -> tuple[str, ContextAssemblyReport]:
+        """Assemble chunks into final context respecting budget.
+
+        Args:
+            chunks: List of ContextChunk objects to assemble
+            budget: ContextBudget constraints for this session
+            model: Optional ModelDescriptor for token counting (defaults to word count)
+
+        Returns:
+            Tuple of (assembled_context_string, assembly_report)
+        """
+        if not chunks:
+            return "", ContextAssemblyReport(
+                total_chunks_received=0,
+                total_tokens_requested=budget.usable_tokens,
+                total_tokens_assembled=0,
+            )
+
+        # Phase 1: Group chunks by section (preserves insertion order)
+        sections_dict: dict[ContextSection, list[ContextChunk]] = defaultdict(list)
+        for chunk in chunks:
+            sections_dict[chunk.section].append(chunk)
+
+        # Phase 2: Process each section in canonical order
+        assembled_sections: dict[ContextSection, str] = {}
+        kept_sections_chunks: dict[ContextSection, list[ContextChunk]] = {}
+        total_tokens = 0
+        sections_trimmed: dict[ContextSection, int] = {}
+
+        for section in sorted(ContextSection):
+            if section not in sections_dict:
+                continue
+
+            section_chunks = sections_dict[section]
+            is_trimmable = section not in self.UNTRIMMED_SECTIONS
+
+            # Apply section-specific trimming logic
+            if section == ContextSection.RETRIEVED_CONTEXT:
+                processed_chunks, tokens_trimmed = self._trim_retrieved_context(
+                    section_chunks, budget.retrieval_allocation, model
+                )
+                if tokens_trimmed > 0:
+                    sections_trimmed[section] = tokens_trimmed
+            elif section == ContextSection.WORKING_MEMORY:
+                processed_chunks, tokens_trimmed = self._trim_working_memory(
+                    section_chunks, budget.working_memory_allocation, model
+                )
+                if tokens_trimmed > 0:
+                    sections_trimmed[section] = tokens_trimmed
+            elif section == ContextSection.REPOSITORY_SNAPSHOT:
+                processed_chunks = self._trim_repository_snapshot(section_chunks)
+            elif is_trimmable:
+                # For other trimmable sections, no special logic yet
+                processed_chunks = section_chunks
+            else:
+                # Never trim SYSTEM_PROMPT, ARCHITECTURAL_DRIFT, CURRENT_PROMPT
+                processed_chunks = section_chunks
+
+            if processed_chunks:
+                section_content = self._render_section(section, processed_chunks)
+                assembled_sections[section] = section_content
+                kept_sections_chunks[section] = list(processed_chunks)
+                total_tokens += sum(c.token_count for c in processed_chunks)
+
+        # Phase 3: Render final context with separators
+        final_context = self._render_assembled_context(assembled_sections)
+
+        # Phase 4: Check if final assembly exceeds budget
+        if model:
+            final_token_count = TokenCounter.count(final_context, model)
+        else:
+            final_token_count = total_tokens
+
+        budget_exceeded = final_token_count > budget.usable_tokens
+
+        if budget_exceeded:
+            logger.warning(
+                f"Assembled context exceeds budget: "
+                f"{final_token_count} > {budget.usable_tokens}; "
+                f"dropping lowest-priority RETRIEVED_CONTEXT chunks"
+            )
+            # Emergency truncation of RETRIEVED_CONTEXT
+            if ContextSection.RETRIEVED_CONTEXT in assembled_sections:
+                del assembled_sections[ContextSection.RETRIEVED_CONTEXT]
+                if ContextSection.RETRIEVED_CONTEXT in kept_sections_chunks:
+                    del kept_sections_chunks[ContextSection.RETRIEVED_CONTEXT]
+                final_context = self._render_assembled_context(assembled_sections)
+
+        report = ContextAssemblyReport(
+            total_chunks_received=len(chunks),
+            total_tokens_requested=budget.usable_tokens,
+            total_tokens_assembled=final_token_count,
+            sections_present=list(assembled_sections.keys()),
+            sections_trimmed=sections_trimmed,
+            chunks_dropped=len(chunks)
+            - sum(len(kept_sections_chunks.get(s, [])) for s in ContextSection),
+            budget_exceeded=budget_exceeded,
+        )
+
+        return final_context, report
+
+    def _trim_retrieved_context(
+        self,
+        chunks: list[ContextChunk],
+        allocation: int,
+        model: ModelDescriptor | None = None,
+    ) -> tuple[list[ContextChunk], int]:
+        """Trim RETRIEVED_CONTEXT to fit allocation.
+
+        Drops lowest trust_score chunks first. Sorts by trust_score descending
+        and keeps chunks until allocation is exceeded.
+        """
+        if sum(c.token_count for c in chunks) <= allocation:
+            return chunks, 0
+
+        # Sort by trust_score descending (keep high-trust chunks first)
+        sorted_chunks = sorted(chunks, key=lambda c: c.trust_score, reverse=True)
+
+        kept_chunks = []
+        tokens_used = 0
+        for chunk in sorted_chunks:
+            if tokens_used + chunk.token_count <= allocation:
+                kept_chunks.append(chunk)
+                tokens_used += chunk.token_count
+
+        tokens_trimmed = sum(c.token_count for c in chunks) - sum(
+            c.token_count for c in kept_chunks
+        )
+        return kept_chunks, tokens_trimmed
+
+    def _trim_working_memory(
+        self,
+        chunks: list[ContextChunk],
+        allocation: int,
+        model: ModelDescriptor | None = None,
+    ) -> tuple[list[ContextChunk], int]:
+        """Trim WORKING_MEMORY to fit allocation.
+
+        Removes oldest turns (first chunks) first to preserve recent context.
+        """
+        if sum(c.token_count for c in chunks) <= allocation:
+            return chunks, 0
+
+        # Keep chunks from the end (most recent turns) first
+        kept_chunks = []
+        tokens_used = 0
+        for chunk in reversed(chunks):
+            if tokens_used + chunk.token_count <= allocation:
+                kept_chunks.insert(0, chunk)
+                tokens_used += chunk.token_count
+
+        tokens_trimmed = sum(c.token_count for c in chunks) - sum(
+            c.token_count for c in kept_chunks
+        )
+        return kept_chunks, tokens_trimmed
+
+    def _trim_repository_snapshot(self, chunks: list[ContextChunk]) -> list[ContextChunk]:
+        """Reduce REPOSITORY_SNAPSHOT to architecture summary if over 2000 tokens.
+
+        Currently returns chunks as-is. Future enhancement: extract and preserve
+        only high-level architecture info when exceeding 2000 tokens.
+        """
+        total_tokens = sum(c.token_count for c in chunks)
+        if total_tokens > 2000:
+            logger.debug(
+                f"REPOSITORY_SNAPSHOT exceeds 2000 tokens ({total_tokens}); "
+                "consider extracting architecture summary only"
+            )
+        return chunks
+
+    def _render_section(self, section: ContextSection, chunks: list[ContextChunk]) -> str:
+        """Render a section with header, content, and footer."""
+        section_name = section.name
+        section_number = section.value
+
+        header = self.SECTION_SEPARATOR.format(name=section_name, number=section_number)
+        content = "\n\n".join(c.content for c in chunks)
+        footer = self.SECTION_END
+
+        return f"{header}\n{content}\n{footer}"
+
+    def _render_assembled_context(self, sections: dict[ContextSection, str]) -> str:
+        """Render all sections in canonical order with separators."""
+        rendered = []
+        for section in sorted(ContextSection):
+            if section in sections:
+                rendered.append(sections[section])
+
+        return "\n\n".join(rendered)

velune/context/budget.py

@@ -0,0 +1,97 @@
+"""Formal context budget allocation system.
+
+Defines token budget constraints for each session mode and allocates
+budget across retrieval, working memory, and system overhead.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from velune.cli.modes import SessionMode
+
+
+@dataclass(frozen=True)
+class ContextBudget:
+    """Token budget allocation across context sections.
+
+    Frozen dataclass enforcing immutability once created. Budget is
+    split proportionally across retrieval, working memory, system overhead,
+    and output reservation based on the active SessionMode.
+    """
+
+    total_tokens: int
+    retrieval_allocation: int
+    working_memory_allocation: int
+    output_reservation: int
+    system_allocation: int = 512
+
+    @classmethod
+    def from_mode(cls, mode: SessionMode, model_context_window: int) -> ContextBudget:
+        """Create a budget for the given session mode and model context.
+
+        Args:
+            mode: SessionMode (OPTIMUS, NORMAL, or GODLY)
+            model_context_window: Model's reported context length in tokens
+
+        Returns:
+            ContextBudget with allocations tuned to the mode's constraints
+        """
+        from velune.cli.modes import SessionMode
+
+        # Phase 1: Determine total usable tokens based on mode
+        if mode == SessionMode.OPTIMUS:
+            total = min(4096, model_context_window)
+        elif mode == SessionMode.NORMAL:
+            total = min(16384, model_context_window)
+        elif mode == SessionMode.GODLY:
+            total = model_context_window
+        else:
+            # Fallback to NORMAL
+            total = min(16384, model_context_window)
+
+        # Phase 2: Reserve output space and system overhead
+        system_alloc = 512
+        output_reserve = min(2048, total // 4)
+        usable = total - output_reserve - system_alloc
+
+        # Phase 3: Allocate remaining budget proportionally
+        # retrieval: 55%, working_memory: 35%
+        retrieval_alloc = int(usable * 0.55)
+        working_memory_alloc = int(usable * 0.35)
+
+        return cls(
+            total_tokens=total,
+            retrieval_allocation=retrieval_alloc,
+            working_memory_allocation=working_memory_alloc,
+            system_allocation=system_alloc,
+            output_reservation=output_reserve,
+        )
+
+    @property
+    def usable_tokens(self) -> int:
+        """Total tokens available for context (minus output and system)."""
+        return self.total_tokens - self.system_allocation - self.output_reservation
+
+    @property
+    def unallocated_tokens(self) -> int:
+        """Tokens not explicitly allocated to retrieval or working memory."""
+        allocated = (
+            self.retrieval_allocation
+            + self.working_memory_allocation
+            + self.system_allocation
+            + self.output_reservation
+        )
+        return self.total_tokens - allocated
+
+    def __str__(self) -> str:
+        """Human-readable budget summary."""
+        return (
+            f"ContextBudget(total={self.total_tokens}, "
+            f"retrieval={self.retrieval_allocation}, "
+            f"working_memory={self.working_memory_allocation}, "
+            f"system={self.system_allocation}, "
+            f"output_reserve={self.output_reservation})"
+        )

velune/context/extractive.py

@@ -0,0 +1,95 @@
+import re
+from collections import Counter
+
+
+def extractive_compress(text: str, target_tokens: int) -> str:
+    """
+    Sentence-importance based extractive compression.
+    Preserves: first sentence, last sentence, high-keyword-density sentences.
+    Drops: boilerplate, repetitive content, low-information sentences.
+    """
+    sentences = _split_sentences(text)
+    if not sentences:
+        return text
+
+    target_chars = target_tokens * 4  # Rough token→char estimate
+    if len(text) <= target_chars:
+        return text  # Already fits
+
+    # Score each sentence
+    word_freq = Counter(re.findall(r"\w+", text.lower()))
+    total_words = sum(word_freq.values())
+
+    scored = []
+    for i, sentence in enumerate(sentences):
+        score = _score_sentence(sentence, word_freq, total_words)
+        # Boost first and last sentences
+        if i == 0:
+            score += 0.5
+        if i == len(sentences) - 1:
+            score += 0.3
+        # Boost sentences with code-like content
+        if any(c in sentence for c in ("def ", "class ", "import ", "() ->", ":=")):
+            score += 0.4
+        scored.append((score, i, sentence))
+
+    # Greedily include highest-scoring sentences until budget exhausted
+    scored.sort(reverse=True)
+    selected = []
+    char_count = 0
+    selected_indices = set()
+
+    for _, idx, sentence in scored:
+        if char_count + len(sentence) <= target_chars:
+            selected.append((idx, sentence))
+            selected_indices.add(idx)
+            char_count += len(sentence)
+
+    # Sort selected sentences back to original order
+    selected.sort(key=lambda x: x[0])
+    result = " ".join(s for _, s in selected)
+
+    if len(text) > len(result) + 50:
+        result += f"\n[COMPRESSED: {len(text)} → {len(result)} chars]"
+
+    return result
+
+
+def _score_sentence(sentence: str, word_freq: Counter, total: int) -> float:
+    words = re.findall(r"\w+", sentence.lower())
+    if not words:
+        return 0.0
+    # TF-IDF inspired: sum of term frequencies of uncommon words
+    score = sum(1.0 / (word_freq[w] + 1) for w in words if len(w) > 3)
+    return score / len(words)
+
+
+def _split_sentences(text: str) -> list[str]:
+    # Split on sentence endings, preserve code blocks intact
+    return re.split(r"(?<=[.!?])\s+", text)
+
+
+def compress_conversation(conversation: list[dict], max_tokens: int) -> list[dict]:
+    """Drop oldest conversation turns until the total fits within max_tokens."""
+    try:
+        from velune.context.window import estimate_tokens
+    except ImportError:
+
+        def estimate_tokens(text: str) -> int:  # type: ignore[misc]
+            return len(text) // 4
+
+    if not conversation:
+        return conversation
+
+    total = sum(estimate_tokens(m.get("content", "")) for m in conversation)
+    if total <= max_tokens:
+        return conversation
+
+    result = list(conversation)
+    while len(result) > 1:
+        total = sum(estimate_tokens(m.get("content", "")) for m in result)
+        if total <= max_tokens:
+            break
+        result = result[1:]
+
+    return result