lionagi 0.16.3__py3-none-any.whl → 0.17.1__py3-none-any.whl

lionagi/__init__.py

@@ -6,9 +6,8 @@ import logging
 
 from pydantic import BaseModel, Field
 
-from . import _types as types
+# Eager imports for commonly used components
 from . import ln as ln
-from .operations.builder import OperationGraphBuilder as Builder
 from .operations.node import Operation
 from .service.imodel import iModel
 from .session.session import Branch, Session
@@ -17,6 +16,29 @@ from .version import __version__
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
+# Module-level lazy loading cache
+_lazy_imports = {}
+
+
+def __getattr__(name: str):
+    """Lazy loading for expensive imports."""
+    if name in _lazy_imports:
+        return _lazy_imports[name]
+
+    if name == "types":
+        from . import _types as types
+
+        _lazy_imports["types"] = types
+        return types
+    elif name == "Builder":
+        from .operations.builder import OperationGraphBuilder as Builder
+
+        _lazy_imports["Builder"] = Builder
+        return Builder
+
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
 __all__ = (
     "Session",
     "Branch",

lionagi/_types.py

@@ -1,3 +1,47 @@
-from .fields import *
-from .models import *
-from .protocols.types import *
+# Lazy loading for heavy type imports to improve startup performance
+_lazy_type_imports = {}
+
+
+def __getattr__(name: str):
+    """Lazy loading for type definitions."""
+    if name in _lazy_type_imports:
+        return _lazy_type_imports[name]
+
+    # Import from fields
+    try:
+        from .fields import __all__ as fields_all
+
+        if name in fields_all:
+            from . import fields
+
+            attr = getattr(fields, name)
+            _lazy_type_imports[name] = attr
+            return attr
+    except (ImportError, AttributeError):
+        pass
+
+    # Import from models
+    try:
+        from .models import __all__ as models_all
+
+        if name in models_all:
+            from . import models
+
+            attr = getattr(models, name)
+            _lazy_type_imports[name] = attr
+            return attr
+    except (ImportError, AttributeError):
+        pass
+
+    # Import from protocols.types
+    try:
+        from .protocols import types as protocol_types
+
+        if hasattr(protocol_types, name):
+            attr = getattr(protocol_types, name)
+            _lazy_type_imports[name] = attr
+            return attr
+    except (ImportError, AttributeError):
+        pass
+
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

lionagi/adapters/_utils.py

@@ -1,12 +1,13 @@
 def check_async_postgres_available():
-    try:
-        import sqlalchemy as sa
-        from pydapter.extras.async_postgres_ import AsyncPostgresAdapter
-        from sqlalchemy.ext.asyncio import create_async_engine
+    from lionagi.utils import is_import_installed
 
+    all_import_present = 0
+    for pkg in ("sqlalchemy", "asyncpg"):
+        if is_import_installed(pkg):
+            all_import_present += 1
+    if all_import_present == 2:
         return True
-    except Exception:
-        return ImportError(
-            "This adapter requires postgres option to be installed. "
-            'Please install them using `uv pip install "lionagi[postgres]"`.'
-        )
+    return ImportError(
+        "This adapter requires postgres option to be installed. "
+        'Please install them using `uv pip install "lionagi[postgres]"`.'
+    )

lionagi/adapters/async_postgres_adapter.py

@@ -13,86 +13,90 @@ from __future__ import annotations
 
 from typing import ClassVar, TypeVar
 
-import sqlalchemy as sa
-from pydapter.extras.async_postgres_ import AsyncPostgresAdapter
-from sqlalchemy.ext.asyncio import create_async_engine
-
-from ._utils import check_async_postgres_available
-
-_ASYNC_POSTGRES_AVAILABLE = check_async_postgres_available()
-
-if isinstance(_ASYNC_POSTGRES_AVAILABLE, ImportError):
-    raise _ASYNC_POSTGRES_AVAILABLE
+from pydapter import AsyncAdapter
 
 T = TypeVar("T")
 
 
-class LionAGIAsyncPostgresAdapter(AsyncPostgresAdapter[T]):
-    """
-    Streamlined async adapter for lionagi Nodes.
-
-    Features:
-    - Auto-creates tables with lionagi schema
-    - Inherits all pydapter v1.0.4+ improvements
-    - No workarounds needed for SQLite or raw SQL
-    """
-
-    obj_key: ClassVar[str] = "lionagi_async_pg"
-
-    @classmethod
-    async def to_obj(
-        cls,
-        subj,
-        /,
-        *,
-        many: bool = True,
-        adapt_meth: str = None,
-        **kw,
-    ):
-        """Write lionagi Node(s) to database with auto-table creation."""
-        # Auto-create table if needed
-        if table := kw.get("table"):
-            if engine_url := (kw.get("dsn") or kw.get("engine_url")):
-                await cls._ensure_table(engine_url, table)
-            elif engine := kw.get("engine"):
-                await cls._ensure_table(engine, table)
-
-        return await super().to_obj(
-            subj, many=many, adapt_meth=adapt_meth, **kw
-        )
-
-    @classmethod
-    async def _ensure_table(cls, engine_or_url, table_name: str):
-        """Create table with lionagi schema if it doesn't exist."""
-        should_dispose = False
-        if isinstance(engine_or_url, str):
-            engine = create_async_engine(engine_or_url, future=True)
-            should_dispose = True
-        else:
-            engine = engine_or_url
-
-        try:
-            async with engine.begin() as conn:
-                # Determine JSON type based on database
-                engine_url = str(engine.url)
-                json_type = (
-                    sa.dialects.postgresql.JSONB
-                    if "postgresql" in engine_url
-                    else sa.JSON
-                )
-
-                # Create table with lionagi schema
-                await conn.run_sync(
-                    lambda sync_conn: sa.Table(
-                        table_name,
-                        sa.MetaData(),
-                        sa.Column("id", sa.String, primary_key=True),
-                        sa.Column("content", json_type),
-                        sa.Column("node_metadata", json_type),
-                        sa.Column("created_at", sa.Float),
-                        sa.Column("embedding", json_type, nullable=True),
-                    ).create(sync_conn, checkfirst=True)
-                )
-        finally:
-            if should_dispose:
-                await engine.dispose()
+def create_lionagi_async_postgres_adapter() -> type[AsyncAdapter]:
+    from pydapter.extras.async_postgres_ import AsyncPostgresAdapter
+
+    class LionAGIAsyncPostgresAdapter(AsyncPostgresAdapter[T]):
+        """
+        Streamlined async adapter for lionagi Nodes.
+
+        Features:
+        - Auto-creates tables with lionagi schema
+        - Inherits all pydapter v1.0.4+ improvements
+        - No workarounds needed for SQLite or raw SQL
+        """
+
+        obj_key: ClassVar[str] = "lionagi_async_pg"
+
+        @classmethod
+        async def to_obj(
+            cls,
+            subj,
+            /,
+            *,
+            many: bool = True,
+            adapt_meth: str = None,
+            **kw,
+        ):
+            """Write lionagi Node(s) to database with auto-table creation."""
+            # Auto-create table if needed
+            if table := kw.get("table"):
+                if engine_url := (kw.get("dsn") or kw.get("engine_url")):
+                    await cls._ensure_table(engine_url, table)
+                elif engine := kw.get("engine"):
+                    await cls._ensure_table(engine, table)
+
+            return await super().to_obj(
+                subj, many=many, adapt_meth=adapt_meth, **kw
+            )
+
+        @classmethod
+        async def _ensure_table(cls, engine_or_url, table_name: str):
+            """Create table with lionagi schema if it doesn't exist."""
+            import sqlalchemy as sa
+            from sqlalchemy.ext.asyncio import create_async_engine
+
+            should_dispose = False
+            if isinstance(engine_or_url, str):
+                engine = create_async_engine(engine_or_url, future=True)
+                should_dispose = True
+            else:
+                engine = engine_or_url
+
+            try:
+                async with engine.begin() as conn:
+                    # Determine JSON type based on database
+                    engine_url = str(engine.url)
+                    json_type = (
+                        sa.dialects.postgresql.JSONB
+                        if "postgresql" in engine_url
+                        else sa.JSON
+                    )
+
+                    # Create table with lionagi schema
+                    await conn.run_sync(
+                        lambda sync_conn: sa.Table(
+                            table_name,
+                            sa.MetaData(),
+                            sa.Column("id", sa.String, primary_key=True),
+                            sa.Column("content", json_type),
+                            sa.Column("node_metadata", json_type),
+                            sa.Column("created_at", sa.Float),
+                            sa.Column("embedding", json_type, nullable=True),
+                        ).create(sync_conn, checkfirst=True)
+                    )
+            finally:
+                if should_dispose:
+                    await engine.dispose()
+
+    return LionAGIAsyncPostgresAdapter
+
+
+LionAGIAsyncPostgresAdapter = create_lionagi_async_postgres_adapter()
+
+__all__ = ("LionAGIAsyncPostgresAdapter",)

lionagi/ln/__init__.py

@@ -1,8 +1,6 @@
 from ._async_call import alcall, bcall
 from ._hash import hash_dict
 from ._json_dump import (
-    DEFAULT_SERIALIZER,
-    DEFAULT_SERIALIZER_OPTION,
     get_orjson_default,
     json_dumpb,
     json_dumps,
@@ -48,8 +46,6 @@ __all__ = (
     "alcall",
     "bcall",
     "hash_dict",
-    "DEFAULT_SERIALIZER",
-    "DEFAULT_SERIALIZER_OPTION",
     "get_orjson_default",
     "json_dumps",
     "make_options",

lionagi/ln/_json_dump.py

@@ -15,8 +15,6 @@ import orjson
 
 __all__ = [
     "get_orjson_default",
-    "DEFAULT_SERIALIZER",
-    "DEFAULT_SERIALIZER_OPTION",
     "make_options",
     "json_dumpb",
     "json_dumps",
@@ -193,10 +191,6 @@ def _cached_default(
 
 # --------- defaults & options -------------------------------------------------
 
-# Compact, no newline, no sorting: neutral default for most use-cases.
-DEFAULT_SERIALIZER_OPTION = 0
-DEFAULT_SERIALIZER = get_orjson_default()
-
 
 def make_options(
     *,

lionagi/operations/__init__.py

@@ -2,11 +2,9 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from .brainstorm.brainstorm import BrainstormOperation, brainstorm
 from .builder import ExpansionStrategy, OperationGraphBuilder
 from .flow import flow
 from .node import BranchOperations, Operation
-from .plan.plan import PlanOperation, plan
 
 Builder = OperationGraphBuilder
 
@@ -17,9 +15,5 @@ __all__ = (
     "flow",
     "BranchOperations",
     "Operation",
-    "plan",
-    "PlanOperation",
-    "brainstorm",
-    "BrainstormOperation",
     "Builder",
 )

lionagi/operations/_visualize_graph.py

@@ -0,0 +1,285 @@
+def visualize_graph(
+    builder,
+    title: str = "Operation Graph",
+    figsize=(14, 10),
+):
+    """Visualization with improved layout for complex graphs."""
+    from lionagi.utils import is_import_installed
+
+    if not is_import_installed("matplotlib"):
+        raise ImportError(
+            "matplotlib is required for visualization. "
+            "Please install it using `pip install matplotlib`."
+        )
+    if not is_import_installed("networkx"):
+        raise ImportError(
+            "networkx is required for visualization. "
+            "Please install it using `pip install networkx`."
+        )
+
+    import matplotlib.pyplot as plt
+    import networkx as nx
+    import numpy as np
+
+    graph = builder.get_graph()
+
+    # Convert to networkx
+    G = nx.DiGraph()
+
+    # Track node positions for hierarchical layout
+    node_levels = {}
+    node_labels = {}
+    node_colors = []
+    node_sizes = []
+
+    # First pass: add nodes and determine levels
+    for node in graph.internal_nodes.values():
+        node_id = str(node.id)[:8]
+        G.add_node(node_id)
+
+        # Determine level based on dependencies
+        in_edges = [
+            e
+            for e in graph.internal_edges.values()
+            if str(e.tail)[:8] == node_id
+        ]
+        if not in_edges:
+            level = 0  # Root nodes
+        else:
+            # Get max level of predecessors + 1
+            pred_levels = []
+            for edge in in_edges:
+                pred_id = str(edge.head)[:8]
+                if pred_id in node_levels:
+                    pred_levels.append(node_levels[pred_id])
+            level = max(pred_levels, default=0) + 1
+
+        node_levels[node_id] = level
+
+        # Create label
+        ref_id = node.metadata.get("reference_id", "")
+        if ref_id:
+            label = f"{node.operation}\n[{ref_id}]"
+        else:
+            label = f"{node.operation}\n{node_id}"
+        node_labels[node_id] = label
+
+        # Color and size based on status and type
+        if node.id in builder._executed:
+            node_colors.append("#90EE90")  # Light green
+            node_sizes.append(4000)
+        elif node.metadata.get("expansion_source"):
+            node_colors.append("#87CEEB")  # Sky blue
+            node_sizes.append(3500)
+        elif node.metadata.get("aggregation"):
+            node_colors.append("#FFD700")  # Gold
+            node_sizes.append(4500)
+        elif node.metadata.get("is_condition_check"):
+            node_colors.append("#DDA0DD")  # Plum
+            node_sizes.append(3500)
+        else:
+            node_colors.append("#E0E0E0")  # Light gray
+            node_sizes.append(3000)
+
+    # Add edges
+    edge_colors = []
+    edge_styles = []
+    edge_widths = []
+    edge_labels = {}
+
+    for edge in graph.internal_edges.values():
+        head_id = str(edge.head)[:8]
+        tail_id = str(edge.tail)[:8]
+        G.add_edge(head_id, tail_id)
+
+        # Style edges based on type
+        edge_label = edge.label[0] if edge.label else ""
+        edge_labels[(head_id, tail_id)] = edge_label
+
+        if "expansion" in edge_label:
+            edge_colors.append("#4169E1")  # Royal blue
+            edge_styles.append("dashed")
+            edge_widths.append(2)
+        elif "aggregate" in edge_label:
+            edge_colors.append("#FF6347")  # Tomato
+            edge_styles.append("dotted")
+            edge_widths.append(2.5)
+        else:
+            edge_colors.append("#808080")  # Gray
+            edge_styles.append("solid")
+            edge_widths.append(1.5)
+
+    # Create improved hierarchical layout
+    pos = {}
+    nodes_by_level = {}
+
+    for node_id, level in node_levels.items():
+        if level not in nodes_by_level:
+            nodes_by_level[level] = []
+        nodes_by_level[level].append(node_id)
+
+    # Position nodes with better spacing algorithm
+    y_spacing = 2.5
+    max_width = 16  # Maximum horizontal spread
+
+    for level, nodes in nodes_by_level.items():
+        num_nodes = len(nodes)
+
+        if num_nodes <= 6:
+            # Normal spacing for small levels
+            x_spacing = 2.5
+            x_offset = -(num_nodes - 1) * x_spacing / 2
+            for i, node_id in enumerate(nodes):
+                pos[node_id] = (x_offset + i * x_spacing, -level * y_spacing)
+        else:
+            # Multi-row layout for large levels
+            nodes_per_row = min(6, int(np.ceil(np.sqrt(num_nodes * 1.5))))
+            rows = int(np.ceil(num_nodes / nodes_per_row))
+
+            for i, node_id in enumerate(nodes):
+                row = i // nodes_per_row
+                col = i % nodes_per_row
+
+                # Calculate row width
+                nodes_in_row = min(
+                    nodes_per_row, num_nodes - row * nodes_per_row
+                )
+                x_spacing = 2.5
+                x_offset = -(nodes_in_row - 1) * x_spacing / 2
+
+                # Add slight y offset for different rows
+                y_offset = row * 0.8
+
+                pos[node_id] = (
+                    x_offset + col * x_spacing,
+                    -level * y_spacing - y_offset,
+                )
+
+    # Create figure
+    plt.figure(figsize=figsize)
+
+    # Draw nodes
+    nx.draw_networkx_nodes(
+        G,
+        pos,
+        node_color=node_colors,
+        node_size=node_sizes,
+        alpha=0.9,
+        linewidths=2,
+        edgecolors="black",
+    )
+
+    # Draw edges with different styles - use curved edges for better visibility
+    for i, (u, v) in enumerate(G.edges()):
+        # Calculate curve based on node positions
+        u_pos = pos[u]
+        v_pos = pos[v]
+
+        # Determine connection style based on relative positions
+        if abs(u_pos[0] - v_pos[0]) > 5:  # Far apart horizontally
+            connectionstyle = "arc3,rad=0.2"
+        else:
+            connectionstyle = "arc3,rad=0.1"
+
+        nx.draw_networkx_edges(
+            G,
+            pos,
+            [(u, v)],
+            edge_color=[edge_colors[i]],
+            style=edge_styles[i],
+            width=edge_widths[i],
+            alpha=0.7,
+            arrows=True,
+            arrowsize=20,
+            arrowstyle="-|>",
+            connectionstyle=connectionstyle,
+        )
+
+    # Draw labels
+    nx.draw_networkx_labels(
+        G,
+        pos,
+        node_labels,
+        font_size=9,
+        font_weight="bold",
+        font_family="monospace",
+    )
+
+    # Draw edge labels (only for smaller graphs)
+    if len(G.edges()) < 20:
+        nx.draw_networkx_edge_labels(
+            G,
+            pos,
+            edge_labels,
+            font_size=7,
+            font_color="darkblue",
+            bbox=dict(
+                boxstyle="round,pad=0.3",
+                facecolor="white",
+                edgecolor="none",
+                alpha=0.7,
+            ),
+        )
+
+    plt.title(title, fontsize=18, fontweight="bold", pad=20)
+    plt.axis("off")
+
+    # Enhanced legend
+    from matplotlib.lines import Line2D
+    from matplotlib.patches import Patch, Rectangle
+
+    legend_elements = [
+        Patch(facecolor="#90EE90", edgecolor="black", label="Executed"),
+        Patch(facecolor="#87CEEB", edgecolor="black", label="Expanded"),
+        Patch(facecolor="#FFD700", edgecolor="black", label="Aggregation"),
+        Patch(facecolor="#DDA0DD", edgecolor="black", label="Condition"),
+        Patch(facecolor="#E0E0E0", edgecolor="black", label="Pending"),
+        Line2D([0], [0], color="#808080", linewidth=2, label="Sequential"),
+        Line2D(
+            [0],
+            [0],
+            color="#4169E1",
+            linewidth=2,
+            linestyle="dashed",
+            label="Expansion",
+        ),
+        Line2D(
+            [0],
+            [0],
+            color="#FF6347",
+            linewidth=2,
+            linestyle="dotted",
+            label="Aggregate",
+        ),
+    ]
+
+    plt.legend(
+        handles=legend_elements,
+        loc="upper left",
+        bbox_to_anchor=(0, 1),
+        frameon=True,
+        fancybox=True,
+        shadow=True,
+        ncol=2,
+    )
+
+    # Add statistics box
+    stats_text = f"Nodes: {len(G.nodes())}\nEdges: {len(G.edges())}\nExecuted: {len(builder._executed)}"
+    if nodes_by_level:
+        max_level = max(nodes_by_level.keys())
+        stats_text += f"\nLevels: {max_level + 1}"
+
+    plt.text(
+        0.98,
+        0.02,
+        stats_text,
+        transform=plt.gca().transAxes,
+        bbox=dict(boxstyle="round,pad=0.5", facecolor="lightgray", alpha=0.8),
+        verticalalignment="bottom",
+        horizontalalignment="right",
+        fontsize=10,
+        fontfamily="monospace",
+    )
+
+    plt.tight_layout()
+    plt.show()