fix: change to pure rank-based RRF for relevance ordering (#4411)

* Fix RRF

* Fix turbopuffer tests

Author: mattzh72

letta/functions/function_sets/base.py

@@ -55,33 +55,31 @@ def conversation_search(
         str: Query result string containing matching messages with timestamps and content.
     """
 
-    import math
-
     from letta.constants import RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE
     from letta.helpers.json_helpers import json_dumps
 
-    if page is None or (isinstance(page, str) and page.lower().strip() == "none"):
-        page = 0
-    try:
-        page = int(page)
-    except:
-        raise ValueError("'page' argument must be an integer")
-    count = RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE
-    # TODO: add paging by page number. currently cursor only works with strings.
-    # original: start=page * count
+    # Use provided limit or default
+    if limit is None:
+        limit = RETRIEVAL_QUERY_DEFAULT_PAGE_SIZE
+
     messages = self.message_manager.list_messages_for_agent(
         agent_id=self.agent_state.id,
         actor=self.user,
         query_text=query,
-        limit=count,
+        roles=roles,
+        limit=limit,
     )
-    total = len(messages)
-    num_pages = math.ceil(total / count) - 1  # 0 index
+
     if len(messages) == 0:
         results_str = "No results found."
     else:
-        results_pref = f"Showing {len(messages)} of {total} results (page {page}/{num_pages}):"
-        results_formatted = [message.content[0].text for message in messages]
+        results_pref = f"Found {len(messages)} results:"
+        results_formatted = []
+        for message in messages:
+            # Extract text content from message
+            text_content = message.content[0].text if message.content else ""
+            result_entry = {"role": message.role, "content": text_content}
+            results_formatted.append(result_entry)
         results_str = f"{results_pref} {json_dumps(results_formatted)}"
     return results_str
 

letta/helpers/tpuf_client.py

@@ -474,8 +474,16 @@ async def query_passages(
                 # for hybrid mode, we get a multi-query response
                 vector_results = self._process_single_query_results(result.results[0], archive_id, tags)
                 fts_results = self._process_single_query_results(result.results[1], archive_id, tags, is_fts=True)
-                # use backwards-compatible wrapper which calls generic RRF
-                return self._reciprocal_rank_fusion(vector_results, fts_results, vector_weight, fts_weight, top_k)
+                # use RRF and return only (passage, score) for backwards compatibility
+                results_with_metadata = self._reciprocal_rank_fusion(
+                    vector_results=[passage for passage, _ in vector_results],
+                    fts_results=[passage for passage, _ in fts_results],
+                    get_id_func=lambda p: p.id,
+                    vector_weight=vector_weight,
+                    fts_weight=fts_weight,
+                    top_k=top_k,
+                )
+                return [(passage, rrf_score) for passage, rrf_score, metadata in results_with_metadata]
             else:
                 # for single queries (vector, fts, timestamp)
                 is_fts = search_mode == "fts"
@@ -499,7 +507,7 @@ async def query_messages(
         fts_weight: float = 0.5,
         start_date: Optional[datetime] = None,
         end_date: Optional[datetime] = None,
-    ) -> List[Tuple[dict, float]]:
+    ) -> List[Tuple[dict, float, dict]]:
         """Query messages from Turbopuffer using vector search, full-text search, or hybrid search.
 
         Args:
@@ -516,7 +524,10 @@ async def query_messages(
             end_date: Optional datetime to filter messages created before this date
 
         Returns:
-            List of (message_dict, score) tuples where message_dict contains id, text, role, created_at
+            List of (message_dict, score, metadata) tuples where:
+            - message_dict contains id, text, role, created_at
+            - score is the final relevance score
+            - metadata contains individual scores and ranking information
         """
         # Check if we should fallback to timestamp-based retrieval
         if query_embedding is None and query_text is None and search_mode not in ["timestamp"]:
@@ -576,28 +587,42 @@ async def query_messages(
             if search_mode == "hybrid":
                 # for hybrid mode, we get a multi-query response
                 vector_results = self._process_message_query_results(result.results[0])
-                fts_results = self._process_message_query_results(result.results[1], is_fts=True)
-                # use generic RRF with lambda to extract ID from dict
-                return self._generic_reciprocal_rank_fusion(
+                fts_results = self._process_message_query_results(result.results[1])
+                # use RRF with lambda to extract ID from dict - returns metadata
+                results_with_metadata = self._reciprocal_rank_fusion(
                     vector_results=vector_results,
                     fts_results=fts_results,
                     get_id_func=lambda msg_dict: msg_dict["id"],
                     vector_weight=vector_weight,
                     fts_weight=fts_weight,
                     top_k=top_k,
                 )
+                # return results with metadata
+                return results_with_metadata
             else:
                 # for single queries (vector, fts, timestamp)
-                is_fts = search_mode == "fts"
-                return self._process_message_query_results(result, is_fts=is_fts)
+                results = self._process_message_query_results(result)
+                # add simple metadata for single search modes
+                results_with_metadata = []
+                for idx, msg_dict in enumerate(results):
+                    metadata = {
+                        "combined_score": 1.0 / (idx + 1),  # Use rank-based score for single mode
+                        "search_mode": search_mode,
+                        f"{search_mode}_rank": idx + 1,  # Add the rank for this search mode
+                    }
+                    results_with_metadata.append((msg_dict, metadata["combined_score"], metadata))
+                return results_with_metadata
 
         except Exception as e:
             logger.error(f"Failed to query messages from Turbopuffer: {e}")
             raise
 
-    def _process_message_query_results(self, result, is_fts: bool = False) -> List[Tuple[dict, float]]:
-        """Process results from a message query into message dicts with scores."""
-        messages_with_scores = []
+    def _process_message_query_results(self, result) -> List[dict]:
+        """Process results from a message query into message dicts.
+
+        For RRF, we only need the rank order - scores are not used.
+        """
+        messages = []
 
         for row in result.rows:
             # Build message dict with key fields
@@ -609,19 +634,9 @@ def _process_message_query_results(self, result, is_fts: bool = False) -> List[T
                 "role": getattr(row, "role", None),
                 "created_at": getattr(row, "created_at", None),
             }
+            messages.append(message_dict)
 
-            # handle score based on search type
-            if is_fts:
-                # for FTS, use the BM25 score directly (higher is better)
-                score = getattr(row, "$score", 0.0)
-            else:
-                # for vector search, convert distance to similarity score
-                distance = getattr(row, "$dist", 0.0)
-                score = 1.0 - distance
-
-            messages_with_scores.append((message_dict, score))
-
-        return messages_with_scores
+        return messages
 
     def _process_single_query_results(
         self, result, archive_id: str, tags: Optional[List[str]], is_fts: bool = False
@@ -663,74 +678,78 @@ def _process_single_query_results(
 
         return passages_with_scores
 
-    def _generic_reciprocal_rank_fusion(
+    def _reciprocal_rank_fusion(
         self,
-        vector_results: List[Tuple[Any, float]],
-        fts_results: List[Tuple[Any, float]],
+        vector_results: List[Any],
+        fts_results: List[Any],
         get_id_func: Callable[[Any], str],
         vector_weight: float,
         fts_weight: float,
         top_k: int,
-    ) -> List[Tuple[Any, float]]:
-        """Generic RRF implementation that works with any object type.
+    ) -> List[Tuple[Any, float, dict]]:
+        """RRF implementation that works with any object type.
 
-        RRF score = vector_weight * (1/(k + vector_rank)) + fts_weight * (1/(k + fts_rank))
+        RRF score = vector_weight * (1/(k + rank)) + fts_weight * (1/(k + rank))
         where k is a constant (typically 60) to avoid division by zero
 
+        This is a pure rank-based fusion following the standard RRF algorithm.
+
         Args:
-            vector_results: List of (item, score) tuples from vector search
-            fts_results: List of (item, score) tuples from FTS
+            vector_results: List of items from vector search (ordered by relevance)
+            fts_results: List of items from FTS (ordered by relevance)
             get_id_func: Function to extract ID from an item
             vector_weight: Weight for vector search results
             fts_weight: Weight for FTS results
             top_k: Number of results to return
 
         Returns:
-            List of (item, score) tuples sorted by RRF score
+            List of (item, score, metadata) tuples sorted by RRF score
+            metadata contains ranks from each result list
         """
-        k = 60  # standard RRF constant
+        k = 60  # standard RRF constant from Cormack et al. (2009)
 
-        # create rank mappings using the get_id_func
-        vector_ranks = {get_id_func(item): rank + 1 for rank, (item, _) in enumerate(vector_results)}
-        fts_ranks = {get_id_func(item): rank + 1 for rank, (item, _) in enumerate(fts_results)}
+        # create rank mappings based on position in result lists
+        # rank starts at 1, not 0
+        vector_ranks = {get_id_func(item): rank + 1 for rank, item in enumerate(vector_results)}
+        fts_ranks = {get_id_func(item): rank + 1 for rank, item in enumerate(fts_results)}
 
-        # combine all unique items
+        # combine all unique items from both result sets
         all_items = {}
-        for item, _ in vector_results:
+        for item in vector_results:
             all_items[get_id_func(item)] = item
-        for item, _ in fts_results:
+        for item in fts_results:
             all_items[get_id_func(item)] = item
 
-        # calculate RRF scores
+        # calculate RRF scores based purely on ranks
         rrf_scores = {}
+        score_metadata = {}
         for item_id in all_items:
-            vector_score = vector_weight / (k + vector_ranks.get(item_id, k + top_k))
-            fts_score = fts_weight / (k + fts_ranks.get(item_id, k + top_k))
-            rrf_scores[item_id] = vector_score + fts_score
+            # RRF formula: sum of 1/(k + rank) across result lists
+            # If item not in a list, we don't add anything (equivalent to rank = infinity)
+            vector_rrf_score = 0.0
+            fts_rrf_score = 0.0
+
+            if item_id in vector_ranks:
+                vector_rrf_score = vector_weight / (k + vector_ranks[item_id])
+            if item_id in fts_ranks:
+                fts_rrf_score = fts_weight / (k + fts_ranks[item_id])
+
+            combined_score = vector_rrf_score + fts_rrf_score
+
+            rrf_scores[item_id] = combined_score
+            score_metadata[item_id] = {
+                "combined_score": combined_score,  # Final RRF score
+                "vector_rank": vector_ranks.get(item_id),
+                "fts_rank": fts_ranks.get(item_id),
+            }
 
-        # sort by RRF score and return top_k
-        sorted_results = sorted([(all_items[iid], score) for iid, score in rrf_scores.items()], key=lambda x: x[1], reverse=True)
+        # sort by RRF score and return with metadata
+        sorted_results = sorted(
+            [(all_items[iid], score, score_metadata[iid]) for iid, score in rrf_scores.items()], key=lambda x: x[1], reverse=True
+        )
 
         return sorted_results[:top_k]
 
-    def _reciprocal_rank_fusion(
-        self,
-        vector_results: List[Tuple[PydanticPassage, float]],
-        fts_results: List[Tuple[PydanticPassage, float]],
-        vector_weight: float,
-        fts_weight: float,
-        top_k: int,
-    ) -> List[Tuple[PydanticPassage, float]]:
-        """Wrapper for backwards compatibility - uses generic RRF for passages."""
-        return self._generic_reciprocal_rank_fusion(
-            vector_results=vector_results,
-            fts_results=fts_results,
-            get_id_func=lambda p: p.id,
-            vector_weight=vector_weight,
-            fts_weight=fts_weight,
-            top_k=top_k,
-        )
-
     @trace_method
     async def delete_passage(self, archive_id: str, passage_id: str) -> bool:
         """Delete a passage from Turbopuffer."""

letta/services/message_manager.py

@@ -1,7 +1,7 @@
 import json
 import uuid
 from datetime import datetime
-from typing import List, Optional, Sequence
+from typing import List, Optional, Sequence, Tuple
 
 from sqlalchemy import delete, exists, func, select, text
 
@@ -1065,7 +1065,7 @@ async def search_messages_async(
         start_date: Optional[datetime] = None,
         end_date: Optional[datetime] = None,
         embedding_config: Optional[EmbeddingConfig] = None,
-    ) -> List[PydanticMessage]:
+    ) -> List[Tuple[PydanticMessage, dict]]:
         """
         Search messages using Turbopuffer if enabled, otherwise fall back to SQL search.
 
@@ -1082,7 +1082,7 @@ async def search_messages_async(
             embedding_config: Optional embedding configuration for generating query embedding
 
         Returns:
-            List of matching messages
+            List of tuples (message, metadata) where metadata contains relevance scores
         """
         from letta.helpers.tpuf_client import TurbopufferClient, should_use_tpuf_for_messages
 
@@ -1133,8 +1133,8 @@ async def search_messages_async(
                     from letta.schemas.letta_message_content import TextContent
                     from letta.schemas.message import Message as PydanticMessage
 
-                    turbopuffer_messages = []
-                    for msg_dict, score in results:
+                    message_tuples = []
+                    for msg_dict, score, metadata in results:
                         # create a message object with the properly extracted text from turbopuffer
                         message = PydanticMessage(
                             id=msg_dict["id"],
@@ -1146,9 +1146,10 @@ async def search_messages_async(
                             created_by_id=actor.id,
                             last_updated_by_id=actor.id,
                         )
-                        turbopuffer_messages.append(message)
+                        # Return tuple of (message, metadata)
+                        message_tuples.append((message, metadata))
 
-                    return turbopuffer_messages
+                    return message_tuples
                 else:
                     return []
 
@@ -1163,7 +1164,16 @@ async def search_messages_async(
                     limit=limit,
                     ascending=False,
                 )
-                return self._combine_assistant_tool_messages(messages)
+                combined_messages = self._combine_assistant_tool_messages(messages)
+                # Add basic metadata for SQL fallback
+                message_tuples = []
+                for message in combined_messages:
+                    metadata = {
+                        "search_mode": "sql_fallback",
+                        "combined_score": None,  # SQL doesn't provide scores
+                    }
+                    message_tuples.append((message, metadata))
+                return message_tuples
         else:
             # use sql-based search
             messages = await self.list_messages_for_agent_async(
@@ -1174,4 +1184,13 @@ async def search_messages_async(
                 limit=limit,
                 ascending=False,
             )
-            return self._combine_assistant_tool_messages(messages)
+            combined_messages = self._combine_assistant_tool_messages(messages)
+            # Add basic metadata for SQL search
+            message_tuples = []
+            for message in combined_messages:
+                metadata = {
+                    "search_mode": "sql",
+                    "combined_score": None,  # SQL doesn't provide scores
+                }
+                message_tuples.append((message, metadata))
+            return message_tuples