Structured
Structured, hierarchical, tabular, and multivector retrieval primitives.
"""Structured, hierarchical, tabular, and multivector retrieval primitives."""
import math
from dataclasses import dataclass
from typing import Dict, Iterable, List, Mapping, Sequence, Set, Tuple
from .text import content_terms
Vector = Sequence[float]
@dataclass(frozen=True)
class TableRow:
"""A table row with stable provenance down to page and row identity."""
table_id: str
row_id: str
cells: Mapping[str, str]
page: int = 0
source: str = ""
def text(self) -> str:
return " | ".join(f"{column}: {value}" for column, value in self.cells.items())
@dataclass(frozen=True)
class TableResult:
"""A ranked row with separate lexical and numeric-match features."""
row: TableRow
score: float
lexical_overlap: float
numeric_overlap: float
@dataclass(frozen=True)
class HierarchyNode:
"""A summary/section node whose leaves retain original evidence IDs."""
identifier: str
text: str
children: Tuple[str, ...] = ()
evidence_ids: Tuple[str, ...] = ()
token_cost: int = 1
def __post_init__(self) -> None:
if not self.identifier or self.token_cost <= 0:
raise ValueError("hierarchy nodes need an identifier and positive token cost")
@dataclass(frozen=True)
class HierarchySelection:
"""Selected nodes and leaf-evidence coverage under a context budget."""
nodes: Tuple[HierarchyNode, ...]
evidence_ids: Tuple[str, ...]
spent: int
def _dot(left: Vector, right: Vector) -> float:
if len(left) != len(right):
raise ValueError("vectors must have the same dimensionality")
return sum(a * b for a, b in zip(left, right))
def cosine(left: Vector, right: Vector) -> float:
"""Cosine similarity with zero vectors mapped to zero similarity."""
numerator = _dot(left, right)
left_norm = math.sqrt(_dot(left, left))
right_norm = math.sqrt(_dot(right, right))
return numerator / (left_norm * right_norm) if left_norm and right_norm else 0.0
def late_interaction_score(query_vectors: Sequence[Vector], document_vectors: Sequence[Vector]) -> float:
"""ColBERT/ColPali-style MaxSim: ``sum_i max_j q_i · d_j``."""
if not query_vectors or not document_vectors:
return 0.0
return sum(max(_dot(query, document) for document in document_vectors) for query in query_vectors)
def pool_vectors(vectors: Sequence[Vector], group_size: int) -> Tuple[Tuple[float, ...], ...]:
"""Mean-pool adjacent page/token vectors to expose the fidelity/storage tradeoff."""
if group_size <= 0:
raise ValueError("group_size must be positive")
if not vectors:
return ()
dimensions = len(vectors[0])
if dimensions == 0 or any(len(vector) != dimensions for vector in vectors):
raise ValueError("vectors must be rectangular and non-empty")
pooled = []
for start in range(0, len(vectors), group_size):
group = vectors[start : start + group_size]
pooled.append(
tuple(sum(vector[index] for vector in group) / len(group) for index in range(dimensions))
)
return tuple(pooled)
def personalized_pagerank(
adjacency: Mapping[str, Mapping[str, float]],
seeds: Mapping[str, float],
damping: float = 0.85,
tolerance: float = 1e-10,
max_iterations: int = 200,
) -> Mapping[str, float]:
"""Weighted Personalized PageRank with dangling-mass redistribution.
Edges with non-positive weights are ignored. The teleport distribution is
normalized from ``seeds`` and also receives mass from dangling nodes. This
is the core propagation mechanism behind many entity/passage graph
retrievers; entity linking and graph construction quality remain separate
concerns.
"""
if not 0.0 < damping < 1.0:
raise ValueError("damping must be between zero and one")
if tolerance <= 0 or max_iterations <= 0:
raise ValueError("tolerance and iteration budget must be positive")
nodes: Set[str] = set(adjacency) | set(seeds)
for edges in adjacency.values():
nodes.update(edges)
if not nodes:
return {}
positive_seeds = {node: max(0.0, seeds.get(node, 0.0)) for node in nodes}
seed_total = sum(positive_seeds.values())
teleport = (
{node: positive_seeds[node] / seed_total for node in nodes}
if seed_total
else {node: 1.0 / len(nodes) for node in nodes}
)
rank = dict(teleport)
normalized_edges: Dict[str, Dict[str, float]] = {}
for node in nodes:
edges = {target: weight for target, weight in adjacency.get(node, {}).items() if weight > 0}
total = sum(edges.values())
normalized_edges[node] = (
{target: weight / total for target, weight in edges.items()} if total else {}
)
for _ in range(max_iterations):
dangling = sum(rank[node] for node in nodes if not normalized_edges[node])
updated = {
node: (1.0 - damping) * teleport[node] + damping * dangling * teleport[node]
for node in nodes
}
for source in nodes:
for target, probability in normalized_edges[source].items():
updated[target] += damping * rank[source] * probability
delta = sum(abs(updated[node] - rank[node]) for node in nodes)
rank = updated
if delta < tolerance:
break
normalizer = sum(rank.values())
return {node: rank[node] / normalizer for node in sorted(nodes)}
def graph_expand(
seed_scores: Mapping[str, float],
adjacency: Mapping[str, Mapping[str, float]],
k: int,
propagation_weight: float = 0.5,
) -> Tuple[Tuple[str, float], ...]:
"""Blend original retrieval scores with graph-propagated relevance."""
if k <= 0:
return ()
if not 0.0 <= propagation_weight <= 1.0:
raise ValueError("propagation_weight must be between zero and one")
propagated = personalized_pagerank(adjacency, seed_scores)
raw_max = max(seed_scores.values(), default=0.0)
normalized_seed = {
node: (score / raw_max if raw_max > 0 else 0.0) for node, score in seed_scores.items()
}
nodes = set(propagated) | set(seed_scores)
blended = {
node: (1.0 - propagation_weight) * normalized_seed.get(node, 0.0)
+ propagation_weight * propagated.get(node, 0.0)
for node in nodes
}
return tuple(sorted(blended.items(), key=lambda item: (-item[1], item[0]))[:k])
def retrieve_table_rows(query: str, rows: Sequence[TableRow], k: int = 5) -> Tuple[TableResult, ...]:
"""Schema-aware lexical/numeric row retrieval with row-level provenance."""
if k <= 0:
return ()
query_terms = set(content_terms(query))
query_numbers = {term for term in query_terms if any(character.isdigit() for character in term)}
results: List[TableResult] = []
for row in rows:
row_terms = set(content_terms(row.text()))
row_numbers = {term for term in row_terms if any(character.isdigit() for character in term)}
lexical = len(query_terms & row_terms) / max(1, len(query_terms))
numeric = (
len(query_numbers & row_numbers) / len(query_numbers) if query_numbers else 0.0
)
header_terms = set(content_terms(" ".join(row.cells.keys())))
schema = len(query_terms & header_terms) / max(1, len(query_terms))
score = 0.65 * lexical + 0.25 * numeric + 0.10 * schema
if score > 0:
results.append(TableResult(row, score, lexical, numeric))
results.sort(key=lambda result: (-result.score, result.row.table_id, result.row.row_id))
return tuple(results[:k])
def select_hierarchy(
nodes: Sequence[HierarchyNode],
relevance: Mapping[str, float],
token_budget: int,
) -> HierarchySelection:
"""Greedily choose hierarchy nodes by new leaf coverage and relevance per token.
A summary node and its descendant leaves can duplicate information. The
marginal coverage term discourages spending the budget twice on identical
leaf evidence while still allowing a highly relevant detailed node.
"""
if token_budget < 0:
raise ValueError("token_budget must be non-negative")
remaining = list(nodes)
selected: List[HierarchyNode] = []
covered: Set[str] = set()
spent = 0
while remaining:
feasible = [node for node in remaining if spent + node.token_cost <= token_budget]
if not feasible:
break
scored = []
for node in feasible:
evidence = set(node.evidence_ids) or {node.identifier}
novelty = len(evidence - covered) / len(evidence)
gain = max(0.0, relevance.get(node.identifier, 0.0)) * (0.25 + 0.75 * novelty)
scored.append((gain / node.token_cost, gain, node.identifier, node))
scored.sort(key=lambda item: (-item[0], -item[1], item[2]))
_, gain, _, chosen = scored[0]
if gain <= 0:
break
selected.append(chosen)
covered.update(chosen.evidence_ids or (chosen.identifier,))
spent += chosen.token_cost
remaining = [node for node in remaining if node.identifier != chosen.identifier]
return HierarchySelection(tuple(selected), tuple(sorted(covered)), spent)