FiDA_Python/src/server/deep_agent/utils/mongodb_util.py

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)