formatting and linting updates

Author: tylerhutcherson

redisvl/extensions/router/schema.py

@@ -64,7 +64,9 @@ class RoutingConfig(BaseModel):
     """The threshold for semantic distance."""
     max_k: int = Field(default=1)
     """The maximum number of top matches to return."""
-    aggregation_method: DistanceAggregationMethod = Field(default=DistanceAggregationMethod.avg)
+    aggregation_method: DistanceAggregationMethod = Field(
+        default=DistanceAggregationMethod.avg
+    )
     """Aggregation method to use to classify queries."""
 
     @validator("max_k")

redisvl/extensions/router/semantic.py

@@ -1,12 +1,17 @@
 import hashlib
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Type
 
 import redis.commands.search.reducers as reducers
 from pydantic.v1 import BaseModel, Field, PrivateAttr
 from redis import Redis
-from redis.commands.search.aggregation import AggregateRequest, AggregateResult
-
-from redisvl.extensions.router.schema import Route, RoutingConfig, RouteMatch, DistanceAggregationMethod
+from redis.commands.search.aggregation import AggregateRequest, AggregateResult, Reducer
+
+from redisvl.extensions.router.schema import (
+    DistanceAggregationMethod,
+    Route,
+    RouteMatch,
+    RoutingConfig,
+)
 from redisvl.index import SearchIndex
 from redisvl.query import RangeQuery
 from redisvl.redis.utils import convert_bytes, make_dict
@@ -18,7 +23,7 @@ class SemanticRouterIndexSchema(IndexSchema):
     """Customized index schema for SemanticRouter."""
 
     @classmethod
-    def from_params(cls, name: str, vector_dims: int) -> 'SemanticRouterIndexSchema':
+    def from_params(cls, name: str, vector_dims: int) -> "SemanticRouterIndexSchema":
         """Create an index schema based on router name and vector dimensions.
 
         Args:
@@ -30,7 +35,7 @@ def from_params(cls, name: str, vector_dims: int) -> 'SemanticRouterIndexSchema'
         """
         return cls(
             index=IndexInfo(name=name, prefix=name),
-            fields=[
+            fields=[  # type: ignore
                 {"name": "route_name", "type": "tag"},
                 {"name": "reference", "type": "text"},
                 {
@@ -87,7 +92,12 @@ def __init__(
             overwrite (bool, optional): Whether to overwrite existing index. Defaults to False.
             **kwargs: Additional arguments.
         """
-        super().__init__(name=name, routes=routes, vectorizer=vectorizer, routing_config=routing_config)
+        super().__init__(
+            name=name,
+            routes=routes,
+            vectorizer=vectorizer,
+            routing_config=routing_config,
+        )
         self._initialize_index(redis_client, redis_url, overwrite)
 
     def _initialize_index(
@@ -130,7 +140,7 @@ def route_names(self) -> List[str]:
         return [route.name for route in self.routes]
 
     @property
-    def route_thresholds(self) -> Dict[str, float]:
+    def route_thresholds(self) -> Dict[str, Optional[float]]:
         """Get the distance thresholds for each route.
 
         Returns:
@@ -168,7 +178,9 @@ def _add_routes(self, routes: List[Route]):
                     }
                 )
                 reference_hash = hashlib.sha256(reference.encode("utf-8")).hexdigest()
-                keys.append(f"{self._index.schema.index.prefix}:{route.name}:{reference_hash}")
+                keys.append(
+                    f"{self._index.schema.index.prefix}:{route.name}:{reference_hash}"
+                )
 
             # set route if does not yet exist client side
             if not self.get(route.name):
@@ -204,7 +216,7 @@ def _build_aggregate_request(
         self,
         vector_range_query: RangeQuery,
         aggregation_method: DistanceAggregationMethod,
-        max_k: int
+        max_k: int,
     ) -> AggregateRequest:
         """Build the Redis aggregation request.
 
@@ -216,6 +228,8 @@ def _build_aggregate_request(
         Returns:
             AggregateRequest: The constructed aggregation request.
         """
+        aggregation_func: Type[Reducer]
+
         if aggregation_method == DistanceAggregationMethod.min:
             aggregation_func = reducers.min
         elif aggregation_method == DistanceAggregationMethod.sum:
@@ -226,7 +240,9 @@ def _build_aggregate_request(
         aggregate_query = str(vector_range_query).split(" RETURN")[0]
         aggregate_request = (
             AggregateRequest(aggregate_query)
-            .group_by("@route_name", aggregation_func("vector_distance").alias("distance"))
+            .group_by(
+                "@route_name", aggregation_func("vector_distance").alias("distance")
+            )
             .sort_by("@distance", max=max_k)
             .dialect(2)
         )
@@ -237,7 +253,7 @@ def _classify(
         self,
         vector: List[float],
         distance_threshold: float,
-        aggregation_method: DistanceAggregationMethod
+        aggregation_method: DistanceAggregationMethod,
     ) -> List[RouteMatch]:
         """Classify a single query vector.
 
@@ -256,16 +272,20 @@ def _classify(
             return_fields=["route_name"],
         )
 
-        aggregate_request = self._build_aggregate_request(vector_range_query, aggregation_method, max_k=1)
-        route_matches: AggregateResult = self._index.client.ft(self._index.name).aggregate(aggregate_request, vector_range_query.params)
+        aggregate_request = self._build_aggregate_request(
+            vector_range_query, aggregation_method, max_k=1
+        )
+        route_matches: AggregateResult = self._index.client.ft(  # type: ignore
+            self._index.name
+        ).aggregate(aggregate_request, vector_range_query.params)
         return [self._process_route(route_match) for route_match in route_matches.rows]
 
     def _classify_many(
         self,
         vector: List[float],
         max_k: int,
         distance_threshold: float,
-        aggregation_method: DistanceAggregationMethod
+        aggregation_method: DistanceAggregationMethod,
     ) -> List[RouteMatch]:
         """Classify multiple query vectors.
 
@@ -284,8 +304,12 @@ def _classify_many(
             distance_threshold=distance_threshold,
             return_fields=["route_name"],
         )
-        aggregate_request = self._build_aggregate_request(vector_range_query, aggregation_method, max_k)
-        route_matches: AggregateResult = self._index.client.ft(self._index.name).aggregate(aggregate_request, vector_range_query.params)
+        aggregate_request = self._build_aggregate_request(
+            vector_range_query, aggregation_method, max_k
+        )
+        route_matches: AggregateResult = self._index.client.ft(  # type: ignore
+            self._index.name
+        ).aggregate(aggregate_request, vector_range_query.params)
         return [self._process_route(route_match) for route_match in route_matches.rows]
 
     def _pass_threshold(self, route_match: Optional[RouteMatch]) -> bool:
@@ -297,7 +321,11 @@ def _pass_threshold(self, route_match: Optional[RouteMatch]) -> bool:
         Returns:
             bool: True if the route match passes the threshold, False otherwise.
         """
-        return route_match is not None and route_match.distance <= route_match.route.distance_threshold
+        if route_match:
+            if route_match.distance is not None and route_match.route is not None:
+                if route_match.route.distance_threshold:
+                    return route_match.distance <= route_match.route.distance_threshold
+        return False
 
     def __call__(
         self,
@@ -320,8 +348,12 @@ def __call__(
                 raise ValueError("Must provide a vector or statement to the router")
             vector = self.vectorizer.embed(statement)
 
-        distance_threshold = distance_threshold or self.routing_config.distance_threshold
-        route_matches = self._classify(vector, distance_threshold, self.routing_config.aggregation_method)
+        distance_threshold = (
+            distance_threshold or self.routing_config.distance_threshold
+        )
+        route_matches = self._classify(
+            vector, distance_threshold, self.routing_config.aggregation_method
+        )
         route_match = route_matches[0] if route_matches else None
 
         if route_match and self._pass_threshold(route_match):
@@ -352,8 +384,16 @@ def route_many(
                 raise ValueError("Must provide a vector or statement to the router")
             vector = self.vectorizer.embed(statement)
 
-        distance_threshold = distance_threshold or self.routing_config.distance_threshold
+        distance_threshold = (
+            distance_threshold or self.routing_config.distance_threshold
+        )
         max_k = max_k or self.routing_config.max_k
-        route_matches = self._classify_many(vector, max_k, distance_threshold, self.routing_config.aggregation_method)
+        route_matches = self._classify_many(
+            vector, max_k, distance_threshold, self.routing_config.aggregation_method
+        )
 
-        return [route_match for route_match in route_matches if self._pass_threshold(route_match)]
+        return [
+            route_match
+            for route_match in route_matches
+            if self._pass_threshold(route_match)
+        ]

redisvl/schema/schema.py

@@ -195,7 +195,7 @@ def validate_and_create_fields(cls, values):
         """
         Validate uniqueness of field names and create valid field instances.
         """
-         # Ensure index is a dictionary for validation
+        # Ensure index is a dictionary for validation
         index = values.get("index")
         if not isinstance(index, IndexInfo):
             index = IndexInfo(**index)

tests/integration/test_semantic_router.py

@@ -1,24 +1,35 @@
 import pytest
 
-from redisvl.extensions.router.schema import Route, RoutingConfig
 from redisvl.extensions.router import SemanticRouter
+from redisvl.extensions.router.schema import Route, RoutingConfig
 
 
 @pytest.fixture
 def routes():
     return [
-        Route(name="greeting", references=["hello", "hi"], metadata={"type": "greeting"}, distance_threshold=0.3),
-        Route(name="farewell", references=["bye", "goodbye"], metadata={"type": "farewell"}, distance_threshold=0.3)
+        Route(
+            name="greeting",
+            references=["hello", "hi"],
+            metadata={"type": "greeting"},
+            distance_threshold=0.3,
+        ),
+        Route(
+            name="farewell",
+            references=["bye", "goodbye"],
+            metadata={"type": "farewell"},
+            distance_threshold=0.3,
+        ),
     ]
 
+
 @pytest.fixture
 def semantic_router(client, routes):
     router = SemanticRouter(
         name="test-router",
         routes=routes,
         routing_config=RoutingConfig(distance_threshold=0.3, max_k=2),
         redis_client=client,
-        overwrite=False
+        overwrite=False,
     )
     yield router
     router._index.delete(drop=True)
@@ -70,6 +81,7 @@ def test_multiple_query(semantic_router):
     assert len(matches) > 0
     assert matches[0].route.name == "greeting"
 
+
 def test_update_routing_config(semantic_router):
     new_config = RoutingConfig(distance_threshold=0.5, max_k=1)
     semantic_router.update_routing_config(new_config)
@@ -85,7 +97,9 @@ def test_vector_query(semantic_router):
 
 
 def test_vector_query_no_match(semantic_router):
-    vector = [0.0] * semantic_router.vectorizer.dims  # Random vector unlikely to match any route
+    vector = [
+        0.0
+    ] * semantic_router.vectorizer.dims  # Random vector unlikely to match any route
     match = semantic_router(vector=vector)
     assert match.route is None
 
@@ -94,7 +108,11 @@ def test_additional_route(semantic_router):
     new_routes = [
         Route(
             name="politics",
-            references=["are you liberal or conservative?", "who will you vote for?", "political speech"],
+            references=[
+                "are you liberal or conservative?",
+                "who will you vote for?",
+                "political speech",
+            ],
             metadata={"type": "greeting"},
         )
     ]