src/utils/caching.py

"""
Caching utilities for QAFD-RAG.

Provides embedding-based caching with similarity matching.
"""

import json
from dataclasses import dataclass
from typing import Optional, Union

import numpy as np

from .logging import logger


def cosine_similarity(v1: np.ndarray, v2: np.ndarray) -> float:
    """
    Compute cosine similarity between two vectors.

    Parameters:
    -----------
    v1, v2 : np.ndarray
        Input vectors

    Returns:
    --------
    float
        Cosine similarity score
    """
    dot_product = np.dot(v1, v2)
    norm1 = np.linalg.norm(v1)
    norm2 = np.linalg.norm(v2)
    return dot_product / (norm1 * norm2)


def quantize_embedding(embedding: np.ndarray, bits: int = 8) -> tuple:
    """
    Quantize an embedding to reduce storage size.

    Parameters:
    -----------
    embedding : np.ndarray
        Float embedding vector
    bits : int
        Bit depth for quantization (default: 8)

    Returns:
    --------
    tuple
        (quantized_array, min_val, max_val)
    """
    min_val = embedding.min()
    max_val = embedding.max()

    scale = (2**bits - 1) / (max_val - min_val)
    quantized = np.round((embedding - min_val) * scale).astype(np.uint8)

    return quantized, min_val, max_val


def dequantize_embedding(
    quantized: np.ndarray, min_val: float, max_val: float, bits: int = 8
) -> np.ndarray:
    """
    Dequantize an embedding back to float values.

    Parameters:
    -----------
    quantized : np.ndarray
        Quantized embedding
    min_val : float
        Original minimum value
    max_val : float
        Original maximum value
    bits : int
        Bit depth used in quantization

    Returns:
    --------
    np.ndarray
        Reconstructed float embedding
    """
    scale = (max_val - min_val) / (2**bits - 1)
    return (quantized * scale + min_val).astype(np.float32)


async def get_best_cached_response(
    hashing_kv,
    current_embedding,
    similarity_threshold: float = 0.95,
    mode: str = "default",
    use_llm_check: bool = False,
    llm_func=None,
    original_prompt: Optional[str] = None,
) -> Union[str, None]:
    """
    Find the best matching cached response based on embedding similarity.

    Parameters:
    -----------
    hashing_kv : BaseKVStorage
        Key-value storage for cache
    current_embedding : np.ndarray
        Embedding of the current query
    similarity_threshold : float
        Minimum similarity score to accept
    mode : str
        Cache mode/namespace
    use_llm_check : bool
        Whether to use LLM for additional similarity verification
    llm_func : callable, optional
        LLM function for similarity checking
    original_prompt : str, optional
        Original prompt for LLM comparison

    Returns:
    --------
    str or None
        Cached response if found, None otherwise
    """
    # Import here to avoid circular import
    from ..prompts import PROMPTS

    mode_cache = await hashing_kv.get_by_id(mode)
    if not mode_cache:
        return None

    best_similarity = -1
    best_response = None
    best_prompt = None
    best_cache_id = None

    for cache_id, cache_data in mode_cache.items():
        if cache_data["embedding"] is None:
            continue

        cached_quantized = np.frombuffer(
            bytes.fromhex(cache_data["embedding"]), dtype=np.uint8
        ).reshape(cache_data["embedding_shape"])
        cached_embedding = dequantize_embedding(
            cached_quantized,
            cache_data["embedding_min"],
            cache_data["embedding_max"],
        )

        similarity = cosine_similarity(current_embedding, cached_embedding)
        if similarity > best_similarity:
            best_similarity = similarity
            best_response = cache_data["return"]
            best_prompt = cache_data["original_prompt"]
            best_cache_id = cache_id

    if best_similarity > similarity_threshold:
        if use_llm_check and llm_func and original_prompt and best_prompt:
            compare_prompt = PROMPTS["similarity_check"].format(
                original_prompt=original_prompt, cached_prompt=best_prompt
            )

            try:
                llm_result = await llm_func(compare_prompt)
                llm_result = llm_result.strip()
                llm_similarity = float(llm_result)

                best_similarity = llm_similarity
                if best_similarity < similarity_threshold:
                    log_data = {
                        "event": "llm_check_cache_rejected",
                        "original_question": original_prompt[:100] + "..."
                        if len(original_prompt) > 100
                        else original_prompt,
                        "cached_question": best_prompt[:100] + "..."
                        if len(best_prompt) > 100
                        else best_prompt,
                        "similarity_score": round(best_similarity, 4),
                        "threshold": similarity_threshold,
                    }
                    logger.info(json.dumps(log_data, ensure_ascii=False))
                    return None
            except Exception as e:
                logger.warning(f"LLM similarity check failed: {e}")
                return None

        prompt_display = (
            best_prompt[:50] + "..." if len(best_prompt) > 50 else best_prompt
        )
        log_data = {
            "event": "cache_hit",
            "mode": mode,
            "similarity": round(best_similarity, 4),
            "cache_id": best_cache_id,
            "original_prompt": prompt_display,
        }
        logger.info(json.dumps(log_data, ensure_ascii=False))
        return best_response
    return None


async def handle_cache(hashing_kv, args_hash: str, prompt: str, mode: str = "default"):
    """
    Handle cache lookup with optional embedding-based matching.

    Parameters:
    -----------
    hashing_kv : BaseKVStorage
        Key-value storage for cache
    args_hash : str
        Hash of the arguments for exact matching
    prompt : str
        The prompt text for embedding-based matching
    mode : str
        Cache mode/namespace

    Returns:
    --------
    tuple
        (cached_response, quantized, min_val, max_val)
    """
    if hashing_kv is None:
        return None, None, None, None

    if mode == "naive":
        mode_cache = await hashing_kv.get_by_id(mode) or {}
        if args_hash in mode_cache:
            return mode_cache[args_hash]["return"], None, None, None
        return None, None, None, None

    embedding_cache_config = hashing_kv.global_config.get(
        "embedding_cache_config",
        {"enabled": False, "similarity_threshold": 0.95, "use_llm_check": False},
    )
    is_embedding_cache_enabled = embedding_cache_config["enabled"]
    use_llm_check = embedding_cache_config.get("use_llm_check", False)

    quantized = min_val = max_val = None
    if is_embedding_cache_enabled:
        embedding_model_func = hashing_kv.global_config["embedding_func"]["func"]
        llm_model_func = hashing_kv.global_config.get("llm_model_func")

        current_embedding = await embedding_model_func([prompt])
        quantized, min_val, max_val = quantize_embedding(current_embedding[0])
        best_cached_response = await get_best_cached_response(
            hashing_kv,
            current_embedding[0],
            similarity_threshold=embedding_cache_config["similarity_threshold"],
            mode=mode,
            use_llm_check=use_llm_check,
            llm_func=llm_model_func if use_llm_check else None,
            original_prompt=prompt if use_llm_check else None,
        )
        if best_cached_response is not None:
            return best_cached_response, None, None, None
    else:
        mode_cache = await hashing_kv.get_by_id(mode) or {}
        if args_hash in mode_cache:
            return mode_cache[args_hash]["return"], None, None, None

    return None, quantized, min_val, max_val


@dataclass
class CacheData:
    """Data structure for cache entries."""

    args_hash: str
    content: str
    prompt: str
    quantized: Optional[np.ndarray] = None
    min_val: Optional[float] = None
    max_val: Optional[float] = None
    mode: str = "default"


async def save_to_cache(hashing_kv, cache_data: CacheData):
    """
    Save data to the cache.

    Parameters:
    -----------
    hashing_kv : BaseKVStorage
        Key-value storage for cache
    cache_data : CacheData
        Data to cache
    """
    if hashing_kv is None or hasattr(cache_data.content, "__aiter__"):
        return

    mode_cache = await hashing_kv.get_by_id(cache_data.mode) or {}

    mode_cache[cache_data.args_hash] = {
        "return": cache_data.content,
        "embedding": cache_data.quantized.tobytes().hex()
        if cache_data.quantized is not None
        else None,
        "embedding_shape": cache_data.quantized.shape
        if cache_data.quantized is not None
        else None,
        "embedding_min": cache_data.min_val,
        "embedding_max": cache_data.max_val,
        "original_prompt": cache_data.prompt,
    }

    await hashing_kv.upsert({cache_data.mode: mode_cache})


__all__ = [
    "cosine_similarity",
    "quantize_embedding",
    "dequantize_embedding",
    "get_best_cached_response",
    "handle_cache",
    "CacheData",
    "save_to_cache",
]