agent-lab-sdk 0.1.35.dev1__py3-none-any.whl → 0.1.37__py3-none-any.whl

agent_lab_sdk/langgraph/checkpoint/agw_saver.py

@@ -1,15 +1,16 @@
 from __future__ import annotations
 
+import orjson
+from random import random
+from langgraph.checkpoint.serde.types import ChannelProtocol
 import asyncio
 import base64
 import logging
 import os
 from contextlib import asynccontextmanager
-from random import random
 from typing import Any, AsyncIterator, Dict, Iterator, Optional, Sequence, Tuple
 
 import httpx
-import orjson
 from langchain_core.runnables import RunnableConfig
 
 from langgraph.checkpoint.base import (
@@ -23,7 +24,7 @@ from langgraph.checkpoint.base import (
 )
 from langgraph.checkpoint.serde.base import SerializerProtocol
 from langgraph.checkpoint.serde.encrypted import EncryptedSerializer
-from langgraph.checkpoint.serde.types import ChannelProtocol
+
 from .serde import Serializer
 
 __all__ = ["AsyncAGWCheckpointSaver"]
@@ -33,17 +34,22 @@ logger = logging.getLogger(__name__)
 TYPED_KEYS = ("type", "blob")
 
 
-def _to_b64(b: bytes | None) -> str | None:
-    return base64.b64encode(b).decode() if b is not None else None
-
-
-def _b64decode_strict(s: str) -> bytes | None:
-    """Возвращает bytes только если строка действительно корректная base64."""
+def _b64decode_strict(value: str) -> bytes | None:
     try:
-        return base64.b64decode(s, validate=True)
+        return base64.b64decode(value, validate=True)
     except Exception:
         return None
 
+# ------------------------------------------------------------------ #
+# helpers for Py < 3.10
+# ------------------------------------------------------------------ #
+try:
+    anext  # type: ignore[name-defined]
+except NameError:  # pragma: no cover
+
+    async def anext(it):
+        return await it.__anext__()
+
 
 class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
     """Persist checkpoints in Agent-Gateway с помощью `httpx` async client."""
@@ -61,15 +67,14 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
     ):
         if not serde:
             base_serde: SerializerProtocol = Serializer()
-            # опционально оборачиваем в AES по ENV
-            _aes_key = (
+            aes_key = (
                 os.getenv("LANGGRAPH_AES_KEY")
                 or os.getenv("AGW_AES_KEY")
                 or os.getenv("AES_KEY")
             )
-            if _aes_key:
+            if aes_key:
                 base_serde = EncryptedSerializer.from_pycryptodome_aes(
-                    base_serde, key=_aes_key
+                    base_serde, key=aes_key
                 )
             serde = base_serde
         super().__init__(serde=serde)
@@ -91,7 +96,7 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
             headers=self.headers,
             timeout=self.timeout,
             verify=verify,
-            trust_env=True,
+            trust_env=True
         )
 
     async def __aenter__(self):  # noqa: D401
@@ -100,63 +105,105 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
     async def __aexit__(self, exc_type, exc, tb):  # noqa: D401
         await self._client.aclose()
 
-    # ----------------------- typed (de)serialize ---------------------
-    def _encode_typed(self, value: Any) -> dict[str, Any]:
-        """value -> {"type": str, "blob": base64str | null}"""
-        t, b = self.serde.dumps_typed(value)
-        return {"type": t, "blob": _to_b64(b)}
-
-    def _decode_typed(self, obj: Any) -> Any:
-        """{type, blob} | [type, blob] | legacy -> python."""
-        # Новый формат: dict с ключами type/blob — только если blob валидная base64 или None
-        if isinstance(obj, dict) and all(k in obj for k in TYPED_KEYS):
-            t = obj.get("type")
-            b64 = obj.get("blob")
-            if b64 is None:
-                return self.serde.loads_typed((t, None))
-            if isinstance(b64, str):
-                b = _b64decode_strict(b64)
-                if b is not None:
-                    return self.serde.loads_typed((t, b))
-            # если невалидно — падаем ниже на общую обработку
-
-        # Допускаем tuple/list вида [type, base64] — только при валидной base64
-        if isinstance(obj, (list, tuple)) and len(obj) == 2 and isinstance(obj[0], str):
-            t, b64 = obj
-            if b64 is None and t == "empty":
-                return self.serde.loads_typed((t, None))
-            if isinstance(b64, str):
-                b = _b64decode_strict(b64)
-                if b is not None:
-                    return self.serde.loads_typed((t, b))
-            # иначе это не typed-пара
-
-        # Если это строка — пробуем как base64 строго, затем как JSON-строку
-        if isinstance(obj, str):
-            b = _b64decode_strict(obj)
-            if b is not None:
-                try:
-                    return self.serde.loads(b)
-                except Exception:
-                    pass
+    # ----------------------- universal dump/load ---------------------
+    # def _safe_dump(self, obj: Any) -> Any:
+    #     """self.serde.dump → гарантированная JSON-строка."""
+    #     dumped = self.serde.dumps(obj)
+    #     if isinstance(dumped, (bytes, bytearray)):
+    #         return base64.b64encode(dumped).decode()  # str
+    #     return dumped  # уже json-совместимо
+
+    def _safe_dump(self, obj: Any) -> Any:
+        """bytes → python-object; fallback base64 для реально бинарных данных."""
+        dumped = self.serde.dumps(obj)
+        if isinstance(dumped, (bytes, bytearray)):
             try:
-                return self.serde.loads(obj.encode())
+                # 1) bytes → str
+                s = dumped.decode()
+                # 2) str JSON → python (list/dict/scalar)
+                return orjson.loads(s)
+            except (UnicodeDecodeError, orjson.JSONDecodeError):
+                # не UTF-8 или не JSON → base64
+                return base64.b64encode(dumped).decode()
+        return dumped
+
+    def _safe_load(self, obj: Any) -> Any:
+        if obj is None:
+            return None
+
+        if isinstance(obj, dict):
+            if all(k in obj for k in TYPED_KEYS):
+                t = obj.get("type")
+                blob = obj.get("blob")
+                if blob is None:
+                    try:
+                        return self.serde.loads_typed((t, None))
+                    except Exception:
+                        return obj
+                if isinstance(blob, str):
+                    payload = _b64decode_strict(blob)
+                    if payload is not None:
+                        try:
+                            return self.serde.loads_typed((t, payload))
+                        except Exception:
+                            # fall back to generic handling below
+                            pass
+            try:
+                return self.serde.loads(orjson.dumps(obj))
             except Exception:
                 return obj
 
-        # dict/list -> считаем это уже JSON и грузим через serde
-        if isinstance(obj, (dict, list)):
+        if isinstance(obj, (list, tuple)):
+            if (
+                len(obj) == 2
+                and isinstance(obj[0], str)
+                and (obj[1] is None or isinstance(obj[1], str))
+            ):
+                blob = obj[1]
+                if blob is None:
+                    try:
+                        return self.serde.loads_typed((obj[0], None))
+                    except Exception:
+                        pass
+                elif isinstance(blob, str):
+                    payload = _b64decode_strict(blob)
+                    if payload is not None:
+                        try:
+                            return self.serde.loads_typed((obj[0], payload))
+                        except Exception:
+                            pass
             try:
-                return self.serde.loads(orjson.dumps(obj))
+                return self.serde.loads(orjson.dumps(list(obj)))
             except Exception:
                 return obj
 
-        # как есть пробуем через serde
+        if isinstance(obj, str):
+            try:
+                return self.serde.loads(obj.encode())
+            except Exception:
+                payload = _b64decode_strict(obj)
+                if payload is not None:
+                    try:
+                        return self.serde.loads(payload)
+                    except Exception:
+                        pass
+                return obj
+
         try:
             return self.serde.loads(obj)
         except Exception:
             return obj
 
+    # def _safe_load(self, obj: Any) -> Any:
+    #     """Обратная операция к _safe_dump."""
+    #     if isinstance(obj, str):
+    #         try:
+    #             return self.serde.load(base64.b64decode(obj))
+    #         except Exception:
+    #             # не base64 — обычная строка
+    #             return self.serde.load(obj)
+    #     return self.serde.load(obj)
+
     # ----------------------- config <-> api --------------------------
     def _to_api_config(self, cfg: RunnableConfig | None) -> Dict[str, Any]:
         if not cfg:
@@ -174,47 +221,44 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
 
     # --------------------- checkpoint (de)ser ------------------------
     def _encode_cp(self, cp: Checkpoint) -> Dict[str, Any]:
-        channel_values = {
-            k: self._encode_typed(v) for k, v in cp.get("channel_values", {}).items()
-        }
-        pending = []
+        pending: list[Any] = []
         for item in cp.get("pending_sends", []) or []:
             try:
                 channel, value = item
-                pending.append({"channel": channel, **self._encode_typed(value)})
             except Exception:
+                pending.append(item)
                 continue
+            pending.append([channel, self._safe_dump(value)])
         return {
             "v": cp["v"],
             "id": cp["id"],
             "ts": cp["ts"],
-            "channelValues": channel_values,
+            "channelValues": {k: self._safe_dump(v) for k, v in cp["channel_values"].items()},
             "channelVersions": cp["channel_versions"],
             "versionsSeen": cp["versions_seen"],
             "pendingSends": pending,
         }
 
     def _decode_cp(self, raw: Dict[str, Any]) -> Checkpoint:
-        cv_raw = raw.get("channelValues") or {}
-        channel_values = {k: self._decode_typed(v) for k, v in cv_raw.items()}
-        ps_raw = raw.get("pendingSends") or []
-        pending_sends = []
-        for obj in ps_raw:
-            # ожидаем {channel, type, blob}
+        pending_sends: list[Tuple[str, Any]] = []
+        for obj in raw.get("pendingSends", []) or []:
             if isinstance(obj, dict) and "channel" in obj:
-                ch = obj["channel"]
-                typed = {k: obj[k] for k in obj.keys() if k in TYPED_KEYS}
-                val = self._decode_typed(typed)
-                pending_sends.append((ch, val))
-            elif isinstance(obj, (list, tuple)) and len(obj) == 2:
-                ch, val = obj
-                pending_sends.append((ch, self._decode_typed(val)))
-
+                channel = obj["channel"]
+                value_payload: Any = obj.get("value")
+                if value_payload is None and all(k in obj for k in TYPED_KEYS):
+                    value_payload = {k: obj[k] for k in TYPED_KEYS}
+                pending_sends.append((channel, self._safe_load(value_payload)))
+            elif isinstance(obj, (list, tuple)) and len(obj) >= 2:
+                channel = obj[0]
+                value_payload = obj[1]
+                pending_sends.append((channel, self._safe_load(value_payload)))
+            else:
+                pending_sends.append(obj)  # сохраняем как есть, если формат неизвестен
         return Checkpoint(
             v=raw["v"],
             id=raw["id"],
             ts=raw["ts"],
-            channel_values=channel_values,
+            channel_values={k: self._safe_load(v) for k, v in raw["channelValues"].items()},
             channel_versions=raw["channelVersions"],
             versions_seen=raw["versionsSeen"],
             pending_sends=pending_sends,
@@ -238,15 +282,22 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
             "thread_id": raw.get("threadId"),
             "thread_ts": raw.get("threadTs"),
             "checkpoint_ns": raw.get("checkpointNs"),
-            "checkpoint_id": raw.get("checkpointId"),
+            "checkpoint_id": raw.get("checkpointId")
         }
 
-    # metadata (de)ser — передаём как есть (JSON-совместимый словарь)
+    # metadata (de)ser
     def _enc_meta(self, md: CheckpointMetadata) -> CheckpointMetadata:
-        return md or {}
+        if not md:
+            return {}
+        out: CheckpointMetadata = {}
+        for k, v in md.items():
+            out[k] = self._enc_meta(v) if isinstance(v, dict) else self._safe_dump(v)  # type: ignore[assignment]
+        return out
 
     def _dec_meta(self, md: Any) -> Any:
-        return md
+        if isinstance(md, dict):
+            return {k: self._dec_meta(v) for k, v in md.items()}
+        return self._safe_load(md)
 
     # ------------------------ HTTP wrapper ---------------------------
     async def _http(self, method: str, path: str, **kw) -> httpx.Response:
@@ -254,36 +305,39 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
             payload = kw.pop("json")
             kw["data"] = orjson.dumps(payload)
             logger.debug("AGW HTTP payload: %s", kw["data"].decode())
+            
         return await self._client.request(method, path, **kw)
 
     # -------------------- api -> CheckpointTuple ----------------------
     def _to_tuple(self, node: Dict[str, Any]) -> CheckpointTuple:
-        pending_writes = None
-        raw_pw = node.get("pendingWrites")
-        if raw_pw:
-            decoded: list[tuple[str, str, Any]] = []
-            for w in raw_pw:
-                if isinstance(w, dict) and "first" in w and "second" in w:
-                    # ожидаем формат, который возвращает бек: first=task_id, second=channel, third=typed
-                    task_id = w["first"]
-                    channel = w["second"]
-                    tv = w.get("third")
-                    value = self._decode_typed(tv)
-                    decoded.append((task_id, channel, value))
+        pending = None
+        if node.get("pendingWrites"):
+            pending = []
+            for w in node["pendingWrites"]:
+                if isinstance(w, dict):
+                    first = w.get("first")
+                    second = w.get("second")
+                    third = w.get("third")
+                    if third is None and isinstance(second, dict) and all(
+                        k in second for k in TYPED_KEYS
+                    ):
+                        third = second
+                    pending.append((first, second, self._safe_load(third)))
                 elif isinstance(w, (list, tuple)):
-                    try:
-                        first, channel, tv = w
-                        decoded.append((first, channel, self._decode_typed(tv)))
-                    except Exception:  # pragma: no cover
+                    if len(w) == 3:
+                        first, second, third = w
+                    elif len(w) == 2:
+                        first, second = w
+                        third = None
+                    else:
                         continue
-            pending_writes = decoded
-
+                    pending.append((first, second, self._safe_load(third)))
         return CheckpointTuple(
             config=self._decode_config(node.get("config")),
             checkpoint=self._decode_cp(node["checkpoint"]),
             metadata=self._dec_meta(node.get("metadata")),
             parent_config=self._decode_config(node.get("parentConfig")),
-            pending_writes=pending_writes,
+            pending_writes=pending,
         )
 
     # =================================================================
@@ -292,7 +346,7 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
     async def aget_tuple(self, cfg: RunnableConfig) -> CheckpointTuple | None:
         cid = get_checkpoint_id(cfg)
         api_cfg = self._to_api_config(cfg)
-        tid = api_cfg.get("threadId")
+        tid = api_cfg["threadId"]
 
         if cid:
             path = f"/checkpoint/{tid}/{cid}"
@@ -304,7 +358,9 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
         resp = await self._http("GET", path, params=params)
         logger.debug("AGW aget_tuple response: %s", resp.text)
 
-        if not resp.text or resp.status_code in (404, 406):
+        if not resp.text:
+            return None
+        if resp.status_code in (404, 406):
             return None
         resp.raise_for_status()
         return self._to_tuple(resp.json())
@@ -339,7 +395,7 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
         payload = {
             "config": self._to_api_config(cfg),
             "checkpoint": self._encode_cp(cp),
-            "metadata": get_checkpoint_metadata(cfg, metadata),
+            "metadata": self._enc_meta(get_checkpoint_metadata(cfg, metadata)),
             "newVersions": new_versions,
         }
         resp = await self._http("POST", "/checkpoint", json=payload)
@@ -354,7 +410,7 @@ class AsyncAGWCheckpointSaver(BaseCheckpointSaver):
         task_id: str,
         task_path: str = "",
     ) -> None:
-        enc = [{"first": ch, "second": self._encode_typed(v)} for ch, v in writes]
+        enc = [{"first": ch, "second": self._safe_dump(v)} for ch, v in writes]
         payload = {
             "config": self._to_api_config(cfg),
             "writes": enc,

{agent_lab_sdk-0.1.35.dev1.dist-info → agent_lab_sdk-0.1.37.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agent-lab-sdk
-Version: 0.1.35.dev1
+Version: 0.1.37
 Summary: SDK для работы с Agent Lab
 Author-email: Andrew Ohurtsov <andermirik@yandex.com>
 License: Proprietary and Confidential — All Rights Reserved

{agent_lab_sdk-0.1.35.dev1.dist-info → agent_lab_sdk-0.1.37.dist-info}/RECORD

@@ -1,6 +1,6 @@
 agent_lab_sdk/__init__.py,sha256=1Dlmv-wuz1QuciymKtYtX7jXzr_fkeGTe7aENfEDl3E,108
 agent_lab_sdk/langgraph/checkpoint/__init__.py,sha256=DnKwR1LwbaQ3qhb124lE-tnojrUIVcCdNzHEHwgpL5M,86
-agent_lab_sdk/langgraph/checkpoint/agw_saver.py,sha256=QeYUAGEldw9SNXup3FRI7gNcGXYQDecKhoF1NyXa7yQ,16365
+agent_lab_sdk/langgraph/checkpoint/agw_saver.py,sha256=K07scbnVV_PT2zbn0kH2lk-Jl36ZVQDng_Nl16sR8Nc,18144
 agent_lab_sdk/langgraph/checkpoint/serde.py,sha256=UTSYbTbhBeL1CAr-XMbaH3SSIx9TeiC7ak22duXvqkw,5175
 agent_lab_sdk/llm/__init__.py,sha256=Yo9MbYdHS1iX05A9XiJGwWN1Hm4IARGav9mNFPrtDeA,376
 agent_lab_sdk/llm/agw_token_manager.py,sha256=_bPPI8muaEa6H01P8hHQOJHiiivaLd8N_d3OT9UT_80,4787
@@ -14,8 +14,8 @@ agent_lab_sdk/schema/input_types.py,sha256=e75nRW7Dz_RHk5Yia8DkFfbqMafsLQsQrJPfz
 agent_lab_sdk/schema/log_message.py,sha256=nadi6lZGRuDSPmfbYs9QPpRJUT9Pfy8Y7pGCvyFF5Mw,638
 agent_lab_sdk/storage/__init__.py,sha256=ik1_v1DMTwehvcAEXIYxuvLuCjJCa3y5qAuJqoQpuSA,81
 agent_lab_sdk/storage/storage.py,sha256=ELpt7GRwFD-aWa6ctinfA_QwcvzWLvKS0Wz8FlxVqAs,2075
-agent_lab_sdk-0.1.35.dev1.dist-info/licenses/LICENSE,sha256=_TRXHkF3S9ilWBPdZcHLI_S-PRjK0L_SeOb2pcPAdV4,417
-agent_lab_sdk-0.1.35.dev1.dist-info/METADATA,sha256=MsteJ1L7FktEWd45-MkhTlgxWHevOduYeL1prnNTjJY,17916
-agent_lab_sdk-0.1.35.dev1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-agent_lab_sdk-0.1.35.dev1.dist-info/top_level.txt,sha256=E1efqkJ89KNmPBWdLzdMHeVtH0dYyCo4fhnSb81_15I,14
-agent_lab_sdk-0.1.35.dev1.dist-info/RECORD,,
+agent_lab_sdk-0.1.37.dist-info/licenses/LICENSE,sha256=_TRXHkF3S9ilWBPdZcHLI_S-PRjK0L_SeOb2pcPAdV4,417
+agent_lab_sdk-0.1.37.dist-info/METADATA,sha256=XTRYyQQVXpDXxU4BkZvuAnsbZJsr-S5RuPKiDP0WKTc,17911
+agent_lab_sdk-0.1.37.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+agent_lab_sdk-0.1.37.dist-info/top_level.txt,sha256=E1efqkJ89KNmPBWdLzdMHeVtH0dYyCo4fhnSb81_15I,14
+agent_lab_sdk-0.1.37.dist-info/RECORD,,