langgraph-cosmosdb-checkpointer 0.0.0__py3-none-any.whl

langgraph_cosmosdb_checkpointer/__init__.py

@@ -0,0 +1,11 @@
+from .cosmosdb import (
+    AsyncCosmosDBSaver,
+    CosmosDBSaver,
+    CosmosDBSettings,
+)
+
+__all__ = [
+    "CosmosDBSettings",
+    "CosmosDBSaver",
+    "AsyncCosmosDBSaver",
+]

langgraph_cosmosdb_checkpointer/cosmosdb.py

@@ -0,0 +1,1074 @@
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import json
+import random
+import zlib
+from contextlib import AbstractAsyncContextManager, AbstractContextManager
+from dataclasses import dataclass
+from typing import Any, AsyncIterator, Iterator, Mapping, Sequence
+
+from azure.core import MatchConditions
+from azure.core.credentials import TokenCredential
+from azure.cosmos import CosmosClient, PartitionKey
+from azure.cosmos.container import ContainerProxy
+from azure.cosmos.database import DatabaseProxy
+from azure.cosmos.exceptions import (
+    CosmosAccessConditionFailedError,
+    CosmosHttpResponseError,
+    CosmosResourceExistsError,
+    CosmosResourceNotFoundError,
+)
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.base import (
+    WRITES_IDX_MAP,
+    BaseCheckpointSaver,
+    ChannelVersions,
+    Checkpoint,
+    CheckpointMetadata,
+    CheckpointTuple,
+    SerializerProtocol,
+    get_checkpoint_id,
+    get_checkpoint_metadata,
+)
+from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
+
+MAX_COSMOS_ITEM_BYTES = 2_000_000
+DEFAULT_INLINE_PAYLOAD_LIMIT_BYTES = 1_500_000
+DEFAULT_CHUNK_SIZE_BYTES = 512_000
+DEFAULT_COMPRESS_ABOVE_BYTES = 16_384
+
+
+@dataclass(slots=True)
+class CosmosDBSettings:
+    database_name: str = "langgraph"
+    checkpoint_container_name: str = "checkpoints"
+    writes_container_name: str = "writes"
+    blob_container_name: str = "checkpoint_blobs"
+    create_containers: bool = True
+    retry_total: int | None = None
+    retry_backoff_max: int | None = None
+    retry_fixed_interval: int | None = None
+    inline_payload_limit_bytes: int = DEFAULT_INLINE_PAYLOAD_LIMIT_BYTES
+    chunk_size_bytes: int = DEFAULT_CHUNK_SIZE_BYTES
+    compress_above_bytes: int = DEFAULT_COMPRESS_ABOVE_BYTES
+    no_response_on_write: bool = False
+
+
+class CosmosDBSaver(BaseCheckpointSaver[str], AbstractContextManager):
+    def __init__(
+        self,
+        checkpoint_container: ContainerProxy,
+        writes_container: ContainerProxy,
+        blob_container: ContainerProxy | None,
+        *,
+        settings: CosmosDBSettings | None = None,
+        serde: SerializerProtocol | None = None,
+        client: CosmosClient | None = None,
+        owns_client: bool = False,
+    ) -> None:
+        super().__init__(serde=serde or JsonPlusSerializer())
+        self._checkpoints = checkpoint_container
+        self._writes = writes_container
+        self._blobs = blob_container
+        self._settings = settings or CosmosDBSettings()
+        self._client = client
+        self._owns_client = owns_client
+
+    @classmethod
+    def from_connection_string(
+        cls,
+        connection_string: str,
+        *,
+        settings: CosmosDBSettings | None = None,
+        credential: str | TokenCredential | Mapping[str, Any] | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> "CosmosDBSaver":
+        cfg = settings or CosmosDBSettings()
+        client_kwargs = _client_kwargs(cfg)
+        client = CosmosClient.from_connection_string(
+            connection_string,
+            credential=credential,
+            **client_kwargs,
+        )
+        checkpoint_container, writes_container, blob_container = _resolve_containers_sync(
+            client=client,
+            settings=cfg,
+        )
+        return cls(
+            checkpoint_container=checkpoint_container,
+            writes_container=writes_container,
+            blob_container=blob_container,
+            settings=cfg,
+            serde=serde,
+            client=client,
+            owns_client=True,
+        )
+
+    @classmethod
+    def from_endpoint(
+        cls,
+        endpoint: str,
+        *,
+        credential: str | TokenCredential | Mapping[str, Any] | None = None,
+        settings: CosmosDBSettings | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> "CosmosDBSaver":
+        cfg = settings or CosmosDBSettings()
+        credential_value = credential
+        if credential_value is None:
+            from azure.identity import DefaultAzureCredential
+
+            credential_value = DefaultAzureCredential()
+        client = CosmosClient(endpoint, credential=credential_value, **_client_kwargs(cfg))
+        checkpoint_container, writes_container, blob_container = _resolve_containers_sync(
+            client=client,
+            settings=cfg,
+        )
+        return cls(
+            checkpoint_container=checkpoint_container,
+            writes_container=writes_container,
+            blob_container=blob_container,
+            settings=cfg,
+            serde=serde,
+            client=client,
+            owns_client=True,
+        )
+
+    def __enter__(self) -> "CosmosDBSaver":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> bool | None:
+        self.close()
+        return None
+
+    def close(self) -> None:
+        if self._owns_client and self._client is not None:
+            self._client.__exit__(None, None, None)
+            self._client = None
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        thread_id, checkpoint_ns, checkpoint_id = _normalize_config(config)
+        if checkpoint_id:
+            doc_id = _checkpoint_doc_id(checkpoint_ns, checkpoint_id)
+            try:
+                doc = self._checkpoints.read_item(item=doc_id, partition_key=thread_id)
+            except CosmosResourceNotFoundError:
+                return None
+            return self._checkpoint_tuple_from_doc(
+                doc=doc,
+                thread_id=thread_id,
+                checkpoint_ns=checkpoint_ns,
+            )
+
+        rows = list(
+            self._checkpoints.query_items(
+                query=(
+                    "SELECT TOP 1 * FROM c "
+                    "WHERE c.doc_type=@doc_type AND c.thread_id=@thread_id AND c.checkpoint_ns=@checkpoint_ns "
+                    "ORDER BY c.checkpoint_id DESC"
+                ),
+                parameters=[
+                    {"name": "@doc_type", "value": "checkpoint"},
+                    {"name": "@thread_id", "value": thread_id},
+                    {"name": "@checkpoint_ns", "value": checkpoint_ns},
+                ],
+                partition_key=thread_id,
+            )
+        )
+        if not rows:
+            return None
+        return self._checkpoint_tuple_from_doc(
+            doc=rows[0],
+            thread_id=thread_id,
+            checkpoint_ns=checkpoint_ns,
+        )
+
+    def list(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[CheckpointTuple]:
+        if config:
+            thread_id, checkpoint_ns, checkpoint_id = _normalize_config(config)
+        else:
+            thread_id, checkpoint_ns, checkpoint_id = None, None, None
+
+        before_checkpoint_id = get_checkpoint_id(before) if before else None
+
+        query = "SELECT * FROM c WHERE c.doc_type=@doc_type"
+        parameters: list[dict[str, Any]] = [{"name": "@doc_type", "value": "checkpoint"}]
+        if thread_id is not None:
+            query += " AND c.thread_id=@thread_id"
+            parameters.append({"name": "@thread_id", "value": thread_id})
+        if checkpoint_ns is not None:
+            query += " AND c.checkpoint_ns=@checkpoint_ns"
+            parameters.append({"name": "@checkpoint_ns", "value": checkpoint_ns})
+        if checkpoint_id is not None:
+            query += " AND c.checkpoint_id=@checkpoint_id"
+            parameters.append({"name": "@checkpoint_id", "value": checkpoint_id})
+        if before_checkpoint_id is not None:
+            query += " AND c.checkpoint_id < @before_checkpoint_id"
+            parameters.append(
+                {"name": "@before_checkpoint_id", "value": before_checkpoint_id}
+            )
+        query += " ORDER BY c.checkpoint_id DESC"
+
+        query_kwargs: dict[str, Any] = {
+            "query": query,
+            "parameters": parameters,
+        }
+        if thread_id is None:
+            query_kwargs["enable_cross_partition_query"] = True
+        else:
+            query_kwargs["partition_key"] = thread_id
+
+        yielded = 0
+        for row in self._checkpoints.query_items(**query_kwargs):
+            tuple_value = self._checkpoint_tuple_from_doc(
+                doc=row,
+                thread_id=row["thread_id"],
+                checkpoint_ns=row["checkpoint_ns"],
+            )
+            if filter and not all(
+                tuple_value.metadata.get(k) == v for k, v in filter.items()
+            ):
+                continue
+            yield tuple_value
+            yielded += 1
+            if limit is not None and yielded >= limit:
+                return
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        thread_id, checkpoint_ns, parent_checkpoint_id = _normalize_config(config)
+        checkpoint_core = checkpoint.copy()
+        channel_values = checkpoint_core.pop("channel_values", {})
+
+        for channel, version in new_versions.items():
+            self._store_channel_blob(
+                thread_id=thread_id,
+                checkpoint_ns=checkpoint_ns,
+                channel=channel,
+                version=version,
+                value=channel_values.get(channel),
+            )
+
+        checkpoint_type, checkpoint_payload, checkpoint_hash = self._serialize_typed(
+            value=checkpoint_core,
+            thread_id=thread_id,
+            payload_scope=f"checkpoint:{checkpoint_ns}:{checkpoint['id']}",
+        )
+        metadata_dict = get_checkpoint_metadata(config, metadata)
+        metadata_type, metadata_payload, metadata_hash = self._serialize_typed(
+            value=metadata_dict,
+            thread_id=thread_id,
+            payload_scope=f"metadata:{checkpoint_ns}:{checkpoint['id']}",
+        )
+
+        doc = {
+            "id": _checkpoint_doc_id(checkpoint_ns, checkpoint["id"]),
+            "doc_type": "checkpoint",
+            "thread_id": thread_id,
+            "checkpoint_ns": checkpoint_ns,
+            "checkpoint_id": checkpoint["id"],
+            "parent_checkpoint_id": parent_checkpoint_id,
+            "checkpoint_type": checkpoint_type,
+            "checkpoint_payload": checkpoint_payload,
+            "checkpoint_hash": checkpoint_hash,
+            "metadata_type": metadata_type,
+            "metadata_payload": metadata_payload,
+            "metadata_hash": metadata_hash,
+            "created_at": checkpoint.get("ts"),
+            "run_id": metadata_dict.get("run_id"),
+        }
+        _assert_item_size(doc, scope=f"checkpoint:{checkpoint['id']}")
+
+        try:
+            self._checkpoints.create_item(doc)
+        except CosmosResourceExistsError:
+            existing = self._checkpoints.read_item(
+                item=doc["id"], partition_key=thread_id
+            )
+            if not _checkpoint_doc_equivalent(existing, doc):
+                raise ValueError(
+                    "Checkpoint conflict for deterministic checkpoint ID."
+                ) from None
+
+        return {
+            "configurable": {
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint["id"],
+            }
+        }
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        thread_id, checkpoint_ns, checkpoint_id = _normalize_config(config)
+        if checkpoint_id is None:
+            raise ValueError("put_writes requires checkpoint_id in config.")
+
+        for idx, (channel, value) in enumerate(writes):
+            mapped_idx = WRITES_IDX_MAP.get(channel, idx)
+            write_type, write_payload, write_hash = self._serialize_typed(
+                value=value,
+                thread_id=thread_id,
+                payload_scope=(
+                    f"write:{checkpoint_ns}:{checkpoint_id}:{task_id}:{mapped_idx}"
+                ),
+            )
+            doc = {
+                "id": _write_doc_id(
+                    checkpoint_ns=checkpoint_ns,
+                    checkpoint_id=checkpoint_id,
+                    task_id=task_id,
+                    idx=mapped_idx,
+                ),
+                "doc_type": "write",
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "checkpoint_id": checkpoint_id,
+                "task_id": task_id,
+                "task_path": task_path,
+                "idx": mapped_idx,
+                "channel": channel,
+                "value_type": write_type,
+                "value_payload": write_payload,
+                "value_hash": write_hash,
+            }
+            _assert_item_size(
+                doc, scope=f"write:{checkpoint_id}:{task_id}:{mapped_idx}"
+            )
+
+            if mapped_idx >= 0:
+                self._create_or_validate_write(doc)
+            else:
+                self._replace_special_write_with_etag(doc)
+
+    def delete_thread(self, thread_id: str) -> None:
+        query = "SELECT c.id FROM c WHERE c.thread_id=@thread_id"
+        params = [{"name": "@thread_id", "value": thread_id}]
+        for container in (self._checkpoints, self._writes, self._blobs):
+            if container is None:
+                continue
+            ids = list(
+                container.query_items(
+                    query=query,
+                    parameters=params,
+                    partition_key=thread_id,
+                )
+            )
+            for row in ids:
+                container.delete_item(item=row["id"], partition_key=thread_id)
+
+    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        return await asyncio.to_thread(self.get_tuple, config)
+
+    async def alist(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        rows = await asyncio.to_thread(
+            lambda: list(self.list(config, filter=filter, before=before, limit=limit))
+        )
+        for row in rows:
+            yield row
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        return await asyncio.to_thread(
+            self.put, config, checkpoint, metadata, new_versions
+        )
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        await asyncio.to_thread(self.put_writes, config, writes, task_id, task_path)
+
+    async def adelete_thread(self, thread_id: str) -> None:
+        await asyncio.to_thread(self.delete_thread, thread_id)
+
+    def get_next_version(self, current: str | None, channel: None) -> str:
+        if current is None:
+            current_v = 0
+        elif isinstance(current, int):
+            current_v = current
+        else:
+            current_v = int(str(current).split(".")[0])
+        return f"{current_v + 1:032}.{random.random():016}"
+
+    def _checkpoint_tuple_from_doc(
+        self,
+        *,
+        doc: Mapping[str, Any],
+        thread_id: str,
+        checkpoint_ns: str,
+    ) -> CheckpointTuple:
+        checkpoint_core = self._deserialize_typed(
+            thread_id=thread_id,
+            payload_type=doc["checkpoint_type"],
+            payload=doc["checkpoint_payload"],
+        )
+        metadata = self._deserialize_typed(
+            thread_id=thread_id,
+            payload_type=doc["metadata_type"],
+            payload=doc["metadata_payload"],
+        )
+        if not isinstance(checkpoint_core, Mapping):
+            raise TypeError("Decoded checkpoint payload is not a mapping.")
+        if not isinstance(metadata, Mapping):
+            raise TypeError("Decoded checkpoint metadata is not a mapping.")
+
+        channel_values = self._load_channel_values(
+            thread_id=thread_id,
+            checkpoint_ns=checkpoint_ns,
+            channel_versions=checkpoint_core.get("channel_versions", {}),
+        )
+        pending_writes = self._load_pending_writes(
+            thread_id=thread_id,
+            checkpoint_ns=checkpoint_ns,
+            checkpoint_id=doc["checkpoint_id"],
+        )
+        checkpoint = dict(checkpoint_core)
+        checkpoint["channel_values"] = channel_values
+
+        parent_checkpoint_id = doc.get("parent_checkpoint_id")
+        parent_config = None
+        if parent_checkpoint_id:
+            parent_config = {
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": checkpoint_ns,
+                    "checkpoint_id": parent_checkpoint_id,
+                }
+            }
+
+        return CheckpointTuple(
+            config={
+                "configurable": {
+                    "thread_id": thread_id,
+                    "checkpoint_ns": checkpoint_ns,
+                    "checkpoint_id": doc["checkpoint_id"],
+                }
+            },
+            checkpoint=checkpoint,  # type: ignore[arg-type]
+            metadata=dict(metadata),  # type: ignore[arg-type]
+            parent_config=parent_config,
+            pending_writes=pending_writes,
+        )
+
+    def _load_pending_writes(
+        self,
+        *,
+        thread_id: str,
+        checkpoint_ns: str,
+        checkpoint_id: str,
+    ) -> list[tuple[str, str, Any]]:
+        query = (
+            "SELECT * FROM c WHERE c.doc_type=@doc_type AND c.thread_id=@thread_id "
+            "AND c.checkpoint_ns=@checkpoint_ns AND c.checkpoint_id=@checkpoint_id "
+            "ORDER BY c.task_id ASC, c.idx ASC"
+        )
+        params = [
+            {"name": "@doc_type", "value": "write"},
+            {"name": "@thread_id", "value": thread_id},
+            {"name": "@checkpoint_ns", "value": checkpoint_ns},
+            {"name": "@checkpoint_id", "value": checkpoint_id},
+        ]
+        rows = self._writes.query_items(
+            query=query,
+            parameters=params,
+            partition_key=thread_id,
+        )
+        output: list[tuple[str, str, Any]] = []
+        for row in rows:
+            value = self._deserialize_typed(
+                thread_id=thread_id,
+                payload_type=row["value_type"],
+                payload=row["value_payload"],
+            )
+            output.append((row["task_id"], row["channel"], value))
+        return output
+
+    def _create_or_validate_write(self, doc: dict[str, Any]) -> None:
+        try:
+            self._writes.create_item(doc)
+            return
+        except CosmosResourceExistsError:
+            existing = self._writes.read_item(item=doc["id"], partition_key=doc["thread_id"])
+        if not _write_doc_equivalent(existing, doc):
+            raise ValueError(
+                "Detected conflicting duplicate write for deterministic write identity."
+            )
+
+    def _replace_special_write_with_etag(self, doc: dict[str, Any]) -> None:
+        try:
+            existing = self._writes.read_item(item=doc["id"], partition_key=doc["thread_id"])
+        except CosmosResourceNotFoundError:
+            self._writes.create_item(doc)
+            return
+
+        if _write_doc_equivalent(existing, doc):
+            return
+
+        try:
+            self._writes.replace_item(
+                item=doc["id"],
+                body=doc,
+                etag=existing.get("_etag"),
+                match_condition=MatchConditions.IfNotModified,
+            )
+        except (CosmosAccessConditionFailedError, CosmosHttpResponseError) as exc:
+            if getattr(exc, "status_code", None) != 412:
+                raise
+            latest = self._writes.read_item(item=doc["id"], partition_key=doc["thread_id"])
+            if not _write_doc_equivalent(latest, doc):
+                raise ValueError(
+                    "Concurrent update conflict while replacing special write."
+                ) from None
+
+    def _store_channel_blob(
+        self,
+        *,
+        thread_id: str,
+        checkpoint_ns: str,
+        channel: str,
+        version: str | int | float,
+        value: Any,
+    ) -> None:
+        if self._blobs is None:
+            if value is not None:
+                raise ValueError("Blob container is required to persist channel values.")
+            return
+
+        doc_id = _channel_blob_doc_id(checkpoint_ns, channel, version)
+        if value is None:
+            doc = {
+                "id": doc_id,
+                "doc_type": "channel_blob",
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "channel": channel,
+                "version": _stringify_version(version),
+                "value_type": "empty",
+                "value_payload": {"mode": "empty"},
+                "value_hash": "empty",
+            }
+        else:
+            value_type, value_payload, value_hash = self._serialize_typed(
+                value=value,
+                thread_id=thread_id,
+                payload_scope=f"channel:{checkpoint_ns}:{channel}:{_stringify_version(version)}",
+            )
+            doc = {
+                "id": doc_id,
+                "doc_type": "channel_blob",
+                "thread_id": thread_id,
+                "checkpoint_ns": checkpoint_ns,
+                "channel": channel,
+                "version": _stringify_version(version),
+                "value_type": value_type,
+                "value_payload": value_payload,
+                "value_hash": value_hash,
+            }
+        _assert_item_size(doc, scope=f"channel_blob:{channel}")
+
+        try:
+            self._blobs.create_item(doc)
+        except CosmosResourceExistsError:
+            existing = self._blobs.read_item(item=doc_id, partition_key=thread_id)
+            if not _channel_blob_equivalent(existing, doc):
+                raise ValueError(
+                    "Conflicting channel blob for deterministic channel version identity."
+                ) from None
+
+    def _load_channel_values(
+        self,
+        *,
+        thread_id: str,
+        checkpoint_ns: str,
+        channel_versions: Mapping[str, Any],
+    ) -> dict[str, Any]:
+        channel_values: dict[str, Any] = {}
+        if self._blobs is None:
+            return channel_values
+
+        for channel, version in channel_versions.items():
+            doc_id = _channel_blob_doc_id(checkpoint_ns, channel, version)
+            try:
+                row = self._blobs.read_item(item=doc_id, partition_key=thread_id)
+            except CosmosResourceNotFoundError:
+                continue
+            if row.get("value_type") == "empty":
+                continue
+            channel_values[channel] = self._deserialize_typed(
+                thread_id=thread_id,
+                payload_type=row["value_type"],
+                payload=row["value_payload"],
+            )
+        return channel_values
+
+    def _serialize_typed(
+        self,
+        *,
+        value: Any,
+        thread_id: str,
+        payload_scope: str,
+    ) -> tuple[str, dict[str, Any], str]:
+        payload_type, raw = self.serde.dumps_typed(value)
+        payload, raw_hash = self._encode_payload(
+            raw=raw,
+            thread_id=thread_id,
+            payload_scope=payload_scope,
+        )
+        return payload_type, payload, raw_hash
+
+    def _deserialize_typed(
+        self,
+        *,
+        thread_id: str,
+        payload_type: str,
+        payload: Mapping[str, Any],
+    ) -> Any:
+        raw = self._decode_payload(thread_id=thread_id, payload=payload)
+        return self.serde.loads_typed((payload_type, raw))
+
+    def _encode_payload(
+        self,
+        *,
+        raw: bytes,
+        thread_id: str,
+        payload_scope: str,
+    ) -> tuple[dict[str, Any], str]:
+        raw_hash = hashlib.sha256(raw).hexdigest()
+        data = raw
+        compressed = False
+        if len(raw) >= self._settings.compress_above_bytes:
+            compressed_candidate = zlib.compress(raw)
+            if len(compressed_candidate) < len(raw):
+                data = compressed_candidate
+                compressed = True
+
+        if len(data) <= self._settings.inline_payload_limit_bytes:
+            return (
+                {
+                    "mode": "inline",
+                    "compressed": compressed,
+                    "raw_hash": raw_hash,
+                    "size_bytes": len(raw),
+                    "data": _b64encode(data),
+                },
+                raw_hash,
+            )
+
+        if self._blobs is None:
+            raise ValueError(
+                "Payload exceeds inline payload limit and no blob container is configured."
+            )
+
+        chunk_ids: list[str] = []
+        for idx in range(0, len(data), self._settings.chunk_size_bytes):
+            chunk = data[idx : idx + self._settings.chunk_size_bytes]
+            chunk_id = _chunk_doc_id(payload_scope, idx // self._settings.chunk_size_bytes)
+            chunk_doc = {
+                "id": chunk_id,
+                "doc_type": "payload_chunk",
+                "thread_id": thread_id,
+                "scope_hash": _hash_token(payload_scope),
+                "chunk_index": idx // self._settings.chunk_size_bytes,
+                "payload": _b64encode(chunk),
+                "payload_hash": hashlib.sha256(chunk).hexdigest(),
+            }
+            _assert_item_size(chunk_doc, scope=f"chunk:{payload_scope}")
+            try:
+                self._blobs.create_item(chunk_doc)
+            except CosmosResourceExistsError:
+                existing = self._blobs.read_item(item=chunk_id, partition_key=thread_id)
+                if not _chunk_doc_equivalent(existing, chunk_doc):
+                    raise ValueError("Payload chunk conflict for deterministic chunk ID.") from None
+            chunk_ids.append(chunk_id)
+
+        return (
+            {
+                "mode": "chunks",
+                "compressed": compressed,
+                "raw_hash": raw_hash,
+                "size_bytes": len(raw),
+                "chunks": chunk_ids,
+            },
+            raw_hash,
+        )
+
+    def _decode_payload(self, *, thread_id: str, payload: Mapping[str, Any]) -> bytes:
+        mode = payload.get("mode")
+        if mode == "empty":
+            return b""
+        if mode == "inline":
+            raw = _b64decode(payload["data"])
+        elif mode == "chunks":
+            if self._blobs is None:
+                raise ValueError("Blob container is required to read chunked payloads.")
+            chunks = payload.get("chunks", [])
+            pieces: list[bytes] = []
+            for chunk_id in chunks:
+                row = self._blobs.read_item(item=chunk_id, partition_key=thread_id)
+                pieces.append(_b64decode(row["payload"]))
+            raw = b"".join(pieces)
+        else:
+            raise ValueError(f"Unknown payload mode: {mode}")
+
+        if payload.get("compressed"):
+            raw = zlib.decompress(raw)
+        if payload.get("raw_hash") and hashlib.sha256(raw).hexdigest() != payload["raw_hash"]:
+            raise ValueError("Payload hash mismatch detected.")
+        return raw
+
+
+class AsyncCosmosDBSaver(BaseCheckpointSaver[str], AbstractAsyncContextManager, AbstractContextManager):
+    def __init__(self, saver: CosmosDBSaver) -> None:
+        super().__init__(serde=saver.serde)
+        self._saver = saver
+
+    @classmethod
+    def from_connection_string(
+        cls,
+        connection_string: str,
+        *,
+        settings: CosmosDBSettings | None = None,
+        credential: str | TokenCredential | Mapping[str, Any] | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> "AsyncCosmosDBSaver":
+        return cls(
+            CosmosDBSaver.from_connection_string(
+                connection_string,
+                settings=settings,
+                credential=credential,
+                serde=serde,
+            )
+        )
+
+    @classmethod
+    def from_endpoint(
+        cls,
+        endpoint: str,
+        *,
+        credential: str | TokenCredential | Mapping[str, Any] | None = None,
+        settings: CosmosDBSettings | None = None,
+        serde: SerializerProtocol | None = None,
+    ) -> "AsyncCosmosDBSaver":
+        return cls(
+            CosmosDBSaver.from_endpoint(
+                endpoint,
+                credential=credential,
+                settings=settings,
+                serde=serde,
+            )
+        )
+
+    def __enter__(self) -> "AsyncCosmosDBSaver":
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> bool | None:
+        self._saver.close()
+        return None
+
+    async def __aenter__(self) -> "AsyncCosmosDBSaver":
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb) -> bool | None:
+        await asyncio.to_thread(self._saver.close)
+        return None
+
+    def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        return self._saver.get_tuple(config)
+
+    def list(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> Iterator[CheckpointTuple]:
+        return self._saver.list(config, filter=filter, before=before, limit=limit)
+
+    def put(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        return self._saver.put(config, checkpoint, metadata, new_versions)
+
+    def put_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        self._saver.put_writes(config, writes, task_id, task_path)
+
+    def delete_thread(self, thread_id: str) -> None:
+        self._saver.delete_thread(thread_id)
+
+    async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
+        return await asyncio.to_thread(self._saver.get_tuple, config)
+
+    async def alist(
+        self,
+        config: RunnableConfig | None,
+        *,
+        filter: dict[str, Any] | None = None,
+        before: RunnableConfig | None = None,
+        limit: int | None = None,
+    ) -> AsyncIterator[CheckpointTuple]:
+        rows = await asyncio.to_thread(
+            lambda: list(self._saver.list(config, filter=filter, before=before, limit=limit))
+        )
+        for row in rows:
+            yield row
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Checkpoint,
+        metadata: CheckpointMetadata,
+        new_versions: ChannelVersions,
+    ) -> RunnableConfig:
+        return await asyncio.to_thread(
+            self._saver.put, config, checkpoint, metadata, new_versions
+        )
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        await asyncio.to_thread(self._saver.put_writes, config, writes, task_id, task_path)
+
+    async def adelete_thread(self, thread_id: str) -> None:
+        await asyncio.to_thread(self._saver.delete_thread, thread_id)
+
+
+def _client_kwargs(settings: CosmosDBSettings) -> dict[str, Any]:
+    kwargs: dict[str, Any] = {
+        "no_response_on_write": settings.no_response_on_write,
+    }
+    if settings.retry_total is not None:
+        kwargs["retry_total"] = settings.retry_total
+    if settings.retry_backoff_max is not None:
+        kwargs["retry_backoff_max"] = settings.retry_backoff_max
+    if settings.retry_fixed_interval is not None:
+        kwargs["retry_fixed_interval"] = settings.retry_fixed_interval
+    return kwargs
+
+
+def _resolve_containers_sync(
+    *,
+    client: CosmosClient,
+    settings: CosmosDBSettings,
+) -> tuple[ContainerProxy, ContainerProxy, ContainerProxy | None]:
+    database: DatabaseProxy
+    if settings.create_containers:
+        database = client.create_database_if_not_exists(id=settings.database_name)
+    else:
+        database = client.get_database_client(settings.database_name)
+
+    if settings.create_containers:
+        checkpoints = database.create_container_if_not_exists(
+            id=settings.checkpoint_container_name,
+            partition_key=PartitionKey(path="/thread_id"),
+            indexing_policy=_checkpoint_indexing_policy(),
+        )
+        writes = database.create_container_if_not_exists(
+            id=settings.writes_container_name,
+            partition_key=PartitionKey(path="/thread_id"),
+            indexing_policy=_writes_indexing_policy(),
+        )
+        blobs = database.create_container_if_not_exists(
+            id=settings.blob_container_name,
+            partition_key=PartitionKey(path="/thread_id"),
+            indexing_policy=_blob_indexing_policy(),
+        )
+    else:
+        checkpoints = database.get_container_client(settings.checkpoint_container_name)
+        writes = database.get_container_client(settings.writes_container_name)
+        blobs = (
+            database.get_container_client(settings.blob_container_name)
+            if settings.blob_container_name
+            else None
+        )
+    return checkpoints, writes, blobs
+
+
+def _normalize_config(config: RunnableConfig) -> tuple[str, str, str | None]:
+    configurable = config.get("configurable") or {}
+    if "thread_ts" in configurable:
+        raise ValueError(
+            "Legacy config key thread_ts is not supported. Use thread_id/checkpoint_ns/checkpoint_id."
+        )
+    thread_id = configurable.get("thread_id")
+    if not thread_id:
+        raise ValueError("config.configurable.thread_id is required.")
+    checkpoint_ns = configurable.get("checkpoint_ns", "")
+    checkpoint_id = get_checkpoint_id(config)
+    return str(thread_id), str(checkpoint_ns), str(checkpoint_id) if checkpoint_id else None
+
+
+def _checkpoint_doc_equivalent(existing: Mapping[str, Any], new_doc: Mapping[str, Any]) -> bool:
+    return (
+        existing.get("checkpoint_hash") == new_doc.get("checkpoint_hash")
+        and existing.get("metadata_hash") == new_doc.get("metadata_hash")
+        and existing.get("parent_checkpoint_id") == new_doc.get("parent_checkpoint_id")
+    )
+
+
+def _write_doc_equivalent(existing: Mapping[str, Any], new_doc: Mapping[str, Any]) -> bool:
+    return (
+        existing.get("channel") == new_doc.get("channel")
+        and existing.get("value_hash") == new_doc.get("value_hash")
+        and existing.get("task_id") == new_doc.get("task_id")
+    )
+
+
+def _channel_blob_equivalent(existing: Mapping[str, Any], new_doc: Mapping[str, Any]) -> bool:
+    return (
+        existing.get("value_type") == new_doc.get("value_type")
+        and existing.get("value_hash") == new_doc.get("value_hash")
+    )
+
+
+def _chunk_doc_equivalent(existing: Mapping[str, Any], new_doc: Mapping[str, Any]) -> bool:
+    return existing.get("payload_hash") == new_doc.get("payload_hash")
+
+
+def _checkpoint_doc_id(checkpoint_ns: str, checkpoint_id: str) -> str:
+    return f"ckpt|{_ns_token(checkpoint_ns)}|{checkpoint_id}"
+
+
+def _write_doc_id(
+    *,
+    checkpoint_ns: str,
+    checkpoint_id: str,
+    task_id: str,
+    idx: int,
+) -> str:
+    return f"w|{_ns_token(checkpoint_ns)}|{checkpoint_id}|{_hash_token(task_id)}|{idx}"
+
+
+def _channel_blob_doc_id(
+    checkpoint_ns: str,
+    channel: str,
+    version: str | int | float,
+) -> str:
+    return (
+        f"b|{_ns_token(checkpoint_ns)}|{_hash_token(channel)}|"
+        f"{_hash_token(_stringify_version(version))}"
+    )
+
+
+def _chunk_doc_id(scope: str, chunk_idx: int) -> str:
+    return f"chunk|{_hash_token(scope)}|{chunk_idx:06d}"
+
+
+def _ns_token(checkpoint_ns: str) -> str:
+    return "root" if not checkpoint_ns else _hash_token(checkpoint_ns)
+
+
+def _hash_token(value: str) -> str:
+    return hashlib.sha256(value.encode("utf-8")).hexdigest()[:16]
+
+
+def _stringify_version(version: str | int | float) -> str:
+    if isinstance(version, str):
+        return version
+    return json.dumps(version, sort_keys=True, separators=(",", ":"))
+
+
+def _assert_item_size(item: Mapping[str, Any], *, scope: str) -> None:
+    item_bytes = len(
+        json.dumps(item, separators=(",", ":"), sort_keys=True, default=str).encode("utf-8")
+    )
+    if item_bytes > MAX_COSMOS_ITEM_BYTES:
+        raise ValueError(
+            f"Cosmos item for {scope} exceeds 2MB limit ({item_bytes} bytes)."
+        )
+
+
+def _b64encode(raw: bytes) -> str:
+    return base64.b64encode(raw).decode("ascii")
+
+
+def _b64decode(data: str) -> bytes:
+    return base64.b64decode(data.encode("ascii"))
+
+
+def _checkpoint_indexing_policy() -> dict[str, Any]:
+    return {
+        "indexingMode": "consistent",
+        "automatic": True,
+        "includedPaths": [{"path": "/*"}],
+        "excludedPaths": [
+            {"path": "/checkpoint_payload/*"},
+            {"path": "/metadata_payload/*"},
+        ],
+    }
+
+
+def _writes_indexing_policy() -> dict[str, Any]:
+    return {
+        "indexingMode": "consistent",
+        "automatic": True,
+        "includedPaths": [{"path": "/*"}],
+        "excludedPaths": [{"path": "/value_payload/*"}],
+    }
+
+
+def _blob_indexing_policy() -> dict[str, Any]:
+    return {
+        "indexingMode": "consistent",
+        "automatic": True,
+        "includedPaths": [
+            {"path": "/thread_id/?"},
+            {"path": "/doc_type/?"},
+            {"path": "/scope_hash/?"},
+            {"path": "/chunk_index/?"},
+            {"path": "/channel/?"},
+            {"path": "/version/?"},
+            {"path": "/id/?"},
+        ],
+        "excludedPaths": [{"path": "/*"}],
+    }

langgraph_cosmosdb_checkpointer-0.0.0.dist-info/METADATA

@@ -0,0 +1,18 @@
+Metadata-Version: 2.4
+Name: langgraph-cosmosdb-checkpointer
+Version: 0.0.0
+Summary: State-of-the-art Azure Cosmos DB checkpointer for LangGraph
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: azure-cosmos>=4.14.6
+Requires-Dist: azure-identity>=1.25.2
+Requires-Dist: langchain>=1.2.10
+Requires-Dist: langgraph>=1.0.9
+Provides-Extra: dev
+Requires-Dist: pytest>=9.0.2; extra == "dev"
+Requires-Dist: pytest-asyncio>=1.3.0; extra == "dev"
+Requires-Dist: ruff>=0.9.0; extra == "dev"
+
+# langgraph-cosmosdb-checkpointer
+
+State-of-the-art Azure Cosmos DB checkpointer for LangGraph.

langgraph_cosmosdb_checkpointer-0.0.0.dist-info/RECORD

@@ -0,0 +1,6 @@
+langgraph_cosmosdb_checkpointer/__init__.py,sha256=8SkbUhhy3suSteK4Ha_7J-O48Tc7RMJQ49N7DrPVveo,177
+langgraph_cosmosdb_checkpointer/cosmosdb.py,sha256=kBZ3Qklg10MBlmY1WBtSCjB-bLPy3LBanHsoCBSbbX0,37835
+langgraph_cosmosdb_checkpointer-0.0.0.dist-info/METADATA,sha256=cIXspDHPg3XV2qyPYLDGhWJGPs8VxpGJKA-MhdqUmIw,605
+langgraph_cosmosdb_checkpointer-0.0.0.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
+langgraph_cosmosdb_checkpointer-0.0.0.dist-info/top_level.txt,sha256=JF8lejbl_uxDEJyPoTnQQwi7JWWXdd2vdLvmPhE-GyQ,32
+langgraph_cosmosdb_checkpointer-0.0.0.dist-info/RECORD,,

langgraph_cosmosdb_checkpointer-0.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

langgraph_cosmosdb_checkpointer-0.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+langgraph_cosmosdb_checkpointer