131 lines
3.9 KiB
Python
131 lines
3.9 KiB
Python
# 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 band_aider.ai_client.completion_client import CompletionClient, OpenAIConfig
|
|
from band_aider.ai_client.token_tracker import record_completion_cache_hit
|
|
|
|
|
|
COMPLETION_CACHE_TTL_SECONDS = 30 * 60
|
|
|
|
|
|
@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 completion responses for this process."""
|
|
|
|
def __init__(self, ttl_seconds: float = COMPLETION_CACHE_TTL_SECONDS) -> None:
|
|
"""Initialize the cache with one absolute TTL."""
|
|
|
|
self.ttl_seconds = ttl_seconds
|
|
self._entries: dict[str, CompletionCacheEntry] = {}
|
|
self._history: deque[CompletionCacheEntry] = deque()
|
|
self._lock = Lock()
|
|
|
|
def get(self, request: str) -> str | None:
|
|
"""Return the cached response for one unexpired request."""
|
|
|
|
with self._lock:
|
|
self._discard_expired()
|
|
entry = self._entries.get(request)
|
|
return entry.response if entry is not None else None
|
|
|
|
def put(self, request: str, response: str) -> None:
|
|
"""Store the latest response for one request."""
|
|
|
|
with self._lock:
|
|
self._discard_expired()
|
|
entry = CompletionCacheEntry(
|
|
request=request,
|
|
response=response,
|
|
created_at=time.monotonic(),
|
|
)
|
|
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 _discard_expired(self) -> None:
|
|
"""Discard entries older than the absolute TTL."""
|
|
|
|
cutoff = time.monotonic() - 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)
|
|
|
|
|
|
class CachedCompletionClient(CompletionClient):
|
|
"""CachedCompletionClient returns recent exact responses before delegating."""
|
|
|
|
def __init__(
|
|
self,
|
|
inner: CompletionClient,
|
|
cache: CompletionCache | None = None,
|
|
) -> None:
|
|
"""Initialize the decorator with its wrapped client."""
|
|
|
|
self._inner = inner
|
|
self._cache = cache or CompletionCache()
|
|
|
|
@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 completion 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:
|
|
configs = self.openai_configs
|
|
config = configs[0] if configs else OpenAIConfig(model="unknown", api_base="unknown")
|
|
record_completion_cache_hit(
|
|
model_name=config.model,
|
|
endpoint=config.api_base,
|
|
)
|
|
if logger is not None:
|
|
logger.info("CachedCompletionClient returned a cached response.")
|
|
return cached
|
|
|
|
response = (
|
|
self.inner.complete(prompt)
|
|
if logger is None
|
|
else self.inner.complete(prompt, logger=logger)
|
|
)
|
|
self._cache.put(prompt, response)
|
|
return response
|