lionagi 0.0.115__py3-none-any.whl → 0.0.204__py3-none-any.whl

lionagi/schema/data_logger.py

@@ -1,6 +1,7 @@
 from collections import deque
-from typing import Optional
-from lionagi.utils import create_path, to_csv
+from typing import Dict, Any
+from ..utils.sys_util import get_timestamp, create_path, as_dict
+from ..utils.io_util import IOUtil
 
 
 class DataLogger:
@@ -38,53 +39,88 @@ class DataLogger:
         self.dir = dir
         self.log = deque(log) if log else deque()
 
-    def __call__(self, entry):
+    def add_entry(self, entry: Dict[str, Any], level: str = "INFO") -> None:
         """
-        Adds a new entry to the log.
+        Adds a new entry to the log with a timestamp and a log level.
 
-        Parameters:
-            entry: The data entry to be added to the log.
-        """        
-        self.log.append(entry)
-
-    def to_csv(self, filename: str, dir: Optional[str] = None, verbose: bool = True, 
-               timestamp: bool = True, dir_exist_ok: bool = True, file_exist_ok: bool = False) -> None:
+        Args:
+            entry (Dict[str, Any]): The data entry to be added to the log.
+            level (str): The log level for the entry (e.g., "INFO", "ERROR"). Defaults to "INFO".
+        """
+        self.log.append({
+            "timestamp": get_timestamp(), "level": level, **as_dict(entry)
+        })
+        
+    def set_dir(self, dir: str) -> None:
         """
-        Exports the logged data to a CSV file and optionally clears the log.
+        Sets the default directory for saving CSV files.
 
         Parameters:
-            filename (str): The name of the CSV file.
+            dir (str): The directory to be set as the default for saving files.
+        """
+        self.dir = dir
 
-            dir (Optional[str]): The directory to save the file. Defaults to the instance's dir attribute.
+    def to_csv(
+        self, filename: str, 
+        file_exist_ok: bool = False,  
+        timestamp = True,
+        time_prefix: bool = False,
+        verbose: bool = True,
+        clear = True
+    ) -> None:
+        """
+        Exports the logged data to a CSV file, using the provided utilities for path creation and timestamping.
 
+        Args:
+            filename (str): The name of the CSV file.
+            file_exist_ok (bool): If True, creates the directory for the file if it does not exist. Defaults to False.
             verbose (bool): If True, prints a message upon completion. Defaults to True.
-
-            timestamp (bool): If True, appends a timestamp to the filename. Defaults to True.
-
-            dir_exist_ok (bool): If True, will not raise an error if the directory already exists. Defaults to True.
-
-            file_exist_ok (bool): If True, overwrites the file if it exists. Defaults to False.
-
-        Side Effects:
-            Clears the log after saving the CSV file.
-
-            Prints a message indicating the save location and number of logs saved if verbose is True.
-        """        
-        dir = dir or self.dir
+            time_prefix (bool): If True, adds the timestamp as a prefix to the filename. Defaults to False.
+        """
+        if not filename.endswith('.csv'):
+            filename += '.csv'
+        
         filepath = create_path(
-            dir=dir, filename=filename, timestamp=timestamp, dir_exist_ok=dir_exist_ok)
-        to_csv(list(self.log), filepath, file_exist_ok=file_exist_ok)
-        n_logs = len(list(self.log))
-        self.log = deque()
+            self.dir, filename, timestamp=timestamp, 
+            dir_exist_ok=file_exist_ok, time_prefix=time_prefix
+        )
+        IOUtil.to_csv(list(self.log), filepath)
+
         if verbose:
-            print(f"{n_logs} logs saved to {filepath}")
-            
-    def set_dir(self, dir: str) -> None:
+            print(f"{len(self.log)} logs saved to {filepath}")
+
+        if clear:
+            self.log.clear()
+
+    def to_jsonl(
+        self, filename: str, 
+        timestamp = False, 
+        time_prefix=False,
+        file_exist_ok: bool = False, 
+        verbose: bool = True, 
+        clear = True
+    ) -> None:
         """
-        Sets the default directory for saving CSV files.
+        Exports the logged data to a JSONL file and optionally clears the log.
 
         Parameters:
-            dir (str): The directory to be set as the default for saving files.
+            filename (str): The name of the JSONL file.
+            file_exist_ok (bool): If True, creates the directory for the file if it does not exist. Defaults to False.
+            verbose (bool): If True, prints a message upon completion. Defaults to True.
         """
-        self.dir = dir
-        
+        if not filename.endswith('.jsonl'):
+            filename += '.jsonl'
+
+        filepath = create_path(
+            self.dir, filename, timestamp=timestamp, 
+            dir_exist_ok=file_exist_ok, time_prefix=time_prefix
+        )
+        
+        for entry in self.log:
+            IOUtil.append_to_jsonl(entry, filepath)
+
+        if verbose:
+            print(f"{len(self.log)} logs saved to {filepath}")
+
+        if clear:
+            self.log.clear()

lionagi/schema/data_node.py

@@ -4,23 +4,74 @@ from typing import Any
 
 class DataNode(BaseNode):
 
-    def to_llama_index(self, **kwargs):
-        # to llama_index textnode
+    def to_llama_index(self, **kwargs) -> Any:
+        """
+        Converts node to llama index format.
+
+        Args:
+            **kwargs: Variable length argument list.
+
+        Returns:
+            The llama index representation of the node.
+
+        Examples:
+            node = DataNode()
+            llama_index = node.to_llama_index()
+        """
         from lionagi.bridge.llama_index import to_llama_index_textnode
         return to_llama_index_textnode(self, **kwargs)
 
-    def to_langchain(self, **kwargs):
-        # to langchain document
+    def to_langchain(self, **kwargs) -> Any:
+        """
+        Converts node to langchain document format.
+
+        Args:
+            **kwargs: Variable length argument list.
+
+        Returns:
+            The langchain document representation of the node.
+
+        Examples:
+            node = DataNode()
+            langchain_doc = node.to_langchain()
+        """
         from lionagi.bridge.langchain import to_langchain_document
         return to_langchain_document(self, **kwargs)
 
     @classmethod
-    def from_llama_index(cls, llama_node: Any, **kwargs):
+    def from_llama_index(cls, llama_node: Any, **kwargs) -> "DataNode":
+        """
+        Creates a DataNode instance from a llama index node.
+
+        Args:
+            llama_node: The llama index node object.
+            **kwargs: Variable length argument list.
+
+        Returns:
+            An instance of DataNode.
+
+        Examples:
+            llama_node = SomeLlamaIndexNode()
+            data_node = DataNode.from_llama_index(llama_node)
+        """
         llama_dict = llama_node.to_dict(**kwargs)
         return cls.from_dict(llama_dict)
 
     @classmethod
-    def from_langchain(cls, lc_doc: Any):
+    def from_langchain(cls, lc_doc: Any) -> "DataNode":
+        """
+        Creates a DataNode instance from a langchain document.
+
+        Args:
+            lc_doc: The langchain document object.
+
+        Returns:
+            An instance of DataNode.
+
+        Examples:
+            lc_doc = SomeLangChainDocument()
+            data_node = DataNode.from_langchain(lc_doc)
+        """
         info_json = lc_doc.to_json()
         info_node = {'lc_id': info_json['id']}
         info_node = {**info_node, **info_json['kwargs']}

lionagi/structures/graph.py

@@ -1,20 +1,56 @@
+from typing import List, Any
 from pydantic import Field
 
-from lionagi.schema.base_node import BaseNode
+from ..utils.call_util import lcall
+from ..schema.base_node import BaseNode
 from .relationship import Relationship
-from lionagi.utils.call_util import lcall
 
 
 class Graph(BaseNode):
+    """
+    Represents a graph structure, consisting of nodes and their relationships.
+    
+    Attributes:
+        nodes (Dict[str, BaseNode]): A dictionary of nodes in the graph.
+        relationships (Dict[str, Relationship]): A dictionary of relationships between nodes in the graph.
+        node_relationships (Dict[str, Dict[str, Dict[str, str]]]): A dictionary tracking the relationships 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
+    """
     nodes: dict = Field(default={})
     relationships: dict = Field(default={})
     node_relationships: dict = Field(default={})
 
-    def add_node(self, node: BaseNode):
+    def add_node(self, node: BaseNode) -> None:
+        """
+        Adds a node to the graph.
+
+        Args:
+            node (BaseNode): The node to add to the graph.
+        """
+
         self.nodes[node.id_] = node
         self.node_relationships[node.id_] = {'in': {}, 'out': {}}
 
-    def add_relationship(self, relationships: Relationship):
+    def add_relationship(self, relationships: Relationship) -> None:
+        """
+        Adds a relationship between nodes in the graph.
+
+        Args:
+            relationship (Relationship): The relationship to add.
+
+        Raises:
+            KeyError: If either the source or target node of the relationship is not found in the graph.
+        """
         if relationships.source_node_id not in self.node_relationships.keys():
             raise KeyError(f'node {relationships.source_node_id} is not found.')
         if relationships.target_node_id not in self.node_relationships.keys():
@@ -24,7 +60,20 @@ class Graph(BaseNode):
         self.node_relationships[relationships.source_node_id]['out'][relationships.id_] = relationships.target_node_id
         self.node_relationships[relationships.target_node_id]['in'][relationships.id_] = relationships.source_node_id
 
-    def get_node_relationships(self, node: BaseNode = None, out_edge=True):
+    def get_node_relationships(self, node: BaseNode = None, out_edge=True) -> List[Relationship]:
+        """
+        Retrieves relationships of a specific node or all relationships in the graph.
+
+        Args:
+            node (Optional[BaseNode]): The node whose relationships to retrieve. If None, retrieves all relationships.
+            out_edge (bool): Whether to retrieve outgoing relationships. If False, retrieves incoming relationships.
+
+        Returns:
+            List[Relationship]: A list of relationships.
+
+        Raises:
+            KeyError: If the specified node is not found in the graph.
+        """
         if node is None:
             return list(self.relationships.values())
 
@@ -40,7 +89,19 @@ class Graph(BaseNode):
             relationships = lcall(relationship_ids, lambda i: self.relationships[i])
             return relationships
 
-    def remove_node(self, node: BaseNode):
+    def remove_node(self, node: BaseNode)-> BaseNode:
+        """
+        Removes a node and its associated relationships from the graph.
+
+        Args:
+            node (BaseNode): The node to remove.
+
+        Returns:
+            BaseNode: The removed node.
+
+        Raises:
+            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')
 
@@ -57,7 +118,19 @@ class Graph(BaseNode):
         self.node_relationships.pop(node.id_)
         return self.nodes.pop(node.id_)
 
-    def remove_relationship(self, relationship: Relationship):
+    def remove_relationship(self, relationship: Relationship) -> Relationship:
+        """
+        Removes a relationship from the graph.
+
+        Args:
+            relationship (Relationship): The relationship to remove.
+
+        Returns:
+            Relationship: The removed relationship.
+
+        Raises:
+            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')
 
@@ -66,19 +139,68 @@ class Graph(BaseNode):
 
         return self.relationships.pop(relationship.id_)
 
-    def node_exists(self, node: BaseNode):
+    def node_exists(self, node: BaseNode) -> bool:
+        """
+        Checks if a node exists in the graph.
+
+        Args:
+            node (BaseNode): The node to check.
+
+        Returns:
+            bool: True if the node exists, False otherwise.
+        """
         if node.id_ in self.nodes.keys():
             return True
         else:
             return False
 
-    def relationship_exists(self, relationship: Relationship):
+    def relationship_exists(self, relationship: Relationship) -> bool:
+        """
+        Checks if a relationship exists in the graph.
+
+        Args:
+            relationship (Relationship): The relationship to check.
+
+        Returns:
+            bool: True if the relationship exists, False otherwise.
+        """
         if relationship.id_ in self.relationships.keys():
             return True
         else:
             return False
 
-    def to_networkx(self, **kwargs):
+    def is_empty(self) -> bool:
+        """
+        Determines if the graph is empty.
+
+        Returns:
+            bool: True if the graph has no nodes, False otherwise.
+        """
+        if self.nodes:
+            return False
+        else:
+            return True
+
+    def clear(self)-> None:
+        """Clears the graph of all nodes and relationships."""
+        self.nodes.clear()
+        self.relationships.clear()
+        self.node_relationships.clear()
+
+    def to_networkx(self, **kwargs) -> Any:
+        """
+        Converts the graph to a NetworkX graph object.
+
+        Args:
+            **kwargs: Additional keyword arguments to pass to the NetworkX DiGraph constructor.
+
+        Returns:
+            Any: A NetworkX directed graph representing the graph.
+
+        Examples:
+            >>> graph = Graph()
+            >>> nx_graph = graph.to_networkx()
+        """
         import networkx as nx
         g = nx.DiGraph(**kwargs)
         for node_id, node in self.nodes.items():

lionagi/structures/relationship.py

@@ -5,15 +5,18 @@ from ..schema.base_node import BaseNode
 
 class Relationship(BaseNode):
     """
-    Relationship class represents a relationship between two nodes in a graph.
-
-    Inherits from BaseNode and adds functionality to manage conditions and relationships 
-    between source and target nodes.
+    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.
+
+    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")
     """
 
     source_node_id: str
@@ -24,8 +27,12 @@ class Relationship(BaseNode):
         """
         Adds a condition to the relationship.
 
-        Parameters:
-            condition (Dict[str, Any]): The condition to be added.
+        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)
 
@@ -33,14 +40,19 @@ class Relationship(BaseNode):
         """
         Removes a condition from the relationship.
 
-        Parameters:
-            condition_key (str): The key of the condition to be removed.
+        Args:
+            condition_key: The key of the condition to be removed.
 
         Returns:
-            Any: The value of the removed condition.
+            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')
@@ -50,11 +62,16 @@ class Relationship(BaseNode):
         """
         Checks if a condition exists in the relationship.
 
-        Parameters:
-            condition_key (str): The key of the condition to check.
+        Args:
+            condition_key: The key of the condition to check.
 
         Returns:
-            bool: True if the condition exists, False otherwise.
+            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
@@ -65,14 +82,21 @@ class Relationship(BaseNode):
         """
         Retrieves a specific condition or all conditions of the relationship.
 
-        Parameters:
-            condition_key (Optional[str]): The key of the specific condition. Defaults to None.
+        Args:
+            condition_key: The key of the specific condition. If None, all conditions are returned.
 
         Returns:
-            Any: The requested condition or all conditions if no key is provided.
+            The requested condition or all conditions if no key is provided.
 
         Raises:
             ValueError: If the specified condition key does not exist.
+
+        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
@@ -85,7 +109,7 @@ class Relationship(BaseNode):
         """
         Checks if the source node exists in a given object.
 
-        Parameters:
+        Args:
             obj (Dict[str, Any]): The object to check.
 
         Returns:
@@ -97,7 +121,7 @@ class Relationship(BaseNode):
         """
         Checks if the target node exists in a given object.
 
-        Parameters:
+        Args:
             obj (Dict[str, Any]): The object to check.
 
         Returns:
@@ -109,7 +133,7 @@ class Relationship(BaseNode):
         """
         Validates the existence of both source and target nodes in a given object.
 
-        Parameters:
+        Args:
             obj (Dict[str, Any]): The object to check.
 
         Returns:
@@ -127,10 +151,24 @@ class Relationship(BaseNode):
             raise ValueError(f"Source node {self.target_node_id} does not exist")
 
     def __str__(self) -> str:
-        """Returns a simple string representation of the Relationship."""
+        """
+        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)'
+        """
         return f"Relationship (id_={self.id_}, from={self.source_node_id}, to={self.target_node_id}, label={self.label})"
 
     def __repr__(self) -> str:
-        """Returns a detailed string representation of the Relationship."""
+        """
+        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)'
+        """
         return f"Relationship(id_={self.id_}, from={self.source_node_id}, to={self.target_node_id}, content={self.content}, " \
                f"metadata={self.metadata}, label={self.label})"

lionagi/structures/structure.py

@@ -1,6 +1,6 @@
 from typing import TypeVar
-from .graph import Graph
 from ..schema import BaseNode
+from .graph import Graph
 from .relationship import Relationship
 
 T = TypeVar('T', bound='BaseNode')
@@ -11,7 +11,7 @@ class Structure(BaseNode):
     """
     Represents the structure of a graph consisting of nodes and relationships.
     """
-    graph: Graph
+    graph: Graph = Graph()
 
     def add_node(self, node: T) -> None:
         """
@@ -31,31 +31,20 @@ class Structure(BaseNode):
         """
         self.graph.add_relationship(relationship)
 
-    # type can be dict or list
-    # @staticmethod
-    # def _typed_return(type: Type[Union[Dict, List]],
-    #                   obj: Optional[Dict[str, Any]] = None
-    #                   ) -> Union[Dict[str, Any], List[Any]]:
-    #     """
-    #     Returns the object in the specified type format.
-    #
-    #     Args:
-    #         type (Type[Union[Dict, List]]): The type to return the object as (dict or list).
-    #
-    #         obj (Optional[Dict[str, Any]]): The object to be converted.
-    #
-    #     Returns:
-    #         Union[Dict[str, Any], List[Any]]: The object in the specified type format.
-    #     """
-    #     if type is list:
-    #         return list(obj.values())
-    #     return obj
-    
     def get_relationships(self) -> list[R]:
         return self.graph.get_node_relationships()
 
-    def get_node_relationships(self, node: T, out_edge=True) -> R:
-        return self.graph.get_node_relationships(node, out_edge)
+    def get_node_relationships(self, node: T, out_edge=True, labels=None) -> R:
+        relationships = self.graph.get_node_relationships(node, out_edge)
+        if labels:
+            if not isinstance(labels, list):
+                labels = [labels]
+            result = []
+            for r in relationships:
+                if r.label in labels:
+                    result.append(r)
+            relationships = result
+        return relationships
 
     def node_exist(self, node: T) -> bool:
         """
@@ -99,4 +88,26 @@ class Structure(BaseNode):
             relationship (R): The relationship instance to be removed.
         """
         return self.graph.remove_relationship(relationship)
-        
+
+    def is_empty(self) -> bool:
+        return self.graph.is_empty()
+
+    # type can be dict or list
+    # @staticmethod
+    # def _typed_return(type: Type[Union[Dict, List]],
+    #                   obj: Optional[Dict[str, Any]] = None
+    #                   ) -> Union[Dict[str, Any], List[Any]]:
+    #     """
+    #     Returns the object in the specified type format.
+    #
+    #     Args:
+    #         type (Type[Union[Dict, List]]): The type to return the object as (dict or list).
+    #
+    #         obj (Optional[Dict[str, Any]]): The object to be converted.
+    #
+    #     Returns:
+    #         Union[Dict[str, Any], List[Any]]: The object in the specified type format.
+    #     """
+    #     if type is list:
+    #         return list(obj.values())
+    #     return obj