Checkpoint

Author: abrookins

redisvl/extensions/llmcache/langcache_api.py

@@ -1,11 +1,14 @@
 import json
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 from langcache import LangCache as LangCacheSDK
+from langcache.models import CacheEntryScope, CacheEntryScopeTypedDict
 
 from redisvl.extensions.llmcache.base import BaseLLMCache
 from redisvl.query.filter import FilterExpression
-from redisvl.utils.utils import current_timestamp, hashify
+from redisvl.utils.utils import current_timestamp
+
+Scope = Optional[Union[CacheEntryScope, CacheEntryScopeTypedDict]]
 
 
 class LangCache(BaseLLMCache):
@@ -20,6 +23,7 @@ def __init__(
         redis_url: str = "redis://localhost:6379",
         connection_kwargs: Dict[str, Any] = {},
         overwrite: bool = False,
+        entry_scope: Scope = None,
         **kwargs,
     ):
         """Initialize a LangCache client.
@@ -32,6 +36,7 @@ def __init__(
             redis_url: URL for Redis connection if no client is provided.
             connection_kwargs: Additional Redis connection parameters.
             overwrite: Whether to overwrite an existing cache with the same name.
+            entry_scope: Optional scope for cache entries.
         """
         # Initialize the base class
         super().__init__(ttl)
@@ -43,7 +48,7 @@ def __init__(
         self._distance_threshold = distance_threshold
         self._ttl = ttl
         self._cache_id = name
-
+        self._entry_scope = entry_scope
         # Initialize LangCache SDK client
         self._api = LangCacheSDK(server_url=redis_url, client=redis_client)
 
@@ -101,12 +106,28 @@ def set_ttl(self, ttl: Optional[int] = None) -> None:
 
     def clear(self) -> None:
         """Clear all entries from the cache while preserving the cache configuration."""
-        self._api.entries.delete_all(cache_id=self._cache_id, attributes={}, scope={})
+        self._api.entries.delete_all(
+            cache_id=self._cache_id,
+            attributes={},
+            scope=(
+                self._entry_scope
+                if self._entry_scope is not None
+                else CacheEntryScope()
+            ),
+        )
 
     async def aclear(self) -> None:
         """Asynchronously clear all entries from the cache."""
-        # Currently using synchronous implementation since langcache doesn't have async API
-        self.clear()
+        # Use the SDK's async delete_all
+        await self._api.entries.delete_all_async(
+            cache_id=self._cache_id,
+            attributes={},
+            scope=(
+                self._entry_scope
+                if self._entry_scope is not None
+                else CacheEntryScope()
+            ),
+        )
 
     def delete(self) -> None:
         """Delete the cache and all its entries."""
@@ -115,8 +136,17 @@ def delete(self) -> None:
 
     async def adelete(self) -> None:
         """Asynchronously delete the cache and all its entries."""
-        # Currently using synchronous implementation since langcache doesn't have async API
-        self.delete()
+        # Clear entries then delete cache asynchronously
+        await self._api.entries.delete_all_async(
+            cache_id=self._cache_id,
+            attributes={},
+            scope=(
+                self._entry_scope
+                if self._entry_scope is not None
+                else CacheEntryScope()
+            ),
+        )
+        await self._api.cache.delete_async(cache_id=self._cache_id)
 
     def drop(
         self, ids: Optional[List[str]] = None, keys: Optional[List[str]] = None
@@ -134,14 +164,13 @@ def drop(
     async def adrop(
         self, ids: Optional[List[str]] = None, keys: Optional[List[str]] = None
     ) -> None:
-        """Asynchronously remove specific entries from the cache.
-
-        Args:
-            ids: List of entry IDs to remove.
-            keys: List of Redis keys to remove.
-        """
-        # Currently using synchronous implementation since langcache doesn't have async API
-        self.drop(ids, keys)
+        """Asynchronously remove specific entries from the cache."""
+        # Use the SDK's async delete for each entry
+        if ids:
+            for entry_id in ids:
+                await self._api.entries.delete_async(
+                    entry_id=entry_id, cache_id=self._cache_id
+                )
 
     def check(
         self,
@@ -151,6 +180,7 @@ def check(
         return_fields: Optional[List[str]] = None,
         filter_expression: Optional[FilterExpression] = None,
         distance_threshold: Optional[float] = None,
+        entry_scope: Optional[Dict[str, Any]] = None,
     ) -> List[Dict[str, Any]]:
         """Check the cache for semantically similar entries.
 
@@ -161,7 +191,7 @@ def check(
             return_fields: Fields to include in the response.
             filter_expression: Optional filter for the search.
             distance_threshold: Override the default distance threshold.
-
+            entry_scope: Optional scope for cache entries.
         Returns:
             List of matching cache entries.
 
@@ -172,18 +202,22 @@ def check(
         if not any([prompt, vector]):
             raise ValueError("Either prompt or vector must be provided")
 
+        _scope = entry_scope or self._entry_scope
+
         if return_fields and not isinstance(return_fields, list):
             raise TypeError("return_fields must be a list")
 
         # Use provided threshold or default
         threshold = distance_threshold or self._distance_threshold
 
-        # Search the cache - note we don't use scope since FilterExpression conversion would be complex
-        # and require proper implementation for CacheEntryScope format
+        # Search the cache - note we don't use scope since FilterExpression conversion
+        # would be complex (impossible?)
         results = self._api.entries.search(
             cache_id=self._cache_id,
             prompt=prompt or "",  # Ensure prompt is never None
             similarity_threshold=threshold,
+            # Type-cast is necessary to handle the scope type correctly
+            scope=_scope,  # type: ignore[arg-type]
         )
 
         # If we need to limit results and have more than requested, slice the list
@@ -235,18 +269,61 @@ async def acheck(
         return_fields: Optional[List[str]] = None,
         filter_expression: Optional[FilterExpression] = None,
         distance_threshold: Optional[float] = None,
+        entry_scope: Scope = None,
     ) -> List[Dict[str, Any]]:
         """Asynchronously check the cache for semantically similar entries."""
-        # Currently using synchronous implementation since langcache doesn't have async API
-        return self.check(
-            prompt,
-            vector,
-            num_results,
-            return_fields,
-            filter_expression,
-            distance_threshold,
+        # Validate inputs
+        if not any([prompt, vector]):
+            raise ValueError("Either prompt or vector must be provided")
+        if return_fields and not isinstance(return_fields, list):
+            raise TypeError("return_fields must be a list")
+
+        # Determine scope to use
+        _scope = entry_scope or self._entry_scope
+
+        # Determine threshold
+        threshold = distance_threshold or self._distance_threshold
+
+        # Perform async search
+        results = await self._api.entries.search_async(
+            cache_id=self._cache_id,
+            prompt=prompt or "",
+            similarity_threshold=threshold,
+            # Type-cast is necessary to handle the scope type correctly
+            scope=_scope,  # type: ignore[arg-type]
         )
 
+        # Limit results
+        if num_results < len(results):
+            results = results[:num_results]
+
+        # Format hits
+        cache_hits: List[Dict[str, Any]] = []
+        for result in results:
+            hit = {
+                "key": result.id,
+                "entry_id": result.id,
+                "prompt": result.prompt,
+                "response": result.response,
+                "vector_distance": result.similarity,
+            }
+            if hasattr(result, "metadata") and result.metadata:
+                try:
+                    metadata_dict = {}
+                    if hasattr(result.metadata, "__dict__"):
+                        metadata_dict = {
+                            k: v
+                            for k, v in result.metadata.__dict__.items()
+                            if not k.startswith("_")
+                        }
+                    hit["metadata"] = metadata_dict
+                except Exception:
+                    hit["metadata"] = {}
+            if return_fields:
+                hit = {k: v for k, v in hit.items() if k in return_fields or k == "key"}
+            cache_hits.append(hit)
+        return cache_hits
+
     def store(
         self,
         prompt: str,
@@ -261,7 +338,7 @@ def store(
         Args:
             prompt: The prompt text.
             response: The response text.
-            vector: Optional vector representation of the prompt.
+            vector: Unused. LangCache manages vectorization internally.
             metadata: Optional metadata to store with the entry.
             filters: Optional filters to associate with the entry.
             ttl: Optional custom TTL for this entry.
@@ -312,84 +389,74 @@ async def astore(
         ttl: Optional[int] = None,
     ) -> str:
         """Asynchronously store a new entry in the cache."""
-        # Currently using synchronous implementation since langcache doesn't have async API
-        return self.store(prompt, response, vector, metadata, filters, ttl)
-
-    def update(self, key: str, **kwargs) -> None:
-        """Update an existing cache entry.
-
-        Args:
-            key: The entry ID to update.
-            **kwargs: Fields to update (prompt, response, metadata, etc.)
-        """
-        # Find the entry to update
-        existing_entries = self._api.entries.search(
-            cache_id=self._cache_id,
-            prompt="",  # Required parameter but we're searching by ID
-            attributes={"id": key},  # Search by ID as an attribute
-            similarity_threshold=1.0,  # We're not doing semantic search
-        )
-
-        if not existing_entries:
-            return
+        # Validate metadata
+        if metadata is not None and not isinstance(metadata, dict):
+            raise ValueError("Metadata must be a dictionary")
 
-        existing_entry = existing_entries[0]
+        # Create entry with optional TTL
+        entry_ttl = ttl if ttl is not None else self._ttl
 
-        # Prepare updated values
-        # CacheEntry objects are Pydantic models, access their attributes directly
-        prompt = kwargs.get(
-            "prompt", existing_entry.prompt if hasattr(existing_entry, "prompt") else ""
-        )
-        response = kwargs.get(
-            "response",
-            existing_entry.response if hasattr(existing_entry, "response") else "",
-        )
+        # Convert ttl to ttl_millis (milliseconds) if provided
+        ttl_millis = entry_ttl * 1000 if entry_ttl is not None else None
 
-        # Prepare attributes for update
+        # Process additional attributes from filters
         attributes = {}
-        if "metadata" in kwargs:
+        if filters:
+            attributes.update(filters)
+
+        # Add metadata to attributes if provided
+        if metadata:
             attributes["metadata"] = (
-                json.dumps(kwargs["metadata"])
-                if isinstance(kwargs["metadata"], dict)
-                else kwargs["metadata"]
+                json.dumps(metadata) if isinstance(metadata, dict) else metadata
             )
 
-        # Convert TTL to milliseconds if provided
-        ttl = kwargs.get("ttl", None)
-        ttl_millis = ttl * 1000 if ttl is not None else None
-
-        # Re-create the entry with updated values
-        self._api.entries.create(
+        # Store the entry and get the response
+        create_response = await self._api.entries.create_async(
             cache_id=self._cache_id,
             prompt=prompt,
             response=response,
             attributes=attributes,
             ttl_millis=ttl_millis,
         )
 
+        # Return the entry ID from the response
+        return create_response.entry_id
+
+    def update(self, key: str, **kwargs) -> None:
+        """Update an existing cache entry.
+
+        Args:
+            key: The entry ID to update.
+            **kwargs: Fields to update (prompt, response, metadata, etc.)
+        """
+        raise NotImplementedError("LangCache SDK does not support update in place")
+
     async def aupdate(self, key: str, **kwargs) -> None:
         """Asynchronously update an existing cache entry."""
-        # Currently using synchronous implementation since langcache doesn't have async API
-        self.update(key, **kwargs)
+        raise NotImplementedError("LangCache SDK does not support update in place")
 
     def disconnect(self) -> None:
-        """Close the Redis connection."""
-        # Redis clients typically don't need explicit disconnection,
-        # as they use connection pooling
-        pass
+        """Close the connection."""
+        if self._api.sdk_configuration.client is not None:
+            self._api.sdk_configuration.client.close()
 
     async def adisconnect(self) -> None:
-        """Asynchronously close the Redis connection."""
-        self.disconnect()
+        """Asynchronously close the connection."""
+        if self._api.sdk_configuration.async_client is not None:
+            await self._api.sdk_configuration.async_client.aclose()
+        if self._api.sdk_configuration.client is not None:
+            self._api.sdk_configuration.client.close()
 
     def __enter__(self):
+        self._api.__enter__()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.disconnect()
+        self._api.__exit__(exc_type, exc_val, exc_tb)
 
     async def __aenter__(self):
+        await self._api.__aenter__()
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        await self.adisconnect()
+        await self._api.__aexit__(exc_type, exc_val, exc_tb)