touchups

Author: yuhongsun96

backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/rerank_documents.py

@@ -21,8 +21,9 @@
 from onyx.configs.agent_configs import AGENT_RERANKING_MAX_QUERY_RETRIEVAL_RESULTS
 from onyx.configs.agent_configs import AGENT_RERANKING_STATS
 from onyx.context.search.models import InferenceSection
-from onyx.context.search.models import SearchRequest
+from onyx.context.search.models import RerankingDetails
 from onyx.context.search.postprocessing.postprocessing import rerank_sections
+from onyx.context.search.postprocessing.postprocessing import reranking_is_runnable
 
 
 def rerank_documents(
@@ -48,27 +49,20 @@ def rerank_documents(
         graph_config.tooling.search_tool
     ), "search_tool must be provided for agentic search"
 
-    search_request = SearchRequest(
-        query=question,
-        persona=graph_config.inputs.search_request.persona,
-        rerank_settings=graph_config.inputs.search_request.rerank_settings,
-    )
+    rerank_settings = graph_config.inputs.search_request.rerank_settings
 
-    if (
-        search_request.rerank_settings
-        and search_request.rerank_settings.rerank_model_name
-        and search_request.rerank_settings.num_rerank > 0
-        and len(verified_documents) > 0
-    ):
+    if reranking_is_runnable(rerank_settings) and len(verified_documents) > 0:
         if len(verified_documents) > 1:
             reranked_documents = rerank_sections(
                 query_str=question,
-                rerank_settings=search_request.rerank_settings,
+                # if runnable, then rerank_settings is not None
+                rerank_settings=cast(RerankingDetails, rerank_settings),
                 sections_to_rerank=verified_documents,
             )
         else:
-            num = "No" if len(verified_documents) == 0 else "One"
-            logger.warning(f"{num} verified document(s) found, skipping reranking")
+            logger.warning(
+                f"{len(verified_documents)} verified document(s) found, skipping reranking"
+            )
             reranked_documents = verified_documents
     else:
         logger.warning("No reranking settings found, using unranked documents")

backend/onyx/context/search/postprocessing/postprocessing.py

@@ -79,7 +79,7 @@ def _remove_metadata_suffix(chunk: InferenceChunkUncleaned) -> str:
 @log_function_time(print_only=True)
 def semantic_reranking(
     query_str: str,
-    rerank_settings: RerankingDetails | None,
+    rerank_settings: RerankingDetails,
     chunks: list[InferenceChunk],
     model_min: int = CROSS_ENCODER_RANGE_MIN,
     model_max: int = CROSS_ENCODER_RANGE_MAX,
@@ -90,9 +90,9 @@ def semantic_reranking(
 
     Note: this updates the chunks in place, it updates the chunk scores which came from retrieval
     """
-    if not rerank_settings or not rerank_settings.rerank_model_name:
-        # Should never reach this part of the flow without reranking settings
-        raise RuntimeError("Reranking flow should not be running")
+    assert (
+        rerank_settings.rerank_model_name
+    ), "Reranking flow cannot run without a specific model"
 
     chunks_to_rerank = chunks[: rerank_settings.num_rerank]
 
@@ -165,9 +165,20 @@ def semantic_reranking(
     return list(ranked_chunks), list(ranked_indices)
 
 
+def reranking_is_runnable(rerank_settings: RerankingDetails | None) -> bool:
+    """Based on the RerankingDetails model, only run rerank if the following conditions are met:
+    - rerank_model_name is not None
+    - num_rerank is greater than 0
+    """
+    if not rerank_settings:
+        return False
+
+    return bool(rerank_settings.rerank_model_name and rerank_settings.num_rerank > 0)
+
+
 def rerank_sections(
     query_str: str,
-    rerank_settings: RerankingDetails | None,
+    rerank_settings: RerankingDetails,
     sections_to_rerank: list[InferenceSection],
     rerank_metrics_callback: Callable[[RerankMetricsContainer], None] | None = None,
 ) -> list[InferenceSection]:
@@ -182,10 +193,6 @@ def rerank_sections(
     """
     chunks_to_rerank = [section.center_chunk for section in sections_to_rerank]
 
-    if not rerank_settings:
-        # Should never reach this part of the flow without reranking settings
-        raise RuntimeError("Reranking settings not found")
-
     ranked_chunks, _ = semantic_reranking(
         query_str=query_str,
         rerank_settings=rerank_settings,
@@ -262,17 +269,13 @@ def search_postprocessing(
 
     rerank_task_id = None
     sections_yielded = False
-    if (
-        search_query.rerank_settings
-        and search_query.rerank_settings.rerank_model_name
-        and search_query.rerank_settings.num_rerank > 0
-    ):
+    if reranking_is_runnable(search_query.rerank_settings):
         post_processing_tasks.append(
             FunctionCall(
                 rerank_sections,
                 (
                     search_query.query,
-                    search_query.rerank_settings,
+                    search_query.rerank_settings,  # Cannot be None here
                     retrieved_sections,
                     rerank_metrics_callback,
                 ),