344 lines
10 KiB
Python
344 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
MMR (Maximal Marginal Relevance) Deduplication for RRF Results
|
|
================================================================
|
|
Re-ranks fused RRF results to maximize relevance while minimizing redundancy.
|
|
|
|
MMR formula:
|
|
MMR = argmax_{d∈R\\S} [ λ·sim(d,q) - (1-λ)·max_{d'∈S} sim(d,d') ]
|
|
|
|
Where:
|
|
R = all candidate results
|
|
S = already-selected results
|
|
q = original query
|
|
sim(d,q) = relevance to query (RRF score)
|
|
sim(d,d') = similarity between documents (embedding cosine)
|
|
λ = trade-off (0.5-0.7 typical; higher = more diverse)
|
|
|
|
Used after _rrf_fuse() in rag_service.py to ensure the final
|
|
top-K results contain minimal duplicate/near-duplicate content.
|
|
"""
|
|
|
|
import contextlib
|
|
import json
|
|
import logging
|
|
import math
|
|
from typing import Any
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def cosine_similarity(vec_a: list[float], vec_b: list[float]) -> float:
|
|
"""Compute cosine similarity between two vectors."""
|
|
if not vec_a or not vec_b or len(vec_a) != len(vec_b):
|
|
return 0.0
|
|
|
|
# Handle different-length vectors by truncating to shorter
|
|
min_len = min(len(vec_a), len(vec_b))
|
|
a = vec_a[:min_len]
|
|
b = vec_b[:min_len]
|
|
|
|
dot = sum(x * y for x, y in zip(a, b, strict=False))
|
|
mag_a = math.sqrt(sum(x * x for x in a))
|
|
mag_b = math.sqrt(sum(x * x for x in b))
|
|
|
|
if mag_a == 0 or mag_b == 0:
|
|
return 0.0
|
|
|
|
return dot / (mag_a * mag_b)
|
|
|
|
|
|
def text_jaccard_similarity(text_a: str, text_b: str) -> float:
|
|
"""
|
|
Token-level Jaccard similarity for text dedup when embeddings unavailable.
|
|
Fast and good enough for surface-level dedup.
|
|
"""
|
|
if not text_a or not text_b:
|
|
return 0.0
|
|
|
|
# Tokenize with simple split + lowercase
|
|
tokens_a = set(text_a.lower().split())
|
|
tokens_b = set(text_b.lower().split())
|
|
|
|
if not tokens_a or not tokens_b:
|
|
return 0.0
|
|
|
|
intersection = tokens_a & tokens_b
|
|
union = tokens_a | tokens_b
|
|
|
|
return len(intersection) / len(union) if union else 0.0
|
|
|
|
|
|
def content_hash_dedup(
|
|
results: list[dict[str, Any]],
|
|
hash_field: str = "doc_id",
|
|
id_field: str = "id",
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Remove exact duplicates based on content hash or ID.
|
|
Fast pre-filter before MMR.
|
|
"""
|
|
seen: set[str] = set()
|
|
deduped = []
|
|
|
|
for r in results:
|
|
# Use doc_id or id as dedup key
|
|
key = r.get(hash_field) or r.get(id_field) or ""
|
|
if not key:
|
|
# Fallback: hash the content
|
|
content = r.get("content", "")[:500]
|
|
key = f"hash:{hash(content)}" if content else ""
|
|
|
|
if key and key in seen:
|
|
continue
|
|
if key:
|
|
seen.add(key)
|
|
deduped.append(r)
|
|
|
|
return deduped
|
|
|
|
|
|
def mmr_rerank(
|
|
results: list[dict[str, Any]],
|
|
query_score_field: str = "rrf_score",
|
|
content_field: str = "content",
|
|
embedding_field: str = "embedding",
|
|
lambda_param: float = 0.6,
|
|
top_k: int = 20,
|
|
similarity_threshold: float = 0.85,
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Maximal Marginal Relevance reranking.
|
|
|
|
Args:
|
|
results: fused results from _rrf_fuse(), sorted by rrf_score desc
|
|
query_score_field: field name for query relevance (rrf_score)
|
|
content_field: field name for text content (used for Jaccard sim)
|
|
embedding_field: field name for vector embeddings (used for cosine sim)
|
|
lambda_param: relevance-diversity trade-off (0=all diverse, 1=all relevant)
|
|
top_k: number of results to return
|
|
similarity_threshold: skip MMR for results below this relevance threshold
|
|
|
|
Returns:
|
|
Re-ranked list with up to top_k results, maximizing relevance + diversity.
|
|
"""
|
|
if not results:
|
|
return []
|
|
|
|
# Pre-filter: remove exact duplicates
|
|
deduped = content_hash_dedup(results)
|
|
if len(deduped) <= 1:
|
|
return deduped[:top_k]
|
|
|
|
# Normalize query relevance scores to [0, 1]
|
|
max_score = max(r.get(query_score_field, 0) for r in deduped)
|
|
if max_score == 0:
|
|
return deduped[:top_k]
|
|
|
|
for r in deduped:
|
|
r["_mmr_relevance"] = r.get(query_score_field, 0) / max_score
|
|
|
|
# Greedy MMR selection
|
|
selected: list[dict[str, Any]] = []
|
|
remaining: list[dict[str, Any]] = list(deduped)
|
|
|
|
# First result: highest relevance
|
|
remaining.sort(key=lambda x: x.get("_mmr_relevance", 0), reverse=True)
|
|
selected.append(remaining.pop(0))
|
|
|
|
while len(selected) < top_k and remaining:
|
|
best_mmr = -1.0
|
|
best_idx = 0
|
|
|
|
for i, candidate in enumerate(remaining):
|
|
relevance = candidate.get("_mmr_relevance", 0)
|
|
|
|
# Skip candidates below threshold
|
|
if relevance < 0.01:
|
|
continue
|
|
|
|
# Maximum similarity to already-selected results
|
|
max_sim = 0.0
|
|
for sel in selected:
|
|
# Try embedding similarity first (most accurate)
|
|
c_emb = candidate.get(embedding_field)
|
|
s_emb = sel.get(embedding_field)
|
|
if c_emb and s_emb and isinstance(c_emb, list) and isinstance(s_emb, list):
|
|
sim = cosine_similarity(c_emb, s_emb)
|
|
else:
|
|
# Fall back to text Jaccard similarity
|
|
c_text = candidate.get(content_field, "") or ""
|
|
s_text = sel.get(content_field, "") or ""
|
|
sim = text_jaccard_similarity(c_text, s_text)
|
|
|
|
max_sim = max(max_sim, sim)
|
|
|
|
# MMR score: λ·relevance - (1-λ)·max_similarity
|
|
mmr_score = lambda_param * relevance - (1 - lambda_param) * max_sim
|
|
|
|
if mmr_score > best_mmr:
|
|
best_mmr = mmr_score
|
|
best_idx = i
|
|
|
|
if best_mmr <= -1.0:
|
|
# No valid candidates left
|
|
break
|
|
|
|
selected.append(remaining.pop(best_idx))
|
|
|
|
# Clean up temporary fields
|
|
for r in deduped:
|
|
r.pop("_mmr_relevance", None)
|
|
|
|
# Add MMR metadata
|
|
for i, r in enumerate(selected):
|
|
r["mmr_rank"] = i + 1
|
|
r["mmr_selected"] = True
|
|
|
|
return selected
|
|
|
|
|
|
async def mmr_dedup_results(
|
|
results: list[dict[str, Any]],
|
|
query_score_field: str = "rrf_score",
|
|
content_field: str = "content",
|
|
lambda_param: float = 0.6,
|
|
top_k: int = 20,
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Async wrapper for MMR deduplication.
|
|
This is the main entry point used by rag_service.py.
|
|
|
|
Attempts to load embeddings from Redis for similarity computation.
|
|
Falls back to text Jaccard if embeddings are unavailable.
|
|
"""
|
|
try:
|
|
# Try to enrich results with embeddings from Redis for better dedup
|
|
enriched = await _enrich_with_embeddings(results)
|
|
return mmr_rerank(
|
|
enriched,
|
|
query_score_field=query_score_field,
|
|
content_field=content_field,
|
|
embedding_field="_embedding",
|
|
lambda_param=lambda_param,
|
|
top_k=top_k,
|
|
)
|
|
except Exception as e:
|
|
# Fall back to text-only dedup
|
|
logger.debug(f"MMR embedding enrichment failed, using text-only: {e}")
|
|
return mmr_rerank(
|
|
results,
|
|
query_score_field=query_score_field,
|
|
content_field=content_field,
|
|
lambda_param=lambda_param,
|
|
top_k=top_k,
|
|
)
|
|
|
|
|
|
async def _enrich_with_embeddings(
|
|
results: list[dict[str, Any]],
|
|
) -> list[dict[str, Any]]:
|
|
"""
|
|
Try to load embedding vectors from Redis for each result.
|
|
Used for more accurate MMR similarity computation.
|
|
Falls back gracefully if Redis/embeddings aren't available.
|
|
"""
|
|
import os
|
|
|
|
import redis.asyncio as aioredis
|
|
|
|
if not results:
|
|
return results
|
|
|
|
# Only enrich if we have doc_ids and collection info
|
|
docs_with_ids = [(i, r) for i, r in enumerate(results) if r.get("doc_id") or r.get("id")]
|
|
|
|
if not docs_with_ids:
|
|
return results
|
|
|
|
r = None
|
|
try:
|
|
r = aioredis.Redis(
|
|
host=os.getenv("REDIS_HOST", "rmi-redis"),
|
|
port=int(os.getenv("REDIS_PORT", "6379")),
|
|
password=os.getenv("REDIS_PASSWORD", ""),
|
|
decode_responses=False, # binary for vectors
|
|
)
|
|
|
|
# Batch fetch embeddings
|
|
pipe = r.pipeline()
|
|
keys_to_fetch = []
|
|
for idx, result in docs_with_ids:
|
|
doc_id = result.get("doc_id") or result.get("id")
|
|
coll = result.get("collection", "wallet_profiles")
|
|
key = f"rag:{coll}:{doc_id}"
|
|
pipe.get(key)
|
|
keys_to_fetch.append((idx, key))
|
|
|
|
raw_docs = await pipe.execute()
|
|
|
|
for (idx, key), data in zip(keys_to_fetch, raw_docs, strict=False):
|
|
if data:
|
|
try:
|
|
doc = json.loads(data)
|
|
embedding = doc.get("vector", [])
|
|
if isinstance(embedding, list) and len(embedding) > 0:
|
|
# Store embedding for MMR similarity computation
|
|
# Use first 128 dims for speed (cosine sim is preserved in high dims)
|
|
results[idx]["_embedding"] = embedding[:128]
|
|
except (json.JSONDecodeError, Exception):
|
|
pass
|
|
|
|
except Exception as e:
|
|
logger.debug(f"Embedding enrichment from Redis failed: {e}")
|
|
finally:
|
|
if r is not None:
|
|
with contextlib.suppress(Exception):
|
|
await r.aclose()
|
|
|
|
return results
|
|
|
|
|
|
def cluster_similar_results(
|
|
results: list[dict[str, Any]],
|
|
content_field: str = "content",
|
|
similarity_threshold: float = 0.7,
|
|
) -> list[list[dict[str, Any]]]:
|
|
"""
|
|
Group results into clusters of similar content.
|
|
Each cluster contains near-duplicate results.
|
|
|
|
Returns list of clusters, where each cluster is a list of similar results.
|
|
The first result in each cluster is the "representative" (highest score).
|
|
"""
|
|
if not results:
|
|
return []
|
|
|
|
clusters: list[list[dict[str, Any]]] = []
|
|
assigned: set[int] = set()
|
|
|
|
for i, result in enumerate(results):
|
|
if i in assigned:
|
|
continue
|
|
|
|
cluster = [result]
|
|
assigned.add(i)
|
|
|
|
for j, other in enumerate(results):
|
|
if j in assigned or j == i:
|
|
continue
|
|
|
|
sim = text_jaccard_similarity(
|
|
result.get(content_field, ""),
|
|
other.get(content_field, ""),
|
|
)
|
|
if sim >= similarity_threshold:
|
|
cluster.append(other)
|
|
assigned.add(j)
|
|
|
|
clusters.append(cluster)
|
|
|
|
# Sort clusters by highest-scoring result
|
|
clusters.sort(key=lambda c: max(r.get("rrf_score", 0) for r in c), reverse=True)
|
|
|
|
return clusters
|