448 lines
16 KiB
Python
448 lines
16 KiB
Python
"""
|
|
Reciprocal Rank Fusion (RRF) utilities for combining search results.
|
|
|
|
RRF is a method to combine multiple ranked lists by computing the reciprocal
|
|
of each item's rank in each list, then summing these reciprocal ranks.
|
|
"""
|
|
|
|
import re
|
|
from typing import Any, TypeVar
|
|
|
|
from sqlalchemy import Select, and_, func, or_, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from src import models
|
|
from src.config import settings
|
|
from src.dependencies import tracked_db
|
|
from src.embedding_client import embedding_client
|
|
from src.exceptions import ValidationException
|
|
from src.models import session_peers_table
|
|
from src.utils.filter import apply_filter
|
|
from src.utils.formatting import ILIKE_ESCAPE_CHAR, escape_ilike_pattern
|
|
from src.vector_store import get_external_vector_store
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def _uses_pgvector_message_search() -> bool:
|
|
"""Return True when semantic message search can stay entirely in Postgres."""
|
|
return (
|
|
settings.VECTOR_STORE.TYPE == "pgvector" or not settings.VECTOR_STORE.MIGRATED
|
|
)
|
|
|
|
|
|
def reciprocal_rank_fusion(*ranked_lists: list[T], k: int = 60, limit: int) -> list[T]:
|
|
"""
|
|
Combine multiple ranked lists using Reciprocal Rank Fusion (RRF).
|
|
|
|
RRF assigns a score to each item based on the formula:
|
|
RRF_score = sum(1 / (k + rank_i)) for all lists where the item appears
|
|
|
|
Where:
|
|
- k is a constant (typically 60) that controls the impact of high-ranked items
|
|
- rank_i is the rank of the item in list i (1-indexed)
|
|
|
|
Args:
|
|
*ranked_lists: Variable number of ranked lists to combine
|
|
k: RRF constant parameter (default: 60)
|
|
limit: Maximum number of results to return
|
|
|
|
Returns:
|
|
list of items ranked by RRF score (highest score first)
|
|
"""
|
|
if not ranked_lists:
|
|
return []
|
|
|
|
# dictionary to store RRF scores for each item
|
|
rrf_scores: dict[T, float] = {}
|
|
|
|
# Process each ranked list
|
|
for ranked_list in ranked_lists:
|
|
for rank, item in enumerate(ranked_list, 1): # 1-indexed ranking
|
|
if item not in rrf_scores:
|
|
rrf_scores[item] = 0.0
|
|
# Add reciprocal rank contribution from this list
|
|
rrf_scores[item] += 1.0 / (k + rank)
|
|
|
|
# Sort items by RRF score (descending order)
|
|
sorted_items = sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True)
|
|
|
|
# Extract just the items (not the scores)
|
|
result = [item for item, _ in sorted_items]
|
|
|
|
return result[:limit]
|
|
|
|
|
|
async def query_external_vector_message_ids(
|
|
workspace_name: str,
|
|
embedding_query: list[float],
|
|
limit: int,
|
|
filters: dict[str, Any] | None = None,
|
|
) -> list[str]:
|
|
"""Query the external vector store and return ordered message IDs."""
|
|
external_vector_store = get_external_vector_store()
|
|
if external_vector_store is None:
|
|
return []
|
|
|
|
namespace = external_vector_store.get_vector_namespace("message", workspace_name)
|
|
|
|
vector_filters: dict[str, Any] = {}
|
|
if filters:
|
|
if "session_id" in filters:
|
|
vector_filters["session_name"] = filters["session_id"]
|
|
if "peer_id" in filters:
|
|
vector_filters["peer_name"] = filters["peer_id"]
|
|
|
|
# Oversample: multiple chunk-level hits can map to the same message,
|
|
# so fetch extra to ensure enough unique messages after deduplication.
|
|
vector_results = await external_vector_store.query(
|
|
namespace,
|
|
embedding_query,
|
|
top_k=limit * 3,
|
|
filters=vector_filters if vector_filters else None,
|
|
)
|
|
|
|
if not vector_results:
|
|
return []
|
|
|
|
seen_message_ids: dict[str, None] = {}
|
|
for result in vector_results:
|
|
message_id = result.metadata.get("message_id")
|
|
if message_id and message_id not in seen_message_ids:
|
|
seen_message_ids[message_id] = None
|
|
|
|
return list(seen_message_ids.keys())
|
|
|
|
|
|
async def fetch_messages_by_ids(
|
|
db: AsyncSession,
|
|
message_ids: list[str],
|
|
filters: dict[str, Any] | None = None,
|
|
) -> list[models.Message]:
|
|
"""Fetch messages by ID and preserve the input ordering."""
|
|
if not message_ids:
|
|
return []
|
|
|
|
stmt = select(models.Message).where(models.Message.public_id.in_(message_ids))
|
|
stmt = apply_filter(stmt, models.Message, filters)
|
|
|
|
result = await db.execute(stmt)
|
|
messages = {msg.public_id: msg for msg in result.scalars().all()}
|
|
|
|
return [messages[msg_id] for msg_id in message_ids if msg_id in messages]
|
|
|
|
|
|
async def _semantic_search_pgvector(
|
|
db: AsyncSession,
|
|
workspace_name: str,
|
|
embedding_query: list[float],
|
|
limit: int,
|
|
filters: dict[str, Any] | None = None,
|
|
) -> list[models.Message]:
|
|
"""
|
|
Perform semantic message search using pgvector in Postgres.
|
|
|
|
Args:
|
|
db: Database session
|
|
workspace_name: Name of the workspace to search in
|
|
embedding_query: Pre-computed embedding for the search query
|
|
limit: Maximum number of results to return
|
|
filters: Optional filters to apply to the message query
|
|
|
|
Returns:
|
|
list of messages ordered by semantic similarity
|
|
"""
|
|
distance_expr = models.MessageEmbedding.embedding.cosine_distance(embedding_query)
|
|
|
|
stmt = (
|
|
select(models.Message)
|
|
.join(
|
|
models.MessageEmbedding,
|
|
models.Message.public_id == models.MessageEmbedding.message_id,
|
|
)
|
|
.where(models.MessageEmbedding.embedding.isnot(None))
|
|
.where(models.MessageEmbedding.workspace_name == workspace_name)
|
|
)
|
|
|
|
if filters:
|
|
internal_filters = filters.copy()
|
|
internal_filters["workspace_id"] = workspace_name
|
|
stmt = apply_filter(stmt, models.Message, internal_filters)
|
|
|
|
# Oversample because a message with multiple embedding chunks can
|
|
# produce duplicate rows; we deduplicate in Python to preserve HNSW
|
|
# index usage (a DISTINCT ON subquery would prevent the index scan).
|
|
stmt = stmt.order_by(distance_expr).limit(limit * 2)
|
|
|
|
result = await db.execute(stmt)
|
|
seen: set[str] = set()
|
|
deduped: list[models.Message] = []
|
|
for msg in result.scalars().all():
|
|
if msg.public_id not in seen:
|
|
seen.add(msg.public_id)
|
|
deduped.append(msg)
|
|
return deduped[:limit]
|
|
|
|
|
|
async def _filter_by_peer_perspective(
|
|
db: AsyncSession,
|
|
messages: list[models.Message],
|
|
workspace_name: str,
|
|
peer_name: str,
|
|
) -> list[models.Message]:
|
|
"""
|
|
Filter messages by peer perspective (temporal session membership).
|
|
|
|
Only keeps messages from sessions where the peer was a member at the time
|
|
the message was created (between joined_at and left_at).
|
|
|
|
Args:
|
|
db: Database session
|
|
messages: List of messages to filter
|
|
workspace_name: Name of the workspace
|
|
peer_name: Name of the peer whose perspective to use
|
|
|
|
Returns:
|
|
Filtered list of messages
|
|
"""
|
|
if not messages:
|
|
return []
|
|
|
|
# Get all session memberships for this peer in this workspace
|
|
session_memberships_query = (
|
|
select(session_peers_table)
|
|
.where(session_peers_table.c.workspace_name == workspace_name)
|
|
.where(session_peers_table.c.peer_name == peer_name)
|
|
)
|
|
result = await db.execute(session_memberships_query)
|
|
memberships = result.all()
|
|
|
|
# Build a lookup of session -> time windows
|
|
session_windows: dict[str, list[tuple[Any, Any]]] = {}
|
|
for membership in memberships:
|
|
session_name = membership.session_name
|
|
if session_name not in session_windows:
|
|
session_windows[session_name] = []
|
|
session_windows[session_name].append((membership.joined_at, membership.left_at))
|
|
|
|
# Filter messages
|
|
filtered_messages: list[models.Message] = []
|
|
for msg in messages:
|
|
if msg.session_name not in session_windows:
|
|
continue
|
|
|
|
# Check if message was created during any of the peer's active windows in this session
|
|
for joined_at, left_at in session_windows[msg.session_name]:
|
|
if msg.created_at >= joined_at and (
|
|
left_at is None or msg.created_at <= left_at
|
|
):
|
|
filtered_messages.append(msg)
|
|
break # Don't add the same message twice
|
|
|
|
return filtered_messages
|
|
|
|
|
|
async def _fulltext_search(
|
|
db: AsyncSession,
|
|
query: str,
|
|
stmt: Select[tuple[models.Message]],
|
|
limit: int,
|
|
) -> list[models.Message]:
|
|
"""
|
|
Perform full-text search using PostgreSQL FTS and ILIKE fallback.
|
|
|
|
Args:
|
|
db: Database session
|
|
query: Search query
|
|
stmt: Base SQL query conditions
|
|
limit: Maximum number of results to return
|
|
|
|
Returns:
|
|
list of messages ordered by text search relevance
|
|
"""
|
|
# Check if query contains special characters that FTS might not handle well
|
|
has_special_chars = bool(
|
|
re.search(r'[~`!@#$%^&*()_+=\[\]{};\':"\\|,.<>/?-]', query)
|
|
)
|
|
|
|
# Escape ILIKE pattern characters to treat user input literally
|
|
escaped_query = escape_ilike_pattern(query)
|
|
|
|
if has_special_chars:
|
|
# For queries with special characters, use exact string matching (ILIKE)
|
|
search_condition = models.Message.content.ilike(
|
|
f"%{escaped_query}%", escape=ILIKE_ESCAPE_CHAR
|
|
)
|
|
fulltext_query = stmt.where(search_condition).order_by(
|
|
models.Message.created_at.desc()
|
|
)
|
|
else:
|
|
# For natural language queries, use full text search with ranking
|
|
fts_condition = func.to_tsvector("english", models.Message.content).op("@@")(
|
|
func.plainto_tsquery("english", query)
|
|
)
|
|
|
|
# Combine FTS with ILIKE as fallback for better coverage
|
|
combined_condition = or_(
|
|
fts_condition,
|
|
models.Message.content.ilike(
|
|
f"%{escaped_query}%", escape=ILIKE_ESCAPE_CHAR
|
|
),
|
|
)
|
|
|
|
fulltext_query = stmt.where(combined_condition).order_by(
|
|
# Order by FTS relevance first, then by creation time
|
|
func.coalesce(
|
|
func.ts_rank(
|
|
func.to_tsvector("english", models.Message.content),
|
|
func.plainto_tsquery("english", query),
|
|
),
|
|
0,
|
|
).desc(),
|
|
models.Message.created_at.desc(),
|
|
)
|
|
|
|
fulltext_query = fulltext_query.limit(limit)
|
|
|
|
result = await db.execute(fulltext_query)
|
|
return list(result.scalars().all())
|
|
|
|
|
|
async def search(
|
|
query: str,
|
|
*,
|
|
filters: dict[str, Any] | None = None,
|
|
limit: int = 10,
|
|
) -> list[models.Message]:
|
|
"""
|
|
Search across message content using a hybrid approach with Reciprocal Rank Fusion (RRF).
|
|
|
|
This function combines semantic search and full-text search results using RRF when both
|
|
are available, providing better search results than either method alone.
|
|
|
|
Args:
|
|
query: Search query to match against message content
|
|
filters: Optional filters to scope search (must include workspace_id for semantic search).
|
|
Special filter 'peer_perspective' will search across all messages from sessions that the peer is/was a member of,
|
|
filtered by the time window when they were actually in the session.
|
|
limit: Maximum number of results to return
|
|
|
|
Returns:
|
|
list of messages that match the search query, ordered by RRF relevance or individual search relevance
|
|
|
|
Raises:
|
|
ValidationException: If query exceeds maximum token limit for embeddings
|
|
"""
|
|
# Base query conditions
|
|
stmt = select(models.Message)
|
|
|
|
# Handle special peer_perspective filter
|
|
peer_perspective_name: str | None = None
|
|
if filters and "peer_perspective" in filters:
|
|
peer_perspective_name = filters["peer_perspective"]
|
|
# Remove from filters dict so apply_filter doesn't try to handle it
|
|
filters = {k: v for k, v in filters.items() if k != "peer_perspective"}
|
|
# Safety: peer_perspective must be scoped to a workspace
|
|
if not filters or (
|
|
"workspace_id" not in filters and "workspace_name" not in filters
|
|
):
|
|
raise ValidationException(
|
|
"peer_perspective requires a workspace scope (workspace_id or workspace_name)."
|
|
)
|
|
|
|
# Join with session_peers_table to get messages from sessions the peer was in
|
|
# Only include messages created during the time window the peer was active
|
|
stmt = stmt.join(
|
|
session_peers_table,
|
|
and_(
|
|
models.Message.session_name == session_peers_table.c.session_name,
|
|
models.Message.workspace_name == session_peers_table.c.workspace_name,
|
|
models.Message.created_at >= session_peers_table.c.joined_at,
|
|
or_(
|
|
session_peers_table.c.left_at.is_(None),
|
|
models.Message.created_at <= session_peers_table.c.left_at,
|
|
),
|
|
),
|
|
).where(session_peers_table.c.peer_name == peer_perspective_name)
|
|
|
|
stmt = apply_filter(stmt, models.Message, filters)
|
|
|
|
workspace_name: str | None = None
|
|
if filters:
|
|
workspace_value = filters.get("workspace_id") or filters.get("workspace_name")
|
|
if isinstance(workspace_value, str):
|
|
workspace_name = workspace_value
|
|
|
|
semantic_limit = limit * 4 if peer_perspective_name else limit * 2
|
|
query_embedding: list[float] | None = None
|
|
semantic_message_ids: list[str] | None = None
|
|
|
|
if settings.EMBED_MESSAGES and isinstance(workspace_name, str):
|
|
try:
|
|
query_embedding = await embedding_client.embed(query)
|
|
except ValueError as e:
|
|
raise ValidationException(
|
|
f"Query exceeds maximum token limit of {settings.EMBEDDING.MAX_INPUT_TOKENS}."
|
|
) from e
|
|
|
|
if not _uses_pgvector_message_search():
|
|
semantic_message_ids = await query_external_vector_message_ids(
|
|
workspace_name=workspace_name,
|
|
embedding_query=query_embedding,
|
|
limit=semantic_limit,
|
|
filters=filters,
|
|
)
|
|
|
|
async def _run_search(active_db: AsyncSession) -> list[models.Message]:
|
|
search_results: list[list[models.Message]] = []
|
|
|
|
if (
|
|
settings.EMBED_MESSAGES
|
|
and isinstance(workspace_name, str)
|
|
and query_embedding is not None
|
|
):
|
|
if _uses_pgvector_message_search():
|
|
semantic_results = await _semantic_search_pgvector(
|
|
db=active_db,
|
|
workspace_name=workspace_name,
|
|
embedding_query=query_embedding,
|
|
limit=semantic_limit,
|
|
filters=filters,
|
|
)
|
|
else:
|
|
semantic_results = await fetch_messages_by_ids(
|
|
db=active_db,
|
|
message_ids=semantic_message_ids or [],
|
|
filters=filters,
|
|
)
|
|
|
|
if peer_perspective_name:
|
|
semantic_results = await _filter_by_peer_perspective(
|
|
active_db,
|
|
semantic_results,
|
|
workspace_name,
|
|
peer_perspective_name,
|
|
)
|
|
|
|
search_results.append(semantic_results)
|
|
|
|
fulltext_results = await _fulltext_search(
|
|
db=active_db,
|
|
query=query,
|
|
stmt=stmt,
|
|
limit=limit * 2,
|
|
)
|
|
search_results.append(fulltext_results)
|
|
|
|
if len(search_results) > 1:
|
|
return reciprocal_rank_fusion(*search_results, limit=limit)
|
|
if len(search_results) == 1:
|
|
return search_results[0][:limit]
|
|
return []
|
|
|
|
async with tracked_db("search.messages") as managed_db:
|
|
combined_results = await _run_search(managed_db)
|
|
for message in combined_results:
|
|
managed_db.expunge(message)
|
|
return combined_results
|