updated tests

Author: staru09

sdks/python/sdk/src/moss/client/models.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
 

sdks/python/sdk/src/moss/client/moss_client.py

@@ -22,7 +22,7 @@
 )
 
 from ..rerankers import get_reranker
-from .models import QueryOptions
+from .models import QueryOptions, RerankOptions
 
 logger = logging.getLogger(__name__)
 
@@ -191,24 +191,35 @@ async def query(
         Otherwise, falls back to the cloud query API.
 
         Args:
-            options: Query options (top_k, alpha, embedding, filter). Example filter:
-                QueryOptions(filter={"$and": [
-                    {"field": "city", "condition": {"$eq": "NYC"}},
-                    {"field": "price", "condition": {"$lt": "50"}},
-                ]})
+            options: Query options (top_k, alpha, embedding, filter, rerank).
+                Reranking is applied client-side after retrieval and works on
+                both the local and cloud paths. Example filter:
+                    QueryOptions(filter={"$and": [
+                        {"field": "city", "condition": {"$eq": "NYC"}},
+                        {"field": "price", "condition": {"$lt": "50"}},
+                    ]})
         """
         is_loaded = await asyncio.to_thread(self._manager.has_index, name)
 
         if is_loaded:
-            return await self._query_local(name, query, options)
+            result = await self._query_local(name, query, options)
+        else:
+            if getattr(options, "filter", None) is not None:
+                logger.warning(
+                    "Metadata filter ignored: filtering is only supported for locally loaded indexes. "
+                    "Call load_index('%s') first.",
+                    name,
+                )
+            result = await self._query_cloud(name, query, options)
 
-        if getattr(options, "filter", None) is not None:
-            logger.warning(
-                "Metadata filter ignored: filtering is only supported for locally loaded indexes. "
-                "Call load_index('%s') first.",
-                name,
-            )
-        return await self._query_cloud(name, query, options)
+        rerank = getattr(options, "rerank", None)
+        if rerank:
+            top_k = getattr(options, "top_k", None)
+            if top_k is None:
+                top_k = 5
+            result = await self._apply_rerank(query, result, rerank, top_k)
+
+        return result
 
     # -- Internal ---------------------------------------------------
 
@@ -218,7 +229,9 @@ async def _query_local(
         query: str,
         options: Optional[QueryOptions],
     ) -> SearchResult:
-        top_k = getattr(options, "top_k", None) or 5
+        top_k = getattr(options, "top_k", None)
+        if top_k is None:
+            top_k = 5
         alpha = getattr(options, "alpha", None)
         if alpha is None:
             alpha = 0.8
@@ -229,7 +242,7 @@ async def _query_local(
         fetch_k = top_k * 4 if rerank else top_k
 
         if query_embedding is not None:
-            result = await asyncio.to_thread(
+            return await asyncio.to_thread(
                 self._manager.query,
                 name,
                 query,
@@ -238,41 +251,46 @@ async def _query_local(
                 alpha,
                 filter,
             )
-        else:
-            try:
-                result = await asyncio.to_thread(
-                    self._manager.query_text,
-                    name,
-                    query,
-                    fetch_k,
-                    alpha,
-                    filter,
-                )
-            except RuntimeError as e:
-                if "requires explicit query embeddings" in str(e):
-                    raise ValueError(
-                        "This index uses custom embeddings. "
-                        "Query embeddings must be provided via QueryOptions.embedding."
-                    ) from e
-                raise
 
-        if rerank:
-            if rerank._instance is None:
-                rerank._instance = get_reranker(
-                    rerank.provider, **rerank.init_kwargs
-                )
-            final_n = rerank.top_n or top_k
-            reranked_docs = await rerank._instance.rerank(
-                query, result.docs, top_k=final_n
-            )
-            result = SearchResult(
-                docs=reranked_docs,
-                query=result.query,
-                index_name=result.index_name,
-                time_taken_ms=result.time_taken_ms,
+        try:
+            return await asyncio.to_thread(
+                self._manager.query_text,
+                name,
+                query,
+                fetch_k,
+                alpha,
+                filter,
             )
+        except RuntimeError as e:
+            if "requires explicit query embeddings" in str(e):
+                raise ValueError(
+                    "This index uses custom embeddings. "
+                    "Query embeddings must be provided via QueryOptions.embedding."
+                ) from e
+            raise
 
-        return result
+    @staticmethod
+    async def _apply_rerank(
+        query: str,
+        result: SearchResult,
+        rerank_opts: RerankOptions,
+        default_top_k: Optional[int],
+    ) -> SearchResult:
+        """Rerank search results. Works on both local and cloud paths."""
+        if rerank_opts._instance is None:
+            rerank_opts._instance = get_reranker(
+                rerank_opts.provider, **rerank_opts.init_kwargs
+            )
+        final_n = rerank_opts.top_n or default_top_k
+        reranked_docs = await rerank_opts._instance.rerank(
+            query, result.docs, top_k=final_n
+        )
+        return SearchResult(
+            docs=reranked_docs,
+            query=result.query,
+            index_name=result.index_name,
+            time_taken_ms=result.time_taken_ms,
+        )
 
     async def _query_cloud(
         self,
@@ -281,15 +299,19 @@ async def _query_cloud(
         options: Optional[QueryOptions],
     ) -> SearchResult:
         """Fallback: query via the cloud API when the index is not loaded locally."""
-        top_k = getattr(options, "top_k", None) or 10
+        top_k = getattr(options, "top_k", None)
+        if top_k is None:
+            top_k = 5
+        rerank = getattr(options, "rerank", None)
+        fetch_k = top_k * 4 if rerank else top_k
         query_embedding = getattr(options, "embedding", None)
 
         request_body: Dict[str, Any] = {
             "query": query,
             "indexName": name,
             "projectId": self._project_id,
             "projectKey": self._project_key,
-            "topK": top_k,
+            "topK": fetch_k,
         }
         if query_embedding is not None:
             request_body["queryEmbedding"] = list(query_embedding)

sdks/python/sdk/src/moss/rerankers/__init__.py

@@ -6,6 +6,8 @@
 
 _REGISTRY: Dict[str, Type[Reranker]] = {}
 
+_MISSING_PROVIDERS: Dict[str, str] = {}
+
 
 def register_reranker(name: str, cls: Type[Reranker]) -> None:
     """Register a reranker class under a provider name.
@@ -24,8 +26,17 @@ def get_reranker(name: str, **kwargs: Any) -> Reranker:
     """Instantiate a reranker by provider name.
 
     Raises:
-        ValueError: If the provider name is not registered.
+        ImportError: If the provider is built-in but its optional dependency
+            isn't installed (e.g. "cohere" without `pip install cohere`).
+        ValueError: If the provider name is not registered and not a known
+            built-in.
     """
+    if name in _MISSING_PROVIDERS:
+        package = _MISSING_PROVIDERS[name]
+        raise ImportError(
+            f"The '{name}' reranker requires the '{package}' package. "
+            f"Install it with: pip install {package}"
+        )
     if name not in _REGISTRY:
         available = list(_REGISTRY) or ["(none registered)"]
         raise ValueError(
@@ -43,4 +54,5 @@ def get_reranker(name: str, **kwargs: Any) -> Reranker:
 
     register_reranker("cohere", CohereReranker)
 except ImportError:
-    pass
+    # `pip install cohere` to enable the Cohere reranker.
+    _MISSING_PROVIDERS["cohere"] = "cohere"

sdks/python/sdk/tests/test_client.py

@@ -354,6 +354,7 @@ async def test_query_with_custom_embedding(self, client):
         opts.alpha = 0.9
         opts.embedding = [0.1, 0.2, 0.3]
         opts.filter = None
+        opts.rerank = None
 
         result = await client.query("idx", "search text", opts)
 
@@ -377,6 +378,7 @@ async def test_query_defaults_top_k_and_alpha(self, client):
         opts.alpha = None
         opts.embedding = [0.5]
         opts.filter = None
+        opts.rerank = None
 
         await client.query("idx", "q", opts)
 
@@ -394,6 +396,8 @@ async def test_query_raises_for_custom_model_without_embedding(self, client):
         opts.embedding = None
         opts.top_k = 5
         opts.alpha = 0.8
+        opts.filter = None
+        opts.rerank = None
 
         with pytest.raises(ValueError, match="custom embeddings"):
             await client.query("idx", "q", opts)
@@ -418,6 +422,7 @@ async def test_query_passes_filter_to_manager(self, client):
         opts.top_k = 5
         opts.alpha = 0.8
         opts.embedding = [0.1]
+        opts.rerank = None
 
         metadata_filter = {"field": "city", "condition": {"$eq": "NYC"}}
         opts.filter = metadata_filter
@@ -443,6 +448,7 @@ async def test_query_passes_none_filter_when_omitted(self, client):
         opts.alpha = 0.8
         opts.embedding = [0.1]
         opts.filter = None
+        opts.rerank = None
 
         await client.query("idx", "q", opts)
 
@@ -464,6 +470,7 @@ async def test_query_passes_complex_and_filter(self, client):
         opts.top_k = 10
         opts.alpha = 0.8
         opts.embedding = [0.5]
+        opts.rerank = None
 
         metadata_filter = {
             "$and": [
@@ -561,6 +568,8 @@ async def test_uses_local_when_index_loaded(self, client):
         opts.top_k = 5
         opts.alpha = 0.8
         opts.embedding = [0.1]
+        opts.filter = None
+        opts.rerank = None
 
         result = await client.query("idx", "q", opts)
 

sdks/python/sdk/tests/test_client_extended.py

@@ -100,6 +100,7 @@ async def test_cloud_fallback_with_custom_embedding(self, unloaded_client):
             opts.top_k = 5
             opts.embedding = [0.1, 0.2, 0.3]
             opts.filter = None
+            opts.rerank = None
 
             result = await unloaded_client.query("idx", "test query", opts)
 
@@ -263,6 +264,7 @@ async def test_query_with_only_top_k(self, client):
         opts.alpha = None
         opts.embedding = None
         opts.filter = None
+        opts.rerank = None
 
         await client.query("idx", "test", opts)
 
@@ -278,6 +280,7 @@ async def test_query_with_only_alpha(self, client):
         opts.alpha = 0.5
         opts.embedding = None
         opts.filter = None
+        opts.rerank = None
 
         await client.query("idx", "test", opts)
 
@@ -293,6 +296,7 @@ async def test_query_alpha_zero_keyword_only(self, client):
         opts.alpha = 0
         opts.embedding = None
         opts.filter = None
+        opts.rerank = None
 
         await client.query("idx", "test", opts)
 
@@ -308,6 +312,7 @@ async def test_query_alpha_one_semantic_only(self, client):
         opts.alpha = 1
         opts.embedding = None
         opts.filter = None
+        opts.rerank = None
 
         await client.query("idx", "test", opts)
 
@@ -333,6 +338,7 @@ async def test_filter_warning_logged_when_unloaded(self, unloaded_client, caplog
 
             opts = MagicMock()
             opts.filter = {"field": "city", "condition": {"$eq": "NYC"}}
+            opts.rerank = None
 
             with caplog.at_level("WARNING"):
                 await unloaded_client.query("idx", "test", opts)

sdks/python/sdk/tests/test_rerankers.py

@@ -73,9 +73,7 @@ def test_multiple_kwargs_forwarded(self):
         opts = RerankOptions(
             provider="cohere", api_key="k", model="rerank-v3.5", top_n=3
         )
-        self.assertEqual(
-            opts.init_kwargs, {"api_key": "k", "model": "rerank-v3.5"}
-        )
+        self.assertEqual(opts.init_kwargs, {"api_key": "k", "model": "rerank-v3.5"})
 
 
 class TestCohereRerankerProtocol(unittest.TestCase):
@@ -132,7 +130,9 @@ async def test_rerank_calls_cohere_sdk(self):
 
         mock_result_1 = type("Result", (), {"index": 1, "relevance_score": 0.95})()
         mock_result_2 = type("Result", (), {"index": 0, "relevance_score": 0.72})()
-        mock_response = type("Response", (), {"results": [mock_result_1, mock_result_2]})()
+        mock_response = type(
+            "Response", (), {"results": [mock_result_1, mock_result_2]}
+        )()
         reranker._client.rerank = AsyncMock(return_value=mock_response)
 
         docs = [
@@ -177,9 +177,7 @@ async def test_rerank_with_top_k(self):
     async def test_rerank_sdk_error(self):
         reranker = CohereReranker(api_key="bad-key")
         reranker._client = AsyncMock()
-        reranker._client.rerank = AsyncMock(
-            side_effect=Exception("Unauthorized")
-        )
+        reranker._client.rerank = AsyncMock(side_effect=Exception("Unauthorized"))
 
         docs = [QueryResultDocumentInfo(id="d1", text="doc", score=0.5)]
 

sdks/python/sdk/tests/test_types.py

@@ -20,6 +20,7 @@ async def test_query_with_filter_option(self, client):
         opts.alpha = 0.8
         opts.embedding = None
         opts.filter = {"field": "city", "condition": {"$eq": "NYC"}}
+        opts.rerank = None
 
         await client.query("idx", "test", opts)
 
@@ -38,6 +39,7 @@ async def test_query_with_custom_embedding_option(self, client):
         opts.alpha = 0.5
         opts.embedding = embedding
         opts.filter = None
+        opts.rerank = None
 
         await client.query("idx", "test", opts)