feat: use default rerank (#1200)

Author: earayu

aperag/flow/runners/rerank.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import logging
-from typing import List, Tuple
+from typing import List, Optional, Tuple
 
 from pydantic import BaseModel, Field
 
@@ -31,9 +31,12 @@
 
 
 class RerankInput(BaseModel):
-    model: str = Field(..., description="Rerank model name")
-    model_service_provider: str = Field(..., description="Model service provider")
-    custom_llm_provider: str = Field(..., description="Custom LLM provider (e.g., 'jina_ai', 'openai')")
+    use_rerank_service: bool = Field(default=True, description="Whether to use rerank service or fallback strategy")
+    model: Optional[str] = Field(default=None, description="Rerank model name")
+    model_service_provider: Optional[str] = Field(default=None, description="Model service provider")
+    custom_llm_provider: Optional[str] = Field(
+        default=None, description="Custom LLM provider (e.g., 'jina_ai', 'openai')"
+    )
     docs: List[DocumentWithScore]
 
 
@@ -49,88 +52,145 @@ class RerankOutput(BaseModel):
 class RerankNodeRunner(BaseNodeRunner):
     async def run(self, ui: RerankInput, si: SystemInput) -> Tuple[RerankOutput, dict]:
         """
-        Run rerank node. ui: user input; si: system input (SystemInput).
-        Returns (output, system_output)
+        Smart rerank node:
+        - use_rerank_service=False: directly use fallback strategy
+        - use_rerank_service=True: try rerank service, fallback on failure
         """
-        query = si.query
         docs = ui.docs
-        result = []
 
         if not docs:
             logger.info("No documents to rerank, returning empty result")
+            return RerankOutput(docs=[]), {}
+
+        # Strategy 1: If not using rerank service, directly use fallback strategy
+        if not ui.use_rerank_service:
+            logger.info("Rerank service disabled, using fallback strategy")
+            result = self._apply_fallback_strategy(docs)
             return RerankOutput(docs=result), {}
 
+        # Strategy 2: Try to use rerank service
         try:
-            # Validate input configuration
-            if not ui.model_service_provider:
-                raise InvalidConfigurationError(
-                    "model_service_provider", ui.model_service_provider, "Model service provider cannot be empty"
-                )
-
-            if not ui.model:
-                raise InvalidConfigurationError("model", ui.model, "Model name cannot be empty")
-
-            if not ui.custom_llm_provider:
-                raise InvalidConfigurationError(
-                    "custom_llm_provider", ui.custom_llm_provider, "Custom LLM provider cannot be empty"
-                )
-
-            # Get API key and base URL from user's model service provider settings
-            api_key = await async_db_ops.query_provider_api_key(ui.model_service_provider, si.user)
-            if not api_key:
-                raise InvalidConfigurationError(
-                    "api_key", api_key, f"API KEY not found for LLM Provider:{ui.model_service_provider}"
-                )
-
-            # Get base_url from LLMProvider
-            try:
-                llm_provider = await async_db_ops.query_llm_provider_by_name(ui.model_service_provider)
-                if not llm_provider:
-                    raise ProviderNotFoundError(ui.model_service_provider, "Rerank")
-                base_url = llm_provider.base_url
-            except Exception as e:
-                logger.error(f"Failed to query LLM provider '{ui.model_service_provider}': {str(e)}")
-                raise ProviderNotFoundError(ui.model_service_provider, "Rerank") from e
-
-            if not base_url:
-                raise InvalidConfigurationError(
-                    "base_url", base_url, f"Base URL not configured for provider '{ui.model_service_provider}'"
-                )
-
-            # Create rerank service with configuration from model service provider
-            rerank_service = RerankService(
-                rerank_provider=ui.custom_llm_provider,
-                rerank_model=ui.model,
-                rerank_service_url=base_url,
-                rerank_service_api_key=api_key,
+            # Check configuration completeness
+            if not self._is_rerank_config_valid(ui):
+                logger.info("Rerank service configuration incomplete, using fallback strategy")
+                result = self._apply_fallback_strategy(docs)
+                return RerankOutput(docs=result), {}
+
+            # Execute actual rerank
+            result = await self._perform_actual_rerank(ui, si)
+            logger.info(f"Successfully reranked {len(result)} documents using rerank service")
+            return RerankOutput(docs=result), {}
+
+        except (InvalidConfigurationError, ProviderNotFoundError) as e:
+            logger.warning(f"Rerank service configuration error, using fallback strategy: {str(e)}")
+            result = self._apply_fallback_strategy(docs)
+            return RerankOutput(docs=result), {}
+
+        except RerankError as e:
+            logger.warning(f"Rerank service operation failed, using fallback strategy: {str(e)}")
+            result = self._apply_fallback_strategy(docs)
+            return RerankOutput(docs=result), {}
+
+        except Exception as e:
+            logger.error(f"Unexpected error during rerank service, using fallback strategy: {str(e)}")
+            result = self._apply_fallback_strategy(docs)
+            return RerankOutput(docs=result), {}
+
+    def _is_rerank_config_valid(self, ui: RerankInput) -> bool:
+        """Check if rerank service configuration is valid"""
+        return (
+            ui.model
+            and ui.model.strip()
+            and ui.model_service_provider
+            and ui.model_service_provider.strip()
+            and ui.custom_llm_provider
+            and ui.custom_llm_provider.strip()
+        )
+
+    async def _perform_actual_rerank(self, ui: RerankInput, si: SystemInput) -> List[DocumentWithScore]:
+        """Execute actual rerank operation"""
+        query = si.query
+        docs = ui.docs
+
+        # Validate configuration
+        if not ui.model_service_provider:
+            raise InvalidConfigurationError(
+                "model_service_provider", ui.model_service_provider, "Model service provider cannot be empty"
             )
 
-            # Validate the service configuration
-            rerank_service.validate_configuration()
+        if not ui.model:
+            raise InvalidConfigurationError("model", ui.model, "Model name cannot be empty")
 
-            logger.info(
-                f"Using rerank service with provider: {ui.model_service_provider}, "
-                f"model: {ui.model}, url: {base_url}, max_docs: {rerank_service.max_documents}"
+        if not ui.custom_llm_provider:
+            raise InvalidConfigurationError(
+                "custom_llm_provider", ui.custom_llm_provider, "Custom LLM provider cannot be empty"
             )
 
-            # Perform reranking
-            result = await rerank_service.async_rerank(query, docs)
-            logger.info(f"Successfully reranked {len(result)} documents")
+        # Get API key and base_url
+        api_key = await async_db_ops.query_provider_api_key(ui.model_service_provider, si.user)
+        if not api_key:
+            raise InvalidConfigurationError(
+                "api_key", api_key, f"API KEY not found for LLM Provider:{ui.model_service_provider}"
+            )
 
-        except (InvalidConfigurationError, ProviderNotFoundError) as e:
-            # Configuration errors - log and return empty result to gracefully degrade
-            logger.error(f"Rerank configuration error: {str(e)}")
-            # For flow execution, we gracefully degrade instead of failing the entire flow
-            result = docs  # Return original documents without reranking
-        except RerankError as e:
-            # Rerank-specific errors - log and return original documents
-            logger.error(f"Rerank operation failed: {str(e)}")
-            # For flow execution, we gracefully degrade instead of failing the entire flow
-            result = docs  # Return original documents without reranking
+        try:
+            llm_provider = await async_db_ops.query_llm_provider_by_name(ui.model_service_provider)
+            if not llm_provider:
+                raise ProviderNotFoundError(ui.model_service_provider, "Rerank")
+            base_url = llm_provider.base_url
         except Exception as e:
-            # Unexpected errors - log and return original documents
-            logger.error(f"Unexpected error during rerank: {str(e)}")
-            # For flow execution, we gracefully degrade instead of failing the entire flow
-            result = docs  # Return original documents without reranking
+            logger.error(f"Failed to query LLM provider '{ui.model_service_provider}': {str(e)}")
+            raise ProviderNotFoundError(ui.model_service_provider, "Rerank") from e
+
+        if not base_url:
+            raise InvalidConfigurationError(
+                "base_url", base_url, f"Base URL not configured for provider '{ui.model_service_provider}'"
+            )
+
+        # Create and execute rerank service
+        rerank_service = RerankService(
+            rerank_provider=ui.custom_llm_provider,
+            rerank_model=ui.model,
+            rerank_service_url=base_url,
+            rerank_service_api_key=api_key,
+        )
+
+        rerank_service.validate_configuration()
+
+        logger.info(
+            f"Using rerank service with provider: {ui.model_service_provider}, "
+            f"model: {ui.model}, url: {base_url}, max_docs: {rerank_service.max_documents}"
+        )
+
+        return await rerank_service.async_rerank(query, docs)
+
+    def _apply_fallback_strategy(self, docs: List[DocumentWithScore]) -> List[DocumentWithScore]:
+        """
+        Apply fallback rerank strategy:
+        1. Graph search results first (better quality, typically 1 result)
+        2. Sort remaining vector and fulltext results by score in descending order
+        """
+        if not docs:
+            return docs
+
+        graph_results = []
+        other_results = []
+
+        for doc in docs:
+            recall_type = doc.metadata.get("recall_type", "")
+            if recall_type == "graph_search":
+                graph_results.append(doc)
+            else:
+                other_results.append(doc)
+
+        # Sort other results by score in descending order
+        other_results.sort(key=lambda x: x.score if x.score is not None else 0.0, reverse=True)
+
+        result = graph_results + other_results
+
+        logger.info(
+            f"Applied fallback rerank strategy: {len(graph_results)} graph results, "
+            f"{len(other_results)} other results sorted by score"
+        )
 
-        return RerankOutput(docs=result), {}
+        return result

aperag/mcp/server.py

@@ -72,6 +72,7 @@ async def search_collection(
     use_vector_index: bool = True,
     use_fulltext_index: bool = False,
     use_graph_index: bool = True,
+    rerank: bool = True,
     topk: int = 5,
     query_keywords: list[str] = None,
 ) -> Dict[str, Any]:
@@ -84,7 +85,8 @@ async def search_collection(
         use_vector_index: Whether to use vector/semantic search (default: True)
         use_fulltext_index: Whether to use full-text keyword search (default: False)
         use_graph_index: Whether to use knowledge graph search (default: True)
-        topk: Maximum number of results to return per search type (default: 10)
+        rerank: Whether to enable reranking of search results for better relevance (default: True)
+        topk: Maximum number of results to return per search type (default: 5)
 
     Returns:
         Search results with relevant documents and metadata (SearchResult format)
@@ -146,7 +148,7 @@ class SearchResult(BaseModel):
         api_key = get_api_key()
 
         # Build search request based on enabled search types
-        search_data = {"query": query}
+        search_data = {"query": query, "rerank": rerank}
 
         # Add search configurations for enabled types
         if use_vector_index:
@@ -346,17 +348,18 @@ async def aperag_usage_guide() -> str:
 1. First, get available collections with essential information: `list_collections()`
 2. Choose a collection from the list
 3. Search the collection: `search_collection(collection_id="abc123", query="your question")`
-   (By default, vector and graph search are enabled for optimal performance)
+   (By default, vector search, graph search, and reranking are enabled for optimal performance)
 
 ## Search Types:
 You can enable/disable any combination of search methods:
 - **Vector search** (use_vector_index): Semantic similarity search using embeddings (default: True)
 - **Full-text search** (use_fulltext_index): Traditional keyword-based text search (default: False)
 - **Graph search** (use_graph_index): Knowledge graph-based search (default: True)
+- **Reranking** (rerank): AI-powered reranking for improved result relevance (default: True)
 
 ⚠️ **Important**: Full-text search can return large amounts of text content which may cause context window overflow with smaller LLM models. Use with caution and consider reducing topk when enabling fulltext search.
 
-By default, vector and graph search are enabled for optimal balance of quality and context size.
+By default, vector search, graph search, and reranking are enabled for optimal balance of quality and context size.
 
 ## Example Workflow:
 ```
@@ -367,13 +370,14 @@ async def aperag_usage_guide() -> str:
 # (collections.items contains collection ID, title, and description)
 collection_id = collections.items[0].id
 
-# Step 3: Search with default methods (vector + graph)
+# Step 3: Search with default methods (vector + graph + rerank)
 results = search_collection(
     collection_id=collection_id,
     query="How to deploy applications?",
     use_vector_index=True,
     use_fulltext_index=False,
     use_graph_index=True,
+    rerank=True,
     topk=5
 )
 
@@ -384,6 +388,7 @@ async def aperag_usage_guide() -> str:
     use_vector_index=True,
     use_fulltext_index=False,
     use_graph_index=False,
+    rerank=True,  # Rerank still enabled for better results
     topk=10
 )
 
@@ -394,6 +399,7 @@ async def aperag_usage_guide() -> str:
     use_vector_index=True,
     use_fulltext_index=True,  # Enable with caution
     use_graph_index=True,
+    rerank=True,  # Rerank for optimal result ordering
     topk=3  # Use smaller topk to manage context size
 )
 ```
@@ -485,6 +491,7 @@ async def aperag_usage_guide() -> str:
     internal_results = search_collection(
         collection_id=collections.items[0].id,
         query="AI developments",
+        rerank=True,  # Default rerank for better results
         topk=5
     )