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

lionagi/__init__.py

@@ -4,7 +4,8 @@ import logging
 
 from dotenv import load_dotenv
 
-from .core.session.types import Branch
+from .core.session.types import Branch, Session
+from .integrations.litellm_.imodel import LiteiModel
 from .protocols.operatives.step import Step
 from .service import iModel
 from .settings import Settings
@@ -20,6 +21,8 @@ __all__ = [
     "LiteiModel",
     "Branch",
     "Step",
+    "Session",
+    "LiteiModel",
 ]
 
 logger = logging.getLogger(__name__)

lionagi/core/session/branch.py

@@ -9,6 +9,7 @@ from pydantic import model_validator
 
 from lionagi.core.generic.types import Component, LogManager, Pile, Progression
 from lionagi.core.typing import ID
+from lionagi.integrations.litellm_.imodel import LiteiModel
 from lionagi.protocols.operatives.instruct import Instruct, OperationInstruct
 from lionagi.service import iModel
 from lionagi.settings import Settings
@@ -24,8 +25,8 @@ class Branch(Component, BranchActionMixin, BranchOperationMixin):
     name: str | None = None
     msgs: MessageManager = None
     acts: ActionManager = None
-    imodel: iModel | None = None
-    parse_imodel: iModel | None = None
+    imodel: iModel | LiteiModel | None = None
+    parse_imodel: iModel | LiteiModel | None = None
 
     @model_validator(mode="before")
     def _validate_data(cls, data: dict) -> dict:

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,11 +2,13 @@
 #
 # 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
 from lionagi.core.typing import ID, Any, BaseModel
-from lionagi.libs.func import alcall
+from lionagi.libs.func.types import alcall
 from lionagi.libs.parse import to_flat_list
 from lionagi.protocols.operatives.instruct import (
     INSTRUCT_FIELD_MODEL,
@@ -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/brainstorm/prompt.py

@@ -1 +1,8 @@
-PROMPT = """Perform a brainstorm session. Fill in {num_instruct} Instruct for the appropriate next step, we will run them separately and concurrently with same external context, but you should supplement each idea with certain amount of uniqueness while adhering to the guidelines and standards of the project. The Instruct should be concisely informational. If you think a particular step requries further extension, you should mention it in the instruct"""
+PROMPT = """Perform a brainstorm session. Generate {num_instruct} concise and distinct instructions (Instruct), each representing a potential next step. We will run them in parallel under the same context. Ensure each idea:
+
+1. Adheres to project guidelines and standards.
+2. Maintains a unique perspective or approach.
+3. Remains succinct yet sufficiently detailed.
+4. Flags any step that needs deeper expansion.
+
+Aim for clarity, practicality, and adherence to the project's core principles."""