Prompt caching & result caching strategies

# openai>=1.0.0, redis>=5.0.0, numpy>=1.26.0
import hashlib, json, os, time, logging, struct
from contextlib import contextmanager
from typing import Optional
import numpy as np
import redis
from openai import OpenAI, RateLimitError, APITimeoutError, APIStatusError

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

client = OpenAI(api_key=os.environ["OPENAI_API_KEY"], timeout=30.0)
r = redis.Redis(
    host=os.environ.get("REDIS_HOST", "localhost"),
    port=int(os.environ.get("REDIS_PORT", 6379)),
    decode_responses=False,  # bytes for embedding storage
)

EXACT_TTL = int(os.environ.get("EXACT_CACHE_TTL", 3600))
SEMAN_TTL = int(os.environ.get("SEMAN_CACHE_TTL", 7200))
SIM_THRESHOLD = float(os.environ.get("SIM_THRESHOLD", 0.93))
COMPLETION_MODEL = os.environ.get("COMPLETION_MODEL", "gpt-4o-mini")
EMBED_MODEL = os.environ.get("EMBED_MODEL", "text-embedding-3-small")

def _normalize(text: str) -> str:
    return " ".join(text.strip().lower().split())

def _exact_key(system: str, user: str, model: str) -> str:
    payload = json.dumps({"model": model, "system": system, "user": _normalize(user)}, sort_keys=True)
    return f"llm:exact:{hashlib.sha256(payload.encode()).hexdigest()}"

def _get_embedding(text: str) -> np.ndarray:
    """Embed text, with one retry on rate limit."""
    for attempt in range(2):
        try:
            resp = client.embeddings.create(model=EMBED_MODEL, input=text)
            return np.array(resp.data[0].embedding, dtype=np.float32)
        except RateLimitError:
            if attempt == 0:
                time.sleep(2)
            else:
                raise

def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9))

def _semantic_lookup(user: str) -> Optional[str]:
    """Scan semantic cache keys and return best match above threshold."""
    query_vec = _get_embedding(user)
    best_score, best_val = 0.0, None
    for key in r.scan_iter("llm:seman:*"):
        raw = r.hgetall(key)
        if not raw:
            continue
        stored_vec = np.frombuffer(raw[b"vec"], dtype=np.float32)
        score = _cosine(query_vec, stored_vec)
        if score > best_score:
            best_score, best_val = score, raw[b"answer"].decode()
    if best_score >= SIM_THRESHOLD:
        log.info("semantic_hit score=%.4f", best_score)
        return best_val
    return None

def _store_semantic(user: str, answer: str) -> None:
    vec = _get_embedding(user)
    key = f"llm:seman:{hashlib.sha256(user.encode()).hexdigest()}"
    pipe = r.pipeline()
    pipe.hset(key, mapping={"vec": vec.tobytes(), "answer": answer.encode()})
    pipe.expire(key, SEMAN_TTL)
    pipe.execute()

def completion_with_cache(
    user_prompt: str,
    system_prompt: str,
    use_semantic: bool = True,
) -> dict:
    start = time.monotonic()
    exact_key = _exact_key(system_prompt, user_prompt, COMPLETION_MODEL)

    # Layer 1: exact-match cache
    cached = r.get(exact_key)
    if cached:
        log.info("cache=exact latency_ms=%.0f", (time.monotonic() - start) * 1000)
        return {"answer": cached.decode(), "cache": "exact", "tokens": 0}

    # Layer 2: semantic cache
    if use_semantic:
        sem_hit = _semantic_lookup(user_prompt)
        if sem_hit:
            log.info("cache=semantic latency_ms=%.0f", (time.monotonic() - start) * 1000)
            return {"answer": sem_hit, "cache": "semantic", "tokens": 0}

    # Layer 3: live LLM call with retries
    last_exc = None
    for attempt in range(3):
        try:
            resp = client.chat.completions.create(
                model=COMPLETION_MODEL,
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_prompt},
                ],
                timeout=25.0,
            )
            answer = resp.choices[0].message.content
            total_tokens = resp.usage.total_tokens
            log.info(
                "cache=miss tokens=%d latency_ms=%.0f attempt=%d",
                total_tokens, (time.monotonic() - start) * 1000, attempt + 1,
            )
            # Store in both cache layers
            r.setex(exact_key, EXACT_TTL, answer.encode())
            if use_semantic:
                _store_semantic(user_prompt, answer)
            return {"answer": answer, "cache": "miss", "tokens": total_tokens}
        except RateLimitError as e:
            last_exc = e
            wait = 2 ** attempt
            log.warning("rate_limit attempt=%d wait=%ds", attempt + 1, wait)
            time.sleep(wait)
        except APITimeoutError as e:
            last_exc = e
            log.warning("timeout attempt=%d", attempt + 1)
            time.sleep(1)
        except APIStatusError as e:
            log.error("api_error status=%d", e.status_code)
            raise

    log.error("all_retries_exhausted")
    raise last_exc