Source code for langchain_couchbase.cache

"""
LangChain Couchbase Caches

Functions "_hash", "_loads_generations" and "_dumps_generations"
are duplicated in this utility from modules:
    - "libs/community/langchain_community/cache.py"
"""

import hashlib
import json
import logging
from typing import Any, Dict, Optional, Union

from couchbase.cluster import Cluster
from langchain_core.caches import RETURN_VAL_TYPE, BaseCache
from langchain_core.embeddings import Embeddings
from langchain_core.load.dump import dumps
from langchain_core.load.load import loads
from langchain_core.outputs import Generation

from langchain_couchbase.vectorstores import CouchbaseVectorStore

logger = logging.getLogger(__file__)


def _hash(_input: str) -> str:
    """Use a deterministic hashing approach."""
    return hashlib.md5(_input.encode()).hexdigest()


def _dumps_generations(generations: RETURN_VAL_TYPE) -> str:
    """
    Serialization for generic RETURN_VAL_TYPE, i.e. sequence of `Generation`

    Args:
        generations (RETURN_VAL_TYPE): A list of language model generations.

    Returns:
        str: a single string representing a list of generations.

    This function (+ its counterpart `_loads_generations`) rely on
    the dumps/loads pair with Reviver, so are able to deal
    with all subclasses of Generation.

    Each item in the list can be `dumps`ed to a string,
    then we make the whole list of strings into a json-dumped.
    """
    return json.dumps([dumps(_item) for _item in generations])


def _loads_generations(generations_str: str) -> Union[RETURN_VAL_TYPE, None]:
    """
    Deserialization of a string into a generic RETURN_VAL_TYPE
    (i.e. a sequence of `Generation`).

    See `_dumps_generations`, the inverse of this function.

    Args:
        generations_str (str): A string representing a list of generations.

    Compatible with the legacy cache-blob format
    Does not raise exceptions for malformed entries, just logs a warning
    and returns none: the caller should be prepared for such a cache miss.

    Returns:
        RETURN_VAL_TYPE: A list of generations.
    """
    try:
        generations = [loads(_item_str) for _item_str in json.loads(generations_str)]
        return generations
    except (json.JSONDecodeError, TypeError):
        # deferring the (soft) handling to after the legacy-format attempt
        pass

    try:
        gen_dicts = json.loads(generations_str)
        # not relying on `_load_generations_from_json` (which could disappear):
        generations = [Generation(**generation_dict) for generation_dict in gen_dicts]
        logger.warning(
            f"Legacy 'Generation' cached blob encountered: '{generations_str}'"
        )
        return generations
    except (json.JSONDecodeError, TypeError):
        logger.warning(
            f"Malformed/unparsable cached blob encountered: '{generations_str}'"
        )
        return None


[docs] class CouchbaseCache(BaseCache): """Couchbase LLM Cache LLM Cache that uses Couchbase as the backend """ PROMPT = "prompt" LLM = "llm" RETURN_VAL = "return_val" def _check_bucket_exists(self) -> bool: """Check if the bucket exists in the linked Couchbase cluster""" bucket_manager = self._cluster.buckets() try: bucket_manager.get_bucket(self._bucket_name) return True except Exception: return False def _check_scope_and_collection_exists(self) -> bool: """Check if the scope and collection exists in the linked Couchbase bucket Raises a ValueError if either is not found""" scope_collection_map: Dict[str, Any] = {} # Get a list of all scopes in the bucket for scope in self._bucket.collections().get_all_scopes(): scope_collection_map[scope.name] = [] # Get a list of all the collections in the scope for collection in scope.collections: scope_collection_map[scope.name].append(collection.name) # Check if the scope exists if self._scope_name not in scope_collection_map.keys(): raise ValueError( f"Scope {self._scope_name} not found in Couchbase " f"bucket {self._bucket_name}" ) # Check if the collection exists in the scope if self._collection_name not in scope_collection_map[self._scope_name]: raise ValueError( f"Collection {self._collection_name} not found in scope " f"{self._scope_name} in Couchbase bucket {self._bucket_name}" ) return True
[docs] def __init__( self, cluster: Cluster, bucket_name: str, scope_name: str, collection_name: str, **kwargs: Dict[str, Any], ) -> None: """Initialize the Couchbase LLM Cache Args: cluster (Cluster): couchbase cluster object with active connection. bucket_name (str): name of the bucket to store documents in. scope_name (str): name of the scope in bucket to store documents in. collection_name (str): name of the collection in the scope to store documents in. """ if not isinstance(cluster, Cluster): raise ValueError( f"cluster should be an instance of couchbase.Cluster, " f"got {type(cluster)}" ) self._cluster = cluster self._bucket_name = bucket_name self._scope_name = scope_name self._collection_name = collection_name # Check if the bucket exists if not self._check_bucket_exists(): raise ValueError( f"Bucket {self._bucket_name} does not exist. " " Please create the bucket before searching." ) try: self._bucket = self._cluster.bucket(self._bucket_name) self._scope = self._bucket.scope(self._scope_name) self._collection = self._scope.collection(self._collection_name) except Exception as e: raise ValueError( "Error connecting to couchbase. " "Please check the connection and credentials." ) from e # Check if the scope and collection exists. Throws ValueError if they don't try: self._check_scope_and_collection_exists() except Exception as e: raise e
[docs] def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]: """Look up from cache based on prompt and llm_string.""" try: doc = self._collection.get( self._generate_key(prompt, llm_string) ).content_as[dict] return _loads_generations(doc[self.RETURN_VAL]) except Exception: return None
def _generate_key(self, prompt: str, llm_string: str) -> str: """Generate the key based on prompt and llm_string.""" return _hash(prompt + llm_string)
[docs] def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None: """Update cache based on prompt and llm_string.""" doc = { self.PROMPT: prompt, self.LLM: llm_string, self.RETURN_VAL: _dumps_generations(return_val), } try: self._collection.upsert( key=self._generate_key(prompt, llm_string), value=doc ) except Exception: logger.error("Error updating cache")
[docs] def clear(self, **kwargs: Any) -> None: """Clear the cache. This will delete all documents in the collection. This requires an index on the collection. """ try: query = f"DELETE FROM `{self._collection_name}`" self._scope.query(query).execute() except Exception: logger.error("Error clearing cache. Please check if you have an index.")
[docs] class CouchbaseSemanticCache(BaseCache, CouchbaseVectorStore): """Couchbase Semantic Cache Cache backed by a Couchbase Server with Vector Store support """ LLM = "llm_string" RETURN_VAL = "return_val"
[docs] def __init__( self, cluster: Cluster, embedding: Embeddings, bucket_name: str, scope_name: str, collection_name: str, index_name: str, score_threshold: Optional[float] = None, ) -> None: """Initialize the Couchbase LLM Cache Args: cluster (Cluster): couchbase cluster object with active connection. embedding (Embeddings): embedding model to use. bucket_name (str): name of the bucket to store documents in. scope_name (str): name of the scope in bucket to store documents in. collection_name (str): name of the collection in the scope to store documents in. index_name (str): name of the Search index to use. score_threshold (float): score threshold to use for filtering results. """ if not isinstance(cluster, Cluster): raise ValueError( f"cluster should be an instance of couchbase.Cluster, " f"got {type(cluster)}" ) self._cluster = cluster self._bucket_name = bucket_name self._scope_name = scope_name self._collection_name = collection_name # Check if the bucket exists if not self._check_bucket_exists(): raise ValueError( f"Bucket {self._bucket_name} does not exist. " " Please create the bucket before searching." ) try: self._bucket = self._cluster.bucket(self._bucket_name) self._scope = self._bucket.scope(self._scope_name) self._collection = self._scope.collection(self._collection_name) except Exception as e: raise ValueError( "Error connecting to couchbase. " "Please check the connection and credentials." ) from e # Check if the scope and collection exists. Throws ValueError if they don't try: self._check_scope_and_collection_exists() except Exception as e: raise e self.score_threshold = score_threshold # Initialize the vector store super().__init__( cluster=cluster, bucket_name=bucket_name, scope_name=scope_name, collection_name=collection_name, embedding=embedding, index_name=index_name, )
[docs] def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]: """Look up from cache based on the semantic similarity of the prompt""" search_results = self.similarity_search_with_score( prompt, k=1, search_options={f"metadata.{self.LLM}": llm_string} ) if search_results: selected_doc, score = search_results[0] else: return None # Check if the score is above the threshold if a threshold is provided if self.score_threshold: if score < self.score_threshold: return None # Note that the llm_string might not match the vector search result. # So if the llm_string does not match, do not return the result. if selected_doc.metadata["llm_string"] != llm_string: return None return _loads_generations(selected_doc.metadata[self.RETURN_VAL])
[docs] def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None: """Update cache based on the prompt and llm_string""" try: self.add_texts( texts=[prompt], metadatas=[ { self.LLM: llm_string, self.RETURN_VAL: _dumps_generations(return_val), } ], ) except Exception: logger.error("Error updating cache")
[docs] def clear(self, **kwargs: Any) -> None: """Clear the cache. This will delete all documents in the collection. This requires an index on the collection. """ try: query = f"DELETE FROM `{self._collection_name}`" self._scope.query(query).execute() except Exception: logger.error("Error clearing cache. Please check if you have an index.")