You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
105 lines
3.5 KiB
Python
105 lines
3.5 KiB
Python
"""DB-backed LLM call cache keyed by SHA-256 of inputs."""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from typing import Any
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from libs.common.logging import get_logger
|
|
from libs.common.time_utils import utc_now
|
|
from libs.llm.exceptions import LLMCacheError
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
_SCHEMA_VERSION = "1.0.0"
|
|
|
|
|
|
def build_cache_key(
|
|
text: str,
|
|
prompt_name: str,
|
|
prompt_version: str,
|
|
model_name: str,
|
|
schema_version: str = _SCHEMA_VERSION,
|
|
) -> str:
|
|
"""Return SHA-256 cache key for this (text, prompt, model, schema) combination."""
|
|
text_hash = hashlib.sha256(text.encode()).hexdigest()
|
|
hint_hash = hashlib.sha256(
|
|
f"{prompt_name}:{prompt_version}:{model_name}:{schema_version}".encode()
|
|
).hexdigest()
|
|
combined = f"{text_hash}:{hint_hash}"
|
|
return hashlib.sha256(combined.encode()).hexdigest()
|
|
|
|
|
|
class LLMCacheStore:
|
|
"""Read/write LLM responses from/to the llm_call_cache table."""
|
|
|
|
async def get(self, session: AsyncSession, cache_key: str) -> dict[str, Any] | None:
|
|
"""Return cached normalized dict if found, else None."""
|
|
from libs.db.models import LLMCallCache # late import to avoid circular deps
|
|
|
|
try:
|
|
result = await session.execute(
|
|
select(LLMCallCache).where(LLMCallCache.cache_key == cache_key)
|
|
)
|
|
row = result.scalar_one_or_none()
|
|
if row is None:
|
|
return None
|
|
logger.debug("llm_cache_hit", cache_key=cache_key[:16])
|
|
return row.normalized_json # type: ignore[return-value]
|
|
except Exception as exc:
|
|
raise LLMCacheError(
|
|
f"Cache read failed: {exc}",
|
|
source="llm_cache",
|
|
context={"cache_key": cache_key[:16]},
|
|
) from exc
|
|
|
|
async def put(
|
|
self,
|
|
session: AsyncSession,
|
|
cache_key: str,
|
|
document_id: str | None,
|
|
model_name: str,
|
|
prompt_version: str,
|
|
schema_version: str,
|
|
raw_prompt: str,
|
|
raw_response: str,
|
|
normalized: dict[str, Any],
|
|
token_usage: dict[str, int],
|
|
elapsed_ms: int,
|
|
) -> None:
|
|
"""Store a new cache entry. Silently skips on duplicate key."""
|
|
from libs.db.models import LLMCallCache # late import
|
|
|
|
try:
|
|
existing = await session.execute(
|
|
select(LLMCallCache).where(LLMCallCache.cache_key == cache_key)
|
|
)
|
|
if existing.scalar_one_or_none() is not None:
|
|
logger.debug("llm_cache_skip_dup", cache_key=cache_key[:16])
|
|
return
|
|
|
|
row = LLMCallCache(
|
|
cache_key=cache_key,
|
|
document_id=document_id,
|
|
model_name=model_name,
|
|
prompt_version=prompt_version,
|
|
schema_version=schema_version,
|
|
raw_prompt=raw_prompt,
|
|
raw_response=raw_response,
|
|
normalized_json=normalized,
|
|
token_usage_json=token_usage,
|
|
elapsed_ms=elapsed_ms,
|
|
created_at_utc=utc_now(),
|
|
)
|
|
session.add(row)
|
|
await session.flush()
|
|
logger.debug("llm_cache_stored", cache_key=cache_key[:16], elapsed_ms=elapsed_ms)
|
|
except Exception as exc:
|
|
raise LLMCacheError(
|
|
f"Cache write failed: {exc}",
|
|
source="llm_cache",
|
|
context={"cache_key": cache_key[:16]},
|
|
) from exc
|