contextual-engine 0.1.0__py3-none-any.whl

contextual/retrieval/context_assembler.py

@@ -0,0 +1,372 @@
+"""Context assembly: Budget-aware formatting of search results for LLM consumption.
+
+Takes search results and assembles them into a context string that fits within
+a token budget. Implements:
+- Token counting (using tiktoken cl100k_base)
+- Budget allocation (15% structural, 85% content per TECHNICAL_SPEC §5)
+- Progressive inclusion with truncation
+- Rich metadata formatting
+
+The assembler ensures that:
+1. Structural context (file paths, imports, class hierarchy) gets 15% of budget
+2. Actual code chunks get 85% of budget
+3. Higher-ranked results are prioritized
+4. Snippets are truncated gracefully if budget is tight
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import tiktoken
+
+from contextual.observability.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from contextual.retrieval.search import SearchResult
+
+
+logger = get_logger(__name__)
+
+# ============================================================================
+# CONSTANTS (from TECHNICAL_SPEC §5)
+# ============================================================================
+
+STRUCTURAL_BUDGET_RATIO = 0.15
+"""Percentage of budget allocated to structural context."""
+
+CONTENT_BUDGET_RATIO = 0.85
+"""Percentage of budget allocated to code chunks."""
+
+DEFAULT_TOKEN_BUDGET = 100_000
+"""Default maximum tokens for assembled context (Claude 3.7 window)."""
+
+# Encoding for token counting (matches Claude's tokenizer)
+TOKENIZER = tiktoken.get_encoding("cl100k_base")
+
+
+# ============================================================================
+# DATA MODELS
+# ============================================================================
+
+
+@dataclass(frozen=True)
+class AssembledContext:
+    """Assembled context ready for LLM consumption.
+
+    Attributes:
+        content: Formatted context string.
+        total_tokens: Total token count.
+        structural_tokens: Tokens used for structural metadata.
+        content_tokens: Tokens used for code chunks.
+        results_included: Number of search results included.
+        results_truncated: Number of results truncated due to budget.
+    """
+
+    content: str
+    total_tokens: int
+    structural_tokens: int
+    content_tokens: int
+    results_included: int
+    results_truncated: int
+
+
+# ============================================================================
+# TOKEN COUNTING
+# ============================================================================
+
+
+def count_tokens(text: str) -> int:
+    """Count tokens in a string using cl100k_base encoding.
+
+    Args:
+        text: String to count tokens for.
+
+    Returns:
+        Number of tokens.
+    """
+    return len(TOKENIZER.encode(text))
+
+
+# ============================================================================
+# STRUCTURAL CONTEXT FORMATTING
+# ============================================================================
+
+
+def _build_structural_context(results: list[SearchResult]) -> str:
+    """Build structural metadata from search results.
+
+    Extracts:
+    - Unique file paths
+    - Language distribution
+    - Symbol types present
+
+    Args:
+        results: Search results.
+
+    Returns:
+        Formatted structural context string.
+    """
+    if not results:
+        return ""
+
+    # Collect unique metadata
+    files = sorted(set(r.path for r in results))
+    languages = sorted(set(r.path.split(".")[-1] for r in results if "." in r.path))
+    symbol_types = sorted(set(r.chunk_type for r in results if r.chunk_type))
+
+    # Format structural context
+    parts = ["# Search Result Context\n"]
+
+    parts.append(f"\n## Files ({len(files)} total)")
+    for file_path in files:
+        parts.append(f"- {file_path}")
+
+    if languages:
+        parts.append("\n## Languages")
+        parts.append(f"{', '.join(languages)}")
+
+    if symbol_types:
+        parts.append("\n## Symbol Types")
+        parts.append(f"{', '.join(symbol_types)}")
+
+    parts.append("\n---\n")
+
+    return "\n".join(parts)
+
+
+# ============================================================================
+# CONTENT FORMATTING
+# ============================================================================
+
+
+def _format_result(result: SearchResult, index: int) -> str:
+    """Format a single search result as a markdown section.
+
+    Args:
+        result: Search result to format.
+        index: Result index (1-based).
+
+    Returns:
+        Formatted result string.
+    """
+    parts = [f"\n## Result {index} (score: {result.score:.3f})"]
+
+    # Metadata header
+    meta_parts = [
+        f"**File:** `{result.path}:{result.start_line}-{result.end_line}`",
+    ]
+
+    if result.symbol_name:
+        meta_parts.append(f"**Symbol:** `{result.symbol_name}` ({result.chunk_type})")
+
+    if result.bm25_rank is not None and result.dense_rank is not None:
+        meta_parts.append(
+            f"**Ranking:** BM25 rank {result.bm25_rank}, "
+            f"Vector rank {result.dense_rank}"
+        )
+
+    parts.append("\n".join(meta_parts))
+
+    # Code snippet
+    parts.append("\n```")
+    parts.append(result.snippet)
+    parts.append("```\n")
+
+    return "\n".join(parts)
+
+
+def _format_result_truncated(result: SearchResult, index: int, max_tokens: int) -> str:
+    """Format a result with snippet truncated to fit token budget.
+
+    Args:
+        result: Search result to format.
+        index: Result index (1-based).
+        max_tokens: Maximum tokens allowed for this result.
+
+    Returns:
+        Formatted result string, potentially with truncated snippet.
+    """
+    # Start with full format
+    full_format = _format_result(result, index)
+    full_tokens = count_tokens(full_format)
+
+    if full_tokens <= max_tokens:
+        return full_format
+
+    # Need to truncate snippet
+    # Calculate overhead from metadata (everything except snippet)
+    metadata_format = full_format.replace(result.snippet, "")
+    metadata_tokens = count_tokens(metadata_format)
+
+    # Tokens available for snippet
+    available_for_snippet = max_tokens - metadata_tokens
+
+    if available_for_snippet < 50:
+        # Not enough room even for truncated snippet - return just metadata
+        return metadata_format.replace("```\n```", "```\n[truncated]\n```")
+
+    # Binary search to find max snippet length that fits
+    snippet = result.snippet
+    left, right = 0, len(snippet)
+    best_length = 0
+
+    while left <= right:
+        mid = (left + right) // 2
+        truncated_snippet = snippet[:mid] + "..."
+        test_format = metadata_format.replace("```\n```", f"```\n{truncated_snippet}\n```")
+        test_tokens = count_tokens(test_format)
+
+        if test_tokens <= max_tokens:
+            best_length = mid
+            left = mid + 1
+        else:
+            right = mid - 1
+
+    if best_length == 0:
+        # Even minimal truncation doesn't fit
+        return metadata_format.replace("```\n```", "```\n[truncated]\n```")
+
+    truncated_snippet = snippet[:best_length] + "..."
+    return metadata_format.replace("```\n```", f"```\n{truncated_snippet}\n```")
+
+
+# ============================================================================
+# PUBLIC API
+# ============================================================================
+
+
+def assemble_context(
+    results: list[SearchResult],
+    token_budget: int = DEFAULT_TOKEN_BUDGET,
+) -> AssembledContext:
+    """Assemble search results into context string with token budget management.
+
+    Process:
+    1. Build structural context (15% of budget)
+    2. Format results in rank order (85% of budget)
+    3. Truncate snippets if needed to fit budget
+    4. Track what was included vs truncated
+
+    Args:
+        results: Search results ordered by relevance.
+        token_budget: Maximum tokens for assembled context.
+
+    Returns:
+        AssembledContext with formatted content and metadata.
+
+    Raises:
+        RetrievalError: If budget is too small to include any results.
+
+    Example:
+        >>> results = hybrid_search(conn, query, embedding, top_k=10)
+        >>> context = assemble_context(results, token_budget=50_000)
+        >>> print(context.content)  # Ready for LLM
+        >>> print(f"Used {context.total_tokens} tokens")
+    """
+    if not results:
+        return AssembledContext(
+            content="No results found.",
+            total_tokens=count_tokens("No results found."),
+            structural_tokens=0,
+            content_tokens=0,
+            results_included=0,
+            results_truncated=0,
+        )
+
+    # Allocate budgets
+    structural_budget = int(token_budget * STRUCTURAL_BUDGET_RATIO)
+    content_budget = int(token_budget * CONTENT_BUDGET_RATIO)
+
+    logger.debug(
+        "Assembling context",
+        total_budget=token_budget,
+        structural_budget=structural_budget,
+        content_budget=content_budget,
+        result_count=len(results),
+    )
+
+    # Build structural context
+    structural_context = _build_structural_context(results)
+    structural_tokens = count_tokens(structural_context)
+
+    if structural_tokens > structural_budget:
+        # Structural context is too large - truncate file list
+        logger.warning(
+            "Structural context exceeds budget",
+            structural_tokens=structural_tokens,
+            structural_budget=structural_budget,
+        )
+        # Fallback: just include counts
+        structural_context = (
+            f"# Search Result Context\n\n"
+            f"{len(results)} results found across "
+            f"{len(set(r.path for r in results))} files\n\n---\n"
+        )
+        structural_tokens = count_tokens(structural_context)
+
+    # Format results with remaining budget
+    remaining_budget = content_budget
+    formatted_results = []
+    results_included = 0
+    results_truncated = 0
+
+    for index, result in enumerate(results, start=1):
+        # Try to format this result
+        formatted = _format_result(result, index)
+        formatted_tokens = count_tokens(formatted)
+
+        if formatted_tokens <= remaining_budget:
+            # Fits fully
+            formatted_results.append(formatted)
+            remaining_budget -= formatted_tokens
+            results_included += 1
+        elif remaining_budget > 100:
+            # Try truncated version
+            truncated = _format_result_truncated(result, index, remaining_budget)
+            truncated_tokens = count_tokens(truncated)
+
+            if truncated_tokens <= remaining_budget:
+                formatted_results.append(truncated)
+                remaining_budget -= truncated_tokens
+                results_included += 1
+                results_truncated += 1
+            else:
+                # Even truncated doesn't fit - stop here
+                break
+        else:
+            # Budget exhausted
+            break
+
+    # Assemble final content
+    content_parts = [structural_context]
+    content_parts.extend(formatted_results)
+
+    if results_included < len(results):
+        omitted = len(results) - results_included
+        content_parts.append(
+            f"\n_({omitted} additional results omitted due to token budget)_\n"
+        )
+
+    final_content = "\n".join(content_parts)
+    total_tokens = count_tokens(final_content)
+
+    logger.info(
+        "Context assembled",
+        total_tokens=total_tokens,
+        structural_tokens=structural_tokens,
+        content_tokens=total_tokens - structural_tokens,
+        results_included=results_included,
+        results_truncated=results_truncated,
+    )
+
+    return AssembledContext(
+        content=final_content,
+        total_tokens=total_tokens,
+        structural_tokens=structural_tokens,
+        content_tokens=total_tokens - structural_tokens,
+        results_included=results_included,
+        results_truncated=results_truncated,
+    )

contextual/retrieval/ranker.py

@@ -0,0 +1,193 @@
+"""Result ranking and diversity filtering.
+
+Provides post-processing for search results:
+- deduplicate_results: Content deduplication
+- apply_mmr: Maximal Marginal Relevance diversity
+- rerank_with_cross_encoder: Cross-encoder reranking (Phase 2)
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from contextual.observability.logging import get_logger
+
+
+if TYPE_CHECKING:
+    from contextual.retrieval.search import SearchResult
+
+logger = get_logger(__name__)
+
+# ============================================================================
+# CONSTANTS
+# ============================================================================
+
+MMR_LAMBDA = 0.7
+"""Default MMR relevance/diversity tradeoff parameter."""
+
+MMR_MAX_PER_FILE = 2
+"""Default maximum number of results per file path."""
+
+MMR_MAX_PER_SYMBOL = 1
+"""Default maximum number of results per symbol name."""
+
+RERANK_MAX_PAIRS = 30
+"""Maximum query-document pairs for cross-encoder reranking."""
+
+
+def deduplicate_results(results: list[SearchResult]) -> list[SearchResult]:
+    """Remove duplicate snippets while preserving overall ordering.
+
+    Duplicates are detected by exact snippet string match. For each unique snippet,
+    the highest-scoring result is retained. The final output follows the input order
+    (stable order) of chosen representatives to avoid introducing rank jitter.
+
+    Args:
+        results: Ranked retrieval results.
+
+    Returns:
+        Deduplicated result list.
+    """
+    if not results:
+        return []
+
+    best_score_by_snippet: dict[str, float] = {}
+    best_chunk_id_by_snippet: dict[str, int] = {}
+
+    # Pass 1: determine representative result for each snippet.
+    for result in results:
+        snippet_key = result.snippet.strip()
+        if snippet_key not in best_score_by_snippet:
+            best_score_by_snippet[snippet_key] = result.score
+            best_chunk_id_by_snippet[snippet_key] = result.chunk_id
+            continue
+
+        if result.score > best_score_by_snippet[snippet_key]:
+            best_score_by_snippet[snippet_key] = result.score
+            best_chunk_id_by_snippet[snippet_key] = result.chunk_id
+
+    # Pass 2: keep only chosen representatives in original order.
+    deduplicated: list[SearchResult] = []
+    emitted_snippets: set[str] = set()
+    for result in results:
+        snippet_key = result.snippet.strip()
+        if snippet_key in emitted_snippets:
+            continue
+
+        chosen_chunk_id = best_chunk_id_by_snippet.get(snippet_key)
+        if chosen_chunk_id == result.chunk_id:
+            deduplicated.append(result)
+            emitted_snippets.add(snippet_key)
+
+    logger.debug(
+        "Deduplicated retrieval results",
+        input_count=len(results),
+        output_count=len(deduplicated),
+        duplicates_removed=len(results) - len(deduplicated),
+    )
+
+    return deduplicated
+
+
+def apply_mmr(
+    results: list[SearchResult],
+    lambda_param: float = MMR_LAMBDA,
+    max_per_file: int = MMR_MAX_PER_FILE,
+    max_per_symbol: int = MMR_MAX_PER_SYMBOL,
+) -> list[SearchResult]:
+    """Apply diversity filtering constraints over ranked results.
+
+    Phase 1 uses deterministic quota-based filtering:
+    - At most `max_per_file` results per file path.
+    - At most `max_per_symbol` results per symbol name.
+
+    Args:
+        results: Ranked retrieval results.
+        lambda_param: Reserved MMR tradeoff parameter for future semantic MMR.
+        max_per_file: Max results allowed for the same file path.
+        max_per_symbol: Max results allowed for the same symbol name.
+
+    Returns:
+        Diversity-filtered results, preserving input ordering.
+    """
+    if not results:
+        return []
+    if max_per_file <= 0:
+        logger.warning("max_per_file must be positive; returning empty result set", value=max_per_file)
+        return []
+    if max_per_symbol <= 0:
+        logger.warning(
+            "max_per_symbol must be positive; returning empty result set",
+            value=max_per_symbol,
+        )
+        return []
+
+    # Phase 1 doesn't yet compute pairwise similarity; keep parameter visible.
+    if not 0.0 <= lambda_param <= 1.0:
+        logger.warning(
+            "lambda_param out of range for MMR; clamping",
+            original=lambda_param,
+        )
+        lambda_param = max(0.0, min(1.0, lambda_param))
+
+    _ = lambda_param  # Reserved for true MMR scoring in Phase 2.
+
+    file_counts: defaultdict[str, int] = defaultdict(int)
+    symbol_counts: defaultdict[str, int] = defaultdict(int)
+    filtered: list[SearchResult] = []
+
+    skipped_file_quota = 0
+    skipped_symbol_quota = 0
+
+    for result in results:
+        file_key = result.path
+        symbol_key = result.symbol_name.strip() if result.symbol_name else ""
+
+        if file_counts[file_key] >= max_per_file:
+            skipped_file_quota += 1
+            continue
+
+        if symbol_key and symbol_counts[symbol_key] >= max_per_symbol:
+            skipped_symbol_quota += 1
+            continue
+
+        filtered.append(result)
+        file_counts[file_key] += 1
+        if symbol_key:
+            symbol_counts[symbol_key] += 1
+
+    logger.debug(
+        "Applied diversity filtering",
+        input_count=len(results),
+        output_count=len(filtered),
+        skipped_file_quota=skipped_file_quota,
+        skipped_symbol_quota=skipped_symbol_quota,
+        max_per_file=max_per_file,
+        max_per_symbol=max_per_symbol,
+    )
+
+    return filtered
+
+
+def rerank_with_cross_encoder(
+    query: str,
+    results: list[SearchResult],
+    top_k: int = RERANK_MAX_PAIRS,
+) -> list[SearchResult]:
+    """Rerank top candidates with a cross-encoder model.
+
+    Args:
+        query: User search query text.
+        results: Candidate retrieval results.
+        top_k: Maximum candidates to rerank with the cross-encoder.
+
+    Returns:
+        Reranked results (Phase 2).
+
+    Raises:
+        NotImplementedError: Always raised in Phase 1.
+    """
+    _ = (query, results, top_k)
+    # TODO(Phase 2): Integrate jina-reranker-v2-base-multilingual over query/chunk pairs.
+    raise NotImplementedError("Cross-encoder reranking coming in Phase 2")