Merge branch 'main' into feat/RAAE-206/vector-dtypes

Author: justin-cechmanek

redisvl/extensions/constants.py

@@ -0,0 +1,29 @@
+"""
+Constants used within the extension classes SemanticCache, BaseSessionManager,
+StandardSessionManager,SemanticSessionManager and SemanticRouter.
+These constants are also used within theses classes corresponding schema.
+"""
+
+# BaseSessionManager
+ID_FIELD_NAME: str = "entry_id"
+ROLE_FIELD_NAME: str = "role"
+CONTENT_FIELD_NAME: str = "content"
+TOOL_FIELD_NAME: str = "tool_call_id"
+TIMESTAMP_FIELD_NAME: str = "timestamp"
+SESSION_FIELD_NAME: str = "session_tag"
+
+# SemanticSessionManager
+SESSION_VECTOR_FIELD_NAME: str = "vector_field"
+
+# SemanticCache
+REDIS_KEY_FIELD_NAME: str = "key"
+ENTRY_ID_FIELD_NAME: str = "entry_id"
+PROMPT_FIELD_NAME: str = "prompt"
+RESPONSE_FIELD_NAME: str = "response"
+CACHE_VECTOR_FIELD_NAME: str = "prompt_vector"
+INSERTED_AT_FIELD_NAME: str = "inserted_at"
+UPDATED_AT_FIELD_NAME: str = "updated_at"
+METADATA_FIELD_NAME: str = "metadata"
+
+# SemanticRouter
+ROUTE_VECTOR_FIELD_NAME: str = "vector"

redisvl/extensions/llmcache/base.py

@@ -1,7 +1,5 @@
 from typing import Any, Dict, List, Optional
 
-from redisvl.redis.utils import hashify
-
 
 class BaseLLMCache:
     def __init__(self, ttl: Optional[int] = None):
@@ -79,14 +77,3 @@ async def astore(
         """Async store the specified key-value pair in the cache along with
         metadata."""
         raise NotImplementedError
-
-    def hash_input(self, prompt: str) -> str:
-        """Hashes the input prompt using SHA256.
-
-        Args:
-            prompt (str): Input string to be hashed.
-
-        Returns:
-            str: Hashed string.
-        """
-        return hashify(prompt)

redisvl/extensions/llmcache/schema.py

@@ -2,6 +2,13 @@
 
 from pydantic.v1 import BaseModel, Field, root_validator, validator
 
+from redisvl.extensions.constants import (
+    CACHE_VECTOR_FIELD_NAME,
+    INSERTED_AT_FIELD_NAME,
+    PROMPT_FIELD_NAME,
+    RESPONSE_FIELD_NAME,
+    UPDATED_AT_FIELD_NAME,
+)
 from redisvl.redis.utils import array_to_buffer, hashify
 from redisvl.schema import IndexSchema
 from redisvl.utils.utils import current_timestamp, deserialize, serialize
@@ -32,7 +39,7 @@ class CacheEntry(BaseModel):
     def generate_id(cls, values):
         # Ensure entry_id is set
         if not values.get("entry_id"):
-            values["entry_id"] = hashify(values["prompt"])
+            values["entry_id"] = hashify(values["prompt"], values.get("filters"))
         return values
 
     @validator("metadata")
@@ -110,12 +117,12 @@ def from_params(cls, name: str, prefix: str, vector_dims: int, dtype: str):
         return cls(
             index={"name": name, "prefix": prefix},  # type: ignore
             fields=[  # type: ignore
-                {"name": "prompt", "type": "text"},
-                {"name": "response", "type": "text"},
-                {"name": "inserted_at", "type": "numeric"},
-                {"name": "updated_at", "type": "numeric"},
+                {"name": PROMPT_FIELD_NAME, "type": "text"},
+                {"name": RESPONSE_FIELD_NAME, "type": "text"},
+                {"name": INSERTED_AT_FIELD_NAME, "type": "numeric"},
+                {"name": UPDATED_AT_FIELD_NAME, "type": "numeric"},
                 {
-                    "name": "prompt_vector",
+                    "name": CACHE_VECTOR_FIELD_NAME,
                     "type": "vector",
                     "attrs": {
                         "dims": vector_dims,

redisvl/extensions/llmcache/semantic.py

@@ -3,6 +3,16 @@
 
 from redis import Redis
 
+from redisvl.extensions.constants import (
+    CACHE_VECTOR_FIELD_NAME,
+    ENTRY_ID_FIELD_NAME,
+    INSERTED_AT_FIELD_NAME,
+    METADATA_FIELD_NAME,
+    PROMPT_FIELD_NAME,
+    REDIS_KEY_FIELD_NAME,
+    RESPONSE_FIELD_NAME,
+    UPDATED_AT_FIELD_NAME,
+)
 from redisvl.extensions.llmcache.base import BaseLLMCache
 from redisvl.extensions.llmcache.schema import (
     CacheEntry,
@@ -15,20 +25,10 @@
 from redisvl.utils.utils import current_timestamp, serialize, validate_vector_dims
 from redisvl.utils.vectorize import BaseVectorizer, HFTextVectorizer
 
-VECTOR_FIELD_NAME = "prompt_vector"
-
 
 class SemanticCache(BaseLLMCache):
     """Semantic Cache for Large Language Models."""
 
-    redis_key_field_name: str = "key"
-    entry_id_field_name: str = "entry_id"
-    prompt_field_name: str = "prompt"
-    response_field_name: str = "response"
-    inserted_at_field_name: str = "inserted_at"
-    updated_at_field_name: str = "updated_at"
-    metadata_field_name: str = "metadata"
-
     _index: SearchIndex
     _aindex: Optional[AsyncSearchIndex] = None
 
@@ -95,12 +95,12 @@ def __init__(
         # Process fields and other settings
         self.set_threshold(distance_threshold)
         self.return_fields = [
-            self.entry_id_field_name,
-            self.prompt_field_name,
-            self.response_field_name,
-            self.inserted_at_field_name,
-            self.updated_at_field_name,
-            self.metadata_field_name,
+            ENTRY_ID_FIELD_NAME,
+            PROMPT_FIELD_NAME,
+            RESPONSE_FIELD_NAME,
+            INSERTED_AT_FIELD_NAME,
+            UPDATED_AT_FIELD_NAME,
+            METADATA_FIELD_NAME,
         ]
 
         # Create semantic cache schema and index
@@ -137,10 +137,10 @@ def __init__(
 
         validate_vector_dims(
             vectorizer.dims,
-            self._index.schema.fields[VECTOR_FIELD_NAME].attrs.dims,  # type: ignore
+            self._index.schema.fields[CACHE_VECTOR_FIELD_NAME].attrs.dims,  # type: ignore
         )
         self._vectorizer = vectorizer
-        self._dtype = self.index.schema.fields[VECTOR_FIELD_NAME].attrs.datatype  # type: ignore[union-attr]
+        self._dtype = self.index.schema.fields[CACHE_VECTOR_FIELD_NAME].attrs.datatype  # type: ignore[union-attr]
 
     def _modify_schema(
         self,
@@ -150,9 +150,7 @@ def _modify_schema(
         """Modify the base cache schema using the provided filterable fields"""
 
         if filterable_fields is not None:
-            protected_field_names = set(
-                self.return_fields + [self.redis_key_field_name]
-            )
+            protected_field_names = set(self.return_fields + [REDIS_KEY_FIELD_NAME])
             for filter_field in filterable_fields:
                 field_name = filter_field["name"]
                 if field_name in protected_field_names:
@@ -305,7 +303,7 @@ async def _avectorize_prompt(self, prompt: Optional[str]) -> List[float]:
     def _check_vector_dims(self, vector: List[float]):
         """Checks the size of the provided vector and raises an error if it
         doesn't match the search index vector dimensions."""
-        schema_vector_dims = self._index.schema.fields[VECTOR_FIELD_NAME].attrs.dims  # type: ignore
+        schema_vector_dims = self._index.schema.fields[CACHE_VECTOR_FIELD_NAME].attrs.dims  # type: ignore
         validate_vector_dims(len(vector), schema_vector_dims)
 
     def check(
@@ -368,7 +366,7 @@ def check(
 
         query = RangeQuery(
             vector=vector,
-            vector_field_name=VECTOR_FIELD_NAME,
+            vector_field_name=CACHE_VECTOR_FIELD_NAME,
             return_fields=self.return_fields,
             distance_threshold=distance_threshold,
             num_results=num_results,
@@ -450,7 +448,7 @@ async def acheck(
 
         query = RangeQuery(
             vector=vector,
-            vector_field_name=VECTOR_FIELD_NAME,
+            vector_field_name=CACHE_VECTOR_FIELD_NAME,
             return_fields=self.return_fields,
             distance_threshold=distance_threshold,
             num_results=num_results,
@@ -485,7 +483,7 @@ def _process_cache_results(
                 cache_hit_dict = {
                     k: v for k, v in cache_hit_dict.items() if k in return_fields
                 }
-            cache_hit_dict[self.redis_key_field_name] = redis_key
+            cache_hit_dict[REDIS_KEY_FIELD_NAME] = redis_key
             cache_hits.append(cache_hit_dict)
         return redis_keys, cache_hits
 
@@ -547,7 +545,7 @@ def store(
         keys = self._index.load(
             data=[cache_entry.to_dict(self._dtype)],
             ttl=ttl,
-            id_field=self.entry_id_field_name,
+            id_field=ENTRY_ID_FIELD_NAME,
         )
         return keys[0]
 
@@ -611,7 +609,7 @@ async def astore(
         keys = await aindex.load(
             data=[cache_entry.to_dict(self._dtype)],
             ttl=ttl,
-            id_field=self.entry_id_field_name,
+            id_field=ENTRY_ID_FIELD_NAME,
         )
         return keys[0]
 
@@ -635,21 +633,19 @@ def update(self, key: str, **kwargs) -> None:
             for k, v in kwargs.items():
 
                 # Make sure the item is in the index schema
-                if k not in set(
-                    self._index.schema.field_names + [self.metadata_field_name]
-                ):
+                if k not in set(self._index.schema.field_names + [METADATA_FIELD_NAME]):
                     raise ValueError(f"{k} is not a valid field within the cache entry")
 
                 # Check for metadata and deserialize
-                if k == self.metadata_field_name:
+                if k == METADATA_FIELD_NAME:
                     if isinstance(v, dict):
                         kwargs[k] = serialize(v)
                     else:
                         raise TypeError(
                             "If specified, cached metadata must be a dictionary."
                         )
 
-            kwargs.update({self.updated_at_field_name: current_timestamp()})
+            kwargs.update({UPDATED_AT_FIELD_NAME: current_timestamp()})
 
             self._index.client.hset(key, mapping=kwargs)  # type: ignore
 
@@ -680,21 +676,19 @@ async def aupdate(self, key: str, **kwargs) -> None:
             for k, v in kwargs.items():
 
                 # Make sure the item is in the index schema
-                if k not in set(
-                    self._index.schema.field_names + [self.metadata_field_name]
-                ):
+                if k not in set(self._index.schema.field_names + [METADATA_FIELD_NAME]):
                     raise ValueError(f"{k} is not a valid field within the cache entry")
 
                 # Check for metadata and deserialize
-                if k == self.metadata_field_name:
+                if k == METADATA_FIELD_NAME:
                     if isinstance(v, dict):
                         kwargs[k] = serialize(v)
                     else:
                         raise TypeError(
                             "If specified, cached metadata must be a dictionary."
                         )
 
-            kwargs.update({self.updated_at_field_name: current_timestamp()})
+            kwargs.update({UPDATED_AT_FIELD_NAME: current_timestamp()})
 
             await aindex.load(data=[kwargs], keys=[key])
 

redisvl/extensions/router/schema.py

@@ -3,6 +3,7 @@
 
 from pydantic.v1 import BaseModel, Field, validator
 
+from redisvl.extensions.constants import ROUTE_VECTOR_FIELD_NAME
 from redisvl.schema import IndexSchema
 
 
@@ -104,7 +105,7 @@ def from_params(cls, name: str, vector_dims: int, dtype: str):
                 {"name": "route_name", "type": "tag"},
                 {"name": "reference", "type": "text"},
                 {
-                    "name": "vector",
+                    "name": ROUTE_VECTOR_FIELD_NAME,
                     "type": "vector",
                     "attrs": {
                         "algorithm": "flat",

redisvl/extensions/router/semantic.py

@@ -8,6 +8,7 @@
 from redis.commands.search.aggregation import AggregateRequest, AggregateResult, Reducer
 from redis.exceptions import ResponseError
 
+from redisvl.extensions.constants import ROUTE_VECTOR_FIELD_NAME
 from redisvl.extensions.router.schema import (
     DistanceAggregationMethod,
     Route,
@@ -28,8 +29,6 @@
 
 logger = get_logger(__name__)
 
-VECTOR_FIELD_NAME = "vector"
-
 
 class SemanticRouter(BaseModel):
     """Semantic Router for managing and querying route vectors."""
@@ -172,7 +171,7 @@ def _add_routes(self, routes: List[Route]):
             reference_vectors = self.vectorizer.embed_many(
                 [reference for reference in route.references],
                 as_buffer=True,
-                dtype=self._index.schema.fields[VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
+                dtype=self._index.schema.fields[ROUTE_VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
             )
             # set route references
             for i, reference in enumerate(route.references):
@@ -246,10 +245,10 @@ def _classify_route(
         """Classify to a single route using a vector."""
         vector_range_query = RangeQuery(
             vector=vector,
-            vector_field_name="vector",
+            vector_field_name=ROUTE_VECTOR_FIELD_NAME,
             distance_threshold=distance_threshold,
             return_fields=["route_name"],
-            dtype=self._index.schema.fields[VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
+            dtype=self._index.schema.fields[ROUTE_VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
         )
 
         aggregate_request = self._build_aggregate_request(
@@ -299,10 +298,10 @@ def _classify_multi_route(
         """Classify to multiple routes, up to max_k (int), using a vector."""
         vector_range_query = RangeQuery(
             vector=vector,
-            vector_field_name="vector",
+            vector_field_name=ROUTE_VECTOR_FIELD_NAME,
             distance_threshold=distance_threshold,
             return_fields=["route_name"],
-            dtype=self._index.schema.fields[VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
+            dtype=self._index.schema.fields[ROUTE_VECTOR_FIELD_NAME].attrs.datatype,  # type: ignore[union-attr]
         )
         aggregate_request = self._build_aggregate_request(
             vector_range_query, aggregation_method, max_k

redisvl/extensions/session_manager/base_session.py

@@ -1,16 +1,15 @@
 from typing import Any, Dict, List, Optional, Union
 
+from redisvl.extensions.constants import (
+    CONTENT_FIELD_NAME,
+    ROLE_FIELD_NAME,
+    TOOL_FIELD_NAME,
+)
 from redisvl.extensions.session_manager.schema import ChatMessage
 from redisvl.utils.utils import create_uuid
 
 
 class BaseSessionManager:
-    id_field_name: str = "entry_id"
-    role_field_name: str = "role"
-    content_field_name: str = "content"
-    tool_field_name: str = "tool_call_id"
-    timestamp_field_name: str = "timestamp"
-    session_field_name: str = "session_tag"
 
     def __init__(
         self,
@@ -107,11 +106,11 @@ def _format_context(
                 context.append(chat_message.content)
             else:
                 chat_message_dict = {
-                    self.role_field_name: chat_message.role,
-                    self.content_field_name: chat_message.content,
+                    ROLE_FIELD_NAME: chat_message.role,
+                    CONTENT_FIELD_NAME: chat_message.content,
                 }
                 if chat_message.tool_call_id is not None:
-                    chat_message_dict[self.tool_field_name] = chat_message.tool_call_id
+                    chat_message_dict[TOOL_FIELD_NAME] = chat_message.tool_call_id
 
                 context.append(chat_message_dict)  # type: ignore