396 lines
14 KiB
Python
396 lines
14 KiB
Python
"""对话媒体存储：将图片/视频/音频二进制与对话 JSON 解耦。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import hashlib
|
||
import json
|
||
import mimetypes
|
||
import os
|
||
import tempfile
|
||
import threading
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
|
||
DEFAULT_MIME_BY_KIND: Dict[str, str] = {
|
||
"image": "image/png",
|
||
"video": "video/mp4",
|
||
"audio": "audio/wav",
|
||
}
|
||
|
||
|
||
def _normalize_kind(value: Optional[str], mime_type: Optional[str] = None) -> str:
|
||
kind = str(value or "").strip().lower()
|
||
if kind in {"image", "video", "audio"}:
|
||
return kind
|
||
mime = str(mime_type or "").strip().lower()
|
||
if mime.startswith("image/"):
|
||
return "image"
|
||
if mime.startswith("video/"):
|
||
return "video"
|
||
if mime.startswith("audio/"):
|
||
return "audio"
|
||
return "binary"
|
||
|
||
|
||
def _normalize_mime(value: Optional[str], kind: str, fallback_path: Optional[str] = None) -> str:
|
||
mime = str(value or "").strip().lower()
|
||
if mime:
|
||
return mime
|
||
if fallback_path:
|
||
guessed, _ = mimetypes.guess_type(str(fallback_path))
|
||
if guessed:
|
||
return guessed
|
||
return DEFAULT_MIME_BY_KIND.get(kind, "application/octet-stream")
|
||
|
||
|
||
def _guess_extension(mime_type: str, kind: str) -> str:
|
||
ext = mimetypes.guess_extension(mime_type or "")
|
||
if ext:
|
||
return ext
|
||
if kind == "image":
|
||
return ".png"
|
||
if kind == "video":
|
||
return ".mp4"
|
||
if kind == "audio":
|
||
return ".wav"
|
||
return ".bin"
|
||
|
||
|
||
def _decode_base64_payload(payload: str) -> bytes:
|
||
data = "".join(str(payload or "").strip().split())
|
||
if not data:
|
||
raise ValueError("base64 数据为空")
|
||
missing_padding = len(data) % 4
|
||
if missing_padding:
|
||
data += "=" * (4 - missing_padding)
|
||
return base64.b64decode(data.encode("utf-8"), validate=False)
|
||
|
||
|
||
def _parse_data_url(data_url: str) -> Tuple[str, str]:
|
||
raw = str(data_url or "").strip()
|
||
if not raw.startswith("data:") or "," not in raw:
|
||
raise ValueError("非法 data URL")
|
||
header, payload = raw.split(",", 1)
|
||
mime = "application/octet-stream"
|
||
if ";" in header:
|
||
mime = header[5:].split(";", 1)[0] or mime
|
||
elif len(header) > 5:
|
||
mime = header[5:]
|
||
if ";base64" not in header.lower():
|
||
raise ValueError("仅支持 base64 data URL")
|
||
return mime, payload
|
||
|
||
|
||
class MediaStore:
|
||
"""媒体存储中心（data/conversations/media_store）。"""
|
||
|
||
INDEX_VERSION = 1
|
||
|
||
def __init__(self, data_dir: Path):
|
||
self.data_dir = Path(data_dir).expanduser().resolve()
|
||
self.conversations_dir = self.data_dir / "conversations"
|
||
self.root_dir = self.conversations_dir / "media_store"
|
||
self.blobs_dir = self.root_dir / "blobs"
|
||
self.index_file = self.root_dir / "index.json"
|
||
self._lock = threading.RLock()
|
||
self._ensure_dirs()
|
||
|
||
def _ensure_dirs(self) -> None:
|
||
self.conversations_dir.mkdir(parents=True, exist_ok=True)
|
||
self.root_dir.mkdir(parents=True, exist_ok=True)
|
||
self.blobs_dir.mkdir(parents=True, exist_ok=True)
|
||
if not self.index_file.exists():
|
||
self._save_index(
|
||
{
|
||
"version": self.INDEX_VERSION,
|
||
"media": {},
|
||
"message_map": {},
|
||
"updated_at": datetime.now().isoformat(),
|
||
}
|
||
)
|
||
|
||
def _atomic_write_json(self, target_file: Path, payload: Dict[str, Any]) -> None:
|
||
target_file.parent.mkdir(parents=True, exist_ok=True)
|
||
temp_path: Optional[Path] = None
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
mode="w",
|
||
encoding="utf-8",
|
||
dir=str(target_file.parent),
|
||
prefix=f".{target_file.name}.",
|
||
suffix=".tmp",
|
||
delete=False,
|
||
) as fh:
|
||
temp_path = Path(fh.name)
|
||
json.dump(payload, fh, ensure_ascii=False, indent=2)
|
||
fh.flush()
|
||
os.fsync(fh.fileno())
|
||
os.replace(str(temp_path), str(target_file))
|
||
finally:
|
||
if temp_path and temp_path.exists():
|
||
try:
|
||
temp_path.unlink()
|
||
except Exception:
|
||
pass
|
||
|
||
def _load_index(self) -> Dict[str, Any]:
|
||
try:
|
||
if not self.index_file.exists():
|
||
return {
|
||
"version": self.INDEX_VERSION,
|
||
"media": {},
|
||
"message_map": {},
|
||
"updated_at": datetime.now().isoformat(),
|
||
}
|
||
data = json.loads(self.index_file.read_text(encoding="utf-8") or "{}")
|
||
if not isinstance(data, dict):
|
||
raise ValueError("index.json 不是对象")
|
||
data.setdefault("version", self.INDEX_VERSION)
|
||
data.setdefault("media", {})
|
||
data.setdefault("message_map", {})
|
||
data.setdefault("updated_at", datetime.now().isoformat())
|
||
if not isinstance(data.get("media"), dict):
|
||
data["media"] = {}
|
||
if not isinstance(data.get("message_map"), dict):
|
||
data["message_map"] = {}
|
||
return data
|
||
except Exception:
|
||
return {
|
||
"version": self.INDEX_VERSION,
|
||
"media": {},
|
||
"message_map": {},
|
||
"updated_at": datetime.now().isoformat(),
|
||
}
|
||
|
||
def _save_index(self, payload: Dict[str, Any]) -> None:
|
||
payload["updated_at"] = datetime.now().isoformat()
|
||
self._atomic_write_json(self.index_file, payload)
|
||
|
||
@staticmethod
|
||
def _build_media_id(sha256_hex: str) -> str:
|
||
return f"sha256:{sha256_hex}"
|
||
|
||
def _blob_path_for(self, sha256_hex: str, ext: str) -> Tuple[Path, str]:
|
||
rel = Path("blobs") / sha256_hex[:2] / f"{sha256_hex}{ext}"
|
||
return self.root_dir / rel, rel.as_posix()
|
||
|
||
def _write_blob_if_missing(self, blob_path: Path, payload: bytes) -> None:
|
||
if blob_path.exists():
|
||
return
|
||
blob_path.parent.mkdir(parents=True, exist_ok=True)
|
||
temp_path: Optional[Path] = None
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
mode="wb",
|
||
dir=str(blob_path.parent),
|
||
prefix=f".{blob_path.name}.",
|
||
suffix=".tmp",
|
||
delete=False,
|
||
) as fh:
|
||
temp_path = Path(fh.name)
|
||
fh.write(payload)
|
||
fh.flush()
|
||
os.fsync(fh.fileno())
|
||
os.replace(str(temp_path), str(blob_path))
|
||
finally:
|
||
if temp_path and temp_path.exists():
|
||
try:
|
||
temp_path.unlink()
|
||
except Exception:
|
||
pass
|
||
|
||
def store_bytes(
|
||
self,
|
||
payload: bytes,
|
||
*,
|
||
kind: Optional[str] = None,
|
||
mime_type: Optional[str] = None,
|
||
source_path: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
if not isinstance(payload, (bytes, bytearray)) or len(payload) == 0:
|
||
raise ValueError("媒体二进制为空")
|
||
sha256_hex = hashlib.sha256(payload).hexdigest()
|
||
media_id = self._build_media_id(sha256_hex)
|
||
normalized_kind = _normalize_kind(kind, mime_type)
|
||
normalized_mime = _normalize_mime(mime_type, normalized_kind, fallback_path=source_path)
|
||
ext = _guess_extension(normalized_mime, normalized_kind)
|
||
blob_path, blob_rel = self._blob_path_for(sha256_hex, ext)
|
||
|
||
with self._lock:
|
||
self._write_blob_if_missing(blob_path, bytes(payload))
|
||
index = self._load_index()
|
||
media = index.setdefault("media", {})
|
||
entry = media.get(media_id) if isinstance(media.get(media_id), dict) else {}
|
||
entry.update(
|
||
{
|
||
"media_id": media_id,
|
||
"sha256": sha256_hex,
|
||
"kind": normalized_kind,
|
||
"mime_type": normalized_mime,
|
||
"size": len(payload),
|
||
"blob_rel_path": blob_rel,
|
||
"created_at": entry.get("created_at") or datetime.now().isoformat(),
|
||
"updated_at": datetime.now().isoformat(),
|
||
}
|
||
)
|
||
media[media_id] = entry
|
||
self._save_index(index)
|
||
return {
|
||
"media_id": media_id,
|
||
"kind": normalized_kind,
|
||
"mime_type": normalized_mime,
|
||
"size": len(payload),
|
||
"sha256": sha256_hex,
|
||
"store_path": f"media_store/{blob_rel}",
|
||
"source_path": source_path,
|
||
}
|
||
|
||
def store_file(
|
||
self,
|
||
file_path: Path,
|
||
*,
|
||
kind: Optional[str] = None,
|
||
mime_type: Optional[str] = None,
|
||
source_path: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
path = Path(file_path).expanduser().resolve()
|
||
payload = path.read_bytes()
|
||
return self.store_bytes(
|
||
payload,
|
||
kind=kind,
|
||
mime_type=mime_type,
|
||
source_path=source_path or str(file_path),
|
||
)
|
||
|
||
def store_base64_data(
|
||
self,
|
||
data_base64: str,
|
||
*,
|
||
kind: Optional[str] = None,
|
||
mime_type: Optional[str] = None,
|
||
source_path: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
payload = _decode_base64_payload(data_base64)
|
||
return self.store_bytes(
|
||
payload,
|
||
kind=kind,
|
||
mime_type=mime_type,
|
||
source_path=source_path,
|
||
)
|
||
|
||
def store_data_url(
|
||
self,
|
||
data_url: str,
|
||
*,
|
||
kind: Optional[str] = None,
|
||
source_path: Optional[str] = None,
|
||
) -> Dict[str, Any]:
|
||
mime_type, base64_payload = _parse_data_url(data_url)
|
||
return self.store_base64_data(
|
||
base64_payload,
|
||
kind=kind,
|
||
mime_type=mime_type,
|
||
source_path=source_path,
|
||
)
|
||
|
||
def _resolve_blob_path(self, media_ref: Dict[str, Any]) -> Optional[Tuple[Path, Dict[str, Any]]]:
|
||
media_id = str(media_ref.get("media_id") or "").strip()
|
||
with self._lock:
|
||
index = self._load_index()
|
||
media = index.get("media") if isinstance(index.get("media"), dict) else {}
|
||
entry = media.get(media_id) if media_id else None
|
||
if isinstance(entry, dict):
|
||
rel = str(entry.get("blob_rel_path") or "").strip()
|
||
if rel:
|
||
path = self.root_dir / rel
|
||
if path.exists() and path.is_file():
|
||
return path, entry
|
||
store_path = str(media_ref.get("store_path") or "").strip()
|
||
if store_path.startswith("media_store/"):
|
||
candidate = self.conversations_dir / store_path
|
||
elif store_path:
|
||
candidate = self.root_dir / store_path
|
||
else:
|
||
candidate = None
|
||
if candidate and candidate.exists() and candidate.is_file():
|
||
if not isinstance(entry, dict):
|
||
entry = {}
|
||
return candidate, entry
|
||
return None
|
||
|
||
def get_media_entry(self, media_id: str) -> Optional[Dict[str, Any]]:
|
||
target_id = str(media_id or "").strip()
|
||
if not target_id:
|
||
return None
|
||
with self._lock:
|
||
index = self._load_index()
|
||
media = index.get("media") if isinstance(index.get("media"), dict) else {}
|
||
entry = media.get(target_id)
|
||
if isinstance(entry, dict):
|
||
return dict(entry)
|
||
return None
|
||
|
||
def load_bytes(self, media_ref: Dict[str, Any]) -> Optional[bytes]:
|
||
resolved = self._resolve_blob_path(media_ref or {})
|
||
if not resolved:
|
||
return None
|
||
blob_path, _entry = resolved
|
||
try:
|
||
return blob_path.read_bytes()
|
||
except Exception:
|
||
return None
|
||
|
||
def load_bytes_by_media_id(self, media_id: str) -> Optional[bytes]:
|
||
target_id = str(media_id or "").strip()
|
||
if not target_id:
|
||
return None
|
||
entry = self.get_media_entry(target_id)
|
||
if not isinstance(entry, dict):
|
||
return None
|
||
ref = {
|
||
"media_id": target_id,
|
||
"store_path": entry.get("blob_rel_path"),
|
||
"kind": entry.get("kind"),
|
||
"mime_type": entry.get("mime_type"),
|
||
}
|
||
return self.load_bytes(ref)
|
||
|
||
def to_data_url(self, media_ref: Dict[str, Any]) -> Optional[str]:
|
||
payload = self.load_bytes(media_ref)
|
||
if payload is None:
|
||
return None
|
||
mime = _normalize_mime(
|
||
media_ref.get("mime_type"),
|
||
_normalize_kind(media_ref.get("kind"), media_ref.get("mime_type")),
|
||
)
|
||
b64 = base64.b64encode(payload).decode("utf-8")
|
||
return f"data:{mime};base64,{b64}"
|
||
|
||
def link_message_media(
|
||
self,
|
||
conversation_id: Optional[str],
|
||
message_id: Optional[str],
|
||
media_refs: List[Dict[str, Any]],
|
||
) -> None:
|
||
conv = str(conversation_id or "").strip()
|
||
msg = str(message_id or "").strip()
|
||
if not conv or not msg or not media_refs:
|
||
return
|
||
media_ids = [str((item or {}).get("media_id") or "").strip() for item in media_refs]
|
||
media_ids = [item for item in media_ids if item]
|
||
if not media_ids:
|
||
return
|
||
with self._lock:
|
||
index = self._load_index()
|
||
message_map = index.setdefault("message_map", {})
|
||
conv_map = message_map.get(conv) if isinstance(message_map.get(conv), dict) else {}
|
||
conv_map[msg] = {
|
||
"media_ids": media_ids,
|
||
"updated_at": datetime.now().isoformat(),
|
||
}
|
||
message_map[conv] = conv_map
|
||
self._save_index(index)
|