diff --git a/docs/modules/serialization.mdx b/docs/modules/serialization.mdx index fc41d38f7..aa2a0d83f 100644 --- a/docs/modules/serialization.mdx +++ b/docs/modules/serialization.mdx @@ -15,11 +15,33 @@ BeeAI framework provides robust serialization capabilities through its built-in - 📦 Snapshots: Create point-in-time captures of component state - 🔧 Reconstruction: Rebuild objects from their serialized representation +The Python example prints the serialized payload, the restored ISO timestamp, and a boolean confirming the round trip. + -{/* */} +{/* */} ```py Python [expandable] -Example coming soon +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from datetime import UTC, datetime + +from beeai_framework.serialization import Serializer + + +def main() -> None: + original = datetime(2024, 1, 1, tzinfo=UTC) + serialized = Serializer.serialize(original) + restored = Serializer.deserialize(serialized, expected_type=datetime) + + print(serialized) + print(restored.isoformat()) + print(restored == original) + + +if __name__ == "__main__": + main() + ``` {/* */} @@ -69,11 +91,37 @@ The serialization process involves: Most BeeAI components can be serialized out of the box. Here's an example using memory: +Running the Python snippet restores the memory, appends an assistant reply, and prints the message count plus the first question. + -{/* */} +{/* */} ```py Python [expandable] -Example coming soon +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +from beeai_framework.backend import AssistantMessage, UserMessage +from beeai_framework.memory import UnconstrainedMemory + + +async def main() -> None: + memory = UnconstrainedMemory() + await memory.add(UserMessage("What is your name?")) + + serialized = memory.serialize() + restored = UnconstrainedMemory.from_serialized(serialized) + + await restored.add(AssistantMessage("Bee")) + + print(len(restored.messages)) + print(restored.messages[0].text) + + +if __name__ == "__main__": + asyncio.run(main()) + ``` {/* */} @@ -107,11 +155,52 @@ If you want to serialize a class that the `Serializer` does not know, you may re You can register external classes with the serializer: +This example registers a lightweight data class and prints the reconstructed token along with its expiry timestamp. + -{/* */} +{/* */} ```py Python [expandable] -Example coming soon +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass +from datetime import UTC, datetime + +from beeai_framework.serialization import Serializer + + +@dataclass +class ApiToken: + value: str + expires_at: datetime + + +Serializer.register( + ApiToken, + to_plain=lambda token: { + "value": token.value, + "expires_at": token.expires_at, + }, + from_plain=lambda payload: ApiToken( + value=payload["value"], + expires_at=payload["expires_at"], + ), +) + + +def main() -> None: + token = ApiToken("example-token", datetime(2025, 1, 1, tzinfo=UTC)) + serialized = Serializer.serialize(token) + restored = Serializer.deserialize(serialized, expected_type=ApiToken) + + print(restored) + print(restored.expires_at.isoformat()) + + +if __name__ == "__main__": + main() + ``` {/* */} @@ -150,11 +239,46 @@ console.info(deserialized); For deeper integration, extend the Serializable class: +The Python variant increments a counter, serializes it, and prints both the live and restored values. + -{/* */} +{/* */} ```py Python [expandable] -Example coming soon +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from beeai_framework.serialization import Serializable + + +class Counter(Serializable[dict[str, int]]): + def __init__(self, value: int = 0) -> None: + self.value = value + + def increment(self) -> None: + self.value += 1 + + def create_snapshot(self) -> dict[str, int]: + return {"value": self.value} + + def load_snapshot(self, snapshot: dict[str, int]) -> None: + self.value = snapshot["value"] + + +def main() -> None: + counter = Counter(3) + counter.increment() + + serialized = counter.serialize() + restored = Counter.from_serialized(serialized) + + print(counter.value) + print(restored.value) + + +if __name__ == "__main__": + main() + ``` {/* */} @@ -199,11 +323,29 @@ Failure to register a class that the `Serializer` does not know will result in t ## Context matters +Deserialize with the `extraClasses` equivalent to make sure message factories are registered; the Python snippet imports both `UserMessage` and `AssistantMessage` so the serializer can hydrate their content and then prints every restored user utterance. + -{/* */} +{/* */} ```py Python [expandable] -Example coming soon +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from beeai_framework.backend import AssistantMessage, UserMessage +from beeai_framework.memory import UnconstrainedMemory + +SERIALIZED_MEMORY = """{"__version":"0.0.0","__root":{"__serializer":true,"__class":"UnconstrainedMemory","__value":{"messages":[{"__serializer":true,"__class":"UserMessage","__value":{"id":null,"meta":{"createdAt":{"__serializer":true,"__class":"datetime","__value":"2025-10-18T19:38:37.859543+00:00"}},"role":"user","content":[{"type":"text","text":"Hello!"}]}},{"__serializer":true,"__class":"AssistantMessage","__value":{"id":null,"meta":{"createdAt":{"__serializer":true,"__class":"datetime","__value":"2025-10-18T19:38:37.859665+00:00"}},"role":"assistant","content":[{"type":"text","text":"Hello, how can I help you?"}]}}]}}}""" + + +def main() -> None: + memory = UnconstrainedMemory.from_serialized(SERIALIZED_MEMORY, extra_classes=[UserMessage, AssistantMessage]) + print([message.text for message in memory.messages]) + + +if __name__ == "__main__": + main() + ``` {/* */} diff --git a/python/beeai_framework/backend/message.py b/python/beeai_framework/backend/message.py index 7ac5bf163..293eb316c 100644 --- a/python/beeai_framework/backend/message.py +++ b/python/beeai_framework/backend/message.py @@ -3,7 +3,6 @@ import enum import json -from abc import ABC from collections.abc import Sequence from datetime import UTC, datetime from enum import Enum @@ -12,6 +11,7 @@ from pydantic import BaseModel, ConfigDict, Field, computed_field from typing_extensions import TypedDict +from beeai_framework.serialization import Serializable from beeai_framework.utils.dicts import exclude_none from beeai_framework.utils.lists import cast_list from beeai_framework.utils.models import to_any_model, to_model @@ -36,7 +36,25 @@ def values(cls) -> set[str]: return {value for key, value in vars(cls).items() if not key.startswith("_") and isinstance(value, str)} -class MessageTextContent(BaseModel): +class SerializableModel(BaseModel, Serializable[dict[str, Any]], auto_register=False): + """Base class for Pydantic models that need serialization support. + + This class combines BaseModel with Serializable to provide consistent + serialization behavior for message content parts. + """ + + def create_snapshot(self) -> dict[str, Any]: + """Create a snapshot of this model for serialization.""" + return self.model_dump() + + def load_snapshot(self, snapshot: dict[str, Any]) -> None: + """Load state from a snapshot dictionary.""" + for field_name, value in snapshot.items(): + object.__setattr__(self, field_name, value) + object.__setattr__(self, "__pydantic_fields_set__", set(snapshot.keys())) + + +class MessageTextContent(SerializableModel): type: Literal["text"] = "text" text: str @@ -47,12 +65,12 @@ class MessageImageContentImageUrl(TypedDict, total=False): format: str -class MessageImageContent(BaseModel): +class MessageImageContent(SerializableModel): type: Literal["image_url"] = "image_url" image_url: MessageImageContentImageUrl -class MessageFileContent(BaseModel): +class MessageFileContent(SerializableModel): """File content part (e.g. PDF or other document) for multimodal user messages. Flattened shape is supported: @@ -78,15 +96,25 @@ def model_post_init(self, __context: Any) -> None: if not (self.file_id or self.file_data): raise ValueError("Either 'file_id' or 'file_data' must be provided for MessageFileContent") + def create_snapshot(self) -> dict[str, Any]: + """Create snapshot including excluded fields for proper serialization.""" + return { + "type": self.type, + "file_id": self.file_id, + "file_data": self.file_data, + "filename": self.filename, + "format": self.format, + } + -class MessageToolResultContent(BaseModel): +class MessageToolResultContent(SerializableModel): type: Literal["tool-result"] = "tool-result" result: Any tool_name: str tool_call_id: str -class MessageToolCallContent(BaseModel): +class MessageToolCallContent(SerializableModel): type: Literal["tool-call"] = "tool-call" id: str tool_name: str @@ -103,7 +131,7 @@ def is_valid(self) -> bool: return False -class Message(ABC, Generic[T]): +class Message(Serializable[dict[str, Any]], Generic[T]): id: str | None role: Role | str content: list[T] @@ -152,6 +180,23 @@ def __str__(self) -> str: def clone(self) -> Self: return type(self)([c.model_copy() for c in self.content], self.meta.copy()) + def create_snapshot(self) -> dict[str, Any]: + """Create a snapshot of this message for serialization.""" + return { + "id": self.id, + "meta": dict(self.meta), + "role": str(self.role), + # Return content as-is - Serializer will handle encoding each item + "content": list(self.content), + } + + def load_snapshot(self, snapshot: dict[str, Any]) -> None: + """Load state from a snapshot dictionary.""" + self.id = snapshot.get("id") + self.meta = dict(snapshot.get("meta") or {}) + # Content is already deserialized by the Serializer as proper objects + self.content = snapshot.get("content", []) + AssistantMessageContent = MessageTextContent | MessageToolCallContent @@ -314,7 +359,7 @@ def from_text(cls, text: str) -> Self: return cls(MessageTextContent(text=text)) -class CustomMessageContent(BaseModel): +class CustomMessageContent(SerializableModel): model_config = ConfigDict(extra="allow") diff --git a/python/beeai_framework/memory/token_memory.py b/python/beeai_framework/memory/token_memory.py index c08cba38d..eca7dcd97 100644 --- a/python/beeai_framework/memory/token_memory.py +++ b/python/beeai_framework/memory/token_memory.py @@ -2,10 +2,11 @@ # SPDX-License-Identifier: Apache-2.0 from math import ceil -from typing import Any +from typing import Any, Self from beeai_framework.backend.message import AnyMessage from beeai_framework.memory.base_memory import BaseMemory +from beeai_framework.serialization import Serializable def simple_estimate(msg: AnyMessage) -> int: @@ -16,12 +17,12 @@ def simple_tokenize(msgs: list[AnyMessage]) -> int: return sum(map(simple_estimate, msgs)) -class TokenMemory(BaseMemory): +class TokenMemory(Serializable[dict[str, Any]], BaseMemory): """Memory implementation that respects token limits.""" def __init__( self, - llm: Any, + llm: Any | None = None, max_tokens: int | None = None, sync_threshold: float = 0.25, capacity_threshold: float = 0.75, @@ -117,8 +118,10 @@ def reset(self) -> None: self._tokens_by_message.clear() async def clone(self) -> "TokenMemory": + llm = self.llm + llm_clone = llm.clone() if llm is not None and hasattr(llm, "clone") else llm cloned = TokenMemory( - self.llm.clone(), + llm_clone, self._max_tokens, self._sync_threshold, self._threshold, @@ -127,3 +130,37 @@ async def clone(self) -> "TokenMemory": cloned._messages = self._messages.copy() cloned._tokens_by_message = self._tokens_by_message.copy() return cloned + + def create_snapshot(self) -> dict[str, Any]: + return { + "messages": list(self._messages), + "max_tokens": self._max_tokens, + "capacity_threshold": self._threshold, + "sync_threshold": self._sync_threshold, + } + + def load_snapshot(self, snapshot: dict[str, Any]) -> None: + self._messages = list(snapshot.get("messages", [])) + self._max_tokens = snapshot.get("max_tokens") + self._threshold = snapshot.get("capacity_threshold", self._threshold) + self._sync_threshold = snapshot.get("sync_threshold", self._sync_threshold) + self._tokens_by_message.clear() + + for message in self._messages: + key = self._get_message_key(message) + estimated_tokens = self.handlers["estimate"](message) + self._tokens_by_message[key] = { + "tokens_count": estimated_tokens, + "dirty": True, + } + + @classmethod + def from_snapshot(cls, snapshot: dict[str, Any]) -> Self: + instance = cls( + llm=None, + max_tokens=snapshot.get("max_tokens"), + sync_threshold=snapshot.get("sync_threshold", 0.25), + capacity_threshold=snapshot.get("capacity_threshold", 0.75), + ) + instance.load_snapshot(snapshot) + return instance diff --git a/python/beeai_framework/memory/unconstrained_memory.py b/python/beeai_framework/memory/unconstrained_memory.py index 0d3147273..98e1a6cf9 100644 --- a/python/beeai_framework/memory/unconstrained_memory.py +++ b/python/beeai_framework/memory/unconstrained_memory.py @@ -1,11 +1,14 @@ # Copyright 2025 © BeeAI a Series of LF Projects, LLC # SPDX-License-Identifier: Apache-2.0 +from typing import Any + from beeai_framework.backend.message import AnyMessage from beeai_framework.memory.base_memory import BaseMemory +from beeai_framework.serialization import Serializable -class UnconstrainedMemory(BaseMemory): +class UnconstrainedMemory(Serializable[dict[str, Any]], BaseMemory): """Simple memory implementation with no constraints.""" def __init__(self) -> None: @@ -33,3 +36,9 @@ async def clone(self) -> "UnconstrainedMemory": cloned = UnconstrainedMemory() cloned._messages = self._messages.copy() return cloned + + def create_snapshot(self) -> dict[str, Any]: + return {"messages": list(self._messages)} + + def load_snapshot(self, snapshot: dict[str, Any]) -> None: + self._messages = list(snapshot.get("messages", [])) diff --git a/python/beeai_framework/serialization/__init__.py b/python/beeai_framework/serialization/__init__.py new file mode 100644 index 000000000..7f00d0b65 --- /dev/null +++ b/python/beeai_framework/serialization/__init__.py @@ -0,0 +1,10 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from beeai_framework.serialization.serializer import Serializable, Serializer, SerializerError + +__all__ = [ + "Serializable", + "Serializer", + "SerializerError", +] diff --git a/python/beeai_framework/serialization/serializer.py b/python/beeai_framework/serialization/serializer.py new file mode 100644 index 000000000..4e4d4e78e --- /dev/null +++ b/python/beeai_framework/serialization/serializer.py @@ -0,0 +1,463 @@ +from __future__ import annotations + +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 +import json +from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping +from dataclasses import dataclass +from datetime import datetime +from itertools import count +from typing import Any, ClassVar, Generic, TypeVar + +T = TypeVar("T") +SerializableT = TypeVar("SerializableT", bound="Serializable[Any]") + + +class SerializerError(RuntimeError): + """Raised when serialization or deserialization fails.""" + + +@dataclass(slots=True) +class _SerializerFactory: + ref: type[Any] + to_plain: Callable[[Any], Any] + from_plain: Callable[[Any, type[Any]], Any] + create_empty: Callable[[], Any] | None = None + update_instance: Callable[[Any, Any], Any] | None = None + + +class Serializer: + """Minimal registry driven serializer compatible with the documentation examples.""" + + _factories: ClassVar[dict[str, _SerializerFactory]] = {} + _type_to_name: ClassVar[dict[type[Any], str]] = {} + version: ClassVar[str] = "0.0.0" + + @staticmethod + def _class_name(ref: type[Any]) -> str: + return ref.__name__ + + @classmethod + def has_factory(cls, name: str) -> bool: + return name in cls._factories + + @classmethod + def has_factory_for_type(cls, ref: type[Any]) -> bool: + return ref in cls._type_to_name + + @classmethod + def register( + cls, + ref: type[Any], + *, + to_plain: Callable[[Any], Any], + from_plain: Callable[[Any, type[Any]], Any], + create_empty: Callable[[], Any] | None = None, + update_instance: Callable[[Any, Any], Any] | None = None, + aliases: Iterable[str] | None = None, + ) -> None: + """Register custom serialization logic for a class.""" + name = cls._class_name(ref) + factory = _SerializerFactory( + ref=ref, + to_plain=to_plain, + from_plain=from_plain, + create_empty=create_empty, + update_instance=update_instance, + ) + + cls._factories[name] = factory + + alias_list = tuple(aliases) if aliases else () + preferred_name = alias_list[0] if alias_list else name + cls._type_to_name[ref] = preferred_name + + for alias in alias_list: + cls._factories[alias] = factory + + @classmethod + def register_serializable( + cls, + ref: type[Serializable[Any]], + *, + aliases: Iterable[str] | None = None, + ) -> None: + """Register classes that implement the Serializable protocol.""" + + if cls.has_factory_for_type(ref): + return + + cls.register( + ref, + to_plain=lambda instance: instance.create_snapshot(), + from_plain=lambda snapshot, factory_ref: factory_ref.from_snapshot(snapshot), + create_empty=lambda: object.__new__(ref), + update_instance=lambda instance, snapshot: instance.load_snapshot(snapshot), + aliases=aliases, + ) + + @classmethod + def deregister(cls, ref: type[Any]) -> None: + """Remove class registration.""" + name = cls._type_to_name.pop(ref, None) + if name is None: + return + + for key in [alias for alias, factory in cls._factories.items() if factory.ref is ref]: + cls._factories.pop(key, None) + + @classmethod + def ensure_registered(cls, ref: type[Any]) -> None: + if cls.has_factory_for_type(ref): + return + + # Avoid circular import by referencing Serializable at runtime + if issubclass(ref, Serializable): + cls.register_serializable(ref) + return + + raise SerializerError(f'Class "{ref.__name__}" is not registered with the serializer.') + + @classmethod + def serialize(cls, raw_data: Any) -> str: + """Serialize python objects into a JSON string.""" + ref_counter = count(1) + seen: dict[int, str] = {} + payload = { + "__version": cls.version, + "__root": cls._encode(raw_data, seen, ref_counter), + } + return json.dumps(payload, separators=(",", ":"), ensure_ascii=False) + + @classmethod + def deserialize( + cls, + data: str | Mapping[str, Any], + *, + expected_type: type[T] | None = None, + extra_classes: Iterable[type[Any]] | None = None, + ) -> T | Any: + """Deserialize objects previously produced by `serialize`.""" + if isinstance(data, str): + try: + payload = json.loads(data) + except json.JSONDecodeError as exc: + raise SerializerError("Invalid serialized payload.") from exc + else: + payload = dict(data) + + if not isinstance(payload, dict) or "__root" not in payload: + raise SerializerError("Serialized payload is missing the '__root' key.") + + if extra_classes: + for extra in extra_classes: + if isinstance(extra, type): + registrar = getattr(extra, "register", None) + if callable(registrar): + registrar() + try: + cls.ensure_registered(extra) + except SerializerError: + # Allow extra classes that register themselves on demand. + continue + + seen: dict[str, Any] = {} + value = cls._decode(payload["__root"], seen) + + if expected_type is not None and not isinstance(value, expected_type): + raise SerializerError( + f"Deserialized value is of type {type(value).__name__}, expected {expected_type.__name__}.", + ) + + return value + + @classmethod + def _find_factory_for_instance(cls, value: Any) -> tuple[_SerializerFactory, str]: + cls.ensure_registered(type(value)) + + for candidate in type(value).__mro__: + name = cls._type_to_name.get(candidate) + if name and name in cls._factories: + return cls._factories[name], name + + raise SerializerError(f'No serializer factory found for "{type(value).__name__}".') + + @classmethod + def _encode(cls, value: Any, seen: dict[int, str], ref_counter: Iterator[int]) -> Any: + if value is None or isinstance(value, bool | int | float | str): + return value + + if isinstance(value, list | tuple | set): + return [cls._encode(item, seen, ref_counter) for item in value] + + if isinstance(value, dict): + return {str(key): cls._encode(item, seen, ref_counter) for key, item in value.items()} + + factory, name = cls._find_factory_for_instance(value) + obj_id = id(value) + if obj_id in seen: + return { + "__serializer": True, + "__ref": seen[obj_id], + } + + ref_id = str(next(ref_counter)) + seen[obj_id] = ref_id + + plain = factory.to_plain(value) + encoded_plain = cls._encode(plain, seen, ref_counter) + return { + "__serializer": True, + "__class": name, + "__ref": ref_id, + "__value": encoded_plain, + } + + @classmethod + def _decode(cls, value: Any, seen: dict[str, Any]) -> Any: + if isinstance(value, list): + return [cls._decode(item, seen) for item in value] + + if isinstance(value, dict): + if value.get("__serializer") is True: + ref_id = value.get("__ref") + class_name = value.get("__class") + + if class_name is None: + if ref_id is None or ref_id not in seen: + raise SerializerError("Encountered reference to unknown object.") + return seen[ref_id] + + if class_name not in cls._factories: + raise SerializerError(f'Class "{class_name}" was not registered.') + + factory = cls._factories[class_name] + payload = value.get("__value") + + if ref_id is not None and ref_id in seen: + instance = seen[ref_id] + if payload is not None and factory.update_instance: + decoded_payload = cls._decode(payload, seen) + factory.update_instance(instance, decoded_payload) + return instance + + if factory.create_empty and factory.update_instance: + instance = factory.create_empty() + if ref_id is not None: + seen[ref_id] = instance + decoded_payload = cls._decode(payload, seen) if payload is not None else None + factory.update_instance(instance, decoded_payload) + return instance + + decoded_payload = cls._decode(payload, seen) if payload is not None else None + instance = factory.from_plain(decoded_payload, factory.ref) + if ref_id is not None: + seen[ref_id] = instance + return instance + + return {key: cls._decode(item, seen) for key, item in value.items()} + + return value + + +class Serializable(ABC, Generic[T]): + """Mixin that provides convenience helpers for registering serializable classes.""" + + def __init_subclass__( + cls, + *, + auto_register: bool = True, + aliases: Iterable[str] | None = None, + **kwargs: Any, + ) -> None: + super().__init_subclass__(**kwargs) + if auto_register: + Serializer.register_serializable(cls, aliases=aliases) + + def serialize(self) -> str: + return Serializer.serialize(self) + + @classmethod + def register(cls, *, aliases: Iterable[str] | None = None) -> None: + Serializer.register_serializable(cls, aliases=aliases) + + @classmethod + def from_serialized( + cls: type[SerializableT], + data: str, + *, + extra_classes: Iterable[type[Any]] | None = None, + ) -> SerializableT: + value = Serializer.deserialize(data, expected_type=cls, extra_classes=extra_classes) + assert isinstance(value, cls) + return value + + @classmethod + def from_snapshot(cls: type[SerializableT], snapshot: T) -> SerializableT: + instance = object.__new__(cls) + instance.load_snapshot(snapshot) + return instance + + @abstractmethod + def create_snapshot(self) -> T | Awaitable[T]: + raise NotImplementedError + + @abstractmethod + def load_snapshot(self, snapshot: T) -> None | Awaitable[None]: + raise NotImplementedError + + +def _parse_datetime(data: Any) -> datetime: + """Parse datetime from various formats including ISO and JavaScript Date strings.""" + if not isinstance(data, str): + raise SerializerError("Encountered malformed payload for Date.") + text = data.replace("Z", "+00:00") if data.endswith("Z") else data + return datetime.fromisoformat(text) + + +def _parse_number(data: Any) -> int | float: + """Parse number from string or numeric value.""" + if data is None: + return 0 + if isinstance(data, str): + number = float(data) + return int(number) if number.is_integer() else number + if isinstance(data, int | float): + return data + raise SerializerError("Encountered malformed payload for Number.") + + +def _parse_list(data: Any) -> list[Any]: + """Parse list from payload.""" + if data is None: + return [] + if isinstance(data, list): + return data.copy() + raise SerializerError("Encountered malformed payload for Array.") + + +def _parse_set(data: Any) -> set[Any]: + """Parse set from payload.""" + if data is None: + return set() + if isinstance(data, list): + return set(data) + raise SerializerError("Encountered malformed payload for Set.") + + +def _parse_dict(data: Any) -> dict[Any, Any]: + """Parse dict from payload (handles both Object and Map formats).""" + if data is None: + return {} + if isinstance(data, Mapping): + return dict(data) + if isinstance(data, list): + # Handle JavaScript Map format: [[key, value], ...] + result: dict[Any, Any] = {} + for pair in data: + if not isinstance(pair, Iterable): + raise SerializerError("Encountered malformed payload for Map.") + try: + key, value = pair + except ValueError as exc: + raise SerializerError("Encountered malformed payload for Map.") from exc + result[key] = value + return result + raise SerializerError("Encountered malformed payload for Object/Map.") + + +# Register datetime with alias for JavaScript Date +Serializer.register( + datetime, + to_plain=lambda value: value.isoformat(), + from_plain=lambda data, ref: _parse_datetime(data), + aliases=("Date",), +) + +# Register built-in types to match TypeScript serializer +Serializer.register( + int, + to_plain=lambda value: str(value), + from_plain=lambda data, ref: _parse_number(data), + aliases=("Number", "BigInt"), +) + +Serializer.register( + float, + to_plain=lambda value: str(value), + from_plain=lambda data, ref: float(data) if data is not None else 0.0, +) + +Serializer.register( + str, + to_plain=lambda value: value, + from_plain=lambda data, ref: str(data) if data is not None else "", + aliases=("String",), +) + +Serializer.register( + bool, + to_plain=lambda value: value, + from_plain=lambda data, ref: ( + data if isinstance(data, bool) else (data.lower() in {"true", "1"} if isinstance(data, str) else bool(data)) + ), + aliases=("Boolean",), +) + +Serializer.register( + list, + to_plain=lambda value: value.copy(), + from_plain=lambda data, ref: _parse_list(data), + create_empty=lambda: [], + update_instance=lambda instance, update: instance.extend(update) if update else None, + aliases=("Array",), +) + +Serializer.register( + set, + to_plain=lambda value: list(value), + from_plain=lambda data, ref: _parse_set(data), + create_empty=lambda: set(), + update_instance=lambda instance, update: instance.update(update) if update else None, + aliases=("Set",), +) + +Serializer.register( + dict, + to_plain=lambda value: dict(value), + from_plain=lambda data, ref: _parse_dict(data), + create_empty=lambda: {}, + update_instance=lambda instance, update: instance.update(update) if update else None, + aliases=("Object", "Map"), +) + + +# Register type(None) for JavaScript Null/Undefined compatibility +class _NoneType: + """Placeholder class for NoneType registration.""" + + +Serializer.register( + _NoneType, + to_plain=lambda value: None, + from_plain=lambda data, ref: None, + aliases=("Null", "Undefined"), +) + + +def _register_pydantic_basemodel() -> None: + """Register Pydantic BaseModel for generic serialization.""" + try: + from pydantic import BaseModel + + Serializer.register( + BaseModel, + to_plain=lambda value: value.model_dump(), + from_plain=lambda data, ref: ref.model_validate(data), + ) + except ImportError: + pass # Pydantic not installed, skip registration + + +_register_pydantic_basemodel() diff --git a/python/examples/serialization/__init__.py b/python/examples/serialization/__init__.py new file mode 100644 index 000000000..32f1f133b --- /dev/null +++ b/python/examples/serialization/__init__.py @@ -0,0 +1,3 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + diff --git a/python/examples/serialization/base.py b/python/examples/serialization/base.py new file mode 100644 index 000000000..651dcb224 --- /dev/null +++ b/python/examples/serialization/base.py @@ -0,0 +1,20 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from datetime import UTC, datetime + +from beeai_framework.serialization import Serializer + + +def main() -> None: + original = datetime(2024, 1, 1, tzinfo=UTC) + serialized = Serializer.serialize(original) + restored = Serializer.deserialize(serialized, expected_type=datetime) + + print(serialized) + print(restored.isoformat()) + print(restored == original) + + +if __name__ == "__main__": + main() diff --git a/python/examples/serialization/context.py b/python/examples/serialization/context.py new file mode 100644 index 000000000..458b96706 --- /dev/null +++ b/python/examples/serialization/context.py @@ -0,0 +1,16 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from beeai_framework.backend import AssistantMessage, UserMessage +from beeai_framework.memory import UnconstrainedMemory + +SERIALIZED_MEMORY = """{"__version":"0.0.0","__root":{"__serializer":true,"__class":"UnconstrainedMemory","__value":{"messages":[{"__serializer":true,"__class":"UserMessage","__value":{"id":null,"meta":{"createdAt":{"__serializer":true,"__class":"datetime","__value":"2025-10-18T19:38:37.859543+00:00"}},"role":"user","content":[{"type":"text","text":"Hello!"}]}},{"__serializer":true,"__class":"AssistantMessage","__value":{"id":null,"meta":{"createdAt":{"__serializer":true,"__class":"datetime","__value":"2025-10-18T19:38:37.859665+00:00"}},"role":"assistant","content":[{"type":"text","text":"Hello, how can I help you?"}]}}]}}}""" + + +def main() -> None: + memory = UnconstrainedMemory.from_serialized(SERIALIZED_MEMORY, extra_classes=[UserMessage, AssistantMessage]) + print([message.text for message in memory.messages]) + + +if __name__ == "__main__": + main() diff --git a/python/examples/serialization/custom_external.py b/python/examples/serialization/custom_external.py new file mode 100644 index 000000000..cad2e0400 --- /dev/null +++ b/python/examples/serialization/custom_external.py @@ -0,0 +1,39 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from dataclasses import dataclass +from datetime import UTC, datetime + +from beeai_framework.serialization import Serializer + + +@dataclass +class ApiToken: + value: str + expires_at: datetime + + +Serializer.register( + ApiToken, + to_plain=lambda token: { + "value": token.value, + "expires_at": token.expires_at, + }, + from_plain=lambda payload, ref: ref( + value=payload["value"], + expires_at=payload["expires_at"], + ), +) + + +def main() -> None: + token = ApiToken("example-token", datetime(2025, 1, 1, tzinfo=UTC)) + serialized = Serializer.serialize(token) + restored = Serializer.deserialize(serialized, expected_type=ApiToken) + + print(restored) + print(restored.expires_at.isoformat()) + + +if __name__ == "__main__": + main() diff --git a/python/examples/serialization/custom_internal.py b/python/examples/serialization/custom_internal.py new file mode 100644 index 000000000..f76fa82e4 --- /dev/null +++ b/python/examples/serialization/custom_internal.py @@ -0,0 +1,33 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from beeai_framework.serialization import Serializable + + +class Counter(Serializable[dict[str, int]]): + def __init__(self, value: int = 0) -> None: + self.value = value + + def increment(self) -> None: + self.value += 1 + + def create_snapshot(self) -> dict[str, int]: + return {"value": self.value} + + def load_snapshot(self, snapshot: dict[str, int]) -> None: + self.value = snapshot["value"] + + +def main() -> None: + counter = Counter(3) + counter.increment() + + serialized = counter.serialize() + restored = Counter.from_serialized(serialized) + + print(counter.value) + print(restored.value) + + +if __name__ == "__main__": + main() diff --git a/python/examples/serialization/memory.py b/python/examples/serialization/memory.py new file mode 100644 index 000000000..b8d21d1c6 --- /dev/null +++ b/python/examples/serialization/memory.py @@ -0,0 +1,24 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +import asyncio + +from beeai_framework.backend import AssistantMessage, UserMessage +from beeai_framework.memory import UnconstrainedMemory + + +async def main() -> None: + memory = UnconstrainedMemory() + await memory.add(UserMessage("What is your name?")) + + serialized = memory.serialize() + restored = UnconstrainedMemory.from_serialized(serialized) + + await restored.add(AssistantMessage("Bee")) + + print(len(restored.messages)) + print(restored.messages[0].text) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/tests/serialization/test_serializer.py b/python/tests/serialization/test_serializer.py new file mode 100644 index 000000000..679ce25e0 --- /dev/null +++ b/python/tests/serialization/test_serializer.py @@ -0,0 +1,172 @@ +# Copyright 2025 © BeeAI a Series of LF Projects, LLC +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from typing import Any, cast + +import pytest + +from beeai_framework.backend import AssistantMessage, UserMessage +from beeai_framework.memory import UnconstrainedMemory +from beeai_framework.serialization import Serializable, Serializer + + +def test_datetime_roundtrip() -> None: + original = datetime(2024, 1, 1, tzinfo=UTC) + + payload = Serializer.serialize(original) + restored = Serializer.deserialize(payload, expected_type=datetime) + + assert restored == original + + +@pytest.mark.asyncio +async def test_unconstrained_memory_roundtrip() -> None: + memory = UnconstrainedMemory() + await memory.add(UserMessage("Hello")) + await memory.add(AssistantMessage("Hi there")) + + payload = memory.serialize() + restored = UnconstrainedMemory.from_serialized(payload) + + assert len(restored.messages) == 2 + assert restored.messages[0].text == "Hello" + + +def test_custom_registration() -> None: + @dataclass + class ApiToken: + value: str + expires_at: datetime + + Serializer.register( + ApiToken, + to_plain=lambda token: { + "value": token.value, + "expires_at": token.expires_at, + }, + from_plain=lambda payload, ref: ref( + value=payload["value"], + expires_at=payload["expires_at"], + ), + ) + + token = ApiToken("secret", datetime(2025, 1, 1, tzinfo=UTC)) + payload = Serializer.serialize(token) + restored = Serializer.deserialize(payload, expected_type=ApiToken) + + assert restored == token + + +def test_serializable_subclass_roundtrip() -> None: + class Counter(Serializable[dict[str, int]]): + def __init__(self, value: int = 0) -> None: + self.value = value + + def create_snapshot(self) -> dict[str, int]: + return {"value": self.value} + + def load_snapshot(self, snapshot: dict[str, int]) -> None: + self.value = snapshot["value"] + + counter = Counter(7) + payload = counter.serialize() + restored = Counter.from_serialized(payload) + + assert restored.value == 7 + + +@pytest.mark.asyncio +async def test_deserialize_with_extra_classes() -> None: + memory = UnconstrainedMemory() + await memory.add(UserMessage("Where are we?")) + + payload = memory.serialize() + + Serializer.deregister(UserMessage) + try: + restored = UnconstrainedMemory.from_serialized(payload, extra_classes=[UserMessage]) + finally: + cast(Any, UserMessage).register() + + assert len(restored.messages) == 1 + assert isinstance(restored.messages[0], UserMessage) + + +def test_shared_reference_roundtrip() -> None: + class Node(Serializable[dict[str, Any]]): + def __init__(self, name: str) -> None: + self.name = name + self.children: list[Node] = [] + + def create_snapshot(self) -> dict[str, Any]: + return {"name": self.name, "children": self.children} + + def load_snapshot(self, snapshot: dict[str, Any]) -> None: + self.name = snapshot["name"] + self.children = snapshot["children"] + + child = Node("child") + root = Node("root") + root.children = [child, child] + + payload = Serializer.serialize(root) + restored = Serializer.deserialize(payload, expected_type=Node) + + assert restored.children[0] is restored.children[1] + + +def test_cyclic_graph_roundtrip() -> None: + class Ring(Serializable[dict[str, Any]]): + def __init__(self, name: str) -> None: + self.name = name + self.next: Ring | None = None + + def create_snapshot(self) -> dict[str, Any]: + return {"name": self.name, "next": self.next} + + def load_snapshot(self, snapshot: dict[str, Any]) -> None: + self.name = snapshot["name"] + self.next = snapshot["next"] + + node = Ring("ring") + node.next = node + + payload = Serializer.serialize(node) + restored = Serializer.deserialize(payload, expected_type=Ring) + + assert restored.next is restored + + +def test_deserialize_typescript_payload() -> None: + class Counter(Serializable[dict[str, int]]): + def __init__(self, value: int = 0) -> None: + self.value = value + + def create_snapshot(self) -> dict[str, int]: + return {"value": self.value} + + def load_snapshot(self, snapshot: dict[str, int]) -> None: + self.value = snapshot["value"] + + payload = ( + '{"__version":"0.0.0","__root":{"__serializer":true,"__class":"Counter","__ref":"1",' + '"__value":{"value":{"__serializer":true,"__class":"Number","__ref":"2","__value":"5"}}}}' + ) + + restored = Serializer.deserialize(payload, expected_type=Counter) + assert isinstance(restored, Counter) + assert restored.value == 5 + + +def test_deserialize_typescript_date() -> None: + payload = ( + '{"__version":"0.0.0","__root":{"__serializer":true,"__class":"Date","__ref":"1",' + '"__value":"2024-01-01T00:00:00.000Z"}}' + ) + + restored = Serializer.deserialize(payload, expected_type=datetime) + assert restored == datetime(2024, 1, 1, tzinfo=UTC)