lionagi 0.0.306__py3-none-any.whl → 0.0.307__py3-none-any.whl

lionagi/core/schema/structure.py

@@ -1,14 +1,14 @@
-from typing import List, Any, Dict
+import time
+from typing import List, Any, Dict, Callable
 from collections import deque
 from pydantic import Field
 
-from lionagi.libs.sys_util import SysUtil
-from lionagi.libs import ln_func_call as func_call
-from lionagi.libs.ln_async import AsyncUtil
+from lionagi.libs import SysUtil, func_call, AsyncUtil
 
-from lionagi.core.schema.base_node import BaseRelatableNode, BaseNode
+from .base_node import BaseRelatableNode, BaseNode, Tool
 from lionagi.core.mail.schema import BaseMail
 
+from lionagi.core.schema.condition import Condition
 
 from lionagi.core.schema.action_node import ActionNode, ActionSelection
 from lionagi.core.schema.base_node import Tool
@@ -19,113 +19,44 @@ class Relationship(BaseRelatableNode):
     Represents a relationship between two nodes in a graph.
 
     Attributes:
-        source_node_id (str): The identifier of the source node.
-        target_node_id (str): The identifier of the target node.
-        condition (Dict[str, Any]): A dictionary representing conditions for the relationship.
+            source_node_id (str): The identifier of the source node.
+            target_node_id (str): The identifier of the target node.
+            condition (Dict[str, Any]): A dictionary representing conditions for the relationship.
 
     Examples:
-        >>> relationship = Relationship(source_node_id="node1", target_node_id="node2")
-        >>> relationship.add_condition({"key": "value"})
-        >>> condition_value = relationship.get_condition("key")
-        >>> relationship.remove_condition("key")
+            >>> relationship = Relationship(source_node_id="node1", target_node_id="node2")
+            >>> relationship.add_condition({"key": "value"})
+            >>> condition_value = relationship.get_condition("key")
+            >>> relationship.remove_condition("key")
     """
 
     source_node_id: str
     target_node_id: str
     bundle: bool = False
-    condition: dict = Field(default={})
-
-    def add_condition(self, condition: Dict[str, Any]) -> None:
-        """
-        Adds a condition to the relationship.
-
-        Args:
-            condition: The condition to be added.
-
-        Examples:
-            >>> relationship = Relationship(source_node_id="node1", target_node_id="node2")
-            >>> relationship.add_condition({"key": "value"})
-        """
-        self.condition.update(condition)
-
-    def remove_condition(self, condition_key: str) -> Any:
-        """
-        Removes a condition from the relationship.
-
-        Args:
-            condition_key: The key of the condition to be removed.
-
-        Returns:
-            The value of the removed condition.
-
-        Raises:
-            KeyError: If the condition key is not found.
-
-        Examples:
-            >>> relationship = Relationship(source_node_id="node1", target_node_id="node2", condition={"key": "value"})
-            >>> relationship.remove_condition("key")
-            'value'
-        """
-        if condition_key not in self.condition.keys():
-            raise KeyError(f"condition {condition_key} is not found")
-        return self.condition.pop(condition_key)
-
-    def condition_exists(self, condition_key: str) -> bool:
-        """
-        Checks if a condition exists in the relationship.
-
-        Args:
-            condition_key: The key of the condition to check.
-
-        Returns:
-            True if the condition exists, False otherwise.
-
-        Examples:
-            >>> relationship = Relationship(source_node_id="node1", target_node_id="node2", condition={"key": "value"})
-            >>> relationship.condition_exists("key")
-            True
-        """
-        if condition_key in self.condition.keys():
-            return True
-        else:
-            return False
-
-    def get_condition(self, condition_key: str | None = None) -> Any:
-        """
-        Retrieves a specific condition or all conditions of the relationship.
+    condition: Callable = None
 
-        Args:
-            condition_key: The key of the specific condition. If None, all conditions are returned.
-
-        Returns:
-            The requested condition or all conditions if no key is provided.
-
-        Raises:
-            ValueError: If the specified condition key does not exist.
+    def add_condition(self, condition: Condition):
+        if not isinstance(condition, Condition):
+            raise ValueError(
+                "Invalid condition type, please use Condition class to build a valid condition"
+            )
+        self.condition = condition
 
-        Examples:
-            >>> relationship = Relationship(source_node_id="node1", target_node_id="node2", condition={"key": "value"})
-            >>> relationship.get_condition("key")
-            'value'
-            >>> relationship.get_condition()
-            {'key': 'value'}
-        """
-        if condition_key is None:
-            return self.condition
-        if self.condition_exists(condition_key=condition_key):
-            return self.condition[condition_key]
-        else:
-            raise ValueError(f"Condition {condition_key} does not exist")
+    def check_condition(self, source_obj):
+        try:
+            return bool(self.condition(source_obj))
+        except:
+            raise ValueError("Invalid relationship condition function")
 
     def _source_existed(self, obj: Dict[str, Any]) -> bool:
         """
         Checks if the source node exists in a given object.
 
         Args:
-            obj (Dict[str, Any]): The object to check.
+                obj (Dict[str, Any]): The object to check.
 
         Returns:
-            bool: True if the source node exists, False otherwise.
+                bool: True if the source node exists, False otherwise.
         """
         return self.source_node_id in obj.keys()
 
@@ -134,10 +65,10 @@ class Relationship(BaseRelatableNode):
         Checks if the target node exists in a given object.
 
         Args:
-            obj (Dict[str, Any]): The object to check.
+                obj (Dict[str, Any]): The object to check.
 
         Returns:
-            bool: True if the target node exists, False otherwise.
+                bool: True if the target node exists, False otherwise.
         """
         return self.target_node_id in obj.keys()
 
@@ -146,13 +77,13 @@ class Relationship(BaseRelatableNode):
         Validates the existence of both source and target nodes in a given object.
 
         Args:
-            obj (Dict[str, Any]): The object to check.
+                obj (Dict[str, Any]): The object to check.
 
         Returns:
-            bool: True if both nodes exist.
+                bool: True if both nodes exist.
 
         Raises:
-            ValueError: If either the source or target node does not exist.
+                ValueError: If either the source or target node does not exist.
         """
         if self._source_existed(obj) and self._target_existed(obj):
             return True
@@ -167,9 +98,9 @@ class Relationship(BaseRelatableNode):
         Returns a simple string representation of the Relationship.
 
         Examples:
-            >>> relationship = Relationship(source_node_id="node1", target_node_id="node2")
-            >>> str(relationship)
-            'Relationship (id_=None, from=node1, to=node2, label=None)'
+                >>> relationship = Relationship(source_node_id="node1", target_node_id="node2")
+                >>> str(relationship)
+                'Relationship (id_=None, from=node1, to=node2, label=None)'
         """
         return (
             f"Relationship (id_={self.id_}, from={self.source_node_id}, to={self.target_node_id}, "
@@ -181,9 +112,9 @@ class Relationship(BaseRelatableNode):
         Returns a detailed string representation of the Relationship.
 
         Examples:
-            >>> relationship = Relationship(source_node_id="node1", target_node_id="node2")
-            >>> repr(relationship)
-            'Relationship(id_=None, from=node1, to=node2, content=None, metadata=None, label=None)'
+                >>> relationship = Relationship(source_node_id="node1", target_node_id="node2")
+                >>> repr(relationship)
+                'Relationship(id_=None, from=node1, to=node2, content=None, metadata=None, label=None)'
         """
         return (
             f"Relationship(id_={self.id_}, from={self.source_node_id}, to={self.target_node_id}, "
@@ -196,20 +127,20 @@ class Graph(BaseRelatableNode):
     Represents a graph structure, consisting of nodes and their relationship.
 
     Attributes:
-        nodes (Dict[str, BaseNode]): A dictionary of nodes in the graph.
-        relationships (Dict[str, Relationship]): A dictionary of relationship between nodes in the graph.
-        node_relationships (Dict[str, Dict[str, Dict[str, str]]]): A dictionary tracking the relationship of each node.
+            nodes (Dict[str, BaseNode]): A dictionary of nodes in the graph.
+            relationships (Dict[str, Relationship]): A dictionary of relationship between nodes in the graph.
+            node_relationships (Dict[str, Dict[str, Dict[str, str]]]): A dictionary tracking the relationship of each node.
 
     Examples:
-        >>> graph = Graph()
-        >>> node = BaseNode(id_='node1')
-        >>> graph.add_node(node)
-        >>> graph.node_exists(node)
-        True
-        >>> relationship = Relationship(id_='rel1', source_node_id='node1', target_node_id='node2')
-        >>> graph.add_relationship(relationship)
-        >>> graph.relationship_exists(relationship)
-        True
+            >>> graph = Graph()
+            >>> node = BaseNode(id_='node1')
+            >>> graph.add_node(node)
+            >>> graph.node_exists(node)
+            True
+            >>> relationship = Relationship(id_='rel1', source_node_id='node1', target_node_id='node2')
+            >>> graph.add_relationship(relationship)
+            >>> graph.relationship_exists(relationship)
+            True
     """
 
     nodes: dict = Field(default={})
@@ -221,7 +152,7 @@ class Graph(BaseRelatableNode):
         Adds a node to the graph.
 
         Args:
-            node (BaseNode): The node to add to the graph.
+                node (BaseNode): The node to add to the graph.
         """
 
         self.nodes[node.id_] = node
@@ -232,10 +163,10 @@ class Graph(BaseRelatableNode):
         Adds a relationship between nodes in the graph.
 
         Args:
-            relationship (Relationship): The relationship to add.
+                relationship (Relationship): The relationship to add.
 
         Raises:
-            KeyError: If either the source or target node of the relationship is not found in the graph.
+                KeyError: If either the source or target node of the relationship is not found in the graph.
         """
         if relationship.source_node_id not in self.node_relationships.keys():
             raise KeyError(f"node {relationship.source_node_id} is not found.")
@@ -257,14 +188,14 @@ class Graph(BaseRelatableNode):
         Retrieves relationship of a specific node or all relationship in the graph.
 
         Args:
-            node (Optional[BaseNode]): The node whose relationship to retrieve. If None, retrieves all relationship.
-            out_edge (bool): Whether to retrieve outgoing relationship. If False, retrieves incoming relationship.
+                node (Optional[BaseNode]): The node whose relationship to retrieve. If None, retrieves all relationship.
+                out_edge (bool): Whether to retrieve outgoing relationship. If False, retrieves incoming relationship.
 
         Returns:
-            List[Relationship]: A list of relationship.
+                List[Relationship]: A list of relationship.
 
         Raises:
-            KeyError: If the specified node is not found in the graph.
+                KeyError: If the specified node is not found in the graph.
         """
         if node is None:
             return list(self.relationships.values())
@@ -285,18 +216,28 @@ class Graph(BaseRelatableNode):
             )
             return relationships
 
+    def get_predecessors(self, node: BaseNode):
+        node_ids = list(self.node_relationships[node.id_]["in"].values())
+        nodes = func_call.lcall(node_ids, lambda i: self.nodes[i])
+        return nodes
+
+    def get_successors(self, node: BaseNode):
+        node_ids = list(self.node_relationships[node.id_]["out"].values())
+        nodes = func_call.lcall(node_ids, lambda i: self.nodes[i])
+        return nodes
+
     def remove_node(self, node: BaseNode) -> BaseNode:
         """
         Removes a node and its associated relationship from the graph.
 
         Args:
-            node (BaseNode): The node to remove.
+                node (BaseNode): The node to remove.
 
         Returns:
-            BaseNode: The removed node.
+                BaseNode: The removed node.
 
         Raises:
-            KeyError: If the node is not found in the graph.
+                KeyError: If the node is not found in the graph.
         """
         if node.id_ not in self.nodes.keys():
             raise KeyError(f"node {node.id_} is not found")
@@ -319,13 +260,13 @@ class Graph(BaseRelatableNode):
         Removes a relationship from the graph.
 
         Args:
-            relationship (Relationship): The relationship to remove.
+                relationship (Relationship): The relationship to remove.
 
         Returns:
-            Relationship: The removed relationship.
+                Relationship: The removed relationship.
 
         Raises:
-            KeyError: If the relationship is not found in the graph.
+                KeyError: If the relationship is not found in the graph.
         """
         if relationship.id_ not in self.relationships.keys():
             raise KeyError(f"relationship {relationship.id_} is not found")
@@ -342,10 +283,10 @@ class Graph(BaseRelatableNode):
         Checks if a node exists in the graph.
 
         Args:
-            node (BaseNode): The node to check.
+                node (BaseNode): The node to check.
 
         Returns:
-            bool: True if the node exists, False otherwise.
+                bool: True if the node exists, False otherwise.
         """
         if node.id_ in self.nodes.keys():
             return True
@@ -357,10 +298,10 @@ class Graph(BaseRelatableNode):
         Checks if a relationship exists in the graph.
 
         Args:
-            relationship (Relationship): The relationship to check.
+                relationship (Relationship): The relationship to check.
 
         Returns:
-            bool: True if the relationship exists, False otherwise.
+                bool: True if the relationship exists, False otherwise.
         """
         if relationship.id_ in self.relationships.keys():
             return True
@@ -372,7 +313,7 @@ class Graph(BaseRelatableNode):
         Determines if the graph is empty.
 
         Returns:
-            bool: True if the graph has no nodes, False otherwise.
+                bool: True if the graph has no nodes, False otherwise.
         """
         if self.nodes:
             return False
@@ -390,14 +331,14 @@ class Graph(BaseRelatableNode):
         Converts the graph to a NetworkX graph object.
 
         Args:
-            **kwargs: Additional keyword arguments to pass to the NetworkX DiGraph constructor.
+                **kwargs: Additional keyword arguments to pass to the NetworkX DiGraph constructor.
 
         Returns:
-            Any: A NetworkX directed graph representing the graph.
+                Any: A NetworkX directed graph representing the graph.
 
         Examples:
-            >>> graph = Graph()
-            >>> nx_graph = graph.to_networkx()
+                >>> graph = Graph()
+                >>> nx_graph = graph.to_networkx()
         """
 
         SysUtil.check_import("networkx")
@@ -408,11 +349,13 @@ class Graph(BaseRelatableNode):
         for node_id, node in self.nodes.items():
             node_info = node.to_dict()
             node_info.pop("node_id")
+            node_info.update({"class_name": node.__class__.__name__})
             g.add_node(node_id, **node_info)
 
         for _, relationship in self.relationships.items():
             relationship_info = relationship.to_dict()
             relationship_info.pop("node_id")
+            relationship_info.update({"class_name": relationship.__class__.__name__})
             source_node_id = relationship_info.pop("source_node_id")
             target_node_id = relationship_info.pop("target_node_id")
             g.add_edge(source_node_id, target_node_id, **relationship_info)
@@ -422,20 +365,37 @@ class Graph(BaseRelatableNode):
 
 class Structure(BaseRelatableNode):
     graph: Graph = Graph()
-    processing_mails: deque = deque()
     pending_ins: dict = {}
     pending_outs: deque = deque()
     execute_stop: bool = False
+    condition_check_result: bool | None = None
 
     def add_node(self, node: BaseNode):
         self.graph.add_node(node)
 
-    # def add_relationship(self, relationship: Relationship):
-    #     self.graph.add_relationship(relationship)
-    def add_relationship(self, from_node: BaseNode, to_node: BaseNode, **kwargs):
+    def add_relationship(
+        self,
+        from_node: BaseNode,
+        to_node: BaseNode,
+        bundle=False,
+        condition=None,
+        **kwargs,
+    ):
+        if isinstance(from_node, Tool) or isinstance(from_node, ActionSelection):
+            raise ValueError(
+                f"type {type(from_node)} should not be the head of the relationship, "
+                f"please switch position and attach it to the tail of the relationship"
+            )
+        if isinstance(to_node, Tool) or isinstance(to_node, ActionSelection):
+            bundle = True
         relationship = Relationship(
-            source_node_id=from_node.id_, target_node_id=to_node.id_, **kwargs
+            source_node_id=from_node.id_,
+            target_node_id=to_node.id_,
+            bundle=bundle,
+            **kwargs,
         )
+        if condition:
+            relationship.add_condition(condition)
         self.graph.add_relationship(relationship)
 
     def get_relationships(self) -> list[Relationship]:
@@ -453,6 +413,12 @@ class Structure(BaseRelatableNode):
             relationships = result
         return relationships
 
+    def get_predecessors(self, node: BaseNode):
+        return self.graph.get_predecessors(node)
+
+    def get_successors(self, node: BaseNode):
+        return self.graph.get_successors(node)
+
     def node_exist(self, node: BaseNode) -> bool:
         return self.graph.node_exist(node)
 
@@ -469,24 +435,12 @@ class Structure(BaseRelatableNode):
         return self.graph.is_empty()
 
     def get_heads(self):
-        heads = deque()
+        heads = []
         for key in self.graph.node_relationships:
             if not self.graph.node_relationships[key]["in"]:
                 heads.append(self.graph.nodes[key])
         return heads
 
-    # def get_next_step(self, current_node: BaseNode):
-    #     next_nodes = deque()
-    #     next_relationships = self.get_node_relationships(current_node)
-    #     for relationship in next_relationships:
-    #         node = self.graph.nodes[relationship.target_node_id]
-    #         next_nodes.append(node)
-    #         further_relationships = self.get_node_relationships(node)
-    #         for f_relationship in further_relationships:
-    #             if f_relationship.bundle:
-    #                 next_nodes.append(self.graph.nodes[f_relationship.target_node_id])
-    #     return next_nodes
-
     @staticmethod
     def parse_to_action(instruction: BaseNode, bundled_nodes: deque):
         action_node = ActionNode(instruction)
@@ -501,12 +455,36 @@ class Structure(BaseRelatableNode):
                 raise ValueError("Invalid bundles nodes")
         return action_node
 
-    def get_next_step(self, current_node: BaseNode):
-        next_nodes = deque()
+    async def check_condition(self, relationship, executable_id):
+        if relationship.condition.source_type == "structure":
+            return self.check_condition_structure(relationship)
+        elif relationship.condition.source_type == "executable":
+            self.send(
+                recipient_id=executable_id, category="condition", package=relationship
+            )
+            while self.condition_check_result is None:
+                await AsyncUtil.sleep(0.1)
+                self.process_relationship_condition(relationship.id_)
+                continue
+            check_result = self.condition_check_result
+            self.condition_check_result = None
+            return check_result
+        else:
+            raise ValueError(f"Invalid source_type.")
+
+    def check_condition_structure(self, relationship):
+        return relationship.condition(self)
+
+    async def get_next_step(self, current_node: BaseNode, executable_id):
+        next_nodes = []
         next_relationships = self.get_node_relationships(current_node)
         for relationship in next_relationships:
             if relationship.bundle:
                 continue
+            if relationship.condition:
+                check = await self.check_condition(relationship, executable_id)
+                if not check:
+                    continue
             node = self.graph.nodes[relationship.target_node_id]
             further_relationships = self.get_node_relationships(node)
             bundled_nodes = deque()
@@ -559,7 +537,21 @@ class Structure(BaseRelatableNode):
         )
         self.pending_outs.append(mail)
 
-    def process(self) -> None:
+    def process_relationship_condition(self, relationship_id):
+        for key in list(self.pending_ins.keys()):
+            skipped_requests = deque()
+            while self.pending_ins[key]:
+                mail = self.pending_ins[key].popleft()
+                if (
+                    mail.category == "condition"
+                    and mail.package["relationship_id"] == relationship_id
+                ):
+                    self.condition_check_result = mail.package["check_result"]
+                else:
+                    skipped_requests.append(mail)
+            self.pending_ins[key] = skipped_requests
+
+    async def process(self) -> None:
         for key in list(self.pending_ins.keys()):
             while self.pending_ins[key]:
                 mail = self.pending_ins[key].popleft()
@@ -573,13 +565,15 @@ class Structure(BaseRelatableNode):
                         raise ValueError(
                             f"Node {mail.package} does not exist in the structure {self.id_}"
                         )
-                    next_nodes = self.get_next_step(self.graph.nodes[mail.package])
+                    next_nodes = await self.get_next_step(
+                        self.graph.nodes[mail.package], mail.sender_id
+                    )
                 elif mail.category == "node" and isinstance(mail.package, BaseNode):
                     if not self.node_exist(mail.package):
                         raise ValueError(
                             f"Node {mail.package} does not exist in the structure {self.id_}"
                         )
-                    next_nodes = self.get_next_step(mail.package)
+                    next_nodes = await self.get_next_step(mail.package, mail.sender_id)
                 else:
                     raise ValueError(f"Invalid mail type for structure")
 
@@ -587,16 +581,24 @@ class Structure(BaseRelatableNode):
                     self.send(
                         recipient_id=mail.sender_id, category="end", package="end"
                     )
-                while next_nodes:
-                    package = next_nodes.popleft()
-                    self.send(
-                        recipient_id=mail.sender_id, category="node", package=package
-                    )
+                else:
+                    if len(next_nodes) == 1:
+                        self.send(
+                            recipient_id=mail.sender_id,
+                            category="node",
+                            package=next_nodes[0],
+                        )
+                    else:
+                        self.send(
+                            recipient_id=mail.sender_id,
+                            category="node_list",
+                            package=next_nodes,
+                        )
 
     async def execute(self, refresh_time=1):
         if not self.acyclic():
             raise ValueError("Structure is not acyclic")
 
         while not self.execute_stop:
-            self.process()
+            await self.process()
             await AsyncUtil.sleep(refresh_time)

lionagi/core/session/__init__.py

@@ -1,5 +1,3 @@
 from .session import Session
-from lionagi.core.branch.branch import Branch
 
-
-__all__ = ["Session", "Branch"]
+__all__ = ["Session"]