elasticsearch_llm_cache.py

"""
Elasticsearch LLM Cache Library
==================================
This library provides an Elasticsearch-based caching mechanism for Language Model (LLM) responses.
Through the ElasticsearchLLMCache class, it facilitates the creation, querying, and updating
of a cache index to store and retrieve LLM responses based on user prompts.

Key Features:
-------------
- Initialize a cache index with specified or default settings.
- Create the cache index with specified mappings if it does not already exist.
- Query the cache for similar prompts using a k-NN (k-Nearest Neighbors) search.
- Update the 'last_hit_date' field of a document when a cache hit occurs.
- Generate a vector for a given prompt using Elasticsearch's text embedding.
- Add new documents (prompts and responses) to the cache.

Requirements:
-------------
- Elasticsearch
- Python 3.6+
- elasticsearch-py library

Usage Example:
--------------
```python
from elasticsearch import Elasticsearch
from elasticsearch_llm_cache import ElasticsearchLLMCache

# Initialize Elasticsearch client
es_client = Elasticsearch()

# Initialize the ElasticsearchLLMCache instance
llm_cache = ElasticsearchLLMCache(es_client)

# Query the cache
prompt_text = "What is the capital of France?"
query_result = llm_cache.query(prompt_text)

# Add to cache
prompt = "What is the capital of France?"
response = "Paris"
add_result = llm_cache.add(prompt, response)
```

This library is covered in depth in the blog post
Elasticsearch as a GenAI Caching Layer
https://www.elastic.co/search-labs/elasticsearch-as-a-genai-caching-layer

Author: Jeff Vestal
Version: 1.0.0

"""

from datetime import datetime
from typing import Dict, List, Optional
from elasticsearch import Elasticsearch
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class ElasticsearchLLMCache:
    def __init__(
        self,
        es_client: Elasticsearch,
        index_name: Optional[str] = None,
        es_model_id: Optional[str] = "sentence-transformers__all-distilroberta-v1",
        create_index=True,
    ):
        """
        Initialize the ElasticsearchLLMCache instance.

        :param es_client: Elasticsearch client object.
        :param index_name: Optional name for the index; defaults to 'llm_cache'.
        :param es_model_id: Model ID for text embedding; defaults to 'sentence-transformers__all-distilroberta-v1'.
        :param create_index: Boolean to determine whether to create a new index; defaults to True.
        """
        self.es = es_client
        self.index_name = index_name or "llm_cache"
        self.es_model_id = es_model_id
        if create_index:
            self.create_index()

    def create_index(self, dims: Optional[int] = 768) -> Dict:
        """
        Create the index if it does not already exist.

        :return: Dictionary containing information about the index creation.
        """
        if not self.es.indices.exists(index=self.index_name):
            mappings = {
                "mappings": {
                    "properties": {
                        "prompt": {"type": "text"},
                        "response": {"type": "text"},
                        "create_date": {"type": "date"},
                        "last_hit_date": {"type": "date"},
                        "prompt_vector": {
                            "type": "dense_vector",
                            "dims": dims,
                            "index": True,
                            "similarity": "dot_product",
                        },
                    }
                }
            }

            self.es.indices.create(index=self.index_name, body=mappings, ignore=400)
            logger.info(f"Index {self.index_name} created.")

            return {"cache_index": self.index_name, "created_new": True}
        else:
            logger.info(f"Index {self.index_name} already exists.")
            return {"cache_index": self.index_name, "created_new": False}

    def update_last_hit_date(self, doc_id: str):
        """
        Update the 'last_hit_date' field of a document to the current datetime.

        :param doc_id: The ID of the document to update.
        """
        update_body = {"doc": {"last_hit_date": datetime.now()}}
        self.es.update(index=self.index_name, id=doc_id, body=update_body)

    def query(
        self,
        prompt_text: str,
        similarity_threshold: Optional[float] = 0.5,
        num_candidates: Optional[int] = 1000,
        create_date_gte: Optional[str] = "now-1y/y",
    ) -> dict:
        """
        Query the index to find similar prompts and update the `last_hit_date` for that document if a hit is found.

        :param prompt_text: The text of the prompt to find similar entries for.
        :param similarity_threshold: The similarity threshold for filtering results; defaults to 0.5.
        :param num_candidates: The number of candidates to consider; defaults to 1000.
        :param create_date_gte: The date range to consider results; defaults to "now-1y/y".
        :return: A dictionary containing the hits or an empty dictionary if no hits are found.
        """
        knn = [
            {
                "field": "prompt_vector",
                "k": 1,
                "num_candidates": num_candidates,
                "similarity": similarity_threshold,
                "query_vector_builder": {
                    "text_embedding": {
                        "model_id": self.es_model_id,
                        "model_text": prompt_text,
                    }
                },
                "filter": {"range": {"create_date": {"gte": create_date_gte}}},
            }
        ]

        fields = ["prompt", "response"]

        resp = self.es.search(
            index=self.index_name, knn=knn, fields=fields, size=1, source=False
        )

        if resp["hits"]["total"]["value"] == 0:
            return {}
        else:
            doc_id = resp["hits"]["hits"][0]["_id"]
            self.update_last_hit_date(doc_id)
            return resp["hits"]["hits"][0]["fields"]

    def _generate_vector(self, prompt: str) -> List[float]:
        """
        Generate a vector for a given prompt using Elasticsearch's text embedding.

        :param prompt: The text prompt to generate a vector for.
        :return: A list of floats representing the vector.
        """
        docs = [{"text_field": prompt}]

        embedding = self.es.ml.infer_trained_model(model_id=self.es_model_id, docs=docs)

        return embedding["inference_results"][0]["predicted_value"]

    def add(self, prompt: str, response: str, source: Optional[str] = None) -> Dict:
        """
        Add a new document to the index.

        :param prompt: The user prompt.
        :param response: The LLM response.
        :param source: Optional source identifier for the LLM.
        :return: A dictionary indicating the successful caching of the new prompt and response.
        """
        prompt_vector = self._generate_vector(prompt=prompt)

        doc = {
            "prompt": prompt,
            "response": response,
            "create_date": datetime.now(),
            "last_hit_date": datetime.now(),
            "prompt_vector": prompt_vector,
            "source": source,  # Optional
        }
        try:
            self.es.index(index=self.index_name, document=doc)
            return {"success": True}
        except Exception as e:
            logger.error(e)
            return {"success": False, "error": e}