lionagi 0.5.3__py3-none-any.whl → 0.5.4__py3-none-any.whl

lionagi/core/session/session.py

@@ -128,9 +128,9 @@ class Session(Component):
             branch: The branch to set as default or its identifier.
         """
         branch = self.branches[branch]
-        if branch and len(branch) == 1:
-            self.default_branch = branch
-        raise ValueError("Session can only have one default branch.")
+        if not isinstance(branch, Branch):
+            raise ValueError("Input value for branch is not a valid branch.")
+        self.default_branch = branch
 
     def to_df(self, branches: ID.RefSeq = None) -> pd.DataFrame:
         out = self.concat_messages(branches=branches)

lionagi/libs/func/async_calls/alcall.py

@@ -148,6 +148,13 @@ async def alcall(
                         ucall(func, i, **kwargs), retry_timeout
                     )
                     return index, result
+
+            except InterruptedError:
+                return index, None
+
+            except asyncio.CancelledError:
+                return index, None
+
             except TimeoutError as e:
                 raise TimeoutError(
                     f"{error_msg or ''} Timeout {retry_timeout} seconds "

lionagi/operations/brainstorm/brainstorm.py

@@ -2,6 +2,8 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import logging
+from typing import Literal
 
 from lionagi.core.session.branch import Branch
 from lionagi.core.session.session import Session
@@ -17,13 +19,41 @@ from lionagi.protocols.operatives.instruct import (
 from ..utils import prepare_instruct, prepare_session
 from .prompt import PROMPT
 
+# ---------------------------------------------------------------------
+# Data Models & Utilities
+# ---------------------------------------------------------------------
+
 
 class BrainstormOperation(BaseModel):
+    """
+    Container for the outcomes of a brainstorming session:
+      1. initial: the initial result of the 'brainstorm' prompt
+      2. brainstorm: the results of auto-run instructions (if auto_run = True)
+      3. explore: the results of exploring those instructions (if auto_explore = True)
+    """
+
     initial: Any
     brainstorm: list[Instruct] | None = None
     explore: list[InstructResponse] | None = None
 
 
+def chunked(iterable, n):
+    """
+    Yield successive n-sized chunks from an iterable.
+
+    Example:
+        >>> list(chunked([1,2,3,4,5], 2))
+        [[1,2],[3,4],[5]]
+    """
+    for i in range(0, len(iterable), n):
+        yield iterable[i : i + n]
+
+
+# ---------------------------------------------------------------------
+# Core Instruction Execution
+# ---------------------------------------------------------------------
+
+
 async def run_instruct(
     ins: Instruct,
     session: Session,
@@ -32,53 +62,50 @@ async def run_instruct(
     verbose: bool = True,
     **kwargs: Any,
 ) -> Any:
-    """Execute an instruction within a brainstorming session.
-
-    Args:
-        ins: The instruction model to run.
-        session: The current session.
-        branch: The branch to operate on.
-        auto_run: Whether to automatically run nested instructions.
-        verbose: Whether to enable verbose output.
-        **kwargs: Additional keyword arguments.
-
-    Returns:
-        The result of the instruction execution.
+    """
+    Execute a single instruction within a brainstorming session.
+    Optionally auto-run any child instructions that result.
     """
 
-    async def run(ins_):
+    async def _run_child_instruction(child_ins: Instruct):
+        """
+        Helper for recursively running child instructions.
+        """
         if verbose:
-            msg_ = (
-                ins_.guidance[:100] + "..."
-                if len(ins_.guidance) > 100
-                else ins_.guidance
+            snippet = (
+                child_ins.guidance[:100] + "..."
+                if len(child_ins.guidance) > 100
+                else child_ins.guidance
             )
-            print(f"\n-----Running instruction-----\n{msg_}")
-        b_ = session.split(branch)
+            print(f"\n-----Running instruction-----\n{snippet}")
+        child_branch = session.split(branch)
         return await run_instruct(
-            ins_, session, b_, False, verbose=verbose, **kwargs
+            child_ins, session, child_branch, False, verbose=verbose, **kwargs
         )
 
+    # Prepare config for the branch operation
     config = {**ins.model_dump(), **kwargs}
-    res = await branch.operate(**config)
+    result = await branch.operate(**config)
     branch.msgs.logger.dump()
-    instructs = []
-
-    if hasattr(res, "instruct_models"):
-        instructs = res.instruct_models
 
-    if auto_run is True and instructs:
-        ress = await alcall(instructs, run)
-        response_ = []
-        for res in ress:
-            if isinstance(res, list):
-                response_.extend(res)
+    # Extract any newly generated instructions
+    instructs = []
+    if hasattr(result, "instruct_models"):
+        instructs = result.instruct_models
+
+    # If we're allowed to auto-run child instructions, handle them
+    if auto_run and instructs:
+        child_results = await alcall(instructs, _run_child_instruction)
+        combined = []
+        for c in child_results:
+            if isinstance(c, list):
+                combined.extend(c)
             else:
-                response_.append(res)
-        response_.insert(0, res)
-        return response_
+                combined.append(c)
+        combined.insert(0, result)
+        return combined
 
-    return res
+    return result
 
 
 async def brainstorm(
@@ -89,116 +116,314 @@ async def brainstorm(
     auto_run: bool = True,
     auto_explore: bool = False,
     explore_kwargs: dict[str, Any] | None = None,
+    explore_strategy: Literal[
+        "concurrent",
+        "sequential",
+        "sequential_concurrent_chunk",
+        "concurrent_sequential_chunk",
+    ] = "concurrent",
     branch_kwargs: dict[str, Any] | None = None,
     return_session: bool = False,
     verbose: bool = False,
+    branch_as_default: bool = True,
     **kwargs: Any,
 ) -> Any:
-    """Perform a brainstorming session.
-
-    Args:
-        instruct: Instruction model or dictionary.
-        num_instruct: Number of instructions to generate.
-        session: Existing session or None to create a new one.
-        branch: Existing branch or reference.
-        auto_run: If True, automatically run generated instructions.
-        branch_kwargs: Additional arguments for branch creation.
-        return_session: If True, return the session with results.
-        verbose: Whether to enable verbose output.
-        **kwargs: Additional keyword arguments.
-
-    Returns:
-        The results of the brainstorming session, optionally with the session.
     """
+    High-level function to perform a brainstorming session.
 
+    Steps:
+      1. Run the initial 'instruct' prompt to generate suggestions.
+      2. Optionally auto-run those suggestions (auto_run=True).
+      3. Optionally explore the resulting instructions (auto_explore=True)
+         using the chosen strategy (concurrent, sequential, etc.).
+    """
+
+    # -----------------------------------------------------------------
+    # Basic Validations and Setup
+    # -----------------------------------------------------------------
     if auto_explore and not auto_run:
         raise ValueError("auto_explore requires auto_run to be True.")
 
     if verbose:
-        print(f"Starting brainstorming...")
+        print("Starting brainstorming...")
 
+    # Make sure the correct field model is present
     field_models: list = kwargs.get("field_models", [])
     if INSTRUCT_FIELD_MODEL not in field_models:
         field_models.append(INSTRUCT_FIELD_MODEL)
-
     kwargs["field_models"] = field_models
+
+    # Prepare session, branch, and the instruction
     session, branch = prepare_session(session, branch, branch_kwargs)
-    instruct = prepare_instruct(
-        instruct, PROMPT.format(num_instruct=num_instruct)
-    )
+    prompt_str = PROMPT.format(num_instruct=num_instruct)
+    instruct = prepare_instruct(instruct, prompt_str)
+
+    # -----------------------------------------------------------------
+    # 1. Initial Brainstorm
+    # -----------------------------------------------------------------
     res1 = await branch.operate(**instruct, **kwargs)
     out = BrainstormOperation(initial=res1)
 
     if verbose:
         print("Initial brainstorming complete.")
 
-    instructs = None
-
-    async def run(ins_):
+    # Helper to run single instructions from the 'brainstorm'
+    async def run_brainstorm_instruction(ins_):
         if verbose:
-            msg_ = (
+            snippet = (
                 ins_.guidance[:100] + "..."
                 if len(ins_.guidance) > 100
                 else ins_.guidance
             )
-            print(f"\n-----Running instruction-----\n{msg_}")
-        b_ = session.split(branch)
+            print(f"\n-----Running instruction-----\n{snippet}")
+        new_branch = session.split(branch)
         return await run_instruct(
-            ins_, session, b_, auto_run, verbose=verbose, **kwargs
+            ins_, session, new_branch, auto_run, verbose=verbose, **kwargs
         )
 
+    # -----------------------------------------------------------------
+    # 2. Auto-run child instructions if requested
+    # -----------------------------------------------------------------
     if not auto_run:
         if return_session:
             return out, session
         return out
 
+    # We run inside the context manager for branching
     async with session.branches:
         response_ = []
+
+        # If the initial result has instructions, run them
         if hasattr(res1, "instruct_models"):
             instructs: list[Instruct] = res1.instruct_models
-            ress = await alcall(instructs, run)
-            ress = to_flat_list(ress, dropna=True)
+            brainstorm_results = await alcall(
+                instructs, run_brainstorm_instruction
+            )
+            brainstorm_results = to_flat_list(brainstorm_results, dropna=True)
 
-            response_ = [
-                res if not isinstance(res, str | dict) else None
-                for res in ress
+            # Filter out plain str/dict responses, keep model-based
+            filtered = [
+                r if not isinstance(r, (str, dict)) else None
+                for r in brainstorm_results
             ]
-            response_ = to_flat_list(response_, unique=True, dropna=True)
+            filtered = to_flat_list(filtered, unique=True, dropna=True)
+
             out.brainstorm = (
-                response_ if isinstance(response_, list) else [response_]
+                filtered if isinstance(filtered, list) else [filtered]
             )
-            response_.insert(0, res1)
+            # Insert the initial result at index 0 for reference
+            filtered.insert(0, res1)
+            response_ = filtered
 
+        # -----------------------------------------------------------------
+        # 3. Explore the results (if auto_explore = True)
+        # -----------------------------------------------------------------
         if response_ and auto_explore:
-
-            async def explore(ins_: Instruct):
-                if verbose:
-                    msg_ = (
-                        ins_.guidance[:100] + "..."
-                        if len(ins_.guidance) > 100
-                        else ins_.guidance
-                    )
-                    print(f"\n-----Exploring Idea-----\n{msg_}")
-                b_ = session.split(branch)
-                res = await b_.instruct(ins_, **(explore_kwargs or {}))
-                return InstructResponse(
-                    instruct=ins_,
-                    response=res,
-                )
-
-            response_ = to_flat_list(
+            # Gather all newly generated instructions
+            all_explore_instructs = to_flat_list(
                 [
-                    i.instruct_models
-                    for i in response_
-                    if hasattr(i, "instruct_models")
+                    r.instruct_models
+                    for r in response_
+                    if hasattr(r, "instruct_models")
                 ],
                 dropna=True,
                 unique=True,
             )
-            res_explore = await alcall(response_, explore)
-            out.explore = res_explore
 
+            # Decide how to explore based on the strategy
+            match explore_strategy:
+                # ---------------------------------------------------------
+                # Strategy A: CONCURRENT
+                # ---------------------------------------------------------
+                case "concurrent":
+
+                    async def explore_concurrently(ins_: Instruct):
+                        if verbose:
+                            snippet = (
+                                ins_.guidance[:100] + "..."
+                                if len(ins_.guidance) > 100
+                                else ins_.guidance
+                            )
+                            print(f"\n-----Exploring Idea-----\n{snippet}")
+                        new_branch = session.split(branch)
+                        resp = await new_branch.instruct(
+                            ins_, **(explore_kwargs or {})
+                        )
+                        return InstructResponse(instruct=ins_, response=resp)
+
+                    res_explore = await alcall(
+                        all_explore_instructs, explore_concurrently
+                    )
+                    out.explore = res_explore
+
+                    # Add messages for logging / auditing
+                    branch.msgs.add_message(
+                        instruction="\n".join(
+                            i.model_dump_json() for i in all_explore_instructs
+                        )
+                    )
+                    branch.msgs.add_message(
+                        assistant_response="\n".join(
+                            i.model_dump_json() for i in res_explore
+                        )
+                    )
+
+                # ---------------------------------------------------------
+                # Strategy B: SEQUENTIAL
+                # ---------------------------------------------------------
+                case "sequential":
+                    explore_results = []
+
+                    # Warn/log if a large number of instructions
+                    if len(all_explore_instructs) > 30:
+                        all_explore_instructs = all_explore_instructs[:30]
+                        logging.warning(
+                            "Maximum number of instructions for sequential exploration is 50. defaulting to 50."
+                        )
+                    if len(all_explore_instructs) > 10:
+                        logging.warning(
+                            "Large number of instructions for sequential exploration. This may take a while."
+                        )
+
+                    for i in all_explore_instructs:
+                        if verbose:
+                            snippet = (
+                                i.guidance[:100] + "..."
+                                if len(i.guidance) > 100
+                                else i.guidance
+                            )
+                            print(f"\n-----Exploring Idea-----\n{snippet}")
+                        seq_res = await branch.instruct(
+                            i, **(explore_kwargs or {})
+                        )
+                        explore_results.append(
+                            InstructResponse(instruct=i, response=seq_res)
+                        )
+
+                    out.explore = explore_results
+
+                # ---------------------------------------------------------
+                # Strategy C: SEQUENTIAL_CONCURRENT_CHUNK
+                # (chunks processed sequentially, each chunk in parallel)
+                # ---------------------------------------------------------
+                case "sequential_concurrent_chunk":
+                    chunk_size = (explore_kwargs or {}).get("chunk_size", 5)
+                    all_responses = []
+
+                    async def explore_concurrent_chunk(
+                        sub_instructs: list[Instruct], base_branch: Branch
+                    ):
+                        """
+                        Explore instructions in a single chunk concurrently.
+                        """
+                        if verbose:
+                            print(
+                                f"\n--- Exploring a chunk of size {len(sub_instructs)} ---\n"
+                            )
+
+                        async def _explore(ins_: Instruct):
+                            child_branch = session.split(base_branch)
+                            child_resp = await child_branch.instruct(
+                                ins_, **(explore_kwargs or {})
+                            )
+                            return InstructResponse(
+                                instruct=ins_, response=child_resp
+                            )
+
+                        # Run all instructions in the chunk concurrently
+                        res_chunk = await alcall(sub_instructs, _explore)
+
+                        # Log messages for debugging / auditing
+                        next_branch = session.split(base_branch)
+                        next_branch.msgs.add_message(
+                            instruction="\n".join(
+                                i.model_dump_json() for i in sub_instructs
+                            )
+                        )
+                        next_branch.msgs.add_message(
+                            assistant_response="\n".join(
+                                i.model_dump_json() for i in res_chunk
+                            )
+                        )
+                        return res_chunk, next_branch
+
+                    # Process each chunk sequentially
+                    for chunk in chunked(all_explore_instructs, chunk_size):
+                        chunk_result, branch = await explore_concurrent_chunk(
+                            chunk, branch
+                        )
+                        all_responses.extend(chunk_result)
+
+                    out.explore = all_responses
+
+                # ---------------------------------------------------------
+                # Strategy D: CONCURRENT_SEQUENTIAL_CHUNK
+                # (all chunks processed concurrently, each chunk sequentially)
+                # ---------------------------------------------------------
+                case "concurrent_sequential_chunk":
+                    chunk_size = (explore_kwargs or {}).get("chunk_size", 5)
+                    all_chunks = list(
+                        chunked(all_explore_instructs, chunk_size)
+                    )
+
+                    async def explore_chunk_sequentially(
+                        sub_instructs: list[Instruct],
+                    ):
+                        """
+                        Explore instructions in a single chunk, one at a time.
+                        """
+                        chunk_results = []
+                        local_branch = session.split(branch)
+
+                        for ins_ in sub_instructs:
+                            if verbose:
+                                snippet = (
+                                    ins_.guidance[:100] + "..."
+                                    if len(ins_.guidance) > 100
+                                    else ins_.guidance
+                                )
+                                print(
+                                    f"\n-----Exploring Idea (sequential in chunk)-----\n{snippet}"
+                                )
+
+                            seq_resp = await local_branch.instruct(
+                                ins_, **(explore_kwargs or {})
+                            )
+                            chunk_results.append(
+                                InstructResponse(
+                                    instruct=ins_, response=seq_resp
+                                )
+                            )
+
+                        return chunk_results
+
+                    # Run all chunks in parallel
+                    all_responses = await alcall(
+                        all_chunks,
+                        explore_chunk_sequentially,
+                        flatten=True,
+                        dropna=True,
+                    )
+                    out.explore = all_responses
+
+                    # Log final messages
+                    branch.msgs.add_message(
+                        instruction="\n".join(
+                            i.model_dump_json() for i in all_explore_instructs
+                        )
+                    )
+                    branch.msgs.add_message(
+                        assistant_response="\n".join(
+                            i.model_dump_json() for i in all_responses
+                        )
+                    )
+
+    if branch_as_default:
+        session.change_default_branch(branch)
+
+    # -----------------------------------------------------------------
+    # 4. Return Results
+    # -----------------------------------------------------------------
     if return_session:
         return out, session
-
     return out

lionagi/operations/plan/plan.py

@@ -2,10 +2,11 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-
 from lionagi.core.session.branch import Branch
 from lionagi.core.session.session import Session
 from lionagi.core.typing import ID, Any, BaseModel, Literal
+from lionagi.libs.func.types import alcall
+from lionagi.libs.parse import to_flat_list
 from lionagi.protocols.operatives.instruct import (
     INSTRUCT_FIELD_MODEL,
     Instruct,
@@ -15,13 +16,45 @@ from lionagi.protocols.operatives.instruct import (
 from ..utils import prepare_instruct, prepare_session
 from .prompt import EXPANSION_PROMPT, PLAN_PROMPT
 
+# ---------------------------------------------------------------------
+# Data Model
+# ---------------------------------------------------------------------
+
 
 class PlanOperation(BaseModel):
+    """
+    Stores all relevant outcomes for a multi-step Plan:
+      * initial: The result of the initial plan prompt
+      * plan: A list of plan steps (Instruct objects) generated from the initial planning
+      * execute: Any responses from executing those plan steps
+    """
+
     initial: Any
     plan: list[Instruct] | None = None
     execute: list[InstructResponse] | None = None
 
 
+# ---------------------------------------------------------------------
+# Utilities
+# ---------------------------------------------------------------------
+
+
+def chunked(iterable, n):
+    """
+    Yield successive n-sized chunks from an iterable.
+    Example:
+        >>> list(chunked([1,2,3,4,5], 2))
+        [[1,2],[3,4],[5]]
+    """
+    for i in range(0, len(iterable), n):
+        yield iterable[i : i + n]
+
+
+# ---------------------------------------------------------------------
+# Single-Step Runner
+# ---------------------------------------------------------------------
+
+
 async def run_step(
     ins: Instruct,
     session: Session,
@@ -29,33 +62,41 @@ async def run_step(
     verbose: bool = True,
     **kwargs: Any,
 ) -> Any:
-    """Execute a single step of the plan.
+    """
+    Execute a single step of the plan with an 'expansion' or guidance prompt.
 
     Args:
         ins: The instruction model for the step.
-        session: The current session.
-        branch: The branch to operate on.
+        session: The current session context.
+        branch: The branch to operate on for this step.
         verbose: Whether to enable verbose output.
-        **kwargs: Additional keyword arguments.
+        **kwargs: Additional keyword arguments passed to the branch operation.
 
     Returns:
-        The result of the branch operation.
+        The result of the branch operation (which may contain more instructions).
     """
     if verbose:
-        instruction = (
+        snippet = (
             ins.instruction[:100] + "..."
             if len(ins.instruction) > 100
             else ins.instruction
         )
-        print(f"Further planning: {instruction}")
+        print(f"Further planning: {snippet}")
 
+    # Incorporate the EXPANSION_PROMPT into guidance
     config = {**ins.model_dump(), **kwargs}
-    guide = config.pop("guidance", "")
-    config["guidance"] = EXPANSION_PROMPT + "\n" + str(guide)
+    guidance_text = config.pop("guidance", "")
+    config["guidance"] = f"{EXPANSION_PROMPT}\n{guidance_text}"
+
+    # Run the step
+    result = await branch.operate(**config)
+    branch.msgs.logger.dump()  # Dump logs if needed
+    return result
 
-    res = await branch.operate(**config)
-    branch.msgs.logger.dump()
-    return res
+
+# ---------------------------------------------------------------------
+# Main Plan Function (with Multiple Execution Strategies)
+# ---------------------------------------------------------------------
 
 
 async def plan(
@@ -65,108 +106,280 @@ async def plan(
     branch: Branch | ID.Ref | None = None,
     auto_run: bool = True,
     auto_execute: bool = False,
-    execution_strategy: Literal["sequential"] = "sequential",
+    execution_strategy: Literal[
+        "sequential",
+        "concurrent",
+        "sequential_concurrent_chunk",
+        "concurrent_sequential_chunk",
+    ] = "sequential",
     execution_kwargs: dict[str, Any] | None = None,
     branch_kwargs: dict[str, Any] | None = None,
     return_session: bool = False,
     verbose: bool = True,
     **kwargs: Any,
-) -> PlanOperation | tuple[list[InstructResponse], Session]:
-    """Create and execute a multi-step plan.
+) -> PlanOperation | tuple[PlanOperation, Session]:
+    """
+    Create and optionally execute a multi-step plan with up to `num_steps`.
+
+    Steps:
+      1. Generate an initial plan with up to `num_steps`.
+      2. Optionally (auto_run=True) expand on each planned step
+         to refine or further clarify them.
+      3. Optionally (auto_execute=True) execute those refined steps
+         according to `execution_strategy`.
 
     Args:
-        instruct: Instruction model or dictionary.
-        num_steps: Number of steps in the plan.
-        session: Existing session or None to create a new one.
-        branch: Existing branch or reference.
-        auto_run: If True, automatically run the steps.
-        branch_kwargs: Additional keyword arguments for branch creation.
-        return_session: If True, return the session along with results.
-        verbose: Whether to enable verbose output.
-        **kwargs: Additional keyword arguments.
+        instruct: Initial instruction or a dict describing it.
+        num_steps: Maximum number of plan steps (must be <= 5).
+        session: An existing Session, or None to create a new one.
+        branch: An existing Branch, or None to create a new one.
+        auto_run: If True, automatically run the intermediate plan steps.
+        auto_execute: If True, automatically execute the fully refined steps.
+        execution_strategy:
+            - "sequential" (default) runs steps one by one
+            - "concurrent" runs all steps in parallel
+            - "sequential_concurrent_chunk" processes chunks sequentially, each chunk in parallel
+            - "concurrent_sequential_chunk" processes all chunks in parallel, each chunk sequentially
+        execution_kwargs: Extra kwargs used during execution calls.
+        branch_kwargs: Extra kwargs for branch/session creation.
+        return_session: Whether to return (PlanOperation, Session) instead of just PlanOperation.
+        verbose: If True, prints verbose logs.
+        **kwargs: Additional arguments for the initial plan operation.
 
     Returns:
-        Results of the plan execution, optionally with the session.
+        A PlanOperation object containing:
+            - initial plan
+            - (optional) plan expansions
+            - (optional) execution responses
+        Optionally returns the session as well, if `return_session=True`.
     """
+
+    # -----------------------------------------------------------------
+    # 0. Basic Validation & Setup
+    # -----------------------------------------------------------------
     if num_steps > 5:
         raise ValueError("Number of steps must be 5 or less")
 
     if verbose:
         print(f"Planning execution with {num_steps} steps...")
 
+    # Ensure the correct field model
     field_models: list = kwargs.get("field_models", [])
     if INSTRUCT_FIELD_MODEL not in field_models:
         field_models.append(INSTRUCT_FIELD_MODEL)
     kwargs["field_models"] = field_models
+
+    # Prepare session/branch
     session, branch = prepare_session(session, branch, branch_kwargs)
-    execute_branch: Branch = session.split(branch)
-    instruct = prepare_instruct(
-        instruct, PLAN_PROMPT.format(num_steps=num_steps)
-    )
+    execute_branch: Branch = session.split(
+        branch
+    )  # a separate branch for execution
 
-    res1 = await branch.operate(**instruct, **kwargs)
-    out = PlanOperation(initial=res1)
+    # -----------------------------------------------------------------
+    # 1. Run the Initial Plan Prompt
+    # -----------------------------------------------------------------
+    plan_prompt = PLAN_PROMPT.format(num_steps=num_steps)
+    instruct = prepare_instruct(instruct, plan_prompt)
+    initial_res = await branch.operate(**instruct, **kwargs)
+
+    # Wrap initial result in the PlanOperation
+    out = PlanOperation(initial=initial_res)
 
     if verbose:
         print("Initial planning complete. Starting step planning...")
 
+    # If we aren't auto-running the steps, just return the initial plan
     if not auto_run:
-        if return_session:
-            return res1, session
-        return res1
+        return (out, session) if return_session else out
 
+    # -----------------------------------------------------------------
+    # 2. Expand Each Step (auto_run=True)
+    # -----------------------------------------------------------------
     results = []
-    if hasattr(res1, "instruct_models"):
-        instructs: list[Instruct] = res1.instruct_models
-        for i, ins in enumerate(instructs, 1):
+    if hasattr(initial_res, "instruct_models"):
+        instructs: list[Instruct] = initial_res.instruct_models
+        for i, step_ins in enumerate(instructs, start=1):
             if verbose:
                 print(f"\n----- Planning step {i}/{len(instructs)} -----")
-            res = await run_step(
-                ins, session, branch, verbose=verbose, **kwargs
+            expanded_res = await run_step(
+                step_ins, session, branch, verbose=verbose, **kwargs
             )
-            results.append(res)
+            results.append(expanded_res)
 
         if verbose:
-            print("\nAll planning completed successfully!")
+            print("\nAll planning steps expanded/refined successfully!")
+
+    # Gather all newly created plan instructions
+    refined_plans = []
+    for step_result in results:
+        if hasattr(step_result, "instruct_models"):
+            for model in step_result.instruct_models:
+                if model and model not in refined_plans:
+                    refined_plans.append(model)
 
-    all_plans = []
-    for res in results:
-        if hasattr(res, "instruct_models"):
-            for i in res.instruct_models:
-                if i and i not in all_plans:
-                    all_plans.append(i)
-    out.plan = all_plans
+    out.plan = refined_plans
 
+    # -----------------------------------------------------------------
+    # 3. Execute the Plan Steps (auto_execute=True)
+    # -----------------------------------------------------------------
     if auto_execute:
         if verbose:
-            print("\nStarting execution of all steps...")
-        results = []
+            print("\nStarting execution of all plan steps...")
+
+        # We now handle multiple strategies:
         match execution_strategy:
+
+            # ---------------------------------------------------------
+            # Strategy A: SEQUENTIAL
+            # ---------------------------------------------------------
             case "sequential":
-                for i, ins in enumerate(all_plans, 1):
+                seq_results = []
+                for i, plan_step in enumerate(refined_plans, start=1):
+                    if verbose:
+                        snippet = (
+                            plan_step.instruction[:100] + "..."
+                            if len(plan_step.instruction) > 100
+                            else plan_step.instruction
+                        )
+                        print(
+                            f"\n------ Executing step {i}/{len(refined_plans)} ------"
+                        )
+                        print(f"Instruction: {snippet}")
+
+                    step_response = await execute_branch.instruct(
+                        plan_step, **(execution_kwargs or {})
+                    )
+                    seq_results.append(
+                        InstructResponse(
+                            instruct=plan_step, response=step_response
+                        )
+                    )
+
+                out.execute = seq_results
+                if verbose:
+                    print("\nAll steps executed successfully (sequential)!")
+
+            # ---------------------------------------------------------
+            # Strategy B: CONCURRENT
+            # ---------------------------------------------------------
+            case "concurrent":
+
+                async def execute_step_concurrently(plan_step: Instruct):
+                    if verbose:
+                        snippet = (
+                            plan_step.instruction[:100] + "..."
+                            if len(plan_step.instruction) > 100
+                            else plan_step.instruction
+                        )
+                        print(f"\n------ Executing step (concurrently) ------")
+                        print(f"Instruction: {snippet}")
+                    local_branch = session.split(execute_branch)
+                    resp = await local_branch.instruct(
+                        plan_step, **(execution_kwargs or {})
+                    )
+                    return InstructResponse(instruct=plan_step, response=resp)
+
+                # Launch all steps in parallel
+                concurrent_res = await alcall(
+                    refined_plans, execute_step_concurrently
+                )
+                out.execute = concurrent_res
+                if verbose:
+                    print("\nAll steps executed successfully (concurrent)!")
+
+            # ---------------------------------------------------------
+            # Strategy C: SEQUENTIAL_CONCURRENT_CHUNK
+            #   - process plan steps in chunks (one chunk after another),
+            #   - each chunk’s steps run in parallel.
+            # ---------------------------------------------------------
+            case "sequential_concurrent_chunk":
+                chunk_size = (execution_kwargs or {}).get("chunk_size", 5)
+                all_exec_responses = []
+
+                async def execute_chunk_concurrently(
+                    sub_steps: list[Instruct],
+                ):
                     if verbose:
                         print(
-                            f"\n------ Executing step {i}/{len(all_plans)} ------"
+                            f"\n--- Executing a chunk of size {len(sub_steps)} concurrently ---"
                         )
-                        msg = (
-                            ins.instruction[:100] + "..."
-                            if len(ins.instruction) > 100
-                            else ins.instruction
+
+                    async def _execute(plan_step: Instruct):
+                        local_branch = session.split(execute_branch)
+                        resp = await local_branch.instruct(
+                            plan_step, **(execution_kwargs or {})
+                        )
+                        return InstructResponse(
+                            instruct=plan_step, response=resp
                         )
-                        print(f"Instruction: {msg}")
-                    res = await execute_branch.instruct(
-                        ins, **(execution_kwargs or {})
+
+                    # run each chunk in parallel
+                    return await alcall(sub_steps, _execute)
+
+                # process each chunk sequentially
+                for chunk in chunked(refined_plans, chunk_size):
+                    chunk_responses = await execute_chunk_concurrently(chunk)
+                    all_exec_responses.extend(chunk_responses)
+
+                out.execute = all_exec_responses
+                if verbose:
+                    print(
+                        "\nAll steps executed successfully (sequential concurrent chunk)!"
                     )
-                    res_ = InstructResponse(instruct=ins, response=res)
-                    results.append(res_)
-                out.execute = results
+
+            # ---------------------------------------------------------
+            # Strategy D: CONCURRENT_SEQUENTIAL_CHUNK
+            #   - split plan steps into chunks,
+            #   - run all chunks in parallel,
+            #   - but each chunk’s steps run sequentially.
+            # ---------------------------------------------------------
+            case "concurrent_sequential_chunk":
+                chunk_size = (execution_kwargs or {}).get("chunk_size", 5)
+                all_chunks = list(chunked(refined_plans, chunk_size))
+
+                async def execute_chunk_sequentially(
+                    sub_steps: list[Instruct],
+                ):
+                    chunk_result = []
+                    local_branch = session.split(execute_branch)
+                    for plan_step in sub_steps:
+                        if verbose:
+                            snippet = (
+                                plan_step.instruction[:100] + "..."
+                                if len(plan_step.instruction) > 100
+                                else plan_step.instruction
+                            )
+                            print(
+                                f"\n--- Executing step (sequential in chunk) ---\nInstruction: {snippet}"
+                            )
+                        resp = await local_branch.instruct(
+                            plan_step, **(execution_kwargs or {})
+                        )
+                        chunk_result.append(
+                            InstructResponse(instruct=plan_step, response=resp)
+                        )
+                    return chunk_result
+
+                # run all chunks in parallel, each chunk sequentially
+                parallel_chunk_results = await alcall(
+                    all_chunks,
+                    execute_chunk_sequentially,
+                    flatten=True,
+                    dropna=True,
+                )
+
+                out.execute = parallel_chunk_results
                 if verbose:
-                    print("\nAll steps executed successfully!")
+                    print(
+                        "\nAll steps executed successfully (concurrent sequential chunk)!"
+                    )
+
             case _:
                 raise ValueError(
                     f"Invalid execution strategy: {execution_strategy}"
                 )
 
-    if return_session:
-        return out, session
-    return out
+    # -----------------------------------------------------------------
+    # 4. Final Return
+    # -----------------------------------------------------------------
+    return (out, session) if return_session else out

lionagi/version.py

@@ -1 +1 @@
-__version__ = "0.5.3"
+__version__ = "0.5.4"

{lionagi-0.5.3.dist-info → lionagi-0.5.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lionagi
-Version: 0.5.3
+Version: 0.5.4
 Summary: An AGentic Intelligence Operating System.
 Author-email: HaiyangLi <quantocean.li@gmail.com>
 License:            Apache License

{lionagi-0.5.3.dist-info → lionagi-0.5.4.dist-info}/RECORD

@@ -1,6 +1,6 @@
 lionagi/__init__.py,sha256=oybfu2VsZc4ElN7ZeaW3KQrz8T8EcSDHPA8lUE-8G2I,537
 lionagi/settings.py,sha256=BOjxRV4N9zQJervvajPhbaHmgZ-nhbCy7AaQJi3Avug,2726
-lionagi/version.py,sha256=tgzuqHKcEdKBaP57F5oXxq4XlW2n9J4Fj8ZGu7nGOZg,22
+lionagi/version.py,sha256=DITpct-LrdIsTgwx2NgH5Ghx5y8Xgz1YMimy1ZV5RTY,22
 lionagi/core/__init__.py,sha256=v8vNyJVIVj8_Oz9RJdVe6ZKUQMYTgDh1VQpnr1KdLaw,112
 lionagi/core/_class_registry.py,sha256=srSWefqCS9EZrMvyA8zCrZ9KFvzAhTIj8g6mJG5KlIc,1982
 lionagi/core/action/__init__.py,sha256=v8vNyJVIVj8_Oz9RJdVe6ZKUQMYTgDh1VQpnr1KdLaw,112
@@ -56,7 +56,7 @@ lionagi/core/models/types.py,sha256=elcUuz_9dx4AhZddnICF-Cs62VJWIBqMET7MiRe4c1I,
 lionagi/core/session/__init__.py,sha256=v8vNyJVIVj8_Oz9RJdVe6ZKUQMYTgDh1VQpnr1KdLaw,112
 lionagi/core/session/branch.py,sha256=r6yNXwTm0oYA-9YOereuvLJtDnhj9IJWJ5fjTyKN88U,4406
 lionagi/core/session/branch_mixins.py,sha256=C9lGHmD1AmXAePw0kGIeuZjZ7SzOrFcIwV6rVaGhoGg,19623
-lionagi/core/session/session.py,sha256=Uup50oMqs00ZH6zgEiqtk0pG_ANByHW_Ael07dfbT1M,5122
+lionagi/core/session/session.py,sha256=cutece_iTs5K_m5soRfU9oTfHmw1icDEvx77E1RelIM,5129
 lionagi/core/session/types.py,sha256=MUGTSa2HWK79p7z-CG22RFP07N5AKnPVNXZwZt_wIvU,202
 lionagi/core/typing/__init__.py,sha256=Y9BK1OUXzjQgIo3epCVwWqUYhFwQQ_ayhRwI1UOmINg,228
 lionagi/core/typing/_concepts.py,sha256=uIzqfwtPBsIREhvT7NDAowc4i-u-69n_DRzLHzvHZO4,3730
@@ -257,7 +257,7 @@ lionagi/libs/func/throttle.py,sha256=iOOmS6i81NGRCClf4WaIhMMXwPt_rZ6WXYBqMJ_weEE
 lionagi/libs/func/types.py,sha256=wdjGNmY82pVACQBuDMUt3XJO0G9_wKHR7VesOeMxo_A,831
 lionagi/libs/func/utils.py,sha256=-LMdUEjksXP6JjjcUBh0XEQiXF30Zmv3xpKbmXjfya8,2674
 lionagi/libs/func/async_calls/__init__.py,sha256=qDHIkH7B-Ka5NJqeeyu_YL3TY36xL8kBwTjtCR4H8AU,495
-lionagi/libs/func/async_calls/alcall.py,sha256=U8-lWwqNwhe50Q07kzyW4YVcbx_LkUN5X1K1zvFg7i0,7283
+lionagi/libs/func/async_calls/alcall.py,sha256=GhPXId0YLvfxpDSYOABwr_3D-nRu8aLWshaYboJQzmQ,7436
 lionagi/libs/func/async_calls/bcall.py,sha256=gwfKpRZkExjVn-1YGZfaboFES8RqUgyaBrdpNtz1IdY,4436
 lionagi/libs/func/async_calls/mcall.py,sha256=9O5gWbBT4iIqXfcdZjgAdqsOSLWrNS4_tt6ou231ozA,4607
 lionagi/libs/func/async_calls/pcall.py,sha256=6u3RJPV-3yOkWxC9iSoVFl1veFfZFJpR0QRyGtfBODI,5831
@@ -317,10 +317,10 @@ lionagi/libs/string_similarity/utils.py,sha256=NdD0qF5tuytWBsm0WMrH0gRBBSxw2p4-m
 lionagi/operations/__init__.py,sha256=Dt7o6DFP7zVR-uxZ4xsGHQcse3XVlF6S8Y9NhaUTn_4,68
 lionagi/operations/utils.py,sha256=kn5SkZRczl1aQ-vJBeVPlMdeyUUD0s5iyuAW4P6KOvQ,1164
 lionagi/operations/brainstorm/__init__.py,sha256=amsoH65wepsx58DfgH-TRTN1wDH5TC24qYI_f02zkVg,61
-lionagi/operations/brainstorm/brainstorm.py,sha256=H2D0ZBCbRA2Ig60IHuXhyH6Hu02jxKhbj028UtS4eik,6279
+lionagi/operations/brainstorm/brainstorm.py,sha256=1Uuc11OH34jEYfDdombX5ui9b-bJTn4bVSLt0jjQUIc,16747
 lionagi/operations/brainstorm/prompt.py,sha256=3a7LsmtiqGAK5mtWoX-2qhsjETBzBx8FxM3cFCBEoOo,497
 lionagi/operations/plan/__init__.py,sha256=SVqoVmlSGz9lsLztm487H2qOLwgyFxSi2yZ8ubn-bgE,43
-lionagi/operations/plan/plan.py,sha256=Ui8fFdRNLiD3RJmnKCwfY48yj0Fxxw2LBye0egYttWE,5611
+lionagi/operations/plan/plan.py,sha256=AX4h_TeyhcXag572ywN9rmvo3fEiuGER7skAKf25_LY,15329
 lionagi/operations/plan/prompt.py,sha256=sKHa_jDahzCJ60oILj1XNYCIlbS-H8ybKRXpf9zd5x0,880
 lionagi/operations/select/__init__.py,sha256=dUd-KS1l404_ueYlIQsVNhS9jAqjn5pJbtUEbbh6KlI,49
 lionagi/operations/select/prompt.py,sha256=wbmuDC96fcQ4LFKjqmeErOQwVRpGWRqcwUeLTWnbeNs,186
@@ -368,7 +368,7 @@ lionagi/strategies/sequential_chunk.py,sha256=jG_WZXG-Ra3yd30CmX4b3XeCNAUrZGA2-i
 lionagi/strategies/sequential_concurrent_chunk.py,sha256=H7GShaqYlD5XxNJMG2GdOR4Vl8JHDhZb5jxNq8zY0hI,3365
 lionagi/strategies/types.py,sha256=fEvE4d1H4SeCcXcd2dz3q4k8jFIBtxYzjxDN7eJRLtI,769
 lionagi/strategies/utils.py,sha256=DX1dvxia8cNRqEJJbssJ3mgRzo7kgWCTA4y5DYLCCZE,1281
-lionagi-0.5.3.dist-info/METADATA,sha256=EDYWV82gJQWvmwXKbRLUaJjEGQ20B2DL-3emvuhdchg,22736
-lionagi-0.5.3.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-lionagi-0.5.3.dist-info/licenses/LICENSE,sha256=VXFWsdoN5AAknBCgFqQNgPWYx7OPp-PFEP961zGdOjc,11288
-lionagi-0.5.3.dist-info/RECORD,,
+lionagi-0.5.4.dist-info/METADATA,sha256=or-K2I6vye-JTu559FMXpy14cZlI5_3xLAgLCURMu1w,22736
+lionagi-0.5.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+lionagi-0.5.4.dist-info/licenses/LICENSE,sha256=VXFWsdoN5AAknBCgFqQNgPWYx7OPp-PFEP961zGdOjc,11288
+lionagi-0.5.4.dist-info/RECORD,,