Merge branch 'main' into feat/RAAE-127/voyage-ai-integration

Author: tylerhutcherson

doctests/data/query_vector.json

@@ -0,0 +1,17 @@
+version: '0.1.0'
+
+index:
+  name: idx:bicycle
+  prefix: bicycle
+  storage_type: json
+
+fields:
+    - name: description
+      type: text
+    - name: description_embeddings
+      type: vector
+      attrs:
+        algorithm: flat
+        dims: 384
+        distance_metric: cosine
+        datatype: float32

doctests/data/query_vector_idx.yaml

@@ -0,0 +1,68 @@
+# EXAMPLE: query_vector
+# HIDE_START
+import json
+import warnings
+import redis
+import numpy as np
+from redisvl.index import SearchIndex
+from redisvl.query import RangeQuery, VectorQuery
+from redisvl.schema import IndexSchema
+from sentence_transformers import SentenceTransformer
+
+
+def embed_text(model, text):
+    return np.array(model.encode(text)).astype(np.float32).tobytes()
+
+r = redis.Redis(decode_responses=True)
+
+warnings.filterwarnings("ignore", category=FutureWarning, message=r".*clean_up_tokenization_spaces.*")
+model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
+# create index
+schema = IndexSchema.from_yaml('data/query_vector_idx.yaml')
+index = SearchIndex(schema, r)
+index.create(overwrite=True, drop=True)
+
+# load data
+with open("data/query_vector.json") as f:
+    bicycles = json.load(f)
+index.load(bicycles)
+# HIDE_END
+
+# STEP_START vector1
+query = "Bike for small kids"
+query_vector = embed_text(model, query)
+print(query_vector[:10]) # >>> b'\x02=c=\x93\x0e\xe0=aC'
+
+vquery = VectorQuery(
+    vector=query_vector,
+    vector_field_name="description_embeddings",
+    num_results=3,
+    return_score=True,
+    return_fields=["description"]
+)
+res = index.query(vquery)
+print(res) # >>> [{'id': 'bicycle:6b702e8b...', 'vector_distance': '0.399111807346', 'description': 'Kids want...
+# REMOVE_START
+assert len(res) == 3
+# REMOVE_END
+# STEP_END
+
+# STEP_START vector2
+vquery = RangeQuery(
+    vector=query_vector,
+    vector_field_name="description_embeddings",
+    distance_threshold=0.5,
+    return_score=True
+).return_fields("description").dialect(2)
+res = index.query(vquery)
+print(res) # >>> [{'id': 'bicycle:6bcb1bb4...', 'vector_distance': '0.399111807346', 'description': 'Kids want...
+# REMOVE_START
+assert len(res) == 2
+# REMOVE_END
+# STEP_END
+
+# REMOVE_START
+# destroy index and data
+index.delete(drop=True)
+# REMOVE_END

doctests/query_vector.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/constants.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/base.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):
         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/schema.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,
@@ -19,15 +29,6 @@
 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"
-    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"
-
     _index: SearchIndex
     _aindex: Optional[AsyncSearchIndex] = None
 
@@ -94,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
@@ -133,7 +134,7 @@ def __init__(
 
         validate_vector_dims(
             vectorizer.dims,
-            self._index.schema.fields[self.vector_field_name].attrs.dims,  # type: ignore
+            self._index.schema.fields[CACHE_VECTOR_FIELD_NAME].attrs.dims,  # type: ignore
         )
         self._vectorizer = vectorizer
 
@@ -145,9 +146,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:
@@ -300,7 +299,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[self.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(
@@ -363,7 +362,7 @@ def check(
 
         query = RangeQuery(
             vector=vector,
-            vector_field_name=self.vector_field_name,
+            vector_field_name=CACHE_VECTOR_FIELD_NAME,
             return_fields=self.return_fields,
             distance_threshold=distance_threshold,
             num_results=num_results,
@@ -444,7 +443,7 @@ async def acheck(
 
         query = RangeQuery(
             vector=vector,
-            vector_field_name=self.vector_field_name,
+            vector_field_name=CACHE_VECTOR_FIELD_NAME,
             return_fields=self.return_fields,
             distance_threshold=distance_threshold,
             num_results=num_results,
@@ -479,7 +478,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
 
@@ -541,7 +540,7 @@ def store(
         keys = self._index.load(
             data=[cache_entry.to_dict()],
             ttl=ttl,
-            id_field=self.entry_id_field_name,
+            id_field=ENTRY_ID_FIELD_NAME,
         )
         return keys[0]
 
@@ -605,7 +604,7 @@ async def astore(
         keys = await aindex.load(
             data=[cache_entry.to_dict()],
             ttl=ttl,
-            id_field=self.entry_id_field_name,
+            id_field=ENTRY_ID_FIELD_NAME,
         )
         return keys[0]
 
@@ -629,21 +628,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
 
@@ -674,21 +671,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/llmcache/semantic.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 IndexInfo, IndexSchema
 
 
@@ -104,7 +105,7 @@ def from_params(cls, name: str, vector_dims: int) -> "SemanticRouterIndexSchema"
                 {"name": "route_name", "type": "tag"},
                 {"name": "reference", "type": "text"},
                 {
-                    "name": "vector",
+                    "name": ROUTE_VECTOR_FIELD_NAME,
                     "type": "vector",
                     "attrs": {
                         "algorithm": "flat",

redisvl/extensions/router/schema.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,
@@ -226,7 +227,7 @@ 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"],
         )
@@ -278,7 +279,7 @@ 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"],
         )