67 lines
2.6 KiB
Python
67 lines
2.6 KiB
Python
"""LLM Cache Key Generation — Deterministic SHA-256 cache key from LLM request parameters."""
|
||
|
||
import hashlib
|
||
import json
|
||
|
||
|
||
def generate_cache_key(
|
||
model: str,
|
||
messages: list[dict[str, str]],
|
||
temperature: float,
|
||
tools: list[dict[str, object]] | None = None,
|
||
tool_choice: str = "auto",
|
||
max_tokens: int = 2000,
|
||
user_id: str | None = None,
|
||
kb_acl_hash: str | None = None,
|
||
) -> str:
|
||
"""Generate a deterministic SHA-256 cache key from LLM request parameters.
|
||
|
||
The key captures ALL inputs that deterministically affect LLM output:
|
||
model, system_prompt (extracted from messages), messages content,
|
||
temperature, tools, tool_choice, and max_tokens.
|
||
|
||
U17 安全扩展：``user_id`` 和 ``kb_acl_hash`` 非 None 时加入 hash 组件，
|
||
实现 per-user namespace 和 ACL-scope 隔离（安全要求 a, b）。为 None 时
|
||
行为与旧版完全一致（向后兼容）。
|
||
|
||
Args:
|
||
model: Model identifier (e.g. "openai/gpt-4o").
|
||
messages: Chat messages list (may include system prompt as first message).
|
||
temperature: Sampling temperature.
|
||
tools: Optional list of tool definitions.
|
||
tool_choice: Tool selection mode ("auto", "none", etc.).
|
||
max_tokens: Maximum response tokens.
|
||
user_id: U17 — 用户 ID，用于 per-user cache namespace 隔离。
|
||
kb_acl_hash: U17 — KB ACL-scope hash，用于 ACL 隔离缓存键。
|
||
|
||
Returns:
|
||
64-character hex SHA-256 hash string.
|
||
"""
|
||
system_prompt = _extract_system_prompt(messages)
|
||
# 单次 SHA-256：用分隔符拼接所有组件，避免逐组件 hash 再 hash 的冗余计算。
|
||
# 分隔符使用长度前缀防止歧义（如 "ab" + "cd" vs "a" + "bcd"）。
|
||
parts = [
|
||
f"m:{model}",
|
||
f"s:{system_prompt}",
|
||
f"msg:{json.dumps(messages, sort_keys=True, ensure_ascii=False)}",
|
||
f"t:{temperature:.2f}",
|
||
f"tools:{json.dumps(tools, sort_keys=True, ensure_ascii=False) if tools is not None else 'null'}",
|
||
f"tc:{tool_choice}",
|
||
f"mt:{max_tokens}",
|
||
]
|
||
# U17 — per-user namespace + ACL scope hash（安全要求 a, b, e）
|
||
if user_id is not None:
|
||
parts.append(f"u:{user_id}")
|
||
if kb_acl_hash is not None:
|
||
parts.append(f"a:{kb_acl_hash}")
|
||
combined = "\x1f".join(parts) # US (Unit Separator) 防止组件内容注入分隔符
|
||
return hashlib.sha256(combined.encode()).hexdigest()
|
||
|
||
|
||
def _extract_system_prompt(messages: list[dict[str, str]]) -> str:
|
||
"""Extract system prompt from messages list."""
|
||
for msg in messages:
|
||
if msg.get("role") == "system":
|
||
return msg.get("content", "")
|
||
return ""
|