Spaces:
Sleeping
Sleeping
File size: 4,582 Bytes
df6cd5e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """
core/task_cache.py — deterministic task cache for repeated ZIP+prompt runs.
The cache is intentionally conservative:
- fingerprint = uploaded ZIP bytes hash + normalized prompt + selected model
- default TTL = 24 hours
- stores both the pipeline state and the exported ZIP artifact (when present)
"""
from __future__ import annotations
import copy
import hashlib
import json
import os
import shutil
import time
from pathlib import Path
from typing import Any, Dict, Optional
CACHE_TTL_SECONDS = 24 * 60 * 60
def _normalize_prompt(prompt: str) -> str:
return " ".join((prompt or "").lower().split())
def _safe_model_name(model_name: Optional[str]) -> str:
return (model_name or "auto").strip().lower() or "auto"
def task_fingerprint(
zip_path: Optional[str],
prompt: str,
model_name: Optional[str] = None,
) -> str:
zip_hash = "nozip"
if zip_path and os.path.exists(zip_path):
h = hashlib.sha256()
with open(zip_path, "rb") as fh:
for chunk in iter(lambda: fh.read(1024 * 1024), b""):
h.update(chunk)
zip_hash = h.hexdigest()[:16]
prompt_norm = _normalize_prompt(prompt)
prompt_hash = hashlib.md5(prompt_norm.encode("utf-8")).hexdigest()[:12]
model_hash = hashlib.md5(_safe_model_name(model_name).encode("utf-8")).hexdigest()[:8]
return f"{zip_hash}-{prompt_hash}-{model_hash}"
class TaskCache:
def __init__(self, workspace: str, ttl_seconds: int = CACHE_TTL_SECONDS) -> None:
self.root = Path(workspace).resolve()
self.entries_dir = self.root / "entries"
self.artifacts_dir = self.root / "artifacts"
self.ttl_seconds = int(ttl_seconds)
self.entries_dir.mkdir(parents=True, exist_ok=True)
self.artifacts_dir.mkdir(parents=True, exist_ok=True)
def _entry_path(self, fingerprint: str) -> Path:
return self.entries_dir / f"{fingerprint}.json"
def _artifact_path(self, fingerprint: str, original_name: str = "cached_result.zip") -> Path:
base = Path(original_name or "cached_result.zip").name
return self.artifacts_dir / f"{fingerprint}__{base}"
def get(self, fingerprint: str) -> Optional[Dict[str, Any]]:
entry_path = self._entry_path(fingerprint)
if not entry_path.exists():
return None
age = time.time() - entry_path.stat().st_mtime
if age > self.ttl_seconds:
try:
entry_path.unlink(missing_ok=True)
except Exception:
pass
return None
try:
payload = json.loads(entry_path.read_text(encoding="utf-8"))
except Exception:
return None
if not isinstance(payload, dict):
return None
return payload
def store(
self,
fingerprint: str,
state: Dict[str, Any],
out_zip_path: Optional[str] = None,
) -> Dict[str, Any]:
payload: Dict[str, Any] = {
"fingerprint": fingerprint,
"cached_at": int(time.time()),
"state": copy.deepcopy(state or {}),
"out_zip_path": None,
}
if out_zip_path and os.path.exists(out_zip_path):
artifact_path = self._artifact_path(fingerprint, os.path.basename(out_zip_path))
shutil.copyfile(out_zip_path, artifact_path)
payload["out_zip_path"] = str(artifact_path)
self._entry_path(fingerprint).write_text(
json.dumps(payload, indent=2, ensure_ascii=False, default=str),
encoding="utf-8",
)
return payload
def restore_result(self, fingerprint: str) -> Optional[tuple[Dict[str, Any], Optional[str]]]:
payload = self.get(fingerprint)
if not payload:
return None
state = copy.deepcopy(payload.get("state") or {})
out_zip_path = payload.get("out_zip_path") or None
if out_zip_path and not os.path.exists(out_zip_path):
out_zip_path = None
state.setdefault("cache", {})
state["cache"].update({
"hit": True,
"fingerprint": fingerprint,
"cached_at": payload.get("cached_at"),
})
state.setdefault("timeline", []).append(
{
"timestamp": time.strftime("%H:%M:%S"),
"agent": "TaskCache",
"status": "OK",
"message": "cache hit — reused prior pipeline result",
}
)
return state, out_zip_path
__all__ = ["TaskCache", "task_fingerprint", "CACHE_TTL_SECONDS"]
|