Add File
This commit is contained in:
562
src/landppt/services/image/cache/image_cache.py
vendored
Normal file
562
src/landppt/services/image/cache/image_cache.py
vendored
Normal file
@@ -0,0 +1,562 @@
|
||||
"""
|
||||
图片缓存管理器
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import shutil
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Any
|
||||
import hashlib
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from ..models import (
|
||||
ImageInfo, ImageCacheInfo, ImageSourceType, ImageProvider
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class ImageCacheManager:
|
||||
"""图片缓存管理器"""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
self.config = config
|
||||
|
||||
# 缓存配置 - 移除过期和清理设置，图片永久有效
|
||||
self.cache_root = Path(config.get('base_dir', config.get('cache_root', 'temp/images_cache')))
|
||||
|
||||
# 缓存大小和清理配置（虽然图片永久有效，但保留配置项）
|
||||
self.max_size_gb = config.get('max_size_gb', 100.0)
|
||||
self.cleanup_interval_hours = config.get('cleanup_interval_hours', 240000)
|
||||
|
||||
# 缓存目录结构
|
||||
self.ai_generated_dir = self.cache_root / 'ai_generated'
|
||||
self.web_search_dir = self.cache_root / 'web_search'
|
||||
self.local_storage_dir = self.cache_root / 'local_storage'
|
||||
self.metadata_dir = self.cache_root / 'metadata'
|
||||
self.thumbnails_dir = self.cache_root / 'thumbnails'
|
||||
|
||||
# 创建目录结构
|
||||
self._create_cache_directories()
|
||||
|
||||
# 缓存索引
|
||||
self._cache_index: Dict[str, ImageCacheInfo] = {}
|
||||
self._load_cache_index()
|
||||
|
||||
def _create_cache_directories(self):
|
||||
"""创建缓存目录结构"""
|
||||
directories = [
|
||||
self.cache_root,
|
||||
self.ai_generated_dir,
|
||||
self.ai_generated_dir / 'dalle',
|
||||
self.ai_generated_dir / 'stable_diffusion',
|
||||
self.ai_generated_dir / 'siliconflow',
|
||||
self.web_search_dir,
|
||||
self.web_search_dir / 'unsplash',
|
||||
self.web_search_dir / 'pixabay',
|
||||
self.local_storage_dir,
|
||||
self.local_storage_dir / 'user_uploads',
|
||||
self.local_storage_dir / 'processed',
|
||||
self.metadata_dir,
|
||||
self.metadata_dir / 'references', # 用于存储同一内容的多个引用
|
||||
self.thumbnails_dir
|
||||
]
|
||||
|
||||
for directory in directories:
|
||||
directory.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def _get_cache_path(self, source_type: ImageSourceType, provider: ImageProvider) -> Path:
|
||||
"""获取缓存路径"""
|
||||
if source_type == ImageSourceType.AI_GENERATED:
|
||||
if provider == ImageProvider.DALLE:
|
||||
return self.ai_generated_dir / 'dalle'
|
||||
elif provider == ImageProvider.STABLE_DIFFUSION:
|
||||
return self.ai_generated_dir / 'stable_diffusion'
|
||||
else:
|
||||
return self.ai_generated_dir
|
||||
|
||||
elif source_type == ImageSourceType.WEB_SEARCH:
|
||||
if provider == ImageProvider.UNSPLASH:
|
||||
return self.web_search_dir / 'unsplash'
|
||||
elif provider == ImageProvider.PIXABAY:
|
||||
return self.web_search_dir / 'pixabay'
|
||||
else:
|
||||
return self.web_search_dir
|
||||
|
||||
elif source_type == ImageSourceType.LOCAL_STORAGE:
|
||||
if provider == ImageProvider.USER_UPLOAD:
|
||||
return self.local_storage_dir / 'user_uploads'
|
||||
else:
|
||||
return self.local_storage_dir / 'processed'
|
||||
|
||||
return self.cache_root
|
||||
|
||||
def _generate_content_hash(self, image_data: bytes) -> str:
|
||||
"""生成图片内容哈希值"""
|
||||
return hashlib.sha256(image_data).hexdigest()
|
||||
|
||||
def _generate_cache_key(self, image_info: ImageInfo, image_data: bytes = None) -> str:
|
||||
"""生成缓存键 - 始终基于内容哈希"""
|
||||
if image_data:
|
||||
return self._generate_content_hash(image_data)
|
||||
|
||||
# 如果没有图片数据，使用图片ID作为临时键（这种情况应该避免）
|
||||
logger.warning(f"Generating cache key without image data for {image_info.image_id}")
|
||||
return hashlib.md5(f"temp_{image_info.image_id}".encode()).hexdigest()
|
||||
|
||||
async def cache_image(self, image_info: ImageInfo, image_data: bytes) -> str:
|
||||
"""缓存图片 - 基于内容去重"""
|
||||
try:
|
||||
# 生成基于内容的缓存键
|
||||
content_hash = self._generate_content_hash(image_data)
|
||||
|
||||
# 检查是否已经缓存了相同内容的图片
|
||||
if content_hash in self._cache_index:
|
||||
existing_cache_info = self._cache_index[content_hash]
|
||||
existing_file_path = Path(existing_cache_info.file_path)
|
||||
|
||||
# 如果文件存在，更新访问信息并保存新的图片元数据引用
|
||||
if existing_file_path.exists():
|
||||
existing_cache_info.update_access()
|
||||
|
||||
# 更新图片信息中的本地路径
|
||||
image_info.local_path = str(existing_file_path)
|
||||
|
||||
# 保存这个新的图片信息作为额外的元数据引用
|
||||
await self._save_image_metadata_reference(content_hash, image_info)
|
||||
|
||||
asyncio.create_task(self._save_cache_index())
|
||||
logger.debug(f"Image content already cached, added new reference: {image_info.image_id}")
|
||||
return content_hash
|
||||
else:
|
||||
# 如果文件不存在，从索引中移除
|
||||
del self._cache_index[content_hash]
|
||||
|
||||
# 选择存储路径 - 优先使用AI生成或网络搜索的路径
|
||||
if image_info.source_type.value in ['ai_generated', 'web_search']:
|
||||
cache_path = self._get_cache_path(image_info.source_type, image_info.provider)
|
||||
else:
|
||||
# 对于本地上传，使用通用路径
|
||||
cache_path = self.local_storage_dir / 'user_uploads'
|
||||
|
||||
# 确定文件扩展名
|
||||
file_extension = Path(image_info.filename).suffix
|
||||
if not file_extension:
|
||||
file_extension = f".{image_info.metadata.format.value}"
|
||||
|
||||
file_path = cache_path / f"{content_hash}{file_extension}"
|
||||
|
||||
# 保存图片文件
|
||||
await asyncio.get_event_loop().run_in_executor(
|
||||
None, self._save_image_file, file_path, image_data
|
||||
)
|
||||
|
||||
# 更新图片信息中的本地路径
|
||||
image_info.local_path = str(file_path)
|
||||
|
||||
# 创建缓存信息 - 移除过期时间，图片永久有效
|
||||
cache_info = ImageCacheInfo(
|
||||
cache_key=content_hash,
|
||||
file_path=str(file_path),
|
||||
file_size=len(image_data),
|
||||
created_at=time.time(),
|
||||
last_accessed=time.time(),
|
||||
access_count=1,
|
||||
expires_at=None # 永不过期
|
||||
)
|
||||
|
||||
# 更新缓存索引
|
||||
self._cache_index[content_hash] = cache_info
|
||||
|
||||
# 保存图片元数据
|
||||
await self._save_image_metadata(content_hash, image_info)
|
||||
|
||||
# 异步保存缓存索引
|
||||
asyncio.create_task(self._save_cache_index())
|
||||
|
||||
logger.debug(f"Image cached successfully: {content_hash}")
|
||||
return content_hash
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to cache image {image_info.image_id}: {e}")
|
||||
raise
|
||||
|
||||
def _save_image_file(self, file_path: Path, image_data: bytes):
|
||||
"""保存图片文件"""
|
||||
with open(file_path, 'wb') as f:
|
||||
f.write(image_data)
|
||||
|
||||
async def get_cached_image(self, cache_key: str) -> Optional[Tuple[ImageInfo, Path]]:
|
||||
"""获取缓存的图片"""
|
||||
try:
|
||||
cache_info = self._cache_index.get(cache_key)
|
||||
if not cache_info:
|
||||
return None
|
||||
|
||||
# 检查文件是否存在
|
||||
file_path = Path(cache_info.file_path)
|
||||
if not file_path.exists():
|
||||
await self.remove_from_cache(cache_key)
|
||||
return None
|
||||
|
||||
# 加载图片元数据
|
||||
image_info = await self._load_image_metadata(cache_key)
|
||||
if not image_info:
|
||||
# 如果没有元数据文件，创建一个基础的ImageInfo对象
|
||||
# 这在PDF转换等场景下是正常的，只需要图片文件本身
|
||||
logger.debug(f"No metadata found for {cache_key}, creating basic ImageInfo")
|
||||
from ..models import ImageInfo, ImageMetadata, ImageFormat, ImageSourceType, ImageProvider
|
||||
|
||||
# 从文件路径推断基本信息
|
||||
file_extension = file_path.suffix.lower()
|
||||
format_map = {
|
||||
'.jpg': ImageFormat.JPEG,
|
||||
'.jpeg': ImageFormat.JPEG,
|
||||
'.png': ImageFormat.PNG,
|
||||
'.webp': ImageFormat.WEBP,
|
||||
'.gif': ImageFormat.GIF
|
||||
}
|
||||
|
||||
image_format = format_map.get(file_extension, ImageFormat.JPEG)
|
||||
|
||||
# 创建基础元数据
|
||||
metadata = ImageMetadata(
|
||||
format=image_format,
|
||||
width=0, # 未知尺寸
|
||||
height=0,
|
||||
file_size=cache_info.file_size
|
||||
)
|
||||
|
||||
# 创建基础ImageInfo
|
||||
image_info = ImageInfo(
|
||||
image_id=cache_key,
|
||||
filename=file_path.name,
|
||||
local_path=str(file_path),
|
||||
source_type=ImageSourceType.LOCAL_STORAGE,
|
||||
provider=ImageProvider.LOCAL_STORAGE,
|
||||
metadata=metadata
|
||||
)
|
||||
|
||||
# 更新访问信息
|
||||
cache_info.update_access()
|
||||
asyncio.create_task(self._save_cache_index())
|
||||
|
||||
logger.debug(f"Cache hit: {cache_key}")
|
||||
return image_info, file_path
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get cached image {cache_key}: {e}")
|
||||
return None
|
||||
|
||||
async def is_cached(self, image_info: ImageInfo, image_data: bytes = None) -> Optional[str]:
|
||||
"""检查图片是否已缓存"""
|
||||
# 如果有图片数据，使用内容哈希检查
|
||||
if image_data:
|
||||
cache_key = self._generate_cache_key(image_info, image_data)
|
||||
cache_info = self._cache_index.get(cache_key)
|
||||
if cache_info:
|
||||
file_path = Path(cache_info.file_path)
|
||||
if file_path.exists():
|
||||
return cache_key
|
||||
|
||||
# 否则使用传统方式检查
|
||||
cache_key = self._generate_cache_key(image_info)
|
||||
cache_info = self._cache_index.get(cache_key)
|
||||
|
||||
if cache_info:
|
||||
file_path = Path(cache_info.file_path)
|
||||
if file_path.exists():
|
||||
return cache_key
|
||||
|
||||
return None
|
||||
|
||||
async def remove_from_cache(self, cache_key: str) -> bool:
|
||||
"""从缓存中移除图片"""
|
||||
try:
|
||||
cache_info = self._cache_index.get(cache_key)
|
||||
if not cache_info:
|
||||
return False
|
||||
|
||||
# 删除图片文件
|
||||
file_path = Path(cache_info.file_path)
|
||||
if file_path.exists():
|
||||
await asyncio.get_event_loop().run_in_executor(None, file_path.unlink)
|
||||
|
||||
# 删除缩略图
|
||||
thumbnail_path = self.thumbnails_dir / f"{cache_key}.jpg"
|
||||
if thumbnail_path.exists():
|
||||
await asyncio.get_event_loop().run_in_executor(None, thumbnail_path.unlink)
|
||||
|
||||
# 删除元数据
|
||||
metadata_path = self.metadata_dir / f"{cache_key}.json"
|
||||
if metadata_path.exists():
|
||||
await asyncio.get_event_loop().run_in_executor(None, metadata_path.unlink)
|
||||
|
||||
# 从索引中移除
|
||||
del self._cache_index[cache_key]
|
||||
|
||||
# 保存索引
|
||||
await self._save_cache_index()
|
||||
|
||||
logger.info(f"Removed from cache: {cache_key}")
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to remove from cache {cache_key}: {e}")
|
||||
return False
|
||||
|
||||
async def _save_image_metadata(self, cache_key: str, image_info: ImageInfo):
|
||||
"""保存图片元数据"""
|
||||
metadata_path = self.metadata_dir / f"{cache_key}.json"
|
||||
metadata = image_info.model_dump()
|
||||
|
||||
def _save():
|
||||
with open(metadata_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
||||
|
||||
await asyncio.get_event_loop().run_in_executor(None, _save)
|
||||
|
||||
async def _save_image_metadata_reference(self, content_hash: str, image_info: ImageInfo):
|
||||
"""保存图片元数据引用 - 为同一内容的图片保存多个引用"""
|
||||
# 创建引用文件名：content_hash_image_id.json
|
||||
reference_filename = f"{content_hash}_{image_info.image_id}.json"
|
||||
reference_path = self.metadata_dir / "references" / reference_filename
|
||||
|
||||
# 确保引用目录存在
|
||||
reference_path.parent.mkdir(exist_ok=True)
|
||||
|
||||
metadata = image_info.model_dump()
|
||||
|
||||
def _save():
|
||||
with open(reference_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(metadata, f, ensure_ascii=False, indent=2)
|
||||
|
||||
await asyncio.get_event_loop().run_in_executor(None, _save)
|
||||
logger.debug(f"Saved metadata reference: {reference_filename}")
|
||||
|
||||
async def _load_image_metadata(self, cache_key: str) -> Optional[ImageInfo]:
|
||||
"""加载图片元数据"""
|
||||
metadata_path = self.metadata_dir / f"{cache_key}.json"
|
||||
if not metadata_path.exists():
|
||||
# 降低日志级别，避免在PDF转换等场景下产生不必要的警告
|
||||
logger.debug(f"Metadata file not found: {metadata_path}")
|
||||
return None
|
||||
|
||||
def _load():
|
||||
with open(metadata_path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
|
||||
try:
|
||||
metadata = await asyncio.get_event_loop().run_in_executor(None, _load)
|
||||
image_info = ImageInfo(**metadata)
|
||||
logger.debug(f"Successfully loaded metadata for cache key: {cache_key}")
|
||||
return image_info
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load image metadata {cache_key}: {e}")
|
||||
logger.error(f"Metadata file path: {metadata_path}")
|
||||
return None
|
||||
|
||||
def _load_cache_index(self):
|
||||
"""加载缓存索引 - 简化版本，直接扫描文件系统"""
|
||||
try:
|
||||
# 清空现有索引
|
||||
self._cache_index.clear()
|
||||
|
||||
# 扫描所有缓存目录
|
||||
cache_dirs = [
|
||||
self.ai_generated_dir,
|
||||
self.web_search_dir,
|
||||
self.local_storage_dir
|
||||
]
|
||||
|
||||
total_files = 0
|
||||
for cache_dir in cache_dirs:
|
||||
if cache_dir.exists():
|
||||
# 递归扫描所有图片文件
|
||||
for file_path in cache_dir.rglob('*'):
|
||||
if file_path.is_file() and file_path.suffix.lower() in ['.jpg', '.jpeg', '.png', '.webp', '.gif']:
|
||||
try:
|
||||
# 使用文件名作为缓存键
|
||||
cache_key = file_path.stem
|
||||
|
||||
# 创建缓存信息
|
||||
stat = file_path.stat()
|
||||
cache_info = ImageCacheInfo(
|
||||
cache_key=cache_key,
|
||||
file_path=str(file_path),
|
||||
file_size=stat.st_size,
|
||||
created_at=stat.st_ctime,
|
||||
last_accessed=stat.st_atime,
|
||||
access_count=1,
|
||||
expires_at=None # 永不过期
|
||||
)
|
||||
|
||||
self._cache_index[cache_key] = cache_info
|
||||
total_files += 1
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to process cache file {file_path}: {e}")
|
||||
|
||||
logger.debug(f"Loaded {total_files} cached images from filesystem")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to load cache index: {e}")
|
||||
# 确保索引不为空，即使出错也继续运行
|
||||
if not self._cache_index:
|
||||
self._cache_index = {}
|
||||
|
||||
async def _save_cache_index(self):
|
||||
"""保存缓存索引 - 简化版本，不再保存JSON文件"""
|
||||
# 不再保存索引文件，直接从文件系统读取
|
||||
logger.debug("Cache index save skipped - using filesystem-based indexing")
|
||||
|
||||
# 移除清理任务相关方法，因为图片永久有效，无需清理
|
||||
|
||||
async def get_cache_size(self) -> int:
|
||||
"""获取缓存总大小"""
|
||||
total_size = 0
|
||||
for cache_info in self._cache_index.values():
|
||||
total_size += cache_info.file_size
|
||||
return total_size
|
||||
|
||||
async def get_cache_stats(self) -> Dict[str, Any]:
|
||||
"""获取缓存统计信息"""
|
||||
total_entries = len(self._cache_index)
|
||||
total_size = await self.get_cache_size()
|
||||
|
||||
# 按来源类型统计
|
||||
source_stats = {}
|
||||
for cache_info in self._cache_index.values():
|
||||
file_path = Path(cache_info.file_path)
|
||||
source_type = file_path.parent.parent.name if file_path.parent.parent.name in ['ai_generated', 'web_search', 'local_storage'] else 'unknown'
|
||||
|
||||
if source_type not in source_stats:
|
||||
source_stats[source_type] = {'count': 0, 'size': 0}
|
||||
|
||||
source_stats[source_type]['count'] += 1
|
||||
source_stats[source_type]['size'] += cache_info.file_size
|
||||
|
||||
# 转换为API需要的格式
|
||||
categories = {}
|
||||
for source_type, stats in source_stats.items():
|
||||
categories[source_type] = stats['count']
|
||||
|
||||
return {
|
||||
'total_entries': total_entries,
|
||||
'total_size_bytes': total_size,
|
||||
'total_size_mb': total_size / (1024 * 1024),
|
||||
'categories': categories,
|
||||
'source_stats': source_stats
|
||||
}
|
||||
|
||||
async def clear_cache(self, source_type: Optional[ImageSourceType] = None) -> int:
|
||||
"""清空缓存"""
|
||||
if source_type:
|
||||
# 清空特定来源的缓存
|
||||
keys_to_remove = []
|
||||
for cache_key, cache_info in self._cache_index.items():
|
||||
file_path = Path(cache_info.file_path)
|
||||
if source_type.value in str(file_path):
|
||||
keys_to_remove.append(cache_key)
|
||||
|
||||
for cache_key in keys_to_remove:
|
||||
await self.remove_from_cache(cache_key)
|
||||
|
||||
return len(keys_to_remove)
|
||||
else:
|
||||
# 清空所有缓存
|
||||
keys_to_remove = list(self._cache_index.keys())
|
||||
|
||||
# 先通过索引删除已知的文件
|
||||
for cache_key in keys_to_remove:
|
||||
await self.remove_from_cache(cache_key)
|
||||
|
||||
# 然后清理可能存在的孤立文件
|
||||
await self._clear_orphaned_files()
|
||||
|
||||
return len(keys_to_remove)
|
||||
|
||||
async def _clear_orphaned_files(self):
|
||||
"""清理孤立的缓存文件"""
|
||||
try:
|
||||
def _clear_directory_contents(directory: Path):
|
||||
"""清理目录中的所有文件，但保留目录结构"""
|
||||
if not directory.exists():
|
||||
return
|
||||
|
||||
for item in directory.iterdir():
|
||||
if item.is_file():
|
||||
item.unlink()
|
||||
elif item.is_dir():
|
||||
# 递归清理子目录
|
||||
_clear_directory_contents(item)
|
||||
|
||||
# 清理所有缓存目录中的文件
|
||||
directories_to_clear = [
|
||||
self.ai_generated_dir,
|
||||
self.web_search_dir,
|
||||
self.local_storage_dir,
|
||||
self.metadata_dir,
|
||||
self.thumbnails_dir
|
||||
]
|
||||
|
||||
for directory in directories_to_clear:
|
||||
await asyncio.get_event_loop().run_in_executor(None, _clear_directory_contents, directory)
|
||||
|
||||
logger.info("Cleared all orphaned cache files")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to clear orphaned files: {e}")
|
||||
# 不抛出异常，因为主要的清理已经完成
|
||||
|
||||
async def deduplicate_cache(self) -> int:
|
||||
"""去重缓存中的重复图片"""
|
||||
try:
|
||||
content_hash_to_keys = {}
|
||||
keys_to_remove = []
|
||||
|
||||
# 按内容哈希分组
|
||||
for cache_key, cache_info in self._cache_index.items():
|
||||
file_path = Path(cache_info.file_path)
|
||||
if not file_path.exists():
|
||||
keys_to_remove.append(cache_key)
|
||||
continue
|
||||
|
||||
try:
|
||||
# 读取文件内容并计算哈希
|
||||
with open(file_path, 'rb') as f:
|
||||
content = f.read()
|
||||
content_hash = hashlib.md5(content).hexdigest()
|
||||
|
||||
if content_hash not in content_hash_to_keys:
|
||||
content_hash_to_keys[content_hash] = []
|
||||
content_hash_to_keys[content_hash].append((cache_key, cache_info))
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to read file for deduplication {file_path}: {e}")
|
||||
keys_to_remove.append(cache_key)
|
||||
|
||||
# 对于每个内容哈希，保留最新的一个，删除其他的
|
||||
for content_hash, key_info_list in content_hash_to_keys.items():
|
||||
if len(key_info_list) > 1:
|
||||
# 按创建时间排序，保留最新的
|
||||
key_info_list.sort(key=lambda x: x[1].created_at, reverse=True)
|
||||
|
||||
# 删除除了第一个（最新的）之外的所有条目
|
||||
for cache_key, cache_info in key_info_list[1:]:
|
||||
keys_to_remove.append(cache_key)
|
||||
logger.info(f"Marking duplicate cache entry for removal: {cache_key}")
|
||||
|
||||
# 删除重复和无效的条目
|
||||
for cache_key in keys_to_remove:
|
||||
await self.remove_from_cache(cache_key)
|
||||
|
||||
logger.info(f"Removed {len(keys_to_remove)} duplicate/invalid cache entries")
|
||||
return len(keys_to_remove)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to deduplicate cache: {e}")
|
||||
return 0
|
||||
Reference in New Issue
Block a user