lionagi 0.5.5__py3-none-any.whl → 0.6.1__py3-none-any.whl

lionagi/protocols/graph/graph.py

@@ -0,0 +1,290 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import deque
+from typing import Any, Literal, Self
+
+from pydantic import Field, model_validator
+
+from lionagi._errors import ItemExistsError, RelationError
+from lionagi.protocols._concepts import Relational
+
+from .._concepts import Relational
+from ..generic.element import ID, Element
+from ..generic.pile import Pile
+from .edge import Edge
+from .node import Node
+
+__all__ = ("Graph",)
+
+
+class Graph(Element, Relational):
+
+    internal_nodes: Pile[Node] = Field(
+        default_factory=lambda: Pile(item_type={Node}, strict_type=False),
+        title="Internal Nodes",
+        description="A collection of nodes in the graph.",
+    )
+    internal_edges: Pile[Edge] = Field(
+        default_factory=lambda: Pile(item_type={Edge}, strict_type=False),
+        title="Internal Edges",
+        description="A collection of edges in the graph.",
+    )
+    node_edge_mapping: dict = Field(default_factory=dict, exclude=True)
+
+    @model_validator(mode="after")
+    def _validate_node_mapping(self) -> Self:
+        self.node_edge_mapping = {}
+        if self.internal_nodes:
+            for node in self.internal_nodes:
+                if node.id not in self.node_edge_mapping:
+                    self.node_edge_mapping[node.id] = {"in": {}, "out": {}}
+
+        if self.internal_edges:
+            for edge in self.internal_edges:
+                self.node_edge_mapping[edge.head]["out"][edge.id] = edge.tail
+                self.node_edge_mapping[edge.tail]["in"][edge.id] = edge.head
+        return self
+
+    def add_node(self, node: Relational) -> None:
+        """Add a node to the graph."""
+        if not isinstance(node, Relational):
+            raise RelationError(
+                "Failed to add node: Invalid node type: "
+                "not a <Relational> entity."
+            )
+        _id = ID.get_id(node)
+        try:
+            self.internal_nodes.insert(len(self.internal_nodes), node)
+            self.node_edge_mapping[_id] = {"in": {}, "out": {}}
+        except ItemExistsError as e:
+            raise RelationError(f"Error adding node: {e}")
+
+    def add_edge(self, edge: Edge, /) -> None:
+        """Add an edge to the graph, linking two existing nodes."""
+        if not isinstance(edge, Edge):
+            raise RelationError("Failed to add edge: Invalid edge type.")
+        if (
+            edge.head not in self.internal_nodes
+            or edge.tail not in self.internal_nodes
+        ):
+            raise RelationError(
+                "Failed to add edge: Either edge head or tail node does"
+                " not exist in the graph."
+            )
+        try:
+            self.internal_edges.insert(len(self.internal_edges), edge)
+            self.node_edge_mapping[edge.head]["out"][edge.id] = edge.tail
+            self.node_edge_mapping[edge.tail]["in"][edge.id] = edge.head
+        except ItemExistsError as e:
+            raise RelationError(f"Error adding node: {e}")
+
+    def remove_node(self, node: ID[Node].Ref, /) -> None:
+        """
+        Remove a node from the graph.
+
+        This method removes a node and all connected edges from the graph.
+
+        Args:
+            node (Node | str): The node or node ID to remove.
+
+        Raises:
+            RelationError: If the node does not exist in the graph.
+        """
+        _id = ID.get_id(node)
+        if _id not in self.internal_nodes:
+            raise RelationError(f"Node {node} not found in the graph nodes.")
+
+        in_edges: dict = self.node_edge_mapping[_id]["in"]
+        for edge_id, node_id in in_edges.items():
+            self.node_edge_mapping[node_id]["out"].pop(edge_id)
+            self.internal_edges.pop(edge_id)
+
+        out_edges: dict = self.node_edge_mapping[_id]["out"]
+        for edge_id, node_id in out_edges.items():
+            self.node_edge_mapping[node_id]["in"].pop(edge_id)
+            self.internal_edges.pop(edge_id)
+
+        self.node_edge_mapping.pop(_id)
+        return self.internal_nodes.pop(_id)
+
+    def remove_edge(self, edge: Edge | str, /) -> None:
+        """
+        Remove an edge from the graph.
+
+        Args:
+            edge (Edge | str): The edge or edge ID to remove.
+
+        Raises:
+            RelationError: If the edge does not exist in the graph.
+        """
+        _id = ID.get_id(edge)
+        if _id not in self.internal_edges:
+            raise RelationError(f"Edge {edge} not found in the graph edges.")
+
+        edge = self.internal_edges[_id]
+        self.node_edge_mapping[edge.head]["out"].pop(_id)
+        self.node_edge_mapping[edge.tail]["in"].pop(_id)
+
+        return self.internal_edges.pop(_id)
+
+    def find_node_edge(
+        self,
+        node: Any,
+        /,
+        direction: Literal["both", "in", "out"] = "both",
+    ) -> list[Edge]:
+        """
+        Find edges associated with a node by direction (in, out, or both).
+
+        Args:
+            node (ID[Node].Ref): The node or node ID.
+            direction: 'in', 'out', or 'both' to filter the edges.
+
+        Returns:
+            list[Edge]: The matching edges.
+
+        Raises:
+            ValueError: If direction is invalid.
+            RelationError: If node is not in the graph.
+        """
+        if direction not in {"both", "in", "out"}:
+            raise ValueError("The direction should be 'both', 'in', or 'out'.")
+
+        _id = ID.get_id(node)
+        if _id not in self.internal_nodes:
+            raise RelationError(f"Node {node} not found in the graph nodes.")
+
+        result = []
+        if direction in {"both", "in"}:
+            for edge_id in self.node_edge_mapping[_id]["in"]:
+                result.append(self.internal_edges[edge_id])
+
+        if direction in {"both", "out"}:
+            for edge_id in self.node_edge_mapping[_id]["out"]:
+                result.append(self.internal_edges[edge_id])
+
+        return Pile(result, item_type={Edge}, strict_type=False)
+
+    def get_heads(self) -> Pile[Node]:
+        """Return nodes with no incoming edges (head nodes)."""
+        result = []
+        for node_id in self.node_edge_mapping.keys():
+            if self.node_edge_mapping[node_id]["in"] == {}:
+                result.append(self.internal_nodes[node_id])
+        return Pile(result, item_type={Node}, strict_type=False)
+
+    def get_predecessors(self, node: Node, /) -> Pile[Node]:
+        """Return all nodes that have outbound edges to the given node."""
+        edges = self.find_node_edge(node, direction="in")
+        result = []
+        for edge in edges:
+            result.append(self.internal_nodes[edge.head])
+        return Pile(result, item_type={Node}, strict_type=False)
+
+    def get_successors(self, node: Node, /) -> Pile[Node]:
+        """
+        Get all successor nodes of a given node.
+
+        Successors are nodes that have incoming edges from the given node.
+
+        Args:
+            node (Node): The node to find successors for.
+
+        Returns:
+            Pile: A Pile containing all successor nodes.
+        """
+        edges = self.find_node_edge(node, direction="out")
+        result = []
+        for edge in edges:
+            result.append(self.internal_nodes[edge.tail])
+        return Pile(result, item_type={Node}, strict_type=False)
+
+    def to_networkx(self, **kwargs) -> Any:
+        """Convert the graph to a NetworkX graph object."""
+        try:
+            import networkx as nx  # type: ignore
+        except ImportError:
+            raise ImportError(
+                "NetworkX is not installed. Please install it with `pip install networkx`."
+            )
+
+        g = nx.DiGraph(**kwargs)
+        for node in self.internal_nodes:
+            node_info = node.to_dict()
+            node_info.pop("id")
+            g.add_node(str(node.id), **node_info)
+
+        for _edge in self.internal_edges:
+            edge_info = _edge.to_dict()
+            edge_info.pop("id")
+            source_node_id = edge_info.pop("head")
+            target_node_id = edge_info.pop("tail")
+            g.add_edge(str(source_node_id), str(target_node_id), **edge_info)
+
+        return g
+
+    def display(
+        self,
+        node_label="lion_class",
+        edge_label="label",
+        draw_kwargs={},
+        **kwargs,
+    ):
+        """Display the graph using NetworkX and Matplotlib."""
+
+        try:
+            import matplotlib.pyplot as plt  # type: ignore
+            import networkx as nx  # type: ignore
+        except ImportError:
+            raise ImportError(
+                "Failed to import Matplotlib. Please install the package with `pip install matplotlib`"
+            )
+
+        g = self.to_networkx(**kwargs)
+        pos = nx.spring_layout(g)
+        nx.draw(
+            g,
+            pos,
+            labels=nx.get_node_attributes(g, node_label),
+            **draw_kwargs,
+        )
+
+        edge_labels = nx.get_edge_attributes(g, edge_label)
+        if edge_labels:
+            nx.draw_networkx_edge_labels(g, pos, edge_labels=edge_labels)
+
+        plt.axis("off")
+        plt.show()
+
+    def is_acyclic(self) -> bool:
+        """Check if the graph is acyclic (contains no cycles)."""
+        node_ids = list(self.internal_nodes.progression)
+        check_deque = deque(node_ids)
+
+        # 0: unvisited, 1: visiting, 2: visited
+        check_dict = {nid: 0 for nid in node_ids}
+
+        def visit(nid):
+            if check_dict[nid] == 2:
+                return True
+            elif check_dict[nid] == 1:
+                return False
+
+            check_dict[nid] = 1
+            for edge_id in self.node_edge_mapping[nid]["out"]:
+                edge: Edge = self.internal_edges[edge_id]
+                if not visit(edge.tail):
+                    return False
+            check_dict[nid] = 2
+            return True
+
+        while check_deque:
+            key = check_deque.pop()
+            if not visit(key):
+                return False
+        return True
+
+    def __contains__(self, item: object) -> bool:
+        return item in self.internal_nodes or item in self.internal_edges

lionagi/protocols/graph/node.py

@@ -0,0 +1,125 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from typing import Any, ClassVar
+
+from pydantic import field_validator
+
+from lionagi._class_registry import LION_CLASS_REGISTRY
+
+from .._concepts import Relational
+from ..adapters.adapter import AdapterRegistry
+from ..adapters.json_adapter import JsonAdapter, JsonFileAdapter
+from ..adapters.pandas_.pd_series_adapter import PandasSeriesAdapter
+from ..generic.element import Element
+
+NODE_DEFAULT_ADAPTERS = (
+    JsonAdapter,
+    JsonFileAdapter,
+    PandasSeriesAdapter,
+)
+
+
+class NodeAdapterRegistry(AdapterRegistry):
+    pass
+
+
+for i in NODE_DEFAULT_ADAPTERS:
+    NodeAdapterRegistry.register(i)
+
+__all__ = ("Node",)
+
+
+class Node(Element, Relational):
+    """
+    A base class for all Nodes in a graph, storing:
+      - Arbitrary content
+      - Metadata as a dict
+      - An optional numeric embedding (list of floats)
+      - Automatic subclass registration
+    """
+
+    _adapter_registry: ClassVar[AdapterRegistry] = NodeAdapterRegistry
+
+    content: Any = None
+    embedding: list[float] | None = None
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
+        """Initialize and register subclasses in the global class registry."""
+        super().__pydantic_init_subclass__(**kwargs)
+        LION_CLASS_REGISTRY[cls.class_name(full=True)] = cls
+
+    @field_validator("embedding", mode="before")
+    def _parse_embedding(
+        cls, value: list[float] | str | None
+    ) -> list[float] | None:
+        if value is None:
+            return None
+        if isinstance(value, str):
+            try:
+                loaded = json.loads(value)
+                if not isinstance(loaded, list):
+                    raise ValueError
+                return [float(x) for x in loaded]
+            except Exception as e:
+                raise ValueError("Invalid embedding string.") from e
+        if isinstance(value, list):
+            try:
+                return [float(x) for x in value]
+            except Exception as e:
+                raise ValueError("Invalid embedding list.") from e
+        raise ValueError(
+            "Invalid embedding type; must be list or JSON-encoded string."
+        )
+
+    def adapt_to(
+        self, obj_key: str, /, many: bool = False, **kwargs: Any
+    ) -> Any:
+        """
+        Convert this Node to another format using a registered adapter.
+        """
+        return self._get_adapter_registry().adapt_to(
+            self, obj_key, many=many, **kwargs
+        )
+
+    @classmethod
+    def adapt_from(
+        cls,
+        obj: Any,
+        obj_key: str,
+        /,
+        many: bool = False,
+        **kwargs: Any,
+    ) -> "Node":
+        """
+        Construct a Node from an external format using a registered adapter.
+        If the adapter returns a dictionary with 'lion_class', we can
+        auto-delegate to the correct subclass via from_dict.
+        """
+        result = cls._get_adapter_registry().adapt_from(
+            cls, obj, obj_key, many=many, **kwargs
+        )
+        # If adapter returned multiple items, choose the first or handle as needed.
+        if isinstance(result, list):
+            result = result[0]
+        return cls.from_dict(result)
+
+    @classmethod
+    def _get_adapter_registry(cls) -> AdapterRegistry:
+        if isinstance(cls._adapter_registry, type):
+            cls._adapter_registry = cls._adapter_registry()
+        return cls._adapter_registry
+
+    @classmethod
+    def register_adapter(cls, adapter: Any) -> None:
+        cls._get_adapter_registry().register(adapter)
+
+    @classmethod
+    def list_adapters(cls) -> list[str]:
+        return cls._get_adapter_registry().list_adapters()
+
+
+# File: protocols/graph/node.py

lionagi/protocols/mail/exchange.py

@@ -0,0 +1,116 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import asyncio
+from typing import Any
+
+from lionagi.protocols.generic.element import IDType
+
+from .._concepts import Communicatable
+from ..generic.element import ID
+from ..generic.pile import Pile
+from .mail import Mail, Package, PackageCategory
+from .mailbox import Mailbox
+
+
+class Exchange:
+
+    def __init__(self, sources: ID[Communicatable].ItemSeq = None):
+        self.sources: Pile[Communicatable] = Pile(
+            item_type={Communicatable}, strict_type=False
+        )
+        self.buffer: dict[IDType, list[Mail]] = {}
+        self.mailboxes: dict[IDType, Mailbox] = {}
+        if sources:
+            self.add_source(sources)
+        self._execute_stop: bool = False
+
+    def add_source(self, sources: ID[Communicatable].ItemSeq):
+        if sources in self.sources:
+            raise ValueError(
+                f"Source {sources} already exists in the mail manager."
+            )
+
+        self.sources.include(sources)
+        for source in sources:
+            self.mailboxes[source.id] = source.mailbox
+        self.buffer.update({source.id: [] for source in sources})
+
+    def delete_source(self, sources: ID[Communicatable].ItemSeq):
+        """this will remove the source from the mail manager and all assosiated pending mails"""
+        if not sources in self.sources:
+            raise ValueError(
+                f"Source {sources} does not exist in the mail manager."
+            )
+
+        self.sources.exclude(sources)
+        for source in sources:
+            self.buffer.pop(source.id)
+            self.mailboxes.pop(source.id)
+
+    @staticmethod
+    def create_mail(
+        sender: ID[Communicatable],
+        recipient: ID[Communicatable],
+        category: PackageCategory | str,
+        item: Any,
+        request_source: Any = None,
+    ) -> Mail:
+        package = Package(
+            category=category, item=item, request_source=request_source
+        )
+        return Mail(sender=sender, recipient=recipient, package=package)
+
+    def collect(self, sender: ID[Communicatable]):
+        """collect all mails from a particular sender"""
+        if sender not in self.sources:
+            raise ValueError(f"Sender source {sender} does not exist.")
+
+        sender_mailbox: Mailbox = self.sources[sender].mailbox
+
+        while sender_mailbox.pending_outs:
+            mail: Mail = sender_mailbox.pile_.popleft()
+            self.buffer[mail.recipient].append(mail)
+
+    def deliver(self, recipient: ID[Communicatable]):
+        if recipient not in self.sources:
+            raise ValueError(f"Recipient source {recipient} does not exist.")
+
+        recipient_mailbox: Mailbox = self.sources[recipient].mailbox
+
+        while self.buffer[recipient]:
+            mail = self.buffer[recipient].pop(0)
+            if mail.recipient != recipient:
+                raise ValueError(
+                    f"Mail recipient {mail.recipient} does not match recipient {recipient}"
+                )
+            if mail.sender not in self.sources:
+                raise ValueError(f"Mail sender {mail.sender} does not exist.")
+            recipient_mailbox.append_in(mail)
+
+    def collect_all(self) -> None:
+        """Collect mail from all sources."""
+        for source in self.sources:
+            self.collect(sender=source.id)
+
+    def deliver_all(self) -> None:
+        """Send mail to all recipients."""
+        for source in self.sources:
+            self.deliver(recipient=source.id)
+
+    async def execute(self, refresh_time: int = 1) -> None:
+        """
+        Execute mail collection and sending process asynchronously.
+
+        This method runs in a loop, collecting and sending mail at
+            regular intervals.
+
+        Args:
+            refresh_time (int, optional): The time to wait between
+                each cycle in seconds. Defaults to 1.
+        """
+        while not self._execute_stop:
+            self.collect_all()
+            self.deliver_all()
+            await asyncio.sleep(refresh_time)

lionagi/protocols/mail/mail.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from pydantic import field_validator
+
+from .._concepts import Sendable
+from ..generic.element import Element
+from ..messages.base import IDType
+from .package import Package, PackageCategory
+
+
+class Mail(Element, Sendable):
+
+    sender: IDType
+    recipient: IDType
+    package: Package
+
+    @field_validator("sender", "recipient")
+    def _validate_sender_recipient(cls, value):
+        return IDType.validate(value)
+
+    @property
+    def category(self) -> PackageCategory:
+        return self.package.category

lionagi/protocols/mail/mailbox.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from lionagi.protocols.generic.element import IDType
+
+from ..generic.pile import Pile, Progression
+from .mail import Mail
+
+__all__ = ("Mailbox",)
+
+
+class Mailbox:
+
+    def __init__(self):
+        self.pile_ = Pile(item_type=Mail, strict_type=True)
+        self.pending_ins: dict[IDType, Progression] = {}
+        self.pending_outs = Progression()
+
+    def __contains__(self, item: Mail) -> bool:
+        return item in self.pile_
+
+    @property
+    def senders(self) -> list[str]:
+        return list(self.pending_ins.keys())
+
+    def append_in(self, item: Mail, /):
+        if item.sender not in self.pending_ins:
+            self.pending_ins[item.sender] = Progression()
+        self.pending_ins[item.sender].include(item)
+        self.pile_.include(item)
+
+    def append_out(self, item: Mail, /):
+        self.pending_outs.include(item)
+        self.pile_.include(item)
+
+    def exclude(self, item: Mail, /):
+        self.pile_.exclude(item)
+        self.pending_outs.exclude(item)
+        for v in self.pending_ins.values():
+            v.exclude(item)
+
+    def __bool__(self) -> bool:
+        return bool(self.pile_)
+
+    def __len__(self) -> int:
+        return len(self.pile_)