Production RAG: Semantic Caching for Sub-50ms LLM Responses

graph TD
    A[User Query] --> B{Layer 1: Lexical Cache Check};
    B -- Hit --> C[Return Cached Response];
    B -- Miss --> D{Layer 2: Semantic Cache Check};
    D -- Hit --> C;
    D -- Miss --> E[Execute Full RAG Pipeline];
    E --> F[LLM Generates Response];
    F --> G{Cache Write Logic};
    G --> H[Return Response to User];
    G --> I[Populate L1 & L2 Caches];

import redis
import hashlib
import json

# Connect to Redis
redis_client = redis.Redis(host='localhost', port=6379, db=0, decode_responses=True)

def get_lexical_cache_key(query: str) -> str:
    """Creates a normalized, order-agnostic key for the query."""
    normalized_query = ''.join(e for e in query.lower() if e.isalnum() or e.isspace())
    # Sorting words makes it resilient to word order changes
    sorted_words = sorted(normalized_query.split())
    canonical_query = " ".join(sorted_words)
    # Use a hash to ensure a fixed-length key
    return f"lexical_cache:{hashlib.sha256(canonical_query.encode()).hexdigest()}"

def get_from_lexical_cache(query: str):
    """Attempt to retrieve a response from the lexical cache."""
    key = get_lexical_cache_key(query)
    cached_result = redis_client.get(key)
    if cached_result:
        print("--- Lexical Cache HIT ---")
        return json.loads(cached_result)
    print("--- Lexical Cache MISS ---")
    return None

def set_in_lexical_cache(query: str, response: dict, ttl_seconds: int = 3600):
    """Store a response in the lexical cache with a TTL."""
    key = get_lexical_cache_key(query)
    redis_client.setex(key, ttl_seconds, json.dumps(response))
    print(f"--- Lexical Cache SET for key: {key} ---")

# Example Usage
query1 = "How can I reset my account password?"
query2 = "reset my password for my account how?"

# First request (miss)
response = get_from_lexical_cache(query1)

# Assume we run the full RAG pipeline and get a response
if not response:
    rag_response = {"answer": "To reset your password, go to the settings page and click 'Reset Password'.", "source_docs": ["doc1.txt"]}
    set_in_lexical_cache(query1, rag_response)

# Second, semantically identical request (now a hit)
response = get_from_lexical_cache(query2)
print(response)

import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
import uuid
import json
import redis

# --- Initialization ---

# 1. Load a fast embedding model
print("Loading embedding model...")
cache_embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
embedding_dim = cache_embedding_model.get_sentence_embedding_dimension()

# 2. Initialize FAISS index
# Using IndexFlatL2 is simple, but for production, IndexHNSWFlat is faster for search.
# We will use IndexIDMap to map FAISS's internal IDs to our own UUIDs.
index = faiss.IndexFlatL2(embedding_dim)
index_map = faiss.IndexIDMap(index)

# 3. Connect to Redis for the Answer Store
redis_client = redis.Redis(host='localhost', port=6379, db=1, decode_responses=True)

# --- Global State (in a real app, this would be managed properly) ---
# In-memory mapping from our integer IDs to string UUIDs for this example
# A production system might use a persistent store or a different FAISS index type
faiss_id_counter = 0
id_to_uuid_map = {}

# --- Core Functions ---

def search_semantic_cache(query: str, threshold: float = 0.98):
    """Search for a semantically similar query in the cache."""
    global index_map, id_to_uuid_map

    if index_map.ntotal == 0:
        print("--- Semantic Cache is empty. MISS ---")
        return None

    query_embedding = cache_embedding_model.encode([query], convert_to_numpy=True)
    # FAISS uses L2 distance, SentenceTransformers embeddings are normalized to length 1.
    # For normalized vectors, L2_distance^2 = 2 * (1 - cosine_similarity).
    # So, we search for a small L2 distance.
    distance_threshold = np.sqrt(2 * (1 - threshold))

    # Search for the 1 nearest neighbor
    distances, ids = index_map.search(query_embedding, k=1)

    if ids[0][0] != -1 and distances[0][0] < distance_threshold:
        print(f"--- Semantic Cache HIT with distance {distances[0][0]} ---")
        matched_faiss_id = ids[0][0]
        matched_uuid = id_to_uuid_map[matched_faiss_id]
        cached_response = redis_client.get(f"semantic_cache:{matched_uuid}")
        return json.loads(cached_response) if cached_response else None
    
    print("--- Semantic Cache MISS ---")
    return None

def add_to_semantic_cache(query: str, response: dict):
    """Add a new query-answer pair to the cache."""
    global faiss_id_counter, index_map, id_to_uuid_map

    query_embedding = cache_embedding_model.encode([query], convert_to_numpy=True)
    new_uuid = str(uuid.uuid4())
    new_faiss_id = faiss_id_counter

    # Add to FAISS index
    index_map.add_with_ids(query_embedding, np.array([new_faiss_id]))

    # Update mappings
    id_to_uuid_map[new_faiss_id] = new_uuid
    faiss_id_counter += 1

    # Add to Redis answer store
    redis_client.set(f"semantic_cache:{new_uuid}", json.dumps(response))
    print(f"--- Semantic Cache SET for UUID: {new_uuid} ---")

# --- Example Flow ---

query1 = "What is the process for submitting an expense report?"
query2 = "How do I file my expenses?"
query3 = "What's the weather like today?"

# 1. First query - miss
response = search_semantic_cache(query1)
if not response:
    # Simulate full RAG pipeline
    rag_response = {"answer": "To submit an expense report, log into the portal, navigate to 'My Expenses', and click 'New Report'.", "source_docs": ["hr_manual_p42.pdf"]}
    add_to_semantic_cache(query1, rag_response)

# 2. Second, semantically similar query - hit!
response = search_semantic_cache(query2)
if response:
    print("Cached response:", response['answer'])

# 3. Third, unrelated query - miss
response = search_semantic_cache(query3)

    // In Redis Answer Store
    "semantic_cache:<uuid>": {
        "answer": "...",
        "source_chunks": ["doc_abc_chunk_3", "doc_abc_chunk_5"]
    }

# In a production FAISS setup
hnsw_index = faiss.IndexHNSWFlat(embedding_dim, 32)
# ... add vectors to hnsw_index

# At search time:
# For a high-priority query where accuracy is key
hnsw_index.hnsw.efSearch = 128 
distances, ids = hnsw_index.search(query_embedding, k=1)

# For a low-priority, high-throughput endpoint
hnsw_index.hnsw.efSearch = 16
distances, ids = hnsw_index.search(query_embedding, k=1)

    import openai
    
    def synthesize_questions_from_chunk(document_chunk: str):
        response = openai.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Given the following text, generate 3 distinct questions that this text could answer. Respond with a JSON list of strings."},
                {"role": "user", "content": document_chunk}
            ]
        )
        try:
            questions = json.loads(response.choices[0].message.content)
            return questions
        except:
            return []

    # chunk = "The SOC 2 Type 2 report is an audit of a service organization's systems..."
    # generated_questions = synthesize_questions_from_chunk(chunk)
    # # Now run these questions through your RAG pipeline to populate the cache.

# --- This is a simulation of the full end-to-end flow ---
import time

# Assume all previous functions (get_lexical_cache_key, etc.) are defined

def full_rag_pipeline(query: str) -> dict:
    """Simulates the slow, expensive RAG pipeline."""
    print("--- EXECUTING FULL RAG PIPELINE (CACHE MISS) ---")
    time.sleep(2.5) # Simulate embedding, vector search, LLM call
    return {"answer": f"This is a freshly generated answer for: '{query}'", "source_docs": ["doc_xyz.pdf"]}

def handle_query_request(query: str):
    """The main entry point for a user query."""
    start_time = time.time()

    # 1. Check Layer 1: Lexical Cache
    response = get_from_lexical_cache(query)
    if response:
        end_time = time.time()
        print(f"Total response time: {(end_time - start_time) * 1000:.2f} ms")
        return response

    # 2. Check Layer 2: Semantic Cache
    response = search_semantic_cache(query, threshold=0.98)
    if response:
        end_time = time.time()
        print(f"Total response time: {(end_time - start_time) * 1000:.2f} ms")
        return response

    # 3. Cache Miss: Execute full pipeline
    response = full_rag_pipeline(query)

    # 4. Populate caches for future requests
    set_in_lexical_cache(query, response)
    add_to_semantic_cache(query, response)

    end_time = time.time()
    print(f"Total response time: {(end_time - start_time) * 1000:.2f} ms")
    return response

# --- Benchmark Run ---

print("\n--- Run 1: Cold query ---")
query_a = "How do I configure the SSO integration?"
handle_query_request(query_a)

print("\n--- Run 2: Identical query (Lexical Hit) ---")
handle_query_request(query_a)

print("\n--- Run 3: Semantically similar query (Semantic Hit) ---")
query_b = "What are the steps to set up single sign-on?"
handle_query_request(query_b)

print("\n--- Run 4: Another cold query ---")
query_c = "What is our data retention policy?"
handle_query_request(query_c)