Python之DeepAgents基于mongodb的检查点

from __future__ import annotations

from datetime import datetime, timezone
from typing import Any, Iterator, AsyncIterator

from bson import Binary
from langchain_core.runnables import RunnableConfig
from langgraph.checkpoint.base import (
    BaseCheckpointSaver,
    Checkpoint,
    CheckpointMetadata,
    CheckpointTuple,
    ChannelVersions,
)
from motor.motor_asyncio import AsyncIOMotorClient
from pymongo import MongoClient


class MongoDBCheckpointSaver(BaseCheckpointSaver):
    def __init__(
        self,
        uri: str = "mongodb://localhost:27017/",
        db_name: str = "langgraph",
        collection_name: str = "checkpoints"
    ):
        super().__init__()
        self.client = MongoClient(uri)
        self.db = self.client[db_name]
        self.collection = self.db[collection_name]

        self.async_client = AsyncIOMotorClient(uri)
        self.async_db = self.async_client[db_name]
        self.async_collection = self.async_db[collection_name]

        self.collection.create_index("thread_id")
        self.collection.create_index([("thread_id", 1), ("checkpoint_id", 1)], unique=True)

    # -------------------------- 真正源码级正确 --------------------------
    def _serialize(self, obj: Any) -> dict:
        # dumps_typed 返回 (type: str, data: bytes)
        typ, data = self.serde.dumps_typed(obj)
        return {"type": typ, "data": Binary(data)}

    def _deserialize(self, data: dict) -> Any:
        return self.serde.loads_typed((data["type"], data["data"]))

    # ------------------------------ 同步方法 ------------------------------
    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
        thread_id = config["configurable"]["thread_id"]
        doc = self.collection.find_one({"thread_id": thread_id}, sort=[("checkpoint_id", -1)])
        if not doc:
            return None

        return CheckpointTuple(
            config=config,
            checkpoint=self._deserialize(doc["checkpoint"]),
            metadata=self._deserialize(doc["metadata"]),
            parent_config=None,
        )

    def list(
        self,
        config: RunnableConfig | None,
        **kwargs: Any
    ) -> Iterator[CheckpointTuple]:
        query = {}
        if config and "configurable" in config:
            thread_id = config["configurable"].get("thread_id")
            if thread_id:
                query["thread_id"] = thread_id

        for doc in self.collection.find(query, sort=[("checkpoint_id", -1)]):
            yield CheckpointTuple(
                config={"configurable": {"thread_id": doc["thread_id"]}},
                checkpoint=self._deserialize(doc["checkpoint"]),
                metadata=self._deserialize(doc["metadata"]),
            )

    def put(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_id = checkpoint["id"]

        self.collection.update_one(
            {"thread_id": thread_id, "checkpoint_id": checkpoint_id},
            {
                "$set": {
                    "thread_id": thread_id,
                    "checkpoint_id": checkpoint_id,
                    "checkpoint": self._serialize(checkpoint),
                    "metadata": self._serialize(metadata),
                    "updated_at": datetime.now(timezone.utc),
                }
            },
            upsert=True,
        )
        return config

    def put_writes(self, *args, **kwargs):
        pass

    # ------------------------------ 异步方法 ------------------------------
    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
        thread_id = config["configurable"]["thread_id"]
        doc = await self.async_collection.find_one(
            {"thread_id": thread_id}, sort=[("checkpoint_id", -1)]
        )
        if not doc:
            return None

        return CheckpointTuple(
            config=config,
            checkpoint=self._deserialize(doc["checkpoint"]),
            metadata=self._deserialize(doc["metadata"]),
        )

    async def alist(
        self, config: RunnableConfig | None, **kwargs: Any
    ) -> AsyncIterator[CheckpointTuple]:
        query = {}
        if config and "configurable" in config:
            thread_id = config["configurable"].get("thread_id")
            if thread_id:
                query["thread_id"] = thread_id

        cursor = self.async_collection.find(query).sort("checkpoint_id", -1)
        async for doc in cursor:
            yield CheckpointTuple(
                config={"configurable": {"thread_id": doc["thread_id"]}},
                checkpoint=self._deserialize(doc["checkpoint"]),
                metadata=self._deserialize(doc["metadata"]),
            )

    async def aput(
        self,
        config: RunnableConfig,
        checkpoint: Checkpoint,
        metadata: CheckpointMetadata,
        new_versions: ChannelVersions,
    ) -> RunnableConfig:
        thread_id = config["configurable"]["thread_id"]
        checkpoint_id = checkpoint["id"]

        await self.async_collection.update_one(
            {"thread_id": thread_id, "checkpoint_id": checkpoint_id},
            {
                "$set": {
                    "thread_id": thread_id,
                    "checkpoint_id": checkpoint_id,
                    "checkpoint": self._serialize(checkpoint),
                    "metadata": self._serialize(metadata),
                    "updated_at": datetime.now(timezone.utc),
                }
            },
            upsert=True,
        )
        return config

    async def aput_writes(self, *args, **kwargs):
        pass

import urllib.parse

from util.ConfigManager import config


mongo_ip = config.get("MONGODB").get("IP")
mongo_port = config.get("MONGODB").get("PORT")
mongo_user = config.get("MONGODB").get("USER")
mongo_pwd = config.get("MONGODB").get("PWD")
mongo_db = config.get("MONGODB").get("DB")
mongo_collection = config.get("MONGODB").get("COLLECTION")


def build_connection_string() -> str:
    if mongo_user and mongo_pwd:
        username = urllib.parse.quote_plus(mongo_user)
        password = urllib.parse.quote_plus(mongo_pwd)
        print(f"mongodb://{username}:{password}@{mongo_ip}:{mongo_port}")
        return f"mongodb://{username}:{password}@{mongo_ip}:{mongo_port}"
    else:
        return f"mongodb://{mongo_ip}:{mongo_port}"

# 实例化一个全局的 MongoDBCheckpointSaver
mongo_db_checkpoint_saver = MongoDBCheckpointSaver(
    uri=build_connection_string(),
    db_name=mongo_db,
    collection_name=mongo_collection
)

session_key = f"{tenant_id}_{session_id}"

if session_key in cls._instances:
    return cls._instances[session_key]

# 设置权限
permissions = [
    FilesystemPermission(
        operations=["read", "write"],
        paths=[f"/workspace/{tenant_id}/{session_id}/**"],
        mode="allow"
    )
]
# 设置后端存储，按租户和会话隔离
backend = CompositeBackend(
    default=StateBackend(),
    routes={
        f"/memories/{tenant_id}/{session_id}/": FilesystemBackend(
            root_dir=f"./memories/{tenant_id}/{session_id}",
            virtual_mode=True
        ),
        f"/workspace/{tenant_id}/{session_id}/": FilesystemBackend(
            root_dir=f"./workspace/{tenant_id}/{session_id}",
            virtual_mode=True
        )
    }
)

# 使用默认系统提示词或自定义提示词
sys_prompt = system_prompt or "You are a research assistant."

# 合并工具列表
tools = custom_tools if custom_tools else [current_data_quality]

# 创建DeepAgent实例
agent = create_deep_agent(
    model=open_ai_client,
    system_prompt=sys_prompt,
    middleware=[
        log_tool_calls
    ],
    tools=tools,
    permissions=permissions,
    backend=backend,
    checkpointer=mongo_db_checkpoint_saver
)

from bson import Binary
from pymongo import MongoClient
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer

from util.MongoPool import build_connection_string, mongo_db

# 固定序列化器
_SERDE = JsonPlusSerializer()


def decode_typed_item(item: dict) -> dict:
    typ = item["type"]
    data = item["data"]
    return _SERDE.loads_typed((typ, data))


def query_chat_messages(
    thread_id: str,
    mongo_uri="mongodb://localhost:27017/",
    db_name="langgraph",
    collection_name="checkpoints"
):
    client = MongoClient(mongo_uri)
    coll = client[db_name][collection_name]

    # 查最新 checkpoint
    doc = coll.find_one({"thread_id": thread_id}, sort=[("checkpoint_id", -1)])
    if not doc:
        return []

    # 解析
    checkpoint = decode_typed_item(doc["checkpoint"])
    messages = checkpoint.get("channel_values", {}).get("messages", [])

    # 格式化：把 LangChain 消息对象 转成 普通字典
    result = []
    for msg in messages:
        try:
            # 消息对象 → 字典
            msg_dict = msg.model_dump()
            # print("-------------------------------------------------")
            # print(msg_dict)
            # print("-------------------------------------------------")
            role = msg_dict.get("type", "ai")
            content = msg_dict.get("content", "")
            result.append({"role": role, "content": content})
        except Exception as e:
            continue

    return result


# ============ 直接使用 ============
if __name__ == "__main__":
    mongodb_url = build_connection_string()
    db_name = mongo_db
    # 替换成你的 thread_id
    thread_id = "sdrowero234249jdoisjfoiet93xxxxx"

    msgs = query_chat_messages(
        thread_id=thread_id,
        mongo_uri=mongodb_url,
        db_name=db_name
    )
    for item in msgs:
        print(f"【{item['role']}】\n{item['content']}\n")