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

lionagi/operations/brainstorm/brainstorm.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import logging
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 
 from pydantic import BaseModel
 
@@ -12,11 +12,13 @@ from lionagi.fields.instruct import (
     Instruct,
     InstructResponse,
 )
-from lionagi.ln import alcall
+from lionagi.ln import alcall, to_list
 from lionagi.protocols.generic.element import ID
-from lionagi.session.branch import Branch
-from lionagi.session.session import Session
-from lionagi.utils import to_list
+
+if TYPE_CHECKING:
+    from lionagi.session.branch import Branch
+    from lionagi.session.session import Session
+
 
 from ..utils import prepare_instruct, prepare_session
 from .prompt import PROMPT
@@ -58,8 +60,8 @@ def chunked(iterable, n):
 
 async def run_instruct(
     ins: Instruct,
-    session: Session,
-    branch: Branch,
+    session: "Session",
+    branch: "Branch",
     auto_run: bool,
     verbose: bool = True,
     **kwargs: Any,
@@ -113,8 +115,8 @@ async def run_instruct(
 async def brainstorm(
     instruct: Instruct | dict[str, Any],
     num_instruct: int = 2,
-    session: Session | None = None,
-    branch: Branch | ID.Ref | None = None,
+    session: "Session" = None,
+    branch: ID["Branch"].Ref | None = None,
     auto_run: bool = True,
     auto_explore: bool = False,
     explore_kwargs: dict[str, Any] | None = None,
@@ -156,8 +158,8 @@ async def brainstorm(
 async def brainstormStream(
     instruct: Instruct | dict[str, Any],
     num_instruct: int = 2,
-    session: Session | None = None,
-    branch: Branch | ID.Ref | None = None,
+    session: "Session" = None,
+    branch: ID["Branch"].Ref | None = None,
     auto_run: bool = True,
     auto_explore: bool = False,
     explore_kwargs: dict[str, Any] | None = None,
@@ -377,7 +379,7 @@ async def brainstormStream(
                     all_responses = []
 
                     async def explore_concurrent_chunk(
-                        sub_instructs: list[Instruct], base_branch: Branch
+                        sub_instructs: list[Instruct], base_branch: "Branch"
                     ):
                         """
                         Explore instructions in a single chunk concurrently.

lionagi/operations/builder.py

@@ -11,9 +11,8 @@ Build → Execute → Expand → Execute → ...
 from enum import Enum
 from typing import Any
 
-from lionagi.operations.node import BranchOperations, Operation
+from lionagi.operations.node import create_operation
 from lionagi.protocols.graph.edge import Edge
-from lionagi.protocols.graph.graph import Graph
 from lionagi.protocols.types import ID
 
 __all__ = (
@@ -62,18 +61,20 @@ class OperationGraphBuilder:
 
     def __init__(self, name: str = "DynamicGraph"):
         """Initialize the incremental graph builder."""
+        from lionagi.protocols.graph.graph import Graph
+
         self.name = name
         self.graph = Graph()
 
         # Track state
-        self._operations: dict[str, Operation] = {}  # All operations by ID
+        self._operations = {}  # All operations by ID
         self._executed: set[str] = set()  # IDs of executed operations
         self._current_heads: list[str] = []  # Current head nodes for linking
         self.last_operation_id: str | None = None
 
     def add_operation(
         self,
-        operation: BranchOperations,
+        operation: str,
         node_id: str | None = None,
         depends_on: list[str] | None = None,
         inherit_context: bool = False,
@@ -95,7 +96,7 @@ class OperationGraphBuilder:
             ID of the created node
         """
         # Create operation node
-        node = Operation(operation=operation, parameters=parameters)
+        node = create_operation(operation=operation, parameters=parameters)
 
         # Store context inheritance strategy
         if inherit_context and depends_on:
@@ -137,7 +138,7 @@ class OperationGraphBuilder:
         self,
         items: list[Any],
         source_node_id: str,
-        operation: BranchOperations,
+        operation: str,
         strategy: ExpansionStrategy = ExpansionStrategy.CONCURRENT,
         inherit_context: bool = False,
         chain_context: bool = False,
@@ -179,7 +180,7 @@ class OperationGraphBuilder:
             params["expanded_from"] = source_node_id
             params["expansion_strategy"] = strategy.value
 
-            node = Operation(
+            node = create_operation(
                 operation=operation,
                 parameters=params,
                 metadata={
@@ -227,7 +228,7 @@ class OperationGraphBuilder:
 
     def add_aggregation(
         self,
-        operation: BranchOperations,
+        operation: str,
         node_id: str | None = None,
         source_node_ids: list[str] | None = None,
         inherit_context: bool = False,
@@ -262,7 +263,7 @@ class OperationGraphBuilder:
             **parameters,
         }
 
-        node = Operation(
+        node = create_operation(
             operation=operation,
             parameters=agg_params,
             metadata={"aggregation": True},
@@ -308,7 +309,7 @@ class OperationGraphBuilder:
         """
         self._executed.update(node_ids)
 
-    def get_unexecuted_nodes(self) -> list[Operation]:
+    def get_unexecuted_nodes(self):
         """
         Get nodes that haven't been executed yet.
 
@@ -323,9 +324,9 @@ class OperationGraphBuilder:
 
     def add_conditional_branch(
         self,
-        condition_check_op: BranchOperations,
-        true_op: BranchOperations,
-        false_op: BranchOperations | None = None,
+        condition_check_op: str,
+        true_op: str,
+        false_op: str | None = None,
         **check_params,
     ) -> dict[str, str]:
         """
@@ -341,7 +342,7 @@ class OperationGraphBuilder:
             Dict with node IDs: {'check': id, 'true': id, 'false': id}
         """
         # Add condition check node
-        check_node = Operation(
+        check_node = create_operation(
             operation=condition_check_op,
             parameters={**check_params, "is_condition_check": True},
         )
@@ -358,7 +359,9 @@ class OperationGraphBuilder:
         result = {"check": check_node.id}
 
         # Add true branch
-        true_node = Operation(operation=true_op, parameters={"branch": "true"})
+        true_node = create_operation(
+            operation=true_op, parameters={"branch": "true"}
+        )
         self.graph.add_node(true_node)
         self._operations[true_node.id] = true_node
         result["true"] = true_node.id
@@ -371,7 +374,7 @@ class OperationGraphBuilder:
 
         # Add false branch if specified
         if false_op:
-            false_node = Operation(
+            false_node = create_operation(
                 operation=false_op, parameters={"branch": "false"}
             )
             self.graph.add_node(false_node)
@@ -389,7 +392,7 @@ class OperationGraphBuilder:
 
         return result
 
-    def get_graph(self) -> Graph:
+    def get_graph(self):
         """
         Get the current graph for execution.
 
@@ -398,7 +401,7 @@ class OperationGraphBuilder:
         """
         return self.graph
 
-    def get_node_by_reference(self, reference_id: str) -> Operation | None:
+    def get_node_by_reference(self, reference_id: str):
         """
         Get a node by its reference ID.
 
@@ -440,292 +443,10 @@ class OperationGraphBuilder:
         }
 
     def visualize(self, title: str = "Operation Graph", figsize=(14, 10)):
+        from ._visualize_graph import visualize_graph
+
         visualize_graph(
             self,
             title=title,
             figsize=figsize,
         )
-
-
-def visualize_graph(
-    builder: OperationGraphBuilder,
-    title: str = "Operation Graph",
-    figsize=(14, 10),
-):
-    """Visualization with improved layout for complex graphs."""
-    from lionagi.protocols.graph.graph import (
-        _MATPLIB_AVAILABLE,
-        _NETWORKX_AVAILABLE,
-    )
-
-    if _MATPLIB_AVAILABLE is not True:
-        raise _MATPLIB_AVAILABLE
-    if _NETWORKX_AVAILABLE is not True:
-        raise _NETWORKX_AVAILABLE
-
-    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()

lionagi/operations/flow.py

@@ -10,16 +10,19 @@ using Events for synchronization and CapacityLimiter for concurrency control.
 """
 
 import os
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from lionagi.ln._async_call import AlcallParams
 from lionagi.ln.concurrency import CapacityLimiter, ConcurrencyEvent
 from lionagi.operations.node import Operation
-from lionagi.protocols.types import EventStatus, Graph
-from lionagi.session.branch import Branch
-from lionagi.session.session import Session
+from lionagi.protocols.types import EventStatus
 from lionagi.utils import to_dict
 
+if TYPE_CHECKING:
+    from lionagi.protocols.graph.graph import Graph
+    from lionagi.session.session import Branch, Session
+
+
 # Maximum concurrency when None is specified (effectively unlimited)
 UNLIMITED_CONCURRENCY = int(os.environ.get("LIONAGI_MAX_CONCURRENCY", "10000"))
 
@@ -29,12 +32,12 @@ class DependencyAwareExecutor:
 
     def __init__(
         self,
-        session: Session,
-        graph: Graph,
+        session: "Session",
+        graph: "Graph",
         context: dict[str, Any] | None = None,
         max_concurrent: int = 5,
         verbose: bool = False,
-        default_branch: Branch | None = None,
+        default_branch: "Branch" = None,
         alcall_params: AlcallParams | None = None,
     ):
         """Initialize the executor.
@@ -402,7 +405,7 @@ class DependencyAwareExecutor:
         branch = self._resolve_branch_for_operation(operation)
         self.operation_branches[operation.id] = branch
 
-    def _resolve_branch_for_operation(self, operation: Operation) -> Branch:
+    def _resolve_branch_for_operation(self, operation: Operation) -> "Branch":
         """Resolve which branch an operation should use - all branches are pre-allocated."""
         # All branches should be pre-allocated
         if operation.id in self.operation_branches:
@@ -501,10 +504,10 @@ class DependencyAwareExecutor:
 
 
 async def flow(
-    session: Session,
-    graph: Graph,
+    session: "Session",
+    graph: "Graph",
     *,
-    branch: Branch | None = None,
+    branch: "Branch" = None,
     context: dict[str, Any] | None = None,
     parallel: bool = True,
     max_concurrent: int = None,
@@ -546,3 +549,94 @@ async def flow(
     )
 
     return await executor.execute()
+
+
+def cleanup_flow_results(
+    result: dict[str, Any], keep_only: list[str] = None
+) -> dict[str, Any]:
+    """
+    Clean up flow execution results to reduce memory usage.
+
+    Args:
+        result: Flow execution result dictionary
+        keep_only: List of operation IDs to keep results for (optional)
+
+    Returns:
+        Modified result dictionary with reduced memory footprint
+    """
+    if not isinstance(result, dict) or "operation_results" not in result:
+        return result
+
+    # If keep_only is specified, only keep those results
+    if keep_only is not None:
+        filtered_results = {
+            op_id: res
+            for op_id, res in result["operation_results"].items()
+            if op_id in keep_only
+        }
+        result["operation_results"] = filtered_results
+        # Update completed_operations to match
+        result["completed_operations"] = [
+            op_id
+            for op_id in result.get("completed_operations", [])
+            if op_id in keep_only
+        ]
+    else:
+        # Clear all results to free memory
+        result["operation_results"] = {}
+        result["completed_operations"] = []
+
+    return result
+
+
+async def flow_with_cleanup(
+    session: "Session",
+    graph: "Graph",
+    context: dict[str, Any] | None = None,
+    parallel: bool = True,
+    max_concurrent: int = 5,
+    verbose: bool = False,
+    branch: "Branch" = None,
+    alcall_params: AlcallParams | None = None,
+    cleanup_results: bool = True,
+    keep_only: list[str] = None,
+) -> dict[str, Any]:
+    """
+    Execute flow with automatic cleanup to prevent memory accumulation.
+
+    Args:
+        session: Session instance for branch management
+        graph: Operation graph to execute
+        context: Initial context data
+        parallel: Execute independent operations in parallel
+        max_concurrent: Max concurrent operations (1 if not parallel)
+        verbose: Enable verbose logging
+        branch: Default branch for operations
+        alcall_params: Parameters for async parallel call execution
+        cleanup_results: Whether to clean up operation results after execution
+        keep_only: List of operation IDs to keep results for (if cleanup_results=True)
+
+    Returns:
+        Execution results (potentially with cleaned up memory footprint)
+    """
+    # Execute the flow normally
+    result = await flow(
+        session=session,
+        graph=graph,
+        context=context,
+        parallel=parallel,
+        max_concurrent=max_concurrent,
+        verbose=verbose,
+        branch=branch,
+        alcall_params=alcall_params,
+    )
+
+    # Clean up session memory
+    if hasattr(session, "cleanup_memory"):
+        session.cleanup_memory()
+
+    # Clean up results if requested
+    if cleanup_results:
+        result = cleanup_flow_results(result, keep_only=keep_only)
+
+    return result

lionagi/operations/node.py

@@ -1,13 +1,15 @@
 import asyncio
 import logging
-from typing import Any, Literal
+from typing import TYPE_CHECKING, Any, Literal
 from uuid import UUID
 
 from anyio import get_cancelled_exc_class
 from pydantic import BaseModel, Field
 
 from lionagi.protocols.types import ID, Event, EventStatus, IDType, Node
-from lionagi.session.branch import Branch
+
+if TYPE_CHECKING:
+    from lionagi.session.branch import Branch
 
 BranchOperations = Literal[
     "chat",
@@ -74,7 +76,7 @@ class Operation(Node, Event):
         """Get the response from the execution."""
         return self.execution.response if self.execution else None
 
-    async def invoke(self, branch: Branch):
+    async def invoke(self, branch: "Branch"):
         meth = branch.get_operation(self.operation)
         if meth is None:
             raise ValueError(f"Unsupported operation type: {self.operation}")
@@ -108,3 +110,12 @@ class Operation(Node, Event):
                 res.append(i)
             return res
         return await meth(**self.request)
+
+
+def create_operation(
+    operation: BranchOperations | str,
+    parameters: dict[str, Any] | BaseModel = None,
+    **kwargs,
+):
+    """Create an Operation node."""
+    return Operation(operation=operation, parameters=parameters, **kwargs)