band-aider-pub/band_aider/ai_client/semantic_completion_cache.py

# SPDX-FileCopyrightText: 2026 Art Sea LLC
#
# SPDX-License-Identifier: AGPL-3.0-only

from __future__ import annotations

from collections import deque
from dataclasses import dataclass
import logging
from threading import Lock
import time
from typing import Callable

from band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
from band_aider.judgement.model2vec_code_parser import Model2VecSimilarityCheck


COMPLETION_CACHE_TTL_SECONDS = 5 * 60
COMPLETION_CACHE_MAX_SEMANTIC_DISTANCE = 0.03


@dataclass(frozen=True)
class CompletionCacheEntry:
    """CompletionCacheEntry stores one cached request and response."""

    request: str
    response: str
    created_at: float


class CompletionCache:
    """CompletionCache stores recent exact and semantically equivalent completions."""

    def __init__(
        self,
        ttl_seconds: float = COMPLETION_CACHE_TTL_SECONDS,
        max_semantic_distance: float = COMPLETION_CACHE_MAX_SEMANTIC_DISTANCE,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        """Initialize the process cache."""

        self.ttl_seconds = ttl_seconds
        self._clock = clock
        self._entries: dict[str, CompletionCacheEntry] = {}
        self._history: deque[CompletionCacheEntry] = deque()
        self._lock = Lock()
        self._semantic_lock = Lock()
        self._semantic_check = Model2VecSimilarityCheck(
            changes=[],
            request_text="",
            searchable_beats=(),
            threshold=1.0 - max_semantic_distance,
        )

    def get(self, request: str) -> tuple[str, bool] | None:
        """Return a cached response and whether the match was semantic."""

        with self._lock:
            self._discard_expired(self._clock())

            exact = self._entries.get(request)
            if exact is not None:
                return exact.response, False

            candidates = tuple(self._entries)

        matched_request = self._semantic_match(request, candidates)
        if matched_request is None:
            return None

        with self._lock:
            self._discard_expired(self._clock())
            matched = self._entries.get(matched_request)
            if matched is None:
                return None
            return matched.response, True

    def put(self, request: str, response: str) -> None:
        """Store the latest response for one request."""

        with self._lock:
            now = self._clock()
            self._discard_expired(now)

            entry = CompletionCacheEntry(
                request=request,
                response=response,
                created_at=now,
            )
            self._entries[request] = entry
            self._history.append(entry)

    def clear(self) -> None:
        """Remove every cached completion."""

        with self._lock:
            self._entries.clear()
            self._history.clear()

    def _semantic_match(
        self,
        request: str,
        candidates: tuple[str, ...],
    ) -> str | None:
        """Return the closest accepted cached request."""

        if not candidates:
            return None

        with self._semantic_lock:
            check = self._semantic_check
            check.changes.clear()
            check.score = None
            check.matched = False
            check.reason = ""
            check.error = ""
            check.best_candidate = ""

            if not check(
                request_text=request,
                searchable_beats=candidates,
            ):
                return None

            return check.best_candidate or None

    def _discard_expired(self, now: float) -> None:
        """Discard entries older than the absolute TTL."""

        cutoff = now - self.ttl_seconds

        while self._history and self._history[0].created_at <= cutoff:
            expired = self._history.popleft()
            current = self._entries.get(expired.request)

            if current is expired:
                self._entries.pop(expired.request, None)


_CACHE = CompletionCache()


class CachedCompletionClient(CompletionClient):
    """CachedCompletionClient returns recent responses before delegating."""

    def __init__(
        self,
        inner: CompletionClient,
        cache: CompletionCache = _CACHE,
    ) -> None:
        """Initialize the decorator."""

        self._inner = inner
        self._cache = cache

    @property
    def inner(self) -> CompletionClient:
        """Return the wrapped completion client."""

        return self._inner

    @property
    def openai_configs(self) -> tuple[OpenAIConfig, ...]:
        """Return the wrapped client's configurations."""

        return self.inner.openai_configs

    def complete(
        self,
        prompt: str,
        logger: logging.Logger | None = None,
    ) -> str:
        """Return a cached response or complete and cache the prompt."""

        cached = self._cache.get(prompt)
        if cached is not None:
            response, semantic = cached
            if logger is not None:
                match_kind = "semantic" if semantic else "exact"
                logger.info(
                    "CachedCompletionClient returned a %s cache hit.",
                    match_kind,
                )
            return response

        response = (
            self.inner.complete(prompt)
            if logger is None
            else self.inner.complete(prompt, logger=logger)
        )
        self._cache.put(prompt, response)
        return response