"""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)