aleksey
a2d238d92e
Make Zed Agent closer to Cursor: deterministic DevOps path index, live Traefik port probe before blind edits, stop-after-edit, and quieter Russian progress.
135 lines
3.3 KiB
Python
135 lines
3.3 KiB
Python
"""Redis cache + in-flight coalescing for identical LiteLLM chat calls."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import json
|
|
import logging
|
|
import os
|
|
from collections.abc import Awaitable, Callable
|
|
from typing import Any
|
|
|
|
log = logging.getLogger("llm_cache")
|
|
|
|
_redis = None
|
|
_redis_tried = False
|
|
_inflight: dict[str, asyncio.Future[str]] = {}
|
|
_inflight_lock = asyncio.Lock()
|
|
|
|
|
|
def _get_redis():
|
|
global _redis, _redis_tried
|
|
if _redis_tried:
|
|
return _redis
|
|
_redis_tried = True
|
|
url = os.environ.get("REDIS_URL", "").strip()
|
|
if not url:
|
|
return None
|
|
try:
|
|
import redis
|
|
|
|
client = redis.from_url(url, decode_responses=True)
|
|
client.ping()
|
|
_redis = client
|
|
except Exception as exc: # noqa: BLE001
|
|
log.warning("llm cache redis unavailable: %s", exc)
|
|
_redis = None
|
|
return _redis
|
|
|
|
|
|
def cache_key(
|
|
*,
|
|
model: str,
|
|
messages: list[dict[str, Any]],
|
|
max_tokens: int,
|
|
temperature: float,
|
|
) -> str:
|
|
blob = json.dumps(
|
|
{
|
|
"model": model,
|
|
"messages": messages,
|
|
"max_tokens": max_tokens,
|
|
"temperature": temperature,
|
|
},
|
|
ensure_ascii=False,
|
|
sort_keys=True,
|
|
separators=(",", ":"),
|
|
)
|
|
digest = hashlib.sha256(blob.encode("utf-8")).hexdigest()
|
|
return f"ai-router:llm-cache:{digest}"
|
|
|
|
|
|
def get_cached(key: str) -> str | None:
|
|
r = _get_redis()
|
|
if not r:
|
|
return None
|
|
try:
|
|
val = r.get(key)
|
|
return str(val) if val is not None else None
|
|
except Exception as exc: # noqa: BLE001
|
|
log.warning("llm cache get failed: %s", exc)
|
|
return None
|
|
|
|
|
|
def set_cached(key: str, value: str, ttl_sec: int) -> None:
|
|
r = _get_redis()
|
|
if not r or ttl_sec <= 0:
|
|
return
|
|
try:
|
|
r.setex(key, ttl_sec, value)
|
|
except Exception as exc: # noqa: BLE001
|
|
log.warning("llm cache set failed: %s", exc)
|
|
|
|
|
|
async def run_cached(
|
|
key: str,
|
|
factory: Callable[[], Awaitable[str]],
|
|
*,
|
|
ttl_sec: int,
|
|
on_hit: Callable[[str], Awaitable[None] | None] | None = None,
|
|
) -> tuple[str, str]:
|
|
"""
|
|
Return (text, source) where source is cache|coalesce|live.
|
|
Coalesce concurrent identical calls onto one upstream request.
|
|
"""
|
|
cached = get_cached(key)
|
|
if cached is not None:
|
|
if on_hit:
|
|
maybe = on_hit("cache")
|
|
if asyncio.iscoroutine(maybe):
|
|
await maybe
|
|
return cached, "cache"
|
|
|
|
async with _inflight_lock:
|
|
existing = _inflight.get(key)
|
|
if existing is not None:
|
|
fut: asyncio.Future[str] = existing
|
|
mine = False
|
|
else:
|
|
fut = asyncio.get_running_loop().create_future()
|
|
_inflight[key] = fut
|
|
mine = True
|
|
|
|
if not mine:
|
|
if on_hit:
|
|
maybe = on_hit("coalesce")
|
|
if asyncio.iscoroutine(maybe):
|
|
await maybe
|
|
return await fut, "coalesce"
|
|
|
|
try:
|
|
text = await factory()
|
|
set_cached(key, text, ttl_sec)
|
|
if not fut.done():
|
|
fut.set_result(text)
|
|
return text, "live"
|
|
except Exception as exc:
|
|
if not fut.done():
|
|
fut.set_exception(exc)
|
|
raise
|
|
finally:
|
|
async with _inflight_lock:
|
|
if _inflight.get(key) is fut:
|
|
_inflight.pop(key, None)
|