Retrievers
Inspectable retrieval algorithms used by the educational pipelines.
"""Inspectable retrieval algorithms used by the educational pipelines.
The implementations intentionally trade scale for clarity. BM25 is a real
sparse ranker. :class:`HashingSemanticRetriever` is an offline semantic proxy
based on synonym canonicalization and signed feature hashing; it demonstrates
the dual-encoder interface but is not a neural embedding model.
"""
import hashlib
import math
import re
from collections import Counter, defaultdict
from typing import Callable, DefaultDict, Dict, Iterable, List, Mapping, Protocol, Sequence, Tuple
from .models import Chunk, SearchResult
from .text import DEFAULT_SYNONYM_GROUPS, semantic_terms, synonym_map, tokenize
class Retriever(Protocol):
"""Minimal retrieval interface shared by all pipeline components."""
def search(self, query: str, k: int = 5) -> List[SearchResult]:
"""Return at most ``k`` results ordered best first."""
class BM25Retriever:
"""Okapi BM25 over title-boosted chunks.
The Robertson/Sparck Jones IDF variant is kept positive with
``log(1 + ratio)``. This is the common Lucene-style form and avoids a term
appearing in most tiny demo documents receiving a negative score.
"""
name = "bm25"
def __init__(
self,
chunks: Sequence[Chunk],
k1: float = 1.2,
b: float = 0.75,
title_boost: int = 2,
) -> None:
if k1 <= 0:
raise ValueError("k1 must be positive")
if not 0 <= b <= 1:
raise ValueError("b must be between zero and one")
self.chunks = tuple(chunks)
self.k1 = k1
self.b = b
self._term_frequencies: List[Counter[str]] = []
self._lengths: List[int] = []
document_frequency: Counter[str] = Counter()
for chunk in self.chunks:
terms = tokenize(((chunk.title + " ") * title_boost) + chunk.text)
frequencies = Counter(terms)
self._term_frequencies.append(frequencies)
self._lengths.append(len(terms))
document_frequency.update(frequencies.keys())
self._average_length = (
sum(self._lengths) / len(self._lengths) if self._lengths else 0.0
)
total = len(self.chunks)
self._idf = {
term: math.log(1.0 + (total - count + 0.5) / (count + 0.5))
for term, count in document_frequency.items()
}
def _score(self, query_terms: Sequence[str], index: int) -> float:
if not self._average_length:
return 0.0
frequencies = self._term_frequencies[index]
length_norm = 1.0 - self.b + self.b * self._lengths[index] / self._average_length
score = 0.0
for term, query_frequency in Counter(query_terms).items():
term_frequency = frequencies.get(term, 0)
if not term_frequency:
continue
saturation = (term_frequency * (self.k1 + 1.0)) / (
term_frequency + self.k1 * length_norm
)
score += self._idf.get(term, 0.0) * saturation * (1.0 + math.log(query_frequency))
return score
def search(self, query: str, k: int = 5) -> List[SearchResult]:
if k <= 0:
return []
query_terms = tokenize(query)
scored = [
(self._score(query_terms, index), chunk)
for index, chunk in enumerate(self.chunks)
]
scored = [(score, chunk) for score, chunk in scored if score > 0.0]
scored.sort(key=lambda item: (-item[0], item[1].id))
return [
SearchResult(
chunk=chunk,
score=score,
rank=rank,
retriever=self.name,
component_scores={self.name: score},
)
for rank, (score, chunk) in enumerate(scored[:k], start=1)
]
class HashingSemanticRetriever:
"""Dependency-free approximation of a semantic bi-encoder.
Both queries and chunks are independently converted to fixed-width vectors.
Explicit synonym groups and character n-grams make the demo robust to some
vocabulary mismatch; signed hashing keeps the representation compact.
A production replacement can preserve ``search`` while swapping in a
trained encoder plus an ANN index.
"""
name = "hash-semantic"
def __init__(
self,
chunks: Sequence[Chunk],
dimensions: int = 512,
synonym_groups: Sequence[Sequence[str]] = DEFAULT_SYNONYM_GROUPS,
minimum_score: float = 0.03,
) -> None:
if dimensions < 64:
raise ValueError("dimensions must be at least 64")
self.chunks = tuple(chunks)
self.dimensions = dimensions
self.minimum_score = minimum_score
self._synonyms = synonym_map(synonym_groups)
raw_features = [self._features(chunk.title + " " + chunk.text) for chunk in self.chunks]
document_frequency: Counter[int] = Counter()
for features in raw_features:
document_frequency.update(features.keys())
total = max(1, len(self.chunks))
self._idf = {
feature: math.log((total + 1.0) / (frequency + 1.0)) + 1.0
for feature, frequency in document_frequency.items()
}
self._vectors = [self._vectorize(features) for features in raw_features]
def _hash(self, feature: str) -> Tuple[int, float]:
digest = hashlib.blake2b(feature.encode("utf-8"), digest_size=8).digest()
value = int.from_bytes(digest, byteorder="big", signed=False)
index = value % self.dimensions
sign = -1.0 if value & (1 << 63) else 1.0
return index, sign
def _features(self, text: str) -> Mapping[int, float]:
features: DefaultDict[int, float] = defaultdict(float)
terms = semantic_terms(text, self._synonyms)
for term in terms:
index, sign = self._hash("w:" + term)
features[index] += sign
if len(term) >= 4:
padded = "^" + term + "$"
for offset in range(len(padded) - 2):
index, sign = self._hash("c:" + padded[offset : offset + 3])
features[index] += 0.18 * sign
for left, right in zip(terms, terms[1:]):
index, sign = self._hash("b:" + left + "_" + right)
features[index] += 0.35 * sign
return features
def _vectorize(self, features: Mapping[int, float]) -> Tuple[float, ...]:
vector = [0.0] * self.dimensions
unseen_idf = math.log(len(self.chunks) + 1.0) + 1.0
for index, frequency in features.items():
magnitude = 1.0 + math.log(abs(frequency)) if abs(frequency) >= 1.0 else abs(frequency)
vector[index] = math.copysign(magnitude, frequency) * self._idf.get(index, unseen_idf)
norm = math.sqrt(sum(value * value for value in vector))
if norm:
vector = [value / norm for value in vector]
return tuple(vector)
@staticmethod
def _cosine(left: Sequence[float], right: Sequence[float]) -> float:
return sum(a * b for a, b in zip(left, right))
def search(self, query: str, k: int = 5) -> List[SearchResult]:
if k <= 0:
return []
query_vector = self._vectorize(self._features(query))
scored = [
(self._cosine(query_vector, vector), chunk)
for vector, chunk in zip(self._vectors, self.chunks)
]
scored = [(score, chunk) for score, chunk in scored if score >= self.minimum_score]
scored.sort(key=lambda item: (-item[0], item[1].id))
return [
SearchResult(
chunk=chunk,
score=score,
rank=rank,
retriever=self.name,
component_scores={self.name: score},
)
for rank, (score, chunk) in enumerate(scored[:k], start=1)
]
class HybridRetriever:
"""Fuse heterogeneous retrievers with weighted reciprocal-rank fusion."""
name = "hybrid-rrf"
def __init__(
self,
retrievers: Sequence[Tuple[str, Retriever, float]],
rrf_constant: int = 60,
candidate_multiplier: int = 4,
) -> None:
if not retrievers:
raise ValueError("at least one retriever is required")
if rrf_constant <= 0:
raise ValueError("rrf_constant must be positive")
self.retrievers = tuple(retrievers)
self.rrf_constant = rrf_constant
self.candidate_multiplier = candidate_multiplier
def search(self, query: str, k: int = 5) -> List[SearchResult]:
if k <= 0:
return []
candidate_k = max(k * self.candidate_multiplier, k)
chunks: Dict[str, Chunk] = {}
fused: DefaultDict[str, float] = defaultdict(float)
components: DefaultDict[str, Dict[str, float]] = defaultdict(dict)
for name, retriever, weight in self.retrievers:
for result in retriever.search(query, candidate_k):
chunk_id = result.chunk.id
chunks[chunk_id] = result.chunk
contribution = weight / (self.rrf_constant + result.rank)
fused[chunk_id] += contribution
components[chunk_id][name] = result.score
components[chunk_id][name + "_rrf"] = contribution
ranked_ids = sorted(fused, key=lambda chunk_id: (-fused[chunk_id], chunk_id))[:k]
return [
SearchResult(
chunk=chunks[chunk_id],
score=fused[chunk_id],
rank=rank,
retriever=self.name,
component_scores=components[chunk_id],
)
for rank, chunk_id in enumerate(ranked_ids, start=1)
]
class MultiQueryRetriever:
"""Retrieve several query views and fuse them, a deterministic RAG-Fusion analogue."""
name = "multi-query"
def __init__(
self,
retriever: Retriever,
rewriter: Callable[[str], Sequence[str]],
rrf_constant: int = 60,
) -> None:
self.retriever = retriever
self.rewriter = rewriter
self.rrf_constant = rrf_constant
def search(self, query: str, k: int = 5) -> List[SearchResult]:
if k <= 0:
return []
queries: List[str] = []
for candidate in (query, *self.rewriter(query)):
if candidate.strip() and candidate not in queries:
queries.append(candidate)
chunks: Dict[str, Chunk] = {}
scores: DefaultDict[str, float] = defaultdict(float)
components: DefaultDict[str, Dict[str, float]] = defaultdict(dict)
candidate_k = max(4 * k, k)
for query_index, rewritten in enumerate(queries):
for result in self.retriever.search(rewritten, candidate_k):
chunk_id = result.chunk.id
chunks[chunk_id] = result.chunk
contribution = 1.0 / (self.rrf_constant + result.rank)
scores[chunk_id] += contribution
components[chunk_id][f"q{query_index}_rrf"] = contribution
ranked_ids = sorted(scores, key=lambda chunk_id: (-scores[chunk_id], chunk_id))[:k]
return [
SearchResult(
chunk=chunks[chunk_id],
score=scores[chunk_id],
rank=rank,
retriever=self.name,
component_scores=components[chunk_id],
)
for rank, chunk_id in enumerate(ranked_ids, start=1)
]
class GraphExpandedRetriever:
"""Expand initially retrieved chunks through shared metadata entities.
This is a transparent local-neighborhood analogue, not Microsoft GraphRAG's
community-report algorithm. It demonstrates why graph expansion can find
logically linked evidence missed by independent similarity ranking.
"""
name = "graph-expanded"
def __init__(
self,
base_retriever: Retriever,
chunks: Sequence[Chunk],
expansion_weight: float = 0.35,
) -> None:
self.base_retriever = base_retriever
self.chunks = tuple(chunks)
self.expansion_weight = expansion_weight
self._by_id = {chunk.id: chunk for chunk in self.chunks}
self._entity_to_chunks: DefaultDict[str, List[str]] = defaultdict(list)
for chunk in self.chunks:
entities = chunk.metadata.get("entities", ())
if isinstance(entities, str):
entities = (entities,)
for entity in entities:
normalized = str(entity).strip().lower()
if normalized:
self._entity_to_chunks[normalized].append(chunk.id)
@staticmethod
def _entities(chunk: Chunk) -> Tuple[str, ...]:
entities = chunk.metadata.get("entities", ())
if isinstance(entities, str):
entities = (entities,)
return tuple(str(entity).strip().lower() for entity in entities if str(entity).strip())
def search(self, query: str, k: int = 5) -> List[SearchResult]:
if k <= 0:
return []
initial = self.base_retriever.search(query, max(k * 3, k))
if not initial:
return []
maximum = max(result.score for result in initial) or 1.0
scores: DefaultDict[str, float] = defaultdict(float)
components: DefaultDict[str, Dict[str, float]] = defaultdict(dict)
query_tokens = set(tokenize(query))
for result in initial:
normalized_score = result.score / maximum
scores[result.chunk.id] = max(scores[result.chunk.id], normalized_score)
components[result.chunk.id]["base"] = normalized_score
for entity in self._entities(result.chunk):
direct_bonus = 0.15 if set(tokenize(entity)) <= query_tokens else 0.0
for neighbor_id in self._entity_to_chunks.get(entity, ()):
if neighbor_id == result.chunk.id:
continue
contribution = normalized_score * self.expansion_weight + direct_bonus
scores[neighbor_id] += contribution
components[neighbor_id]["graph"] = (
components[neighbor_id].get("graph", 0.0) + contribution
)
ranked_ids = sorted(scores, key=lambda chunk_id: (-scores[chunk_id], chunk_id))[:k]
return [
SearchResult(
chunk=self._by_id[chunk_id],
score=scores[chunk_id],
rank=rank,
retriever=self.name,
component_scores=components[chunk_id],
)
for rank, chunk_id in enumerate(ranked_ids, start=1)
]
class AdaptiveRetriever:
"""Route identifier, multi-hop, and ordinary semantic queries differently."""
name = "adaptive"
MULTIHOP_CUES = re.compile(
r"\b(compare|difference|relationship|before|after|versus|vs|combine|both|evolve)\b|\band\b",
re.IGNORECASE,
)
IDENTIFIER = re.compile(r"(?:[A-Z]{2,}(?:-[A-Z0-9]+)*|\b\d{4}\b|\"[^\"]+\")")
def __init__(
self,
sparse: Retriever,
hybrid: Retriever,
graph: Retriever,
) -> None:
self.sparse = sparse
self.hybrid = hybrid
self.graph = graph
self.last_route = "hybrid"
def route_for(self, query: str) -> str:
if self.MULTIHOP_CUES.search(query):
return "graph"
if self.IDENTIFIER.search(query):
return "sparse"
return "hybrid"
def search(self, query: str, k: int = 5) -> List[SearchResult]:
route = self.route_for(query)
self.last_route = route
retriever = {"sparse": self.sparse, "hybrid": self.hybrid, "graph": self.graph}[route]
results = retriever.search(query, k)
return [
SearchResult(
chunk=result.chunk,
score=result.score,
rank=result.rank,
retriever=f"{self.name}:{route}",
component_scores=result.component_scores,
)
for result in results
]
def simple_query_variants(query: str) -> Sequence[str]:
"""Create cheap query views without pretending to be an LLM rewriter."""
terms = tokenize(query)
content = [
term
for term in terms
if term not in {"what", "which", "who", "when", "where", "why", "how", "does", "did"}
]
variants = [" ".join(content)]
if "difference" in terms or "compare" in terms:
variants.append(" ".join(term for term in content if term not in {"difference", "compare"}))
return tuple(variant for variant in variants if variant and variant != query.lower())