flock-core 0.3.23__py3-none-any.whl → 0.3.31__py3-none-any.whl

flock/core/serialization/serializable.py

@@ -1,12 +1,33 @@
-"""Module for serializable objects in the system."""
-
+# src/flock/core/serialization/serializable.py
 import json
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any, TypeVar
 
-import cloudpickle
-import msgpack
+# Use yaml if available, otherwise skip yaml methods
+try:
+    import yaml
+
+    YAML_AVAILABLE = True
+except ImportError:
+    YAML_AVAILABLE = False
+
+# Use msgpack if available
+try:
+    import msgpack
+
+    MSGPACK_AVAILABLE = True
+except ImportError:
+    MSGPACK_AVAILABLE = False
+
+# Use cloudpickle
+try:
+    import cloudpickle
+
+    PICKLE_AVAILABLE = True
+except ImportError:
+    PICKLE_AVAILABLE = False
+
 
 T = TypeVar("T", bound="Serializable")
 
@@ -15,79 +36,280 @@ class Serializable(ABC):
     """Base class for all serializable objects in the system.
 
     Provides methods for serializing/deserializing objects to various formats.
+    Subclasses MUST implement to_dict and from_dict.
     """
 
     @abstractmethod
     def to_dict(self) -> dict[str, Any]:
-        """Convert instance to dictionary representation."""
+        """Convert instance to a dictionary representation suitable for serialization.
+        This method should handle converting nested Serializable objects and callables.
+        """
         pass
 
     @classmethod
     @abstractmethod
     def from_dict(cls: type[T], data: dict[str, Any]) -> T:
-        """Create instance from dictionary representation."""
+        """Create instance from a dictionary representation.
+        This method should handle reconstructing nested Serializable objects and callables.
+        """
         pass
 
-    def to_json(self) -> str:
+    # --- JSON Methods ---
+    def to_json(self, indent: int | None = 2) -> str:
         """Serialize to JSON string."""
+        # Import encoder locally to avoid making it a hard dependency if JSON isn't used
+        from .json_encoder import FlockJSONEncoder
+
         try:
-            return json.dumps(self.to_dict())
-        except Exception:
-            raise
+            # Note: to_dict should ideally prepare the structure fully.
+            # FlockJSONEncoder is a fallback for types missed by to_dict.
+            return json.dumps(
+                self.to_dict(), cls=FlockJSONEncoder, indent=indent
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to serialize {self.__class__.__name__} to JSON: {e}"
+            ) from e
 
     @classmethod
     def from_json(cls: type[T], json_str: str) -> T:
         """Create instance from JSON string."""
         try:
-            return cls.from_dict(json.loads(json_str))
-        except Exception:
+            data = json.loads(json_str)
+            return cls.from_dict(data)
+        except json.JSONDecodeError as e:
+            raise ValueError(f"Invalid JSON string: {e}") from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to deserialize {cls.__name__} from JSON: {e}"
+            ) from e
+
+    # --- YAML Methods ---
+    def to_yaml(self, sort_keys=False, default_flow_style=False) -> str:
+        """Serialize to YAML string."""
+        if not YAML_AVAILABLE:
+            raise NotImplementedError(
+                "YAML support requires PyYAML: pip install pyyaml"
+            )
+        try:
+            # to_dict should prepare a structure suitable for YAML dumping
+            return yaml.dump(
+                self.to_dict(),
+                sort_keys=sort_keys,
+                default_flow_style=default_flow_style,
+                allow_unicode=True,
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to serialize {self.__class__.__name__} to YAML: {e}"
+            ) from e
+
+    @classmethod
+    def from_yaml(cls: type[T], yaml_str: str) -> T:
+        """Create instance from YAML string."""
+        if not YAML_AVAILABLE:
+            raise NotImplementedError(
+                "YAML support requires PyYAML: pip install pyyaml"
+            )
+        try:
+            data = yaml.safe_load(yaml_str)
+            if not isinstance(data, dict):
+                raise TypeError(
+                    f"YAML did not yield a dictionary for {cls.__name__}"
+                )
+            return cls.from_dict(data)
+        except yaml.YAMLError as e:
+            raise ValueError(f"Invalid YAML string: {e}") from e
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to deserialize {cls.__name__} from YAML: {e}"
+            ) from e
+
+    def to_yaml_file(self, path: Path | str, **yaml_dump_kwargs) -> None:
+        """Serialize to YAML file."""
+        if not YAML_AVAILABLE:
+            raise NotImplementedError(
+                "YAML support requires PyYAML: pip install pyyaml"
+            )
+        path = Path(path)
+        try:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            yaml_str = self.to_yaml(**yaml_dump_kwargs)
+            path.write_text(yaml_str, encoding="utf-8")
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to write {self.__class__.__name__} to YAML file {path}: {e}"
+            ) from e
+
+    @classmethod
+    def from_yaml_file(cls: type[T], path: Path | str) -> T:
+        """Create instance from YAML file."""
+        if not YAML_AVAILABLE:
+            raise NotImplementedError(
+                "YAML support requires PyYAML: pip install pyyaml"
+            )
+        path = Path(path)
+        try:
+            yaml_str = path.read_text(encoding="utf-8")
+            return cls.from_yaml(yaml_str)
+        except FileNotFoundError:
             raise
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to read {cls.__name__} from YAML file {path}: {e}"
+            ) from e
 
-    def to_msgpack(self, path: Path | None = None) -> bytes:
+    # --- MsgPack Methods ---
+    def to_msgpack(self) -> bytes:
         """Serialize to msgpack bytes."""
+        if not MSGPACK_AVAILABLE:
+            raise NotImplementedError(
+                "MsgPack support requires msgpack: pip install msgpack"
+            )
         try:
-            msgpack_bytes = msgpack.packb(self.to_dict())
-            if path:
-                path.write_bytes(msgpack_bytes)
-            return msgpack_bytes
-        except Exception:
-            raise
+            # Use default hook for complex types if needed, or rely on to_dict
+            return msgpack.packb(self.to_dict(), use_bin_type=True)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to serialize {self.__class__.__name__} to MsgPack: {e}"
+            ) from e
 
     @classmethod
     def from_msgpack(cls: type[T], msgpack_bytes: bytes) -> T:
         """Create instance from msgpack bytes."""
+        if not MSGPACK_AVAILABLE:
+            raise NotImplementedError(
+                "MsgPack support requires msgpack: pip install msgpack"
+            )
         try:
-            return cls.from_dict(msgpack.unpackb(msgpack_bytes))
-        except Exception:
-            raise
+            # Use object_hook if custom deserialization is needed beyond from_dict
+            data = msgpack.unpackb(msgpack_bytes, raw=False)
+            if not isinstance(data, dict):
+                raise TypeError(
+                    f"MsgPack did not yield a dictionary for {cls.__name__}"
+                )
+            return cls.from_dict(data)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to deserialize {cls.__name__} from MsgPack: {e}"
+            ) from e
+
+    def to_msgpack_file(self, path: Path | str) -> None:
+        """Serialize to msgpack file."""
+        if not MSGPACK_AVAILABLE:
+            raise NotImplementedError(
+                "MsgPack support requires msgpack: pip install msgpack"
+            )
+        path = Path(path)
+        try:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            msgpack_bytes = self.to_msgpack()
+            path.write_bytes(msgpack_bytes)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to write {self.__class__.__name__} to MsgPack file {path}: {e}"
+            ) from e
 
     @classmethod
-    def from_msgpack_file(cls: type[T], path: Path) -> T:
+    def from_msgpack_file(cls: type[T], path: Path | str) -> T:
         """Create instance from msgpack file."""
+        if not MSGPACK_AVAILABLE:
+            raise NotImplementedError(
+                "MsgPack support requires msgpack: pip install msgpack"
+            )
+        path = Path(path)
         try:
-            return cls.from_msgpack(path.read_bytes())
-        except Exception:
+            msgpack_bytes = path.read_bytes()
+            return cls.from_msgpack(msgpack_bytes)
+        except FileNotFoundError:
             raise
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to read {cls.__name__} from MsgPack file {path}: {e}"
+            ) from e
 
+    # --- Pickle Methods (Use with caution due to security risks) ---
     def to_pickle(self) -> bytes:
-        """Serialize to pickle bytes."""
+        """Serialize to pickle bytes using cloudpickle."""
+        if not PICKLE_AVAILABLE:
+            raise NotImplementedError(
+                "Pickle support requires cloudpickle: pip install cloudpickle"
+            )
         try:
             return cloudpickle.dumps(self)
-        except Exception:
-            raise
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to serialize {self.__class__.__name__} to Pickle: {e}"
+            ) from e
 
     @classmethod
-    def from_pickle(cls, pickle_bytes: bytes) -> T:
-        """Create instance from pickle bytes."""
+    def from_pickle(cls: type[T], pickle_bytes: bytes) -> T:
+        """Create instance from pickle bytes using cloudpickle."""
+        if not PICKLE_AVAILABLE:
+            raise NotImplementedError(
+                "Pickle support requires cloudpickle: pip install cloudpickle"
+            )
         try:
-            return cloudpickle.loads(pickle_bytes)
-        except Exception:
-            raise
+            instance = cloudpickle.loads(pickle_bytes)
+            if not isinstance(instance, cls):
+                raise TypeError(
+                    f"Deserialized object is not of type {cls.__name__}"
+                )
+            return instance
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to deserialize {cls.__name__} from Pickle: {e}"
+            ) from e
+
+    def to_pickle_file(self, path: Path | str) -> None:
+        """Serialize to pickle file using cloudpickle."""
+        if not PICKLE_AVAILABLE:
+            raise NotImplementedError(
+                "Pickle support requires cloudpickle: pip install cloudpickle"
+            )
+        path = Path(path)
+        try:
+            path.parent.mkdir(parents=True, exist_ok=True)
+            pickle_bytes = self.to_pickle()
+            path.write_bytes(pickle_bytes)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to write {self.__class__.__name__} to Pickle file {path}: {e}"
+            ) from e
 
     @classmethod
-    def from_pickle_file(cls: type[T], path: Path) -> T:
-        """Create instance from pickle file."""
+    def from_pickle_file(cls: type[T], path: Path | str) -> T:
+        """Create instance from pickle file using cloudpickle."""
+        if not PICKLE_AVAILABLE:
+            raise NotImplementedError(
+                "Pickle support requires cloudpickle: pip install cloudpickle"
+            )
+        path = Path(path)
         try:
-            return cls.from_pickle(path.read_bytes())
-        except Exception:
+            pickle_bytes = path.read_bytes()
+            return cls.from_pickle(pickle_bytes)
+        except FileNotFoundError:
             raise
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to read {cls.__name__} from Pickle file {path}: {e}"
+            ) from e
+
+    # _filter_none_values remains unchanged
+    @staticmethod
+    def _filter_none_values(data: Any) -> Any:
+        """Filter out None values from dictionaries and lists recursively."""
+        if isinstance(data, dict):
+            return {
+                k: Serializable._filter_none_values(v)
+                for k, v in data.items()
+                if v is not None
+            }
+        elif isinstance(data, list):
+            # Filter None from list items AND recursively filter within items
+            return [
+                Serializable._filter_none_values(item)
+                for item in data
+                if item is not None
+            ]
+        return data

flock/core/serialization/serialization_utils.py

@@ -0,0 +1,199 @@
+# src/flock/core/serialization/serialization_utils.py
+"""Utilities for recursive serialization/deserialization with callable handling."""
+
+import importlib
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel
+
+# Use TYPE_CHECKING to avoid circular imports
+if TYPE_CHECKING:
+    pass
+
+from flock.core.logging.logging import get_logger
+
+logger = get_logger("serialization.utils")
+
+# Remove this line to avoid circular import at module level
+# FlockRegistry = get_registry()  # Get singleton instance
+
+# --- Serialization Helper ---
+
+
+def serialize_item(item: Any) -> Any:
+    """Recursively prepares an item for serialization (e.g., to dict for YAML/JSON).
+    Converts known callables to their path strings using FlockRegistry.
+    Converts Pydantic models using model_dump.
+    """
+    # Import the registry lazily when needed
+    from flock.core.flock_registry import get_registry
+
+    FlockRegistry = get_registry()
+
+    if isinstance(item, BaseModel):
+        dumped = item.model_dump(mode="json", exclude_none=True)
+        return serialize_item(dumped)
+    elif callable(item) and not isinstance(item, type):
+        path_str = FlockRegistry.get_callable_path_string(
+            item
+        )  # Use registry helper
+        if path_str:
+            return {"__callable_ref__": path_str}
+        else:
+            logger.warning(
+                f"Could not get path string for callable {item}, storing as string."
+            )
+            return str(item)
+    elif isinstance(item, Mapping):
+        return {key: serialize_item(value) for key, value in item.items()}
+    elif isinstance(item, Sequence) and not isinstance(item, str):
+        return [serialize_item(sub_item) for sub_item in item]
+    elif isinstance(
+        item, type
+    ):  # Handle type objects themselves (e.g. if stored directly)
+        type_name = FlockRegistry.get_component_type_name(
+            item
+        )  # Check components first
+        if type_name:
+            return {"__component_ref__": type_name}
+        type_name = FlockRegistry._get_path_string(
+            item
+        )  # Check regular types/classes by path
+        if type_name:
+            return {"__type_ref__": type_name}
+        logger.warning(
+            f"Could not serialize type object {item}, storing as string."
+        )
+        return str(item)
+    else:
+        # Return basic types as is
+        return item
+
+
+# --- Deserialization Helper ---
+
+
+def deserialize_item(item: Any) -> Any:
+    """Recursively processes a deserialized item (e.g., from YAML/JSON dict).
+    Converts reference dicts back to actual callables or types using FlockRegistry.
+    Handles nested lists and dicts.
+    """
+    # Import the registry lazily when needed
+    from flock.core.flock_registry import get_registry
+
+    FlockRegistry = get_registry()
+
+    if isinstance(item, Mapping):
+        if "__callable_ref__" in item and len(item) == 1:
+            path_str = item["__callable_ref__"]
+            try:
+                return FlockRegistry.get_callable(path_str)
+            except KeyError:
+                logger.error(
+                    f"Callable reference '{path_str}' not found during deserialization."
+                )
+                return None
+        elif "__component_ref__" in item and len(item) == 1:
+            type_name = item["__component_ref__"]
+            try:
+                return FlockRegistry.get_component(type_name)
+            except KeyError:
+                logger.error(
+                    f"Component reference '{type_name}' not found during deserialization."
+                )
+                return None
+        elif "__type_ref__" in item and len(item) == 1:
+            type_name = item["__type_ref__"]
+            try:
+                # For general types, use get_type or fallback to dynamic import like get_callable
+                # Using get_type for now, assuming it needs registration
+                return FlockRegistry.get_type(type_name)
+            except KeyError:
+                # Attempt dynamic import as fallback if get_type fails (similar to get_callable)
+                try:
+                    if "." not in type_name:  # Builtins?
+                        mod = importlib.import_module("builtins")
+                    else:
+                        module_name, class_name = type_name.rsplit(".", 1)
+                        mod = importlib.import_module(module_name)
+                    type_obj = getattr(mod, class_name)
+                    if isinstance(type_obj, type):
+                        FlockRegistry.register_type(
+                            type_obj, type_name
+                        )  # Cache it
+                        return type_obj
+                    else:
+                        raise TypeError()
+                except Exception:
+                    logger.error(
+                        f"Type reference '{type_name}' not found in registry or via dynamic import."
+                    )
+                    return None
+
+        else:
+            # Recursively deserialize dictionary values
+            return {key: deserialize_item(value) for key, value in item.items()}
+    elif isinstance(item, Sequence) and not isinstance(item, str):
+        return [deserialize_item(sub_item) for sub_item in item]
+    else:
+        # Return basic types as is
+        return item
+
+
+# --- Component Deserialization Helper ---
+def deserialize_component(
+    data: dict | None, expected_base_type: type
+) -> Any | None:
+    """Deserializes a component (Module, Evaluator, Router) from its dict representation.
+    Uses the 'type' field to find the correct class via FlockRegistry.
+    """
+    # Import the registry and COMPONENT_BASE_TYPES lazily when needed
+    from flock.core.flock_registry import COMPONENT_BASE_TYPES, get_registry
+
+    FlockRegistry = get_registry()
+
+    if data is None:
+        return None
+    if not isinstance(data, dict):
+        logger.error(
+            f"Expected dict for component deserialization, got {type(data)}"
+        )
+        return None
+
+    type_name = data.get(
+        "type"
+    )  # Assuming 'type' key holds the class name string
+    if not type_name:
+        logger.error(f"Component data missing 'type' field: {data}")
+        return None
+
+    try:
+        ComponentClass = FlockRegistry.get_component(type_name)  # Use registry
+        # Optional: Keep the base type check
+        if COMPONENT_BASE_TYPES and not issubclass(
+            ComponentClass, expected_base_type
+        ):
+            raise TypeError(
+                f"Deserialized class {type_name} is not a subclass of {expected_base_type.__name__}"
+            )
+
+        # Recursively deserialize the data *before* passing to Pydantic constructor
+        # This handles nested callables/types within the component's config/data
+        deserialized_data_for_init = {}
+        for k, v in data.items():
+            # Don't pass the 'type' field itself to the constructor if it matches class name
+            if k == "type" and v == ComponentClass.__name__:
+                continue
+            deserialized_data_for_init[k] = deserialize_item(v)
+
+        # Use Pydantic constructor directly. Assumes keys match field names.
+        # from_dict could be added to components for more complex logic if needed.
+        return ComponentClass(**deserialized_data_for_init)
+
+    except (KeyError, TypeError, Exception) as e:
+        logger.error(
+            f"Failed to deserialize component of type '{type_name}': {e}",
+            exc_info=True,
+        )
+        return None

flock/evaluators/declarative/declarative_evaluator.py

@@ -9,6 +9,8 @@ from flock.core.mixin.prompt_parser import PromptParserMixin
 
 
 class DeclarativeEvaluatorConfig(FlockEvaluatorConfig):
+    """Configuration for the DeclarativeEvaluator."""
+
     agent_type_override: str | None = None
     model: str | None = "openai/gpt-4o"
     use_cache: bool = True

flock/modules/memory/memory_module.py

@@ -6,8 +6,12 @@ from typing import Any, Literal
 from pydantic import Field
 from tqdm import tqdm
 
-from flock.core import FlockAgent, FlockModule, FlockModuleConfig
 from flock.core.context.context import FlockContext
+
+# if TYPE_CHECKING:
+#     from flock.core import FlockAgent
+from flock.core.flock_agent import FlockAgent
+from flock.core.flock_module import FlockModule, FlockModuleConfig
 from flock.core.logging.logging import get_logger
 from flock.modules.memory.memory_parser import MemoryMappingParser
 from flock.modules.memory.memory_storage import FlockMemoryStore, MemoryEntry
@@ -282,7 +286,10 @@ class MemoryModule(FlockModule):
         return set(concept_list)
 
     async def _summarize_mode(
-        self, agent: FlockAgent, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
     ) -> str:
         """Extract information chunks using summary mode."""
         split_signature = agent.create_dspy_signature_class(
@@ -300,7 +307,10 @@ class MemoryModule(FlockModule):
         return "\n".join(split_result.chunks)
 
     async def _semantic_splitter_mode(
-        self, agent: FlockAgent, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
     ) -> str | list[dict[str, str]]:
         """Extract information chunks using semantic mode."""
         split_signature = agent.create_dspy_signature_class(
@@ -318,7 +328,10 @@ class MemoryModule(FlockModule):
         return split_result.chunks
 
     async def _character_splitter_mode(
-        self, agent: FlockAgent, inputs: dict[str, Any], result: dict[str, Any]
+        self,
+        agent: FlockAgent,
+        inputs: dict[str, Any],
+        result: dict[str, Any],
     ) -> list[str]:
         """Extract information chunks by splitting text into fixed character lengths."""
         full_text = json.dumps(inputs) + (json.dumps(result) if result else "")

flock/modules/output/output_module.py

@@ -3,11 +3,13 @@
 import json
 import os
 from datetime import datetime
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from pydantic import Field
 
-from flock.core import FlockAgent
+if TYPE_CHECKING:
+    from flock.core import FlockAgent
+
 from flock.core.context.context import FlockContext
 from flock.core.flock_module import FlockModule, FlockModuleConfig
 from flock.core.logging.formatters.themed_formatter import (
@@ -17,6 +19,10 @@ from flock.core.logging.formatters.themes import OutputTheme
 from flock.core.logging.logging import get_logger
 from flock.core.serialization.json_encoder import FlockJSONEncoder
 
+# from flock.core.logging.formatters.themes import OutputTheme
+# from flock.core.logging.logging import get_logger
+# from flock.core.serialization.json_encoder import FlockJSONEncoder
+
 logger = get_logger("module.output")
 
 
@@ -168,7 +174,7 @@ class OutputModule(FlockModule):
 
     async def post_evaluate(
         self,
-        agent: FlockAgent,
+        agent: "FlockAgent",
         inputs: dict[str, Any],
         result: dict[str, Any],
         context: FlockContext | None = None,

flock/workflow/activities.py

@@ -8,9 +8,9 @@ from temporalio import activity
 from flock.core.context.context import FlockContext
 from flock.core.context.context_vars import FLOCK_CURRENT_AGENT, FLOCK_MODEL
 from flock.core.flock_agent import FlockAgent
+from flock.core.flock_registry import get_registry
 from flock.core.flock_router import HandOffRequest
 from flock.core.logging.logging import get_logger
-from flock.core.registry.agent_registry import Registry
 from flock.core.util.input_resolver import resolve_inputs
 
 logger = get_logger("activities")
@@ -26,7 +26,7 @@ async def run_agent(context: FlockContext) -> dict:
     """
     # Start a top-level span for the entire run_agent activity.
     with tracer.start_as_current_span("run_agent") as span:
-        registry = Registry()
+        registry = get_registry()
         previous_agent_name = ""
         if isinstance(context, dict):
             context = FlockContext.from_dict(context)