cache/query_cache.py

# COST: ZERO OpenAI tokens.
# All embedding uses sentence-transformers 'all-MiniLM-L6-v2' running locally.
# Similarity scoring is pure-Python cosine distance against diskcache entries.
# No network call is made during cache lookup or save.

from datetime import datetime, timezone
from pathlib import Path
import math
import uuid

import diskcache

STORE_PATH = Path(__file__).parent / "query_store"
_cache = diskcache.Cache(str(STORE_PATH))

# --- model singleton ---
_model = None


def _get_model():
    global _model
    if _model is None:
        from sentence_transformers import SentenceTransformer
        _model = SentenceTransformer("all-MiniLM-L6-v2")
    return _model


# --- math helpers ---

def _cosine_similarity(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    mag_a = math.sqrt(sum(x * x for x in a))
    mag_b = math.sqrt(sum(y * y for y in b))
    if mag_a == 0 or mag_b == 0:
        return 0.0
    return dot / (mag_a * mag_b)


# --- public API ---

def embed_query_local(query: str) -> list[float]:
    return _get_model().encode(query, convert_to_numpy=True).tolist()


def find_similar_cached(query: str, threshold: float = 0.92) -> dict | None:
    query_emb = embed_query_local(query)

    best_key = None
    best_score = -1.0

    for key in _cache:
        entry = _cache[key]
        score = _cosine_similarity(query_emb, entry["query_embedding"])
        if score > best_score:
            best_score = score
            best_key = key

    if best_key is None or best_score < threshold:
        return None

    entry = _cache[best_key]
    entry["hit_count"] += 1
    _cache[best_key] = entry
    return entry


def save_to_cache(query: str, answer: str, sources: list[str]) -> None:
    entry = {
        "query": query,
        "query_embedding": embed_query_local(query),
        "answer": answer,
        "sources": sources,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "hit_count": 0,
    }
    _cache[str(uuid.uuid4())] = entry


def get_cache_stats() -> dict:
    entries = [_cache[k] for k in _cache]
    total_hits = sum(e["hit_count"] for e in entries)
    most_asked = [
        e["query"]
        for e in sorted(entries, key=lambda e: e["hit_count"], reverse=True)[:5]
    ]
    return {
        "total_entries": len(entries),
        "total_hits": total_hits,
        "most_asked": most_asked,
    }