145 lines
4.9 KiB
Python
Executable File
145 lines
4.9 KiB
Python
Executable File
from typing import Optional
|
||
from pymongo import MongoClient
|
||
from pymongo.collection import Collection
|
||
from pymongo.errors import PyMongoError
|
||
import logging
|
||
|
||
from datetime import datetime
|
||
|
||
from src.core.config import MONGO_URI
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class ThreadImageMinIOStore:
|
||
"""
|
||
根據 thread_id 存取/更新 current_image 的 MinIO 物件路徑（不存 binary）
|
||
|
||
儲存格式範例：
|
||
{
|
||
"thread_id": "thread_abc123",
|
||
"current_image_path": "images/2025/03/thread_abc123_latest.png",
|
||
"updated_at": ISODate,
|
||
"metadata": {"format": "png", "desc": "生成的貓圖", "size_bytes": 512345}
|
||
}
|
||
|
||
使用方式：
|
||
store = ThreadImageMinIOStore("mongodb://localhost:27017/", "deepagents_db")
|
||
store.save_image_path("thread_abc123", "images/cat/001.png", "https://minio.example.com/bucket/images/cat/001.png")
|
||
path_info = store.get_image_path("thread_abc123")
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
mongo_uri: str,
|
||
db_name: str = "deepagents_db",
|
||
collection_name: str = "agent_image_paths",
|
||
connect_timeout_ms: int = 5000,
|
||
server_selection_timeout_ms: int = 5000,
|
||
):
|
||
self.client = MongoClient(
|
||
mongo_uri,
|
||
connectTimeoutMS=connect_timeout_ms,
|
||
serverSelectionTimeoutMS=server_selection_timeout_ms,
|
||
retryWrites=True,
|
||
retryReads=True,
|
||
)
|
||
self.db = self.client[db_name]
|
||
self.collection: Collection = self.db[collection_name]
|
||
|
||
# 建立唯一索引
|
||
self.collection.create_index("thread_id", unique=True)
|
||
|
||
def save_image_path(
|
||
self,
|
||
thread_id: str,
|
||
object_path: list, # MinIO 中的相對路徑，例如 "test/123.png" 或 "images/20250320/abc.png"
|
||
metadata: Optional[dict] = None
|
||
) -> bool:
|
||
"""
|
||
保存或更新某個 thread 的 current_image MinIO 路徑
|
||
|
||
Args:
|
||
thread_id: 對話執行緒 ID
|
||
object_path: MinIO bucket 內的物件路徑 (不含 bucket 名稱)
|
||
metadata: 額外資訊，例如 {"prompt": "...", "format": "png", "width": 1024}
|
||
|
||
Returns:
|
||
bool: 是否成功
|
||
"""
|
||
document = {
|
||
"thread_id": thread_id,
|
||
"current_image_path": object_path,
|
||
"updated_at": datetime.now(),
|
||
"metadata": metadata or {}
|
||
}
|
||
|
||
try:
|
||
result = self.collection.update_one(
|
||
{"thread_id": thread_id},
|
||
{"$set": document},
|
||
upsert=True
|
||
)
|
||
action = "updated" if result.modified_count > 0 else "inserted"
|
||
logger.info(f"Image path for thread {thread_id} {action}: {object_path}")
|
||
return True
|
||
except PyMongoError as e:
|
||
logger.error(f"Failed to save image path for thread {thread_id}: {e}")
|
||
return False
|
||
|
||
def get_image_path(self, thread_id: str) -> Optional[dict]:
|
||
"""
|
||
取得某 thread 的 current_image MinIO 資訊
|
||
|
||
Returns:
|
||
{
|
||
"current_image_path": str,
|
||
"updated_at": datetime,
|
||
"metadata": dict
|
||
} 或 None
|
||
"""
|
||
try:
|
||
doc = self.collection.find_one({"thread_id": thread_id})
|
||
if not doc:
|
||
return None
|
||
|
||
return {
|
||
"current_image_path": doc.get("current_image_path"),
|
||
"updated_at": doc.get("updated_at"),
|
||
"metadata": doc.get("metadata", {})
|
||
}
|
||
except PyMongoError as e:
|
||
logger.error(f"Failed to get image path for thread {thread_id}: {e}")
|
||
return None
|
||
|
||
def get_object_path_only(self, thread_id: str) -> Optional[str]:
|
||
"""只取 MinIO 相對路徑，方便直接給 MinIO client 使用"""
|
||
info = self.get_image_path(thread_id)
|
||
return info["current_image_path"] if info else None
|
||
|
||
def delete_image_path(self, thread_id: str) -> bool:
|
||
"""刪除某 thread 的記錄（不影響 MinIO 實際檔案）"""
|
||
try:
|
||
result = self.collection.delete_one({"thread_id": thread_id})
|
||
if result.deleted_count > 0:
|
||
logger.info(f"Image path record for thread {thread_id} deleted")
|
||
return True
|
||
return False
|
||
except PyMongoError as e:
|
||
logger.error(f"Failed to delete image path for thread {thread_id}: {e}")
|
||
return False
|
||
|
||
def close(self):
|
||
self.client.close()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
image_store = ThreadImageMinIOStore(MONGO_URI, "agent_tool_generate_db")
|
||
success = image_store.save_image_path(
|
||
thread_id="121233",
|
||
object_path=["test/123.png"],
|
||
metadata={"prompt": "prompt", "generated_at": str(datetime.now())})
|
||
print(success)
|
||
info = image_store.get_image_path("121233")
|
||
print(info)
|