Build a Semantic Cache That Cuts an AI Agent's LLM Bill by 53%

Customer question
       |
       v
Embed locally with all-MiniLM-L6-v2
       |
       v
Find nearest cached question in ChromaDB
       |
       +-- Similarity >= 0.82, current policy, not expired
       |             |
       |             v
       |      Return cached answer
       |
       +-- Similarity below threshold, stale, or not found
                     |
                     v
                  Call Gemini
                     |
                     v
             Save question + answer
                     |
                     v
                 Return answer

uv init semantic-cache --python 3.12
cd semantic-cache
uv add chromadb==1.5.9 google-genai==2.29.0 matplotlib==3.11.2 python-dotenv==1.2.4
mkdir src

GEMINI_API_KEY=your_key_here

.env
.venv/
__pycache__/
*.pyc
.cache/
*.png

"""Settings for the semantic cache demo.

Inputs: GEMINI_API_KEY from .env.
Outputs: constants used by every other module (model, prices, threshold, paths).
"""

import os
import sys
from pathlib import Path

from dotenv import load_dotenv

ROOT = Path(__file__).resolve().parent.parent
load_dotenv(ROOT / ".env")

# Some Windows consoles default to a non-UTF-8 encoding.
sys.stdout.reconfigure(encoding="utf-8")

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY", "")
MODEL = "gemini-3.5-flash-lite"

# The example spaces requests to respect its stated free-tier rate limit.
MIN_INTERVAL_S = 4.2

# Illustrative paid-tier USD prices per 1M tokens.
PRICE_IN_PER_M = 0.30
PRICE_OUT_PER_M = 2.50

# Only reuse cached answers when similarity clears this value.
THRESHOLD = 0.82

# Expire entries after one day.
TTL_SECONDS = 24 * 60 * 60

CACHE_DIR = ROOT / ".cache"
COLLECTION = "answers"

"""The store policy and a day of customer questions.

Inputs: none (all data is inline).
Outputs: POLICY and TRAFFIC, a list of (question, intent) pairs in arrival
order. Intent tags grade cache accuracy; the agent never sees them.
"""

POLICY = """Kumkum Home - customer support policy (version {version})

Returns: unused items can be returned within {return_days} days of delivery.
Custom-stitched items cannot be returned.
Refunds: prepaid orders are refunded to the original payment method within
5-7 working days of the return reaching our warehouse. COD orders are
refunded by bank transfer (NEFT) within 7-10 working days, after you share
your bank details.
Cancellation: an order can be cancelled free of charge until it is dispatched.
After dispatch, refuse the parcel at the door instead.
Subscriptions: the monthly candle box subscription can be cancelled anytime
from My Account > Subscriptions; the current month's box still ships.
Cash on delivery: available for orders up to Rs 5,000. A Rs 49 COD fee applies.
Delivery: metro cities in 2-4 days, other locations in 5-7 days. Free shipping
above Rs 999.
Warranty: brass and copper items carry a 1-year warranty against manufacturing
defects.
Exchange: size or colour exchange is free, once per order, within the return
window.
Damaged items: report damage within 48 hours of delivery with a photo of the
parcel.
Gift wrap: available for Rs 50 per item at checkout.
Contact: support@kumkumhome.in, 10am-7pm IST, Monday to Saturday.
"""

# (question, intent). Some questions are paraphrases; others sound similar
# but have different answers (for example, order vs subscription cancellation).
TRAFFIC: list[tuple[str, str]] = [
    ("How many days do I have to return an item?", "returns_window"),
    ("Can I cancel my order?", "cancel_order"),
    ("What is the return window?", "returns_window"),
    ("How long does delivery take?", "delivery_time"),
    ("Is cash on delivery available?", "cod_limit"),
    ("When will I get my refund for a prepaid order?", "refund_prepaid"),
    ("How many days to return a product?", "returns_window"),
    ("How do I cancel my candle subscription?", "cancel_subscription"),
    ("Is there a warranty on brass items?", "warranty"),
    ("I want to cancel my order, how?", "cancel_order"),
    ("How long will delivery take to my city?", "delivery_time"),
    ("When do I get my refund for a COD order?", "refund_cod"),
    ("Can I return something after a week?", "returns_window"),
    ("Do you charge extra for cash on delivery?", "cod_fee"),
    ("Is shipping free?", "free_shipping"),
    ("My item arrived damaged, what do I do?", "damaged"),
    ("can i cancel the order.", "cancel_order"),
    ("Can I exchange for a different size?", "exchange"),
    ("What is the return window for items?", "returns_window"),
    ("Do brass products come with a warranty?", "warranty"),
    ("How long does a prepaid refund take?", "refund_prepaid"),
    ("Do you offer gift wrapping?", "gift_wrap"),
    ("How do I cancel my order?", "cancel_order"),
    ("How can I contact customer support?", "contact"),
    ("How many days for delivery?", "delivery_time"),
    ("What is the maximum order value for cash on delivery?", "cod_limit"),
    ("How many days do I get to return an item?", "returns_window"),
    ("The product came broken. What should I do?", "damaged"),
    ("How do I stop my subscription?", "cancel_subscription"),
    ("What is the delivery time?", "delivery_time"),
    ("Can I return a custom stitched item?", "returns_custom"),
    ("Is there a fee for COD?", "cod_fee"),
    ("I want to cancel my order", "cancel_order"),
    ("Is shipping free on my order?", "free_shipping"),
    ("How long does it take to refund a COD order?", "refund_cod"),
    ("What's your return window?", "returns_window"),
    ("Can I exchange an item for a different colour?", "exchange"),
    ("How long will delivery take?", "delivery_time"),
    ("Is there a warranty on copper items?", "warranty"),
    ("Can I cancel my subscription?", "cancel_subscription"),
    ("When will I receive my refund for a prepaid order?", "refund_prepaid"),
    ("What are your customer support hours?", "contact"),
    ("Can I cancel my order?", "cancel_order"),
    ("Item arrived damaged, what should I do?", "damaged"),
    ("Do you have gift wrap?", "gift_wrap"),
    ("How many days do I have to return an item?", "returns_window"),
    ("Is cash on delivery available for my order?", "cod_limit"),
    ("Delivery time to a small town?", "delivery_time"),
    ("Is there any extra charge for cash on delivery?", "cod_fee"),
    ("When will my refund come for a COD order?", "refund_cod"),
    ("Can I exchange it for another size?", "exchange"),
    ("Is shipping free?", "free_shipping"),
    ("Can I return an item after a week?", "returns_window"),
    ("How do I reach customer support?", "contact"),
    ("How can I cancel my order?", "cancel_order"),
    ("Do brass items have a warranty?", "warranty"),
    ("How long does delivery take?", "delivery_time"),
    ("Can I return custom stitched items?", "returns_custom"),
    ("What is the return window?", "returns_window"),
    ("How long does a prepaid refund take?", "refund_prepaid"),
]

PYTHONPATH=. uv run python -c "from src.traffic import TRAFFIC; print(len(TRAFFIC), len({i for _, i in TRAFFIC}))"

60 15

"""One call to Gemini, with tokens, cost, and latency tracked.

Inputs: a customer question and the policy text to answer from.
Outputs: an Answer with reply, token counts, dollar cost, and latency.
"""

import asyncio
import time
from dataclasses import dataclass

from google import genai
from google.genai import errors, types

from src.config import (
    GEMINI_API_KEY,
    MIN_INTERVAL_S,
    MODEL,
    PRICE_IN_PER_M,
    PRICE_OUT_PER_M,
)

_client = genai.Client(api_key=GEMINI_API_KEY) if GEMINI_API_KEY else None

INSTRUCTIONS = (
    "You are the support agent for Kumkum Home. Answer the customer in one or two "
    "short sentences, using only the policy below.

"
)


@dataclass
class Answer:
    """A reply from the model and what it cost."""

    text: str
    in_tokens: int
    out_tokens: int
    cost: float
    latency_ms: float


def price(in_tokens: int, out_tokens: int) -> float:
    """Return the paid-tier dollar cost of one call."""
    return (in_tokens * PRICE_IN_PER_M + out_tokens * PRICE_OUT_PER_M) / 1_000_000


_last_call = 0.0


async def _pace() -> None:
    """Wait so calls start at least MIN_INTERVAL_S apart."""
    global _last_call
    wait = _last_call + MIN_INTERVAL_S - time.monotonic()
    if wait > 0:
        await asyncio.sleep(wait)
    _last_call = time.monotonic()


async def _generate(
    question: str, policy: str
) -> tuple[types.GenerateContentResponse, float]:
    """Send one request; retry on 429/503. Return response and latency in ms.

    Only a successful attempt is timed, so time spent pacing or waiting for
    a retry does not count as model latency in the comparison.
    """
    if _client is None:
        raise ValueError("GEMINI_API_KEY is not set - add it to .env")

    config = types.GenerateContentConfig(
        system_instruction=INSTRUCTIONS + policy,
        thinking_config=types.ThinkingConfig(thinking_level="MINIMAL"),
    )

    for attempt in range(4):
        await _pace()
        start = time.perf_counter()
        try:
            response = await _client.aio.models.generate_content(
                model=MODEL, contents=question, config=config
            )
            return response, (time.perf_counter() - start) * 1000
        except errors.APIError as err:
            if err.code not in (429, 503) or attempt == 3:
                raise
            await asyncio.sleep(15)


async def answer(question: str, policy: str) -> Answer:
    """Ask the model one question and measure tokens, cost, and latency."""
    response, latency_ms = await _generate(question, policy)

    usage = response.usage_metadata
    in_tokens = usage.prompt_token_count or 0
    # Thinking tokens are also counted at the output price in this example.
    out_tokens = (usage.candidates_token_count or 0) + (usage.thoughts_token_count or 0)
    return Answer(
        text=(response.text or "").strip(),
        in_tokens=in_tokens,
        out_tokens=out_tokens,
        cost=price(in_tokens, out_tokens),
        latency_ms=latency_ms,
    )

PYTHONPATH=. uv run python -c "import asyncio; from src.llm import answer; from src.traffic import POLICY; print(asyncio.run(answer('What is the return window?', POLICY.format(version=1, return_days=7))))"

"""Semantic cache: reuse an answer when a new question means the same thing.

Inputs: questions embedded locally by the MiniLM ONNX model.
Outputs: a Hit with the cached answer and its similarity, or None on a miss.
Entries carry the policy version and a timestamp so stale ones can be dropped.
"""

import time
import uuid
from dataclasses import dataclass

import chromadb
from chromadb.utils.embedding_functions import ONNXMiniLM_L6_V2

from src.config import CACHE_DIR, COLLECTION, THRESHOLD, TTL_SECONDS

# Reuse a single embedder throughout the process.
EMBEDDER = ONNXMiniLM_L6_V2()
# Embed individual questions without fixed-size padding.
EMBEDDER.tokenizer.no_padding()

@dataclass
class Hit:
    """A cached answer that is close enough to reuse."""

    answer: str
    question: str
    similarity: float
    age_s: float


class SemanticCache:
    """Question-to-answer store, looked up by meaning instead of exact text."""

    def __init__(
        self, threshold: float = THRESHOLD, ttl_s: float = TTL_SECONDS
    ) -> None:
        """Open (or create) the cache on disk."""
        self.threshold = threshold
        self.ttl_s = ttl_s
        client = chromadb.PersistentClient(path=str(CACHE_DIR))
        self.col = client.get_or_create_collection(
            COLLECTION,
            embedding_function=EMBEDDER,
            configuration={"hnsw": {"space": "cosine"}},
        )

    def lookup(self, question: str, policy_version: int) -> Hit | None:
        """Return the closest fresh entry if it clears the threshold, else None."""
        if self.col.count() == 0:
            return None

        result = self.col.query(
            query_texts=[question],
            n_results=1,
            where={"policy_version": policy_version},
        )
        if not result["ids"][0]:
            return None

        similarity = 1 - result["distances"][0][0]
        meta = result["metadatas"][0][0]
        age_s = time.time() - meta["created_at"]
        if age_s > self.ttl_s:
            # Expired; evict so the next answer can be stored freshly.
            self.col.delete(ids=result["ids"][0])
            return None
        if similarity < self.threshold:
            return None
        return Hit(meta["answer"], result["documents"][0][0], similarity, age_s)

    def store(self, question: str, answer: str, policy_version: int) -> None:
        """Save a fresh answer under the question it was generated for."""
        self.col.add(
            ids=[uuid.uuid4().hex],
            documents=[question],
            metadatas=[{
                "answer": answer,
                "policy_version": policy_version,
                "created_at": time.time(),
            }],
        )

    def invalidate_similar(
        self, topic: str, min_similarity: float = 0.5
    ) -> list[str]:
        """Delete entries whose questions are close to a topic; return them."""
        if self.col.count() == 0:
            return []
        result = self.col.query(
            query_texts=[topic], n_results=self.col.count()
        )
        doomed = [
            (entry_id, question)
            for entry_id, question, distance in zip(
                result["ids"][0], result["documents"][0], result["distances"][0]
            )
            if 1 - distance >= min_similarity
        ]
        if doomed:
            self.col.delete(ids=[entry_id for entry_id, _ in doomed])
        return [question for _, question in doomed]

    def clear(self) -> None:
        """Remove every entry."""
        ids = self.col.get()["ids"]
        if ids:
            self.col.delete(ids=ids)

PYTHONPATH=. uv run python -c "
from src.cache import SemanticCache
c = SemanticCache(threshold=0.0); c.clear()
c.store('Can I cancel my order?', 'Until dispatch', 1)
for q in ['How do I cancel my order?', 'Can I cancel my subscription?', 'How do I stop my subscription?']:
    print(round(c.lookup(q, 1).similarity, 3), q)
c.clear()
"

0.895 How do I cancel my order?
0.693 Can I cancel my subscription?
0.413 How do I stop my subscription?

"""The support agent: check the cache first; call the LLM only on a miss.

Inputs: a question, the cache, and the current policy (text + version).
Outputs: a Record describing the hit/miss, cost, and latency.
"""

import time
from dataclasses import dataclass

from src.cache import SemanticCache
from src.llm import answer


@dataclass
class Record:
    """One answered question and what it cost."""

    question: str
    answer: str
    hit: bool
    cost: float
    latency_ms: float
    similarity: float = 0.0
    matched: str = ""


async def ask(
    question: str,
    policy: str,
    version: int,
    cache: SemanticCache | None,
) -> Record:
    """Answer one question, from the cache when possible."""
    lookup_ms = 0.0
    if cache is not None:
        start = time.perf_counter()
        hit = cache.lookup(question, version)
        lookup_ms = (time.perf_counter() - start) * 1000
        if hit is not None:
            return Record(
                question, hit.answer, True, 0.0, lookup_ms,
                hit.similarity, hit.question
            )

    reply = await answer(question, policy)
    if cache is not None:
        cache.store(question, reply.text, version)

    # A miss pays for both the lookup and the model call.
    return Record(
        question, reply.text, False, reply.cost,
        lookup_ms + reply.latency_ms
    )

"""Pick a similarity threshold from data instead of guessing one.

Inputs: TRAFFIC, with questions and their known intents.
Outputs: cache hits and false hits for each threshold, including the most
similar question pair with different intents.

The sweep embeds questions locally and replays the cache in memory.
"""

from dataclasses import dataclass

import numpy as np

from src.cache import EMBEDDER
from src.traffic import TRAFFIC


@dataclass
class SweepPoint:
    """The result of replaying the traffic at one threshold."""

    threshold: float
    hits: int
    false_hits: int
    example: str = ""


def similarity_matrix(questions: list[str]) -> np.ndarray:
    """Embed questions and return pairwise cosine similarities."""
    # Embed one at a time: lengths vary when tokenizer padding is disabled.
    vectors = np.array([EMBEDDER([question])[0] for question in questions])
    vectors = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)
    return vectors @ vectors.T


def replay(sims: np.ndarray, threshold: float) -> SweepPoint:
    """Replay traffic through an empty cache; count hits and false hits."""
    intents = [intent for _, intent in TRAFFIC]
    cached: list[int] = []
    point = SweepPoint(threshold, 0, 0)
    for i, (question, intent) in enumerate(TRAFFIC):
        best = max(cached, key=lambda j: sims[i, j], default=None)
        if best is None or sims[i, best] < threshold:
            cached.append(i)
            continue
        point.hits += 1
        if intents[best] != intent:
            point.false_hits += 1
            point.example = point.example or (
                f'"{question}" got the answer to "{TRAFFIC[best][0]}" '
                f'({sims[i, best]:.2f})'
            )
    return point


def closest_wrong_pair(sims: np.ndarray) -> tuple[float, str, str]:
    """Find the most similar pair of questions asking different things."""
    best = (0.0, "", "")
    for i, (qi, ii) in enumerate(TRAFFIC):
        for j, (qj, ij) in enumerate(TRAFFIC[:i]):
            if ii != ij and sims[i, j] > best[0]:
                best = (float(sims[i, j]), qi, qj)
    return best


def sweep(
    thresholds: list[float],
) -> tuple[list[SweepPoint], tuple[float, str, str]]:
    """Replay at each threshold; return results and closest wrong pair."""
    sims = similarity_matrix([question for question, _ in TRAFFIC])
    return [replay(sims, t) for t in thresholds], closest_wrong_pair(sims)

"""Turn two replay runs into numbers and charts.

Inputs: Records from one run without the cache and another with it.
Outputs: a printed totals table and cache_report.png (cost and latency).
"""

import statistics
from pathlib import Path

import matplotlib

matplotlib.use("Agg")  # Write PNG files; never open an interactive window.
import matplotlib.pyplot as plt

from src.agent import Record

NO_CACHE = "#eb6834"
CACHE = "#2a78d6"
INK = "#52514e"
GRID = "#e4e3df"


def totals(records: list[Record]) -> dict:
    """Sum one run: LLM calls, dollars, and median latency."""
    return {
        "calls": sum(not r.hit for r in records),
        "cost": sum(r.cost for r in records),
        "p50_ms": statistics.median(r.latency_ms for r in records),
        "hit_p50_ms": statistics.median(
            [r.latency_ms for r in records if r.hit] or [0]
        ),
    }


def print_totals(plain: list[Record], cached: list[Record]) -> float:
    """Print the side-by-side table; return the percentage of cost saved."""
    a, b = totals(plain), totals(cached)
    saved = 100 * (1 - b["cost"] / a["cost"])
    per_100k = 100_000 / len(plain)

    print(f"
{'':22}{'no cache':>14}{'with cache':>14}")
    print(f"{'LLM calls':22}{a['calls']:>14}{b['calls']:>14}")
    print(
        f"{'cost (paid tier)':22}"
        f"{('$' + format(a['cost'], '.5f')):>14}"
        f"{('$' + format(b['cost'], '.5f')):>14}"
    )
    print(
        f"{'cost per 100k questions':22}"
        f"{('$' + format(a['cost'] * per_100k, '.2f')):>14}"
        f"{('$' + format(b['cost'] * per_100k, '.2f')):>14}"
    )
    print(
        f"{'median latency':22}"
        f"{a['p50_ms']:>12.0f}ms{b['p50_ms']:>12.0f}ms"
    )
    print(
        f"{'median cache hit':22}"
        f"{'-':>14}{b['hit_p50_ms']:>12.0f}ms"
    )
    print(
        f"
Cost saved: {saved:.0f}% "
        f"({len(plain) - b['calls']} of {len(plain)} answered from cache)"
    )
    return saved


def _style(ax: plt.Axes, title: str, ylabel: str) -> None:
    """Recessive grid and axes, left-aligned title."""
    ax.set_title(title, loc="left", fontsize=12, color="#0b0b0b")
    ax.set_xlabel("question # (arrival order)", color=INK)
    ax.set_ylabel(ylabel, color=INK)
    ax.grid(axis="y", color=GRID, linewidth=0.8)
    ax.tick_params(colors=INK)
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)


def plot_cost(
    ax: plt.Axes, plain: list[Record], cached: list[Record]
) -> None:
    """Cumulative model spend, question by question, with and without cache."""
    xs = range(1, len(plain) + 1)
    for records, color, label in (
        (plain, NO_CACHE, "no cache"),
        (cached, CACHE, "with cache"),
    ):
        running, total = [], 0.0
        for r in records:
            total += r.cost
            running.append(total * 1000)  # milli-dollars are easier to read
        ax.plot(xs, running, color=color, linewidth=2, label=label)
        ax.annotate(
            f"{label}  ${total:.4f}", (len(records), running[-1]),
            xytext=(-4, 6), textcoords="offset points", ha="right", color=INK
        )
    _style(ax, "Cumulative LLM cost", "cost (milli-dollars, paid tier)")
    ax.legend(frameon=False, loc="upper left")


def plot_latency(
    ax: plt.Axes, plain: list[Record], cached: list[Record]
) -> None:
    """Latency of each request; cache hits sit far lower on a log scale."""
    xs = range(1, len(plain) + 1)
    ax.scatter(
        xs, [r.latency_ms for r in plain],
        s=18, color=NO_CACHE, label="no cache"
    )
    ax.scatter(
        xs, [r.latency_ms for r in cached],
        s=18, color=CACHE, label="with cache"
    )
    ax.set_yscale("log")
    _style(ax, "Latency per question", "milliseconds (log scale)")
    ax.legend(frameon=False, loc="center right")


def save_chart(
    plain: list[Record], cached: list[Record], saved: float, path: Path
) -> None:
    """Draw both panels side by side and save a PNG."""
    fig, (left, right) = plt.subplots(
        1, 2, figsize=(13, 5), facecolor="#fcfcfb"
    )
    plot_cost(left, plain, cached)
    plot_latency(right, plain, cached)
    fig.suptitle(
        f"Semantic cache: {saved:.0f}% cheaper on {len(plain)} support questions",
        x=0.01, ha="left", fontsize=15, fontweight="bold"
    )
    fig.tight_layout()
    fig.savefig(path, dpi=150)
    plt.close(fig)

"""Command line for the semantic cache demo.

Commands:
    ask "question"     answer one question through the persistent cache
    replay             run traffic without, then with, the cache; totals + chart
    sweep              hits vs false hits at each threshold (free, no LLM calls)
    invalidate-demo    show a stale answer, then three ways to get rid of it
    clear              empty the cache
"""

import argparse
import asyncio

from src.agent import Record, ask
from src.cache import SemanticCache
from src.config import ROOT, THRESHOLD
from src.report import print_totals, save_chart
from src.sweep import sweep
from src.traffic import POLICY, TRAFFIC

POLICY_V1 = POLICY.format(version=1, return_days=7)


def show(n: int, record: Record) -> None:
    """Print one line per question: what happened, how long it took."""
    if record.hit:
        how = f'HIT  {record.similarity:.2f} ~ "{record.matched}"'
    else:
        how = "MISS -> LLM"
    print(f"{n:>3}  {record.latency_ms:>7.0f} ms  {how:<58} {record.question}")


async def run_traffic(cache: SemanticCache | None) -> list[Record]:
    """Send every question in TRAFFIC through the agent, in order."""
    records = []
    for n, (question, _) in enumerate(TRAFFIC, start=1):
        record = await ask(question, POLICY_V1, 1, cache)
        show(n, record)
        records.append(record)
    return records


async def cmd_replay() -> None:
    """Same traffic, cache off vs cache on (starting empty)."""
    print("=== Pass 1: no cache - every question calls the LLM ===")
    plain = await run_traffic(None)

    print(f"
=== Pass 2: semantic cache, threshold {THRESHOLD} ===")
    cache = SemanticCache()
    cache.clear()
    cached = await run_traffic(cache)
    saved = print_totals(plain, cached)
    path = ROOT / "cache_report.png"
    save_chart(plain, cached, saved, path)
    print(f"Chart saved to {path.name}")


def cmd_sweep() -> None:
    """Show what each threshold buys and what it risks."""
    thresholds = [round(0.50 + 0.05 * i, 2) for i in range(10)]