tokenpack-rag 0.1.0__py3-none-any.whl

tokenpack/__init__.py

@@ -0,0 +1,28 @@
+"""Context-window budget optimization for retrieval-augmented LLM workflows."""
+
+from tokenpack.chunk_profiles import ChunkSizeConfig, resolve_chunk_size_config
+from tokenpack.chunking import SemanticThresholdChunker, StructureAwareChunker
+from tokenpack.compression import CompressionConfig, CompressionResult, compress_chunks
+from tokenpack.dataset import GoldRecord
+from tokenpack.embeddings import DEFAULT_EMBEDDING_MODEL, SentenceTransformerEmbedder, make_embedder
+from tokenpack.reranking import CrossEncoderReranker, apply_reranker, blend_reranker_scores
+from tokenpack.selectors import select_chunks
+
+__all__ = [
+    "CrossEncoderReranker",
+    "DEFAULT_EMBEDDING_MODEL",
+    "GoldRecord",
+    "SemanticThresholdChunker",
+    "SentenceTransformerEmbedder",
+    "StructureAwareChunker",
+    "ChunkSizeConfig",
+    "CompressionConfig",
+    "CompressionResult",
+    "compress_chunks",
+    "apply_reranker",
+    "blend_reranker_scores",
+    "make_embedder",
+    "resolve_chunk_size_config",
+    "select_chunks",
+]
+

tokenpack/benchmark.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+
+from tokenpack.dataset import GoldRecord, propose_gold_records
+from tokenpack.embeddings import Embedder, cosine
+from tokenpack.index import ChunkIndex
+from tokenpack.models import ScoredChunk, SelectionResult
+from tokenpack.scoring import DEFAULT_SCORING_PROFILE, score_chunks
+from tokenpack.selectors import select_chunks
+
+
+STRATEGIES = [
+    "document-prefix",
+    "full-document",
+    "top-k",
+    "production-rag",
+    "budget-top-k",
+    "greedy-value",
+    "greedy-density",
+    "mmr",
+    "knapsack",
+    "knapsack-redundancy",
+    "knapsack-coverage",
+]
+
+
+def synthetic_queries(index: ChunkIndex, sample_size: int = 12) -> list[dict]:
+    return [
+        {
+            "query": record.query,
+            "evidence_chunk_id": record.evidence_chunk_ids[0],
+            "source_path": record.source_path,
+        }
+        for record in propose_gold_records(index, sample_size=sample_size)
+    ]
+
+
+def run_benchmark(
+    index: ChunkIndex,
+    embedder: Embedder,
+    budget: int,
+    reserve_output: int,
+    sample_size: int = 12,
+    candidate_pool: int = 250,
+    scoring: str = DEFAULT_SCORING_PROFILE,
+) -> dict:
+    """Developer smoke benchmark using auto-proposed single-evidence queries."""
+
+    records = propose_gold_records(index, sample_size=sample_size)
+    payload = run_gold_benchmark(
+        index=index,
+        embedder=embedder,
+        records=records,
+        budgets=[budget],
+        reserve_output=reserve_output,
+        candidate_pool=candidate_pool,
+        scoring=scoring,
+    )
+    return payload["budgets"][0] | {
+        "mode": "smoke",
+        "scoring": scoring,
+        "query_count": len(records),
+        "queries": payload["budgets"][0]["queries"],
+    }
+
+
+def run_gold_benchmark(
+    index: ChunkIndex,
+    embedder: Embedder,
+    records: list[GoldRecord],
+    budgets: list[int],
+    reserve_output: int,
+    candidate_pool: int = 250,
+    strategies: list[str] | None = None,
+    redundancy_penalty: float = 0.35,
+    scoring: str = DEFAULT_SCORING_PROFILE,
+) -> dict:
+    strategy_names = strategies or STRATEGIES
+    budget_runs = [
+        _run_gold_for_budget(
+            index=index,
+            embedder=embedder,
+            records=records,
+            budget=budget,
+            reserve_output=reserve_output,
+            candidate_pool=candidate_pool,
+            strategies=strategy_names,
+            redundancy_penalty=redundancy_penalty,
+            scoring=scoring,
+        )
+        for budget in budgets
+    ]
+    return {
+        "mode": "gold",
+        "scoring": scoring,
+        "query_count": len(records),
+        "reserve_output": reserve_output,
+        "strategies": strategy_names,
+        "budgets": budget_runs,
+    }
+
+
+def save_benchmark(payload: dict, path: str | Path) -> None:
+    target = Path(path)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    target.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+def _run_gold_for_budget(
+    index: ChunkIndex,
+    embedder: Embedder,
+    records: list[GoldRecord],
+    budget: int,
+    reserve_output: int,
+    candidate_pool: int,
+    strategies: list[str],
+    redundancy_penalty: float,
+    scoring: str,
+) -> dict:
+    effective_budget = max(0, budget - reserve_output)
+    totals = {
+        strategy: {
+            "recall": 0.0,
+            "precision": 0.0,
+            "coverage": 0.0,
+            "used_tokens": 0,
+            "budget_utilization": 0.0,
+            "total_value": 0.0,
+            "redundancy": 0.0,
+            "latency": 0.0,
+            "runs": 0,
+        }
+        for strategy in strategies
+    }
+    per_query = []
+    for record in records:
+        query_embedding = embedder.embed([record.query])[0]
+        scored = score_chunks(
+            query_embedding,
+            index.chunks,
+            index.embeddings,
+            scoring=scoring,
+            query_text=record.query,
+        )
+        scored_redundant = score_chunks(
+            query_embedding,
+            index.chunks,
+            index.embeddings,
+            redundancy_penalty=redundancy_penalty,
+            scoring=scoring,
+            query_text=record.query,
+        )
+        query_result = {
+            "query": record.query,
+            "evidence_chunk_ids": record.evidence_chunk_ids,
+            "strategies": {},
+        }
+        for strategy in strategies:
+            source_scores = scored_redundant if strategy == "knapsack-redundancy" else scored
+            started = time.perf_counter()
+            result = select_chunks(
+                source_scores,
+                strategy=strategy,
+                budget=effective_budget,
+                candidate_pool=candidate_pool,
+                embeddings=index.embeddings,
+                coverage_query=record.query,
+            )
+            elapsed = time.perf_counter() - started
+            metrics = _selection_metrics(result, record.evidence_chunk_ids, effective_budget)
+            metrics["latency_seconds"] = elapsed
+            _accumulate(totals[strategy], metrics, result)
+            query_result["strategies"][strategy] = metrics | {
+                "selected_count": len(result.selected),
+                "selected_chunk_ids": [item.chunk.id for item in result.selected],
+            }
+        per_query.append(query_result)
+
+    summary = {strategy: _summarize(values) for strategy, values in totals.items()}
+    return {
+        "budget": budget,
+        "reserve_output": reserve_output,
+        "effective_budget": effective_budget,
+        "query_count": len(records),
+        "summary": summary,
+        "queries": per_query,
+    }
+
+
+def _selection_metrics(result: SelectionResult, evidence_ids: list[str], budget: int) -> dict:
+    selected_ids = {item.chunk.id for item in result.selected}
+    evidence_set = set(evidence_ids)
+    matched = selected_ids & evidence_set
+    recall = len(matched) / max(1, len(evidence_set))
+    precision = len(matched) / max(1, len(selected_ids))
+    return {
+        "evidence_recall_at_budget": recall,
+        "evidence_precision": precision,
+        "coverage_ratio": 1.0 if evidence_set and evidence_set.issubset(selected_ids) else 0.0,
+        "used_tokens": result.used_tokens,
+        "budget_utilization": result.used_tokens / max(1, budget),
+        "over_budget": result.used_tokens > budget,
+        "over_budget_tokens": max(0, result.used_tokens - budget),
+        "total_value": result.total_value,
+        "value_density": result.total_value / max(1, result.used_tokens),
+        "redundancy_score": redundancy_score(result.selected),
+    }
+
+
+def redundancy_score(selected: list[ScoredChunk]) -> float:
+    embeddings = [item.embedding for item in selected if item.embedding is not None]
+    if len(embeddings) < 2:
+        return 0.0
+    total = 0.0
+    comparisons = 0
+    for left_index, left in enumerate(embeddings):
+        for right in embeddings[left_index + 1 :]:
+            total += max(0.0, cosine(left, right))
+            comparisons += 1
+    return total / max(1, comparisons)
+
+
+def _accumulate(total: dict, metrics: dict, result: SelectionResult) -> None:
+    total["recall"] += metrics["evidence_recall_at_budget"]
+    total["precision"] += metrics["evidence_precision"]
+    total["coverage"] += metrics["coverage_ratio"]
+    total["used_tokens"] += result.used_tokens
+    total["budget_utilization"] += metrics["budget_utilization"]
+    total["over_budget"] = total.get("over_budget", 0) + int(metrics["over_budget"])
+    total["over_budget_tokens"] = total.get("over_budget_tokens", 0) + metrics["over_budget_tokens"]
+    total["total_value"] += result.total_value
+    total["redundancy"] += metrics["redundancy_score"]
+    total["latency"] += metrics["latency_seconds"]
+    total["runs"] += 1
+
+
+def _summarize(values: dict) -> dict:
+    runs = max(1, values["runs"])
+    avg_tokens = values["used_tokens"] / runs
+    avg_value = values["total_value"] / runs
+    return {
+        "evidence_recall_at_budget": values["recall"] / runs,
+        "evidence_precision": values["precision"] / runs,
+        "coverage_ratio": values["coverage"] / runs,
+        "avg_used_tokens": avg_tokens,
+        "budget_utilization": values["budget_utilization"] / runs,
+        "over_budget_rate": values.get("over_budget", 0) / runs,
+        "avg_over_budget_tokens": values.get("over_budget_tokens", 0) / runs,
+        "avg_total_value": avg_value,
+        "value_density": avg_value / max(1.0, avg_tokens),
+        "redundancy_score": values["redundancy"] / runs,
+        "latency_seconds": values["latency"] / runs,
+    }

tokenpack/chunk_profiles.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class ChunkSizeConfig:
+    target_tokens: int
+    min_tokens: int
+    max_tokens: int
+
+
+CHUNK_SIZE_PRESETS: dict[str, ChunkSizeConfig] = {
+    "default": ChunkSizeConfig(target_tokens=650, min_tokens=120, max_tokens=900),
+    "low-budget": ChunkSizeConfig(target_tokens=250, min_tokens=40, max_tokens=320),
+}
+
+
+def resolve_chunk_size_config(
+    preset: str,
+    target_tokens: int,
+    min_tokens: int,
+    max_tokens: int,
+) -> ChunkSizeConfig:
+    if preset == "manual":
+        return ChunkSizeConfig(
+            target_tokens=target_tokens,
+            min_tokens=min_tokens,
+            max_tokens=max_tokens,
+        )
+    try:
+        return CHUNK_SIZE_PRESETS[preset]
+    except KeyError as exc:
+        choices = ", ".join(["manual", *CHUNK_SIZE_PRESETS])
+        raise ValueError(f"Unknown chunk size preset: {preset}. Choose one of: {choices}.") from exc

tokenpack/chunking.py

@@ -0,0 +1,340 @@
+from __future__ import annotations
+
+import hashlib
+import re
+from collections.abc import Sequence
+from typing import Any
+
+from tokenpack.embeddings import cosine
+from tokenpack.models import Chunk, TextBlock
+from tokenpack.tokenization import TokenCounter
+
+SENTENCE_RE = re.compile(r"(?<=[.!?])\s+")
+
+
+class _ChunkGroupBase:
+    """Shared token-bounded grouping helpers for structure-aware chunkers."""
+
+    def __init__(
+        self,
+        target_tokens: int = 650,
+        min_tokens: int = 120,
+        max_tokens: int = 900,
+        token_counter: TokenCounter | None = None,
+    ) -> None:
+        if min_tokens > target_tokens or target_tokens > max_tokens:
+            raise ValueError("Expected min_tokens <= target_tokens <= max_tokens.")
+        self.target_tokens = target_tokens
+        self.min_tokens = min_tokens
+        self.max_tokens = max_tokens
+        self.token_counter = token_counter or TokenCounter()
+
+    def _split_large_block(self, block_id: int, block: TextBlock) -> list[Chunk]:
+        units = self._split_block_units(block)
+        chunks: list[Chunk] = []
+        current_units: list[str] = []
+        current_tokens = 0
+        split_offset = 0
+
+        def flush() -> None:
+            nonlocal current_units, current_tokens, split_offset
+            if not current_units:
+                return
+            separator = "\n" if block.metadata.get("content_type") == "code" else " "
+            text = separator.join(current_units).strip()
+            if text:
+                chunks.append(
+                    self._make_chunk(
+                        [(block_id, block, self.token_counter.count(text))],
+                        text_override=text,
+                        suffix=f"split-{split_offset}",
+                        char_start=block.char_start,
+                        char_end=block.char_start + len(text),
+                    )
+                )
+            split_offset += 1
+            current_units = []
+            current_tokens = 0
+
+        for unit in units:
+            unit_tokens = max(1, self.token_counter.count(unit))
+            if unit_tokens > self.max_tokens:
+                flush()
+                for piece in self._split_oversized_unit(unit):
+                    chunks.append(
+                        self._make_chunk(
+                            [(block_id, block, self.token_counter.count(piece))],
+                            text_override=piece,
+                            suffix=f"split-{split_offset}",
+                            char_start=block.char_start,
+                            char_end=block.char_start + len(piece),
+                        )
+                    )
+                    split_offset += 1
+                continue
+            if current_units and current_tokens + unit_tokens > self.max_tokens:
+                flush()
+            current_units.append(unit)
+            current_tokens += unit_tokens
+            if current_tokens >= self.target_tokens:
+                flush()
+        flush()
+        return chunks
+
+    def _split_block_units(self, block: TextBlock) -> list[str]:
+        text = block.text.strip()
+        if not text:
+            return []
+        if block.metadata.get("content_type") == "code":
+            units = [line.rstrip() for line in text.splitlines() if line.strip()]
+            return units or [text]
+        paragraphs = [item.strip() for item in re.split(r"\n\s*\n+", text) if item.strip()]
+        units: list[str] = []
+        for paragraph in paragraphs or [text]:
+            sentences = [sentence.strip() for sentence in SENTENCE_RE.split(paragraph) if sentence.strip()]
+            units.extend(sentences or [paragraph])
+        return units
+
+    def _split_oversized_unit(self, unit: str) -> list[str]:
+        words = unit.split()
+        pieces: list[str] = []
+        current: list[str] = []
+        for word in words:
+            current.append(word)
+            if self.token_counter.count(" ".join(current)) >= self.target_tokens:
+                pieces.append(" ".join(current))
+                current = []
+        if current:
+            pieces.append(" ".join(current))
+        return pieces or [unit]
+
+    def _flush(self, items: list[tuple[int, TextBlock, int]]) -> list[Chunk]:
+        if not items:
+            return []
+        return [self._make_chunk(items)]
+
+    def _make_chunk(
+        self,
+        items: list[tuple[int, TextBlock, int]],
+        text_override: str | None = None,
+        suffix: str = "",
+        char_start: int | None = None,
+        char_end: int | None = None,
+    ) -> Chunk:
+        blocks = [item[1] for item in items]
+        text = text_override if text_override is not None else "\n\n".join(block.text for block in blocks)
+        token_count = max(1, self.token_counter.count(text))
+        start_page = next((block.page for block in blocks if block.page is not None), None)
+        end_page = next((block.page for block in reversed(blocks) if block.page is not None), start_page)
+        digest_input = (
+            f"{blocks[0].source_path}:{blocks[0].paragraph_index}:"
+            f"{blocks[-1].paragraph_index}:{suffix}:{text[:80]}"
+        )
+        digest = hashlib.sha1(digest_input.encode("utf-8", errors="replace")).hexdigest()[:12]
+        metadata = self._chunk_metadata(blocks)
+        return Chunk(
+            id=f"chunk-{digest}",
+            text=text,
+            source_path=blocks[0].source_path,
+            document_index=blocks[0].document_index,
+            start_page=start_page,
+            end_page=end_page,
+            start_paragraph=blocks[0].paragraph_index,
+            end_paragraph=blocks[-1].paragraph_index,
+            char_start=blocks[0].char_start if char_start is None else char_start,
+            char_end=blocks[-1].char_end if char_end is None else char_end,
+            token_count=token_count,
+            block_ids=[item[0] for item in items],
+            metadata=metadata,
+        )
+
+    def _chunk_metadata(self, blocks: list[TextBlock]) -> dict[str, Any]:
+        metadata: dict[str, Any] = {"bbox": [block.bbox for block in blocks if block.bbox is not None]}
+        all_keys = {key for block in blocks for key in block.metadata}
+        for key in all_keys:
+            values = [block.metadata.get(key) for block in blocks if key in block.metadata]
+            unique_values = {repr(value): value for value in values}
+            if len(unique_values) == 1:
+                metadata[key] = values[0]
+            elif key == "content_type":
+                metadata[key] = "mixed"
+        start_lines = [block.metadata.get("start_line") for block in blocks if isinstance(block.metadata.get("start_line"), int)]
+        end_lines = [block.metadata.get("end_line") for block in blocks if isinstance(block.metadata.get("end_line"), int)]
+        if start_lines:
+            metadata["start_line"] = min(start_lines)
+        if end_lines:
+            metadata["end_line"] = max(end_lines)
+        return metadata
+
+
+class StructureAwareChunker(_ChunkGroupBase):
+    """Chunk documents and code using structural metadata plus semantic drift."""
+
+    def __init__(
+        self,
+        target_tokens: int = 650,
+        min_tokens: int = 120,
+        max_tokens: int = 900,
+        token_counter: TokenCounter | None = None,
+        block_embeddings: Sequence[list[float]] | None = None,
+        semantic_threshold: float = 0.35,
+    ) -> None:
+        super().__init__(
+            target_tokens=target_tokens,
+            min_tokens=min_tokens,
+            max_tokens=max_tokens,
+            token_counter=token_counter,
+        )
+        self.block_embeddings = block_embeddings
+        self.semantic_threshold = semantic_threshold
+
+    def chunk(self, blocks: Sequence[TextBlock]) -> list[Chunk]:
+        if self.block_embeddings is not None and len(self.block_embeddings) != len(blocks):
+            raise ValueError("StructureAwareChunker requires one embedding per text block when semantic boundaries are enabled.")
+
+        chunks: list[Chunk] = []
+        current: list[tuple[int, TextBlock, int]] = []
+        current_tokens = 0
+
+        def flush() -> None:
+            nonlocal current, current_tokens
+            chunks.extend(self._flush(current))
+            current = []
+            current_tokens = 0
+
+        for block_id, block in enumerate(blocks):
+            block_tokens = max(1, self.token_counter.count(block.text))
+            content_type = block.metadata.get("content_type", "document")
+            is_symbol = bool(block.metadata.get("symbol_name") or block.metadata.get("symbol_kind"))
+
+            if content_type == "code" and is_symbol:
+                flush()
+                if block_tokens > self.max_tokens:
+                    chunks.extend(self._split_large_block(block_id, block))
+                else:
+                    chunks.extend(self._flush([(block_id, block, block_tokens)]))
+                continue
+
+            if block_tokens > self.max_tokens:
+                flush()
+                chunks.extend(self._split_large_block(block_id, block))
+                continue
+
+            if current and self._structural_boundary(current[-1][1], block):
+                flush()
+
+            would_exceed = current_tokens + block_tokens > self.max_tokens
+            topic_shift = self._semantic_boundary(block_id, block, current)
+            good_enough = current_tokens >= self.min_tokens
+            if current and good_enough and (would_exceed or topic_shift):
+                flush()
+
+            current.append((block_id, block, block_tokens))
+            current_tokens += block_tokens
+
+            if current_tokens >= self.target_tokens:
+                flush()
+
+        flush()
+        for chunk in chunks:
+            chunk.metadata["chunker"] = "structure-aware"
+            if self.block_embeddings is not None:
+                chunk.metadata["semantic_threshold"] = self.semantic_threshold
+        return chunks
+
+    def _structural_boundary(self, previous: TextBlock, current: TextBlock) -> bool:
+        if previous.source_path != current.source_path:
+            return True
+        previous_type = previous.metadata.get("content_type", "document")
+        current_type = current.metadata.get("content_type", "document")
+        if previous_type != current_type:
+            return True
+        if previous.metadata.get("section_hint") != current.metadata.get("section_hint"):
+            return bool(previous.metadata.get("section_hint") or current.metadata.get("section_hint"))
+        return False
+
+    def _semantic_boundary(self, block_id: int, block: TextBlock, current: list[tuple[int, TextBlock, int]]) -> bool:
+        if self.block_embeddings is None or not current:
+            return False
+        previous_id, previous, _ = current[-1]
+        if previous.source_path != block.source_path:
+            return False
+        if previous.metadata.get("content_type", "document") != "document":
+            return False
+        if block.metadata.get("content_type", "document") != "document":
+            return False
+        similarity = cosine(self.block_embeddings[previous_id], self.block_embeddings[block_id])
+        return similarity < self.semantic_threshold
+
+
+class SemanticThresholdChunker(_ChunkGroupBase):
+    """Start a new chunk when adjacent block embeddings indicate topic drift."""
+
+    def __init__(
+        self,
+        block_embeddings: Sequence[list[float]],
+        similarity_threshold: float = 0.35,
+        target_tokens: int = 650,
+        min_tokens: int = 120,
+        max_tokens: int = 900,
+        token_counter: TokenCounter | None = None,
+    ) -> None:
+        super().__init__(
+            target_tokens=target_tokens,
+            min_tokens=min_tokens,
+            max_tokens=max_tokens,
+            token_counter=token_counter,
+        )
+        self.block_embeddings = block_embeddings
+        self.similarity_threshold = similarity_threshold
+
+    def chunk(self, blocks: Sequence[TextBlock]) -> list[Chunk]:
+        if len(self.block_embeddings) != len(blocks):
+            raise ValueError("SemanticThresholdChunker requires one embedding per text block.")
+
+        chunks: list[Chunk] = []
+        current: list[tuple[int, TextBlock, int]] = []
+        current_tokens = 0
+
+        for block_id, block in enumerate(blocks):
+            block_tokens = max(1, self.token_counter.count(block.text))
+            if block_tokens > self.max_tokens:
+                chunks.extend(self._flush(current))
+                current = []
+                current_tokens = 0
+                chunks.extend(self._split_large_block(block_id, block))
+                continue
+
+            starts_new_doc = current and block.source_path != current[-1][1].source_path
+            would_exceed = current_tokens + block_tokens > self.max_tokens
+            topic_shift = self._topic_shift(block_id, current)
+            good_enough = current_tokens >= self.min_tokens
+
+            if current and (starts_new_doc or (would_exceed and good_enough) or (topic_shift and good_enough)):
+                chunks.extend(self._flush(current))
+                current = []
+                current_tokens = 0
+
+            current.append((block_id, block, block_tokens))
+            current_tokens += block_tokens
+
+            if current_tokens >= self.target_tokens:
+                chunks.extend(self._flush(current))
+                current = []
+                current_tokens = 0
+
+        chunks.extend(self._flush(current))
+        for chunk in chunks:
+            chunk.metadata["chunker"] = "semantic-threshold"
+            chunk.metadata["similarity_threshold"] = self.similarity_threshold
+        return chunks
+
+    def _topic_shift(self, block_id: int, current: list[tuple[int, TextBlock, int]]) -> bool:
+        if not current:
+            return False
+        previous_id = current[-1][0]
+        if current[-1][1].source_path != current[0][1].source_path:
+            return False
+        similarity = cosine(self.block_embeddings[previous_id], self.block_embeddings[block_id])
+        return similarity < self.similarity_threshold
+