Merge branch '0.6.0' into feat/RAAE-769-add-langcache-wrapper

Author: abrookins

docs/user_guide/08_semantic_router.ipynb

@@ -8,7 +8,7 @@ routes:
   metadata:
     category: tech
     priority: 1
-  distance_threshold: 1.0
+  distance_threshold: 0.71
 - name: sports
   references:
   - who won the game last night?
@@ -19,7 +19,7 @@ routes:
   metadata:
     category: sports
     priority: 2
-  distance_threshold: 0.5
+  distance_threshold: 0.72
 - name: entertainment
   references:
   - what are the top movies right now?

docs/user_guide/router.yaml

@@ -100,6 +100,7 @@ def from_params(cls, name: str, vector_dims: int, dtype: str):
         return cls(
             index={"name": name, "prefix": name},  # type: ignore
             fields=[  # type: ignore
+                {"name": "reference_id", "type": "tag"},
                 {"name": "route_name", "type": "tag"},
                 {"name": "reference", "type": "text"},
                 {

redisvl/extensions/router/schema.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type, Union
 
 import redis.commands.search.reducers as reducers
 import yaml
@@ -8,6 +8,7 @@
 from redis.commands.search.aggregation import AggregateRequest, AggregateResult, Reducer
 from redis.exceptions import ResponseError
 
+from redisvl.exceptions import RedisModuleVersionError
 from redisvl.extensions.constants import ROUTE_VECTOR_FIELD_NAME
 from redisvl.extensions.router.schema import (
     DistanceAggregationMethod,
@@ -17,10 +18,17 @@
     SemanticRouterIndexSchema,
 )
 from redisvl.index import SearchIndex
-from redisvl.query import VectorRangeQuery
+from redisvl.query import FilterQuery, VectorRangeQuery
+from redisvl.query.filter import Tag
+from redisvl.redis.connection import RedisConnectionFactory
 from redisvl.redis.utils import convert_bytes, make_dict
 from redisvl.utils.log import get_logger
-from redisvl.utils.utils import deprecated_argument, hashify, model_to_dict
+from redisvl.utils.utils import (
+    deprecated_argument,
+    hashify,
+    model_to_dict,
+    scan_by_pattern,
+)
 from redisvl.utils.vectorize.base import BaseVectorizer
 from redisvl.utils.vectorize.text.huggingface import HFTextVectorizer
 
@@ -98,9 +106,41 @@ def __init__(
             routes=routes,
             vectorizer=vectorizer,
             routing_config=routing_config,
+            redis_url=redis_url,
+            redis_client=redis_client,
         )
+
         self._initialize_index(redis_client, redis_url, overwrite, **connection_kwargs)
 
+        self._index.client.json().set(f"{self.name}:route_config", f".", self.to_dict())  # type: ignore
+
+    @classmethod
+    def from_existing(
+        cls,
+        name: str,
+        redis_client: Optional[Redis] = None,
+        redis_url: str = "redis://localhost:6379",
+        **kwargs,
+    ) -> "SemanticRouter":
+        """Return SemanticRouter instance from existing index."""
+        try:
+            if redis_url:
+                redis_client = RedisConnectionFactory.get_redis_connection(
+                    redis_url=redis_url,
+                    **kwargs,
+                )
+            elif redis_client:
+                RedisConnectionFactory.validate_sync_redis(redis_client)
+        except RedisModuleVersionError as e:
+            raise RedisModuleVersionError(
+                f"Loading from existing index failed. {str(e)}"
+            )
+
+        router_dict = redis_client.json().get(f"{name}:route_config")  # type: ignore
+        return cls.from_dict(
+            router_dict, redis_url=redis_url, redis_client=redis_client
+        )
+
     @deprecated_argument("dtype")
     def _initialize_index(
         self,
@@ -111,9 +151,11 @@ def _initialize_index(
         **connection_kwargs,
     ):
         """Initialize the search index and handle Redis connection."""
+
         schema = SemanticRouterIndexSchema.from_params(
             self.name, self.vectorizer.dims, self.vectorizer.dtype  # type: ignore
         )
+
         self._index = SearchIndex(
             schema=schema,
             redis_client=redis_client,
@@ -174,10 +216,10 @@ def update_route_thresholds(self, route_thresholds: Dict[str, Optional[float]]):
             if route.name in route_thresholds:
                 route.distance_threshold = route_thresholds[route.name]  # type: ignore
 
-    def _route_ref_key(self, route_name: str, reference: str) -> str:
+    @staticmethod
+    def _route_ref_key(index: SearchIndex, route_name: str, reference_hash: str) -> str:
         """Generate the route reference key."""
-        reference_hash = hashify(reference)
-        return f"{self._index.prefix}:{route_name}:{reference_hash}"
+        return f"{index.prefix}:{route_name}:{reference_hash}"
 
     def _add_routes(self, routes: List[Route]):
         """Add routes to the router and index.
@@ -195,14 +237,18 @@ def _add_routes(self, routes: List[Route]):
             )
             # set route references
             for i, reference in enumerate(route.references):
+                reference_hash = hashify(reference)
                 route_references.append(
                     {
+                        "reference_id": reference_hash,
                         "route_name": route.name,
                         "reference": reference,
                         "vector": reference_vectors[i],
                     }
                 )
-                keys.append(self._route_ref_key(route.name, reference))
+                keys.append(
+                    self._route_ref_key(self._index, route.name, reference_hash)
+                )
 
             # set route if does not yet exist client side
             if not self.get(route.name):
@@ -438,7 +484,7 @@ def remove_route(self, route_name: str) -> None:
         else:
             self._index.drop_keys(
                 [
-                    self._route_ref_key(route.name, reference)
+                    self._route_ref_key(self._index, route.name, hashify(reference))
                     for reference in route.references
                 ]
             )
@@ -596,3 +642,155 @@ def to_yaml(self, file_path: str, overwrite: bool = True) -> None:
         with open(fp, "w") as f:
             yaml_data = self.to_dict()
             yaml.dump(yaml_data, f, sort_keys=False)
+
+    # reference methods
+    def add_route_references(
+        self,
+        route_name: str,
+        references: Union[str, List[str]],
+    ) -> List[str]:
+        """Add a reference(s) to an existing route.
+
+        Args:
+            router_name (str): The name of the router.
+            references (Union[str, List[str]]): The reference or list of references to add.
+
+        Returns:
+            List[str]: The list of added references keys.
+        """
+
+        if isinstance(references, str):
+            references = [references]
+
+        route_references: List[Dict[str, Any]] = []
+        keys: List[str] = []
+
+        # embed route references as a single batch
+        reference_vectors = self.vectorizer.embed_many(references, as_buffer=True)
+
+        # set route references
+        for i, reference in enumerate(references):
+            reference_hash = hashify(reference)
+
+            route_references.append(
+                {
+                    "reference_id": reference_hash,
+                    "route_name": route_name,
+                    "reference": reference,
+                    "vector": reference_vectors[i],
+                }
+            )
+            keys.append(self._route_ref_key(self._index, route_name, reference_hash))
+
+        keys = self._index.load(route_references, keys=keys)
+
+        route = self.get(route_name)
+        if not route:
+            raise ValueError(f"Route {route_name} not found in the SemanticRouter")
+        route.references.extend(references)
+        self._update_router_state()
+        return keys
+
+    @staticmethod
+    def _make_filter_queries(ids: List[str]) -> List[FilterQuery]:
+        """Create a filter query for the given ids."""
+
+        queries = []
+
+        for id in ids:
+            fe = Tag("reference_id") == id
+            fq = FilterQuery(
+                return_fields=["reference_id", "route_name", "reference"],
+                filter_expression=fe,
+            )
+            queries.append(fq)
+
+        return queries
+
+    def get_route_references(
+        self,
+        route_name: str = "",
+        reference_ids: List[str] = [],
+        keys: List[str] = [],
+    ) -> List[Dict[str, Any]]:
+        """Get references for an existing route route.
+
+        Args:
+            router_name (str): The name of the router.
+            references (Union[str, List[str]]): The reference or list of references to add.
+
+        Returns:
+            List[Dict[str, Any]]]: Reference objects stored
+        """
+
+        if reference_ids:
+            queries = self._make_filter_queries(reference_ids)
+        elif route_name:
+            if not keys:
+                keys = scan_by_pattern(
+                    self._index.client, f"{self._index.prefix}:{route_name}:*"  # type: ignore
+                )
+
+            queries = self._make_filter_queries(
+                [key.split(":")[-1] for key in convert_bytes(keys)]
+            )
+        else:
+            raise ValueError(
+                "Must provide a route name, reference ids, or keys to get references"
+            )
+
+        res = self._index.batch_query(queries)
+
+        return [r[0] for r in res if len(r) > 0]
+
+    def delete_route_references(
+        self,
+        route_name: str = "",
+        reference_ids: List[str] = [],
+        keys: List[str] = [],
+    ) -> int:
+        """Get references for an existing semantic router route.
+
+        Args:
+            router_name Optional(str): The name of the router.
+            reference_ids Optional(List[str]]): The reference or list of references to delete.
+            keys Optional(List[str]]): List of fully qualified keys (prefix:router:reference_id) to delete.
+
+        Returns:
+            int: Number of objects deleted
+        """
+
+        if reference_ids and not keys:
+            queries = self._make_filter_queries(reference_ids)
+            res = self._index.batch_query(queries)
+            keys = [r[0]["id"] for r in res if len(r) > 0]
+        elif not keys:
+            keys = scan_by_pattern(
+                self._index.client, f"{self._index.prefix}:{route_name}:*"  # type: ignore
+            )
+
+        if not keys:
+            raise ValueError(f"No references found for route {route_name}")
+
+        to_be_deleted = []
+        for key in keys:
+            route_name = key.split(":")[-2]
+            to_be_deleted.append(
+                (route_name, convert_bytes(self._index.client.hgetall(key)))  # type: ignore
+            )
+
+        deleted = self._index.drop_keys(keys)
+
+        for route_name, delete in to_be_deleted:
+            route = self.get(route_name)
+            if not route:
+                raise ValueError(f"Route {route_name} not found in the SemanticRouter")
+            route.references.remove(delete["reference"])
+
+        self._update_router_state()
+
+        return deleted
+
+    def _update_router_state(self) -> None:
+        """Update the router configuration in Redis."""
+        self._index.client.json().set(f"{self.name}:route_config", f".", self.to_dict())  # type: ignore

redisvl/extensions/router/semantic.py

@@ -14,6 +14,7 @@
     Iterable,
     List,
     Optional,
+    Sequence,
     Tuple,
     Union,
 )
@@ -833,7 +834,7 @@ def search(self, *args, **kwargs) -> "Result":
             raise RedisSearchError(f"Error while searching: {str(e)}") from e
 
     def batch_query(
-        self, queries: List[BaseQuery], batch_size: int = 10
+        self, queries: Sequence[BaseQuery], batch_size: int = 10
     ) -> List[List[Dict[str, Any]]]:
         """Execute a batch of queries and process results."""
         results = self.batch_search(

redisvl/index/index.py

@@ -18,9 +18,9 @@ def _generate_run_router(test_data: List[LabeledData], router: SemanticRouter) -
         run_dict[td.id] = {}
         route_match = router(td.query)
         if route_match and route_match.name == td.query_match:
-            run_dict[td.id][td.query_match] = 1
+            run_dict[td.id][td.query_match] = np.int64(1)
         else:
-            run_dict[td.id][NULL_RESPONSE_KEY] = 1
+            run_dict[td.id][NULL_RESPONSE_KEY] = np.int64(1)
 
     return Run(run_dict)
 

redisvl/utils/optimize/router.py

@@ -1,5 +1,6 @@
 from typing import List
 
+import numpy as np
 from ranx import Qrels
 
 from redisvl.utils.optimize.schema import LabeledData
@@ -13,10 +14,10 @@ def _format_qrels(test_data: List[LabeledData]) -> Qrels:
 
     for td in test_data:
         if td.query_match:
-            qrels_dict[td.id] = {td.query_match: 1}
+            qrels_dict[td.id] = {td.query_match: np.int64(1)}
         else:
             # This is for capturing true negatives from test set
-            qrels_dict[td.id] = {NULL_RESPONSE_KEY: 1}
+            qrels_dict[td.id] = {NULL_RESPONSE_KEY: np.int64(1)}
 
     return Qrels(qrels_dict)