Merge pull request #592 from aurelio-labs/james/async-limit

fix: add limit to acall and fix async hybrid router

Author: jamescalam

semantic_router/routers/base.py

@@ -639,6 +639,8 @@ def _pass_routes(
                         route.llm = self.llm
                 # call dynamic route to generate the function_call content
                 route_choice = route(query=text)
+                if route_choice is not None and route_choice.similarity_score is None:
+                    route_choice.similarity_score = total_score
                 passed_routes.append(route_choice)
             elif passed and route is not None and simulate_static:
                 passed_routes.append(
@@ -757,6 +759,7 @@ async def acall(
         self,
         text: Optional[str] = None,
         vector: Optional[List[float] | np.ndarray] = None,
+        limit: int | None = 1,
         simulate_static: bool = False,
         route_filter: Optional[List[str]] = None,
     ) -> RouteChoice | list[RouteChoice]:
@@ -796,7 +799,7 @@ async def acall(
             scored_routes=scored_routes,
             simulate_static=simulate_static,
             text=text,
-            limit=1,
+            limit=limit,
         )
 
     def _index_ready(self) -> bool:

semantic_router/routers/hybrid.py

@@ -319,7 +319,7 @@ def __call__(
         if not self.index.is_ready():
             raise ValueError("Index is not ready.")
         if self.sparse_encoder is None:
-            raise ValueError
+            raise ValueError("Sparse encoder is not set.")
         potential_sparse_vector: List[SparseEmbedding] | None = None
         # if no vector provided, encode text to get vector
         if vector is None:
@@ -360,6 +360,70 @@ def __call__(
         )
         return route_choices
 
+    async def acall(
+        self,
+        text: Optional[str] = None,
+        vector: Optional[List[float] | np.ndarray] = None,
+        limit: int | None = 1,
+        simulate_static: bool = False,
+        route_filter: Optional[List[str]] = None,
+        sparse_vector: dict[int, float] | SparseEmbedding | None = None,
+    ) -> RouteChoice | list[RouteChoice]:
+        """Asynchronously call the router to get a route choice.
+
+        :param text: The text to route.
+        :type text: Optional[str]
+        :param vector: The vector to route.
+        :type vector: Optional[List[float] | np.ndarray]
+        :param simulate_static: Whether to simulate a static route (ie avoid dynamic route
+            LLM calls during fit or evaluate).
+        :type simulate_static: bool
+        :param route_filter: The route filter to use.
+        :type route_filter: Optional[List[str]]
+        :param sparse_vector: The sparse vector to use.
+        :type sparse_vector: dict[int, float] | SparseEmbedding | None
+        :return: The route choice.
+        :rtype: RouteChoice
+        """
+        if not self.index.is_ready():
+            # TODO: need async version for qdrant
+            raise ValueError("Index is not ready.")
+        if self.sparse_encoder is None:
+            raise ValueError("Sparse encoder is not set.")
+        potential_sparse_vector: List[SparseEmbedding] | None = None
+        # if no vector provided, encode text to get vector
+        if vector is None:
+            if text is None:
+                raise ValueError("Either text or vector must be provided")
+            vector, potential_sparse_vector = await self._async_encode(
+                text=[text], input_type="queries"
+            )
+        # convert to numpy array if not already
+        vector = xq_reshape(xq=vector)
+        if sparse_vector is None:
+            if text is None:
+                raise ValueError("Either text or sparse_vector must be provided")
+            sparse_vector = (
+                potential_sparse_vector[0] if potential_sparse_vector else None
+            )
+        # get scores and routes
+        scores, routes = await self.index.aquery(
+            vector=vector[0],
+            top_k=self.top_k,
+            route_filter=route_filter,
+            sparse_vector=sparse_vector,
+        )
+        query_results = [
+            {"route": d, "score": s.item()} for d, s in zip(routes, scores)
+        ]
+        scored_routes = self._score_routes(query_results=query_results)
+        return await self._async_pass_routes(
+            scored_routes=scored_routes,
+            simulate_static=simulate_static,
+            text=text,
+            limit=limit,
+        )
+
     def _convex_scaling(
         self, dense: np.ndarray, sparse: list[SparseEmbedding]
     ) -> tuple[np.ndarray, list[SparseEmbedding]]:

tests/unit/test_router.py

@@ -24,7 +24,7 @@
 from semantic_router.llms import BaseLLM, OpenAILLM
 from semantic_router.route import Route
 from semantic_router.routers import HybridRouter, RouterConfig, SemanticRouter
-from semantic_router.schema import SparseEmbedding
+from semantic_router.schema import RouteChoice, SparseEmbedding
 
 PINECONE_SLEEP = 8
 RETRY_COUNT = 10
@@ -206,6 +206,16 @@ def routes_4():
     ]
 
 
+@pytest.fixture
+def routes_5():
+    return [
+        Route(name="Route 1", utterances=["Hello", "Hi"], metadata={"type": "default"}),
+        Route(name="Route 2", utterances=["Goodbye", "Bye", "Au revoir"]),
+        Route(name="Route 3", utterances=["Hello", "Hi"]),
+        Route(name="Route 4", utterances=["Goodbye", "Bye", "Au revoir"]),
+    ]
+
+
 @pytest.fixture
 def route_single_utterance():
     return [
@@ -643,3 +653,188 @@ async def test_async_encode(
                 assert sparse_encode_queries_spy.called
         else:
             assert sparse_call_spy.called
+
+
+@pytest.mark.parametrize(
+    "router_cls",
+    [
+        HybridRouter,
+        SemanticRouter,
+    ],
+)
+class TestRouter:
+    def test_query_parameter(self, router_cls, routes_5, mocker):
+        """Test that we return expected values in RouteChoice objects."""
+        # Create router with mock encoders
+        dense_encoder = MockSymmetricDenseEncoder(name="Dense Encoder")
+        if router_cls == HybridRouter:
+            sparse_encoder = MockSymmetricSparseEncoder(name="Sparse Encoder")
+            router = router_cls(
+                encoder=dense_encoder,
+                sparse_encoder=sparse_encoder,
+                routes=routes_5,
+                auto_sync="local",
+            )
+        else:
+            router = router_cls(
+                encoder=dense_encoder,
+                routes=routes_5,
+                auto_sync="local",
+            )
+
+        # Setup a mock for the similarity calculation method
+        _ = mocker.patch.object(
+            router,
+            "_score_routes",
+            return_value=[
+                ("Route 1", 0.9, [0.1, 0.2, 0.3]),
+                ("Route 2", 0.8, [0.4, 0.5, 0.6]),
+                ("Route 3", 0.7, [0.7, 0.8, 0.9]),
+                ("Route 4", 0.6, [1.0, 1.1, 1.2]),
+            ],
+        )
+
+        # Test without limit (should return only the top match)
+        result = router("test query")
+        assert result is not None
+        assert isinstance(result, RouteChoice)
+
+        # Confirm we have Route 1 and sim score
+        assert result.name == "Route 1"
+        assert result.similarity_score == 0.9
+        assert result.function_call is None
+
+    @pytest.mark.asyncio
+    async def test_async_query_parameter(self, router_cls, routes_5, mocker):
+        """Test that we return expected values in RouteChoice objects."""
+        # Create router with mock encoders
+        dense_encoder = MockSymmetricDenseEncoder(name="Dense Encoder")
+        if router_cls == HybridRouter:
+            sparse_encoder = MockSymmetricSparseEncoder(name="Sparse Encoder")
+            router = router_cls(
+                encoder=dense_encoder,
+                sparse_encoder=sparse_encoder,
+                routes=routes_5,
+                auto_sync="local",
+            )
+        else:
+            router = router_cls(
+                encoder=dense_encoder,
+                routes=routes_5,
+                auto_sync="local",
+            )
+
+        # Setup a mock for the similarity calculation method
+        _ = mocker.patch.object(
+            router,
+            "_score_routes",
+            return_value=[
+                ("Route 1", 0.9, [0.1, 0.2, 0.3]),
+                ("Route 2", 0.8, [0.4, 0.5, 0.6]),
+                ("Route 3", 0.7, [0.7, 0.8, 0.9]),
+                ("Route 4", 0.6, [1.0, 1.1, 1.2]),
+            ],
+        )
+
+        # Test without limit (should return only the top match)
+        result = await router.acall("test query")
+        assert result is not None
+        assert isinstance(result, RouteChoice)
+
+        # Confirm we have Route 1 and sim score
+        assert result.name == "Route 1"
+        assert result.similarity_score == 0.9
+        assert result.function_call is None
+
+    def test_limit_parameter(self, router_cls, routes_5, mocker):
+        """Test that the limit parameter works correctly for sync router calls."""
+        # Create router with mock encoders
+        dense_encoder = MockSymmetricDenseEncoder(name="Dense Encoder")
+        if router_cls == HybridRouter:
+            sparse_encoder = MockSymmetricSparseEncoder(name="Sparse Encoder")
+            router = router_cls(
+                encoder=dense_encoder,
+                sparse_encoder=sparse_encoder,
+                routes=routes_5,
+                auto_sync="local",
+            )
+        else:
+            router = router_cls(
+                encoder=dense_encoder,
+                routes=routes_5,
+                auto_sync="local",
+            )
+
+        # Setup a mock for the similarity calculation method
+        _ = mocker.patch.object(
+            router,
+            "_score_routes",
+            return_value=[
+                ("Route 1", 0.9, [0.1, 0.2, 0.3]),
+                ("Route 2", 0.8, [0.4, 0.5, 0.6]),
+                ("Route 3", 0.7, [0.7, 0.8, 0.9]),
+                ("Route 4", 0.6, [1.0, 1.1, 1.2]),
+            ],
+        )
+
+        # Test without limit (should return only the top match)
+        result = router("test query")
+        assert result is not None
+        assert isinstance(result, RouteChoice)
+
+        # Test with limit=2 (should return top 2 matches)
+        result = router("test query", limit=2)
+        assert result is not None
+        assert len(result) == 2
+
+        # Test with limit=None (should return all matches)
+        result = router("test query", limit=None)
+        assert result is not None
+        assert len(result) == 4  # Should return all matches
+
+    @pytest.mark.asyncio
+    async def test_async_limit_parameter(self, router_cls, routes_5, mocker):
+        """Test that the limit parameter works correctly for async router calls."""
+        # Create router with mock encoders
+        dense_encoder = MockSymmetricDenseEncoder(name="Dense Encoder")
+        if router_cls == HybridRouter:
+            sparse_encoder = MockSymmetricSparseEncoder(name="Sparse Encoder")
+            router = router_cls(
+                encoder=dense_encoder,
+                sparse_encoder=sparse_encoder,
+                routes=routes_5,
+                auto_sync="local",
+            )
+        else:
+            router = router_cls(
+                encoder=dense_encoder,
+                routes=routes_5,
+                auto_sync="local",
+            )
+
+        # Setup a mock for the async similarity calculation method
+        _ = mocker.patch.object(
+            router,
+            "_score_routes",
+            return_value=[
+                ("Route 1", 0.9, [0.1, 0.2, 0.3]),
+                ("Route 2", 0.8, [0.4, 0.5, 0.6]),
+                ("Route 3", 0.7, [0.7, 0.8, 0.9]),
+                ("Route 4", 0.6, [1.0, 1.1, 1.2]),
+            ],
+        )
+
+        # Test without limit (should return only the top match)
+        result = await router.acall("test query")
+        assert result is not None
+        assert isinstance(result, RouteChoice)
+
+        # Test with limit=2 (should return top 2 matches)
+        result = await router.acall("test query", limit=2)
+        assert result is not None
+        assert len(result) == 2
+
+        # Test with limit=None (should return all matches)
+        result = await router.acall("test query", limit=None)
+        assert result is not None
+        assert len(result) == 4  # Should return all matches