evoagentx/optimizers/engine/registry.py

import re
from typing import Any, Callable, Dict, List
from ...prompts.template import PromptTemplate
import copy
import warnings
_INDEX_RE = re.compile(r'^(.*?)\[(.*?)\]$')

# _PATH_RE = re.compile(r"""
#     ([a-zA-Z_]\w*)       |  # attr name
#     \[(\d+)\]            |  # list index
#     \[['"](.+?)['"]\]       # dict key
# """, re.VERBOSE)

_PATH_RE = re.compile(r"""
    ([a-zA-Z_]\w*)               |  # attr name
    \[\s*(-?\d+)\s*\]            |  # list index (allow negative)
    \[\s*['"]([^'\"]+)['"]\s*\]     # dict key (inside quotes)
""", re.VERBOSE)

class OptimizableField:
    """
    Represents a parameter that can be optimized.

    This class encapsulates a runtime attribute using dynamic getter and setter
    functions. It allows the parameter to be exposed and manipulated by an external
    optimizer. An initial snapshot of the field can be stored and later used to reset
    the field to its original value.
    """

    def __init__(
        self,
        name: str,
        getter: Callable[[], Any],
        setter: Callable[[Any], None]
    ):
        """
        Initialize an OptimizableField instance.

        Parameters
        ----------
        name : str
            The alias used to register the field in the registry.
        getter : Callable[[], Any]
            A function that returns the current value of the field.
        setter : Callable[[Any], None]
            A function that sets a new value to the field.
        """
        self.name = name
        self._get = getter
        self._set = setter
        self._initial_value = None

    def get(self) -> Any:
        """
        Retrieve the current value of the field.

        Returns
        -------
        Any
            The current value of the field.
        """
        return self._get()

    def set(self, value: Any) -> None:
        """
        Update the field with a new value.

        Parameters
        ----------
        value : Any
            The new value to assign to the field.
        """
        self._set(value)
    
    def init_snapshot(self) -> None:
        """
        Capture a snapshot of the current field value.

        This method stores a deep copy of the current field value so that it
        can be restored later using `reset()`.
        """
        current = self.get()
        self._initial_value = safe_deepcopy(current)
    
    def reset(self) -> None:
        """
        Reset the field to its initial value.

        If the current value object defines a `__reset__()` method, it will be
        called to perform the reset. Otherwise, the field is reset to the deep-copied
        initial value stored by `init_snapshot()`.

        Raises
        ------
        ValueError
            If `init_snapshot()` has not been called before `reset()`.
        """
        current = self.get()

        if self._initial_value is None:
            raise ValueError(f"Field '{self.name}' has no snapshot. Call init_snapshot() first.")

        if hasattr(current, "__reset__") and callable(current.__reset__):
            current.__reset__()
        else:
            self.set(safe_deepcopy(self._initial_value))


class ParamRegistry:
    """
    Central registry for all parameters that can be exposed to optimization.

    Allows dynamic binding and tracking of runtime attributes via dot-paths,
    dictionary keys, or list indices. Provides getter/setter access to all
    registered parameters for optimizers.
    """
    def __init__(self) -> None:
        """Initialize an empty registry of optimizable fields."""
        self.fields: Dict[str, OptimizableField] = {}

    def register_field(self, field: OptimizableField):
        """Manually register an OptimizableField with its alias name."""
        field.init_snapshot()
        self.fields[field.name] = field

    def get(self, name: str) -> Any:
        """Retrieve the current value of a registered field by name."""
        return self.fields[name].get()
    
    def get_field(self, name: str) -> OptimizableField:
        """Retrieve the OptimizableField object by name."""
        if name not in self.fields:
            raise ValueError(f"Field '{name}' is not registered.")
        else:
            return self.fields[name]

    def set(self, name: str, value: Any):
        """Set the value of a registered field by name."""
        self.fields[name].set(value)

    def names(self) -> List[str]:
        """Return a list of all registered field names (aliases)."""
        return list(self.fields.keys())
    
    def reset(self):
        """Roll back all registered fields to their initial values."""
        for field in self.fields.values():
            field.reset()
    
    def reset_field(self, name: str):
        """Roll back a registered field to its initial value."""
        self.fields[name].reset()

    def track(self, root_or_obj: Any, path_or_attr: str, *, name: str | None = None):
        """
        Register a parameter to be optimized. Supports both nested paths and direct attributes.

        Parameters:
        - root_or_obj (Any): the base object or container
        - path_or_attr (str): a path like 'prompt.template' or a direct attribute like 'template'
        - name (str | None): optional alias for this parameter

        Supported formats:
        - registry.track(program, "prompt.template")              # nested attribute
        - registry.track(program, "metadata['style']")           # dictionary key
        - registry.track(program, "components[2].prefix")        # list index
        - registry.track(program.prompt, "template")             # direct object + attribute
        - registry.track([
            (program, "prompt.template"),
            (program, "metadata['style']", "style"),
            (program.prompt, "prefix", "prompt_prefix")
          ])                                                    # batch registration
        - registry.track(program, "prompt.template").track(program, "prompt.prefix")  # chained calls
        
        - registry.track(program, "prompt_template_obj")  # register a prompt_template instance

        Returns:
        - self (PromptRegistry): for chaining
        """
        if isinstance(root_or_obj, list | tuple):
            # batch mode: track([(obj, path), (obj, path, alias), ...])
            # Example:
            # registry.track([
            #     (program, "prompt.template"),
            #     (program, "metadata['style']", "style"),
            #     (program.prompt, "prefix", "prompt_prefix")
            # ])
            for item in root_or_obj:
                if len(item) == 2:
                    self.track(item[0], item[1])
                elif len(item) == 3:
                    self.track(item[0], item[1], name=item[2])
            return self

        if "." in path_or_attr or "[" in path_or_attr:
            return self._track_path(root_or_obj, path_or_attr, name)
        else:
            key = name or path_or_attr

            def getter():
                return getattr(root_or_obj, path_or_attr)

            def setter(v):
                setattr(root_or_obj, path_or_attr, v)

            field = OptimizableField(key, getter, setter)
            if key in self.fields:
                import warnings
                warnings.warn(f"Field '{key}' is already registered. Overwriting.")
            self.register_field(field)
            return self

    def _track_path(self, root: Any, path: str, name: str | None = None):
        """
        Internal helper that registers a nested field (via dot path, index, or key)
        as an OptimizableField by dynamically creating getter and setter functions.

        Parameters:
        - root (Any): the root object to start walking from
        - path (str): dot-separated path supporting list/dict access
        - name (Optional[str]): alias for the parameter (defaults to last path segment)

        Returns:
        - self
        """
        key = name if name is not None else path
        parent, leaf = self._walk(root, path)

        def getter():
            return parent[leaf] if isinstance(parent, (list, dict)) else getattr(parent, leaf)

        def setter(v):
            if isinstance(parent, (list, dict)):
                parent[leaf] = v
            else:
                setattr(parent, leaf, v)

        field = OptimizableField(key, getter, setter)
        self.register_field(field)
        return self
    

    def _walk(self, root, path: str):
        """
        Internal helper to resolve a dot-separated path string into its parent container
        and the leaf attribute/key/index for assignment or retrieval.

        Supports:
        - Nested attributes: e.g. "a.b.c"
        - Dict key access: e.g. "config['key']"
        - List index access: e.g. "layers[0]"

        Parameters:
        - root (Any): root object to walk from
        - path (str): path string to resolve
        - create_missing (bool): unused placeholder for future extensions

        Returns:
        - (parent, leaf): where parent[leaf] or getattr(parent, leaf) is the target
        """
        cur = root
        parts = []
        for match in _PATH_RE.finditer(path):
            attr, idx, key = match.groups()
            if attr:
                parts.append(attr)
            elif idx:
                parts.append(int(idx))
            elif key:
                parts.append(key)

        for part in parts[:-1]:
            if isinstance(part, int):
                cur = cur[part]
            else:
                cur = getattr(cur, part) if hasattr(cur, part) else cur[part]

        leaf = parts[-1]
        parent = cur
        return parent, leaf

    def _walk_old(self, root, path: str):
        """
        Unused Function
        Internal helper to resolve a dot-separated path string into its parent container
        and the leaf attribute/key/index for assignment or retrieval.

        Supports:
        - Nested attributes: e.g. "a.b.c"
        - Dict key access: e.g. "config['key']"
        - List index access: e.g. "layers[0]"

        Parameters:
        - root (Any): root object to walk from
        - path (str): path string to resolve
        - create_missing (bool): unused placeholder for future extensions

        Returns:
        - (parent, leaf): where parent[leaf] or getattr(parent, leaf) is the target
        """
        cur = root
        parts = path.split(".")
        for part in parts[:-1]:
            m = _INDEX_RE.match(part)
            if m:
                attr, idx = m.groups()
                cur = getattr(cur, attr) if attr else cur
                idx = idx.strip()
                if (idx.startswith("'") and idx.endswith("'")) or (idx.startswith('"') and idx.endswith('"')):
                    idx = idx[1:-1]
                elif idx.isdigit():
                    idx = int(idx)
                cur = cur[idx]
            else:
                cur = getattr(cur, part)

        leaf = parts[-1]
        m = _INDEX_RE.match(leaf)
        if m:
            attr, idx = m.groups()
            parent = getattr(cur, attr) if attr else cur
            idx = idx.strip()
            if (idx.startswith("'") and idx.endswith("'")) or (idx.startswith('"') and idx.endswith('"')):
                idx = idx[1:-1]
            elif idx.isdigit():
                idx = int(idx)
            return parent, idx
        return cur, leaf


def safe_deepcopy(obj):
    """
    Safely attempt to deep copy any Python object, with graceful fallback behavior.

    This function performs a standard `copy.deepcopy` when possible. If that fails
    (e.g., due to the presence of uncopyable components such as file handles, threads,
    or custom classes that don't support deep copying), it falls back to a more resilient strategy:

    1. Attempts to create a blank instance of the object's class using `__new__`.
    2. Recursively copies all attributes found in the object's `__dict__`, using:
    - `safe_deepcopy` for deep recursive copy,
    - `copy.copy` as a shallow fallback,
    - or the original reference as a last resort.
    3. If the object has no `__dict__` or cannot be instantiated, returns the original object.

    Parameters:
        obj (Any): The object to be deep copied.

    Returns:
        Any: A deep copy of the input object if possible, or a best-effort fallback copy.
    
    Warnings:
        Issues a `warnings.warn()` message whenever:
        - The deep copy fails and fallback mechanisms are used.
        - An attribute copy fails and falls back to a shallower or direct reference.
        - The class cannot be re-instantiated and the original reference is returned.

    Notes:
        - This function is intended for robust copying in systems where user-defined objects,
        templates, or agents may not support strict deep copying.
        - It is not guaranteed to preserve identity semantics or copy objects with `__slots__`.
        - For critical correctness or mutation isolation, ensure your objects are deepcopy-compatible.

    Example:
        >>> obj = CustomObject()
        >>> obj_copy = safe_deepcopy(obj)
    """
    try:
        return copy.deepcopy(obj)
    except Exception:
        warnings.warn(f"Failed to deepcopy {obj.__class__.__name__}. Falling back to advanced handling.")
        pass  # fallback to custom handling

    try:
        # Try to create a blank instance of the same class
        new_instance = obj.__class__.__new__(obj.__class__)
    except Exception:
        warnings.warn(f"Failed to create a blank instance of {obj.__class__.__name__}. Falling back to reference.")
        return obj  # can't even make a blank instance, fallback to reference

    for attr, value in getattr(obj, "__dict__", {}).items():
        try:
            setattr(new_instance, attr, safe_deepcopy(value))  # recursive copy
        except Exception:
            try:
                warnings.warn(f"Failed to copy {attr} of {obj.__class__.__name__}. Falling back to shallow copy.")
                setattr(new_instance, attr, copy.copy(value))  # shallow copy
            except Exception:
                warnings.warn(f"Failed to copy {attr} of {obj.__class__.__name__}. Falling back to reference.")
                setattr(new_instance, attr, value)  # fallback to reference

    return new_instance
    

class PromptTemplateRegister(ParamRegistry):
    """
    Unused Class
    Enhanced parameter registry that supports directly registering PromptTemplate instances
    or prompt strings as a single optimizable object.
    """

    def track(self, root_or_obj: Any, path_or_attr: str, *, name: str | None = None):
        if isinstance(root_or_obj, (list, tuple)):
            for item in root_or_obj:
                if len(item) == 2:
                    self.track(item[0], item[1])
                elif len(item) == 3:
                    self.track(item[0], item[1], name=item[2])
            return self
        
        if '.' in path_or_attr or '[' in path_or_attr:
            return self._track_path(root_or_obj, path_or_attr, name)
        else:
            key = name or path_or_attr

        try:
            value = getattr(root_or_obj, path_or_attr)
        except AttributeError:
            return super().track(root_or_obj, path_or_attr, name=name)

        if isinstance(value, (str, PromptTemplate)):
            # Register the entire prompt or PromptTemplate object
            field = OptimizableField(
                key,
                getter=lambda: getattr(root_or_obj, path_or_attr),
                setter=lambda v: setattr(root_or_obj, path_or_attr, v)
            )
            self.register_field(field)
            return self

        # Fall back to original path-based tracking if not str/template
        return super().track(root_or_obj, path_or_attr, name=name)