lionagi 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

lionagi/core/session/session.py

@@ -1,985 +1,311 @@
-from collections import deque
-from typing import Tuple
-
-from pathlib import Path
-from lionagi.libs import BaseService, convert, dataframe
-
-from lionagi.core.generic import DataLogger
-from lionagi.core.tool import ToolManager, Tool, TOOL_TYPE
+"""
+Copyright 2024 HaiyangLi
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+from lionagi.core.collections import (
+    Pile,
+    Progression,
+    progression,
+    pile,
+    iModel,
+)
+from lionagi.core.message import System
+from typing import Any
+from lionagi.core.action.tool_manager import ToolManager
+
+from lionagi.libs import SysUtil
+from lionagi.core.session.branch import Branch
+from lionagi.core.collections import pile, Pile, Exchange
+from lionagi.core.collections.abc import get_lion_id
+from lionagi.core.collections.util import to_list_type
 from lionagi.core.mail.mail_manager import MailManager
-from lionagi.core.messages.schema import System, Instruction
-from lionagi.core.branch.branch import Branch
-from lionagi.core.flow.polyflow import PolyChat
 
 
 class Session:
     """
-    Represents a session for managing conversations and branches.
-
-    A `Session` encapsulates the state and behavior for managing conversations and their branches.
-    It provides functionality for initializing and managing conversation sessions, including setting up default
-    branches, configuring language learning models, managing tools, and handling session data logging.
+    A session for managing branches, mail transfer, and interactions with a model.
 
     Attributes:
-            branches (dict[str, Branch]): A dictionary of branch instances associated with the session.
-            service (BaseService]): The external service instance associated with the | Nonesession.
-            mail_manager (BranchManager): The manager for handling branches within the session.
-            datalogger (Optional[Any]): The datalogger instance for session data logging.
+        ln_id (str): The unique identifier for the session.
+        timestamp (str): The timestamp when the session was created.
+        system (System): The default system message for the session.
+        system_sender (str): The sender of the system message.
+        branches (Pile[Branch]): The pile of branches in the session.
+        mail_transfer (Exchange): The exchange for managing mail transfer.
+        mail_manager (MailManager): The manager for handling mail.
+        imodel (iModel): The model associated with the session.
+        user (str): The user associated with the session.
+        default_branch (Branch): The default branch of the session.
     """
 
     def __init__(
         self,
-        system: dict | list | System | None = None,
-        sender: str | None = None,
-        llmconfig: dict[str, str | int | dict] | None = None,
-        service: BaseService | None = None,
-        branches: dict[str, Branch] | None = None,
-        default_branch: Branch | None = None,
-        default_branch_name: str | None = None,
-        tools: TOOL_TYPE | None = None,
-        # instruction_sets: Optional[List[Instruction]] = None,
-        tool_manager: ToolManager | None = None,
-        messages: dataframe.ln_DataFrame | None = None,
-        datalogger: None | DataLogger = None,
-        persist_path: Path | str | None = None,
+        system=None,  # the default system message for the session
+        branches: Any | None = None,
+        system_sender: str | None = None,
+        user: str | None = None,
+        imodel=None,
+        tools=None,
     ):
-        """Initialize a new session with optional configuration for managing conversations.
-
-        Args:
-                system (Optional[Union[str, System]]): The system message.
-                sender (str | None): the default sender name for default branch
-                llmconfig (dict[str, Any] | None): Configuration for language learning models.
-                service (BaseService]): External service  | Nonenstance.
-                branches (dict[str, Branch] | None): dictionary of branch instances.
-                default_branch (Branch | None): The default branch for the session.
-                default_branch_name (str | None): The name of the default branch.
-                tools (TOOL_TYPE | None): List of tools available for the session.
-                instruction_sets (Optional[List[Instruction]]): List of instruction sets.
-                tool_manager (Optional[Any]): Manager for handling tools.
-                messages (Optional[List[dict[str, Any]]]): Initial list of messages.
-                datalogger (Optional[Any]): Logger instance for the session.
-                persist_path (str | None): Directory path for saving session data.
-
-        Examples:
-                >>> session = Session(system="you are a helpful assistant", sender="researcher")
-        """
-        self.branches = branches if isinstance(branches, dict) else {}
-        self.service = service
-        self.setup_default_branch(
-            system=system,
-            sender=sender,
-            default_branch=default_branch,
-            default_branch_name=default_branch_name,
-            messages=messages,
-            # instruction_sets=instruction_sets,
-            tool_manager=tool_manager,
-            service=service,
-            llmconfig=llmconfig,
-            tools=tools,
-            persist_path=persist_path,
-            datalogger=datalogger,
-        )
-        self.mail_manager = MailManager(self.branches)
-        self.datalogger = self.default_branch.datalogger
-        for key, branch in self.branches.items():
-            branch.name = key
-
-    # --- default branch methods ---- #
-
-    @property
-    def messages(self):
-        return self.default_branch.messages
-
-    @property
-    def messages_describe(self):
-        """
-        Provides a descriptive summary of all messages in the branch.
-
-        Returns:
-                dict[str, Any]: A dictionary containing summaries of messages by role and sender, total message count,
-                instruction sets, registered tools, and message details.
-
-        Examples:
-                >>> session.messages_describe
-                {'total_messages': 100, 'by_sender': {'User123': 60, 'Bot': 40}}
-        """
-        return self.default_branch.messages_describe
-
-    @property
-    def has_tools(self) -> bool:
-        """
-        Checks if there are any tools registered in the tool manager.
-
-        Returns:
-                bool: True if there are tools registered, False otherwise.
-
-        Examples:
-                >>> session.has_tools
-                True
-        """
-        return self.default_branch.has_tools
-
-    @property
-    def last_message(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves the last message from the conversation.
-
-        Returns:
-                pd.Series: The last message as a pandas Series.
-        """
-        return self.default_branch.last_message
-
-    @property
-    def first_system(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves the first system message from the conversation.
-
-        Returns:
-                pd.Series: The first system message as a pandas Series.
-        """
-        return self.default_branch.first_system
-
-    @property
-    def last_response(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves the last response message from the conversation.
-
-        Returns:
-                pd.Series: The last response message as a pandas Series.
-        """
-        return self.default_branch.last_response
-
-    @property
-    def last_response_content(self) -> dict:
-        """
-        Retrieves the content of the last response message from the conversation.
-
-        Returns:
-                dict: The content of the last response message as a dictionary
-        """
-        return self.default_branch.last_response_content
-
-    @property
-    def tool_request(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves all tool request messages from the conversation.
-
-        Returns:
-                dataframe.ln_DataFrame: A DataFrame containing all tool request messages.
-        """
-        return self.default_branch.tool_request
-
-    @property
-    def tool_response(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves all tool response messages from the conversation.
-
-        Returns:
-                dataframe.ln_DataFrame: A DataFrame containing all tool response messages.
-        """
-        return self.default_branch.tool_response
-
-    @property
-    def responses(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves all response messages from the conversation.
-
-        Returns:
-                dataframe.ln_DataFrame: A DataFrame containing all response messages.
-        """
-        return self.default_branch.responses
-
-    @property
-    def assistant_responses(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves all assistant responses from the conversation, excluding tool requests and responses.
-
-        Returns:
-                dataframe.ln_DataFrame: A DataFrame containing assistant responses excluding tool requests and responses.
-        """
-        return self.default_branch.assistant_responses
-
-    @property
-    def info(self) -> dict[str, int]:
-        """
-        Get a summary of the conversation messages categorized by role.
-
-        Returns:
-                dict[str, int]: A dictionary with keys as message roles and values as counts.
-        """
-
-        return self.default_branch.info
-
-    @property
-    def sender_info(self) -> dict[str, int]:
-        """
-        Provides a descriptive summary of the conversation, including total message count and summary by sender.
-
-        Returns:
-                dict[str, Any]: A dictionary containing the total number of messages and a summary categorized by sender.
-        """
-        return self.default_branch.sender_info
-
-    def register_tools(self, tools):
-        self.default_branch.register_tools(tools)
-
-    @classmethod
-    def from_csv(
-        cls,
-        filepath: Path | str,
-        system: dict | list | System | None = None,
-        sender: str | None = None,
-        llmconfig: dict[str, str | int | dict] | None = None,
-        service: BaseService = None,
-        default_branch_name: str = "main",
-        tools: TOOL_TYPE = False,  # instruction_sets=None,
-        tool_manager=None,
-        **kwargs,
-    ) -> "Session":
-        """
-        Creates a Session instance from a CSV file containing messages.
-
-        Args:
-                filepath (str): Path to the CSV file.
-                name (str | None): Name of the branch, default is None.
-                instruction_sets (Optional[dict[str, InstructionSet]]): Instruction sets, default is None.
-                tool_manager (Optional[ToolManager]): Tool manager for the branch, default is None.
-                service (BaseService]): External service for the branch, default | Noneis None.
-                llmconfig (Optional[dict]): Configuration for language learning models, default is None.
-                tools (TOOL_TYPE | None): Initial list of tools to register, default is None.
-                **kwargs: Additional keyword arguments for pd.read_csv().
-
-        Returns:
-                Branch: A new Branch instance created from the CSV data.
-
-        Examples:
-                >>> branch = Branch.from_csv("path/to/messages.csv", name="ImportedBranch")
-        """
-        df = dataframe.read_csv(filepath, **kwargs)
-
-        return cls(
-            system=system,
-            sender=sender,
-            llmconfig=llmconfig,
-            service=service,
-            default_branch_name=default_branch_name,
-            tools=tools,
-            tool_manager=tool_manager,
-            messages=df,
-            **kwargs,
-        )
+        self.ln_id = SysUtil.create_id()
+        self.timestamp = SysUtil.get_timestamp(sep=None)[:-6]
+        system = system or "You are a helpful assistant, let's think step by step"
+        self.system = System(system=system, sender=system_sender)
+        self.system_sender = system_sender
+        self.branches: Pile[Branch] = self._validate_branches(branches)
+        self.mail_transfer = Exchange()
+        self.mail_manager = MailManager([self.mail_transfer])
+        self.imodel = imodel or iModel()
+        self.user = user
+        self.default_branch = None
+        if self.branches.size() == 0:
+            self.new_branch(system=self.system.clone())
+        else:
+            self.default_branch = self.branches[0]
+        if tools:
+            self.default_branch.tool_manager.register_tools(tools)
 
-    @classmethod
-    def from_json(
-        cls,
-        filepath: Path | str,
-        system: dict | list | System | None = None,
-        sender: str | None = None,
-        llmconfig: dict[str, str | int | dict] | None = None,
-        service: BaseService = None,
-        default_branch_name: str = "main",
-        tools: TOOL_TYPE = False,  # instruction_sets=None,
-        tool_manager=None,
-        **kwargs,
-    ) -> "Session":
+    def _validate_branches(self, value):
         """
-        Creates a Branch instance from a JSON file containing messages.
+        Validates and converts the branches input to a Pile of Branch objects.
 
         Args:
-                filepath (str): Path to the JSON file.
-                name (str | None): Name of the branch, default is None.
-                instruction_sets (Optional[dict[str, InstructionSet]]): Instruction sets, default is None.
-                tool_manager (Optional[ToolManager]): Tool manager for the branch, default is None.
-                service (BaseService]): External service for the branch, default | Noneis None.
-                llmconfig (Optional[dict]): Configuration for language learning models, default is None.
-                **kwargs: Additional keyword arguments for pd.read_json().
+            value (Any): The input value to validate and convert.
 
         Returns:
-                Branch: A new Branch instance created from the JSON data.
-
-        Examples:
-                >>> branch = Branch.from_json_string("path/to/messages.json", name="JSONBranch")
-        """
-        df = dataframe.read_json(filepath, **kwargs)
-        return cls(
-            system=system,
-            sender=sender,
-            llmconfig=llmconfig,
-            service=service,
-            default_branch_name=default_branch_name,
-            tools=tools,  # instruction_sets=instruction_sets,
-            tool_manager=tool_manager,
-            messages=df,
-            **kwargs,
-        )
-
-    def to_csv_file(
-        self,
-        filename: str = "messages.csv",
-        dir_exist_ok: bool = True,
-        timestamp: bool = True,
-        time_prefix: bool = False,
-        verbose: bool = True,
-        clear: bool = True,
-        **kwargs,
-    ):
-        """
-        Saves the branch's messages to a CSV file.
-
-        Args:
-                filename (str): The name of the output CSV file, default is 'messages.csv'.
-                dir_exist_ok (bool): If True, does not raise an error if the directory already exists, default is True.
-                timestamp (bool): If True, appends a timestamp to the filename, default is True.
-                time_prefix (bool): If True, adds a timestamp prefix to the filename, default is False.
-                verbose (bool): If True, prints a message upon successful save, default is True.
-                clear (bool): If True, clears the messages after saving, default is True.
-                **kwargs: Additional keyword arguments for DataFrame.to_csv().
-
-        Examples:
-                >>> branch.to_csv_file("exported_messages.csv")
-                >>> branch.to_csv_file("timed_export.csv", timestamp=True, time_prefix=True)
-        """
-        for name, branch in self.branches.items():
-            f_name = f"{name}_{filename}"
-            branch.to_csv_file(
-                filename=f_name,
-                dir_exist_ok=dir_exist_ok,
-                timestamp=timestamp,
-                time_prefix=time_prefix,
-                verbose=verbose,
-                clear=clear,
-                **kwargs,
-            )
-
-    def to_json_file(
-        self,
-        filename: str = "messages.json",
-        dir_exist_ok: bool = False,
-        timestamp: bool = True,
-        time_prefix: bool = False,
-        verbose: bool = True,
-        clear: bool = True,
-        **kwargs,
-    ):
-        """
-        Saves the branch's messages to a JSON file.
-
-        Args:
-                filename (str): The name of the output JSON file, default is 'messages.json'.
-                dir_exist_ok (bool): If True, does not raise an error if the directory already exists, default is True.
-                timestamp (bool): If True, appends a timestamp to the filename, default is True.
-                time_prefix (bool): If True, adds a timestamp prefix to the filename, default is False.
-                verbose (bool): If True, prints a message upon successful save, default is True.
-                clear (bool): If True, clears the messages after saving, default is True.
-                **kwargs: Additional keyword arguments for DataFrame.to_json().
-
-        Examples:
-                >>> branch.to_json_file("exported_messages.json")
-                >>> branch.to_json_file("timed_export.json", timestamp=True, time_prefix=True)
-        """
-
-        for name, branch in self.branches.items():
-            f_name = f"{name}_{filename}"
-            branch.to_json_file(
-                filename=f_name,
-                dir_exist_ok=dir_exist_ok,
-                timestamp=timestamp,
-                time_prefix=time_prefix,
-                verbose=verbose,
-                clear=clear,
-                **kwargs,
-            )
-
-    def log_to_csv(
-        self,
-        filename: str = "log.csv",
-        dir_exist_ok: bool = True,
-        timestamp: bool = True,
-        time_prefix: bool = False,
-        verbose: bool = True,
-        clear: bool = True,
-        **kwargs,
-    ):
-        """
-        Saves the branch's log data to a CSV file.
-
-        This method is designed to export log data, potentially including operations and intertools,
-        to a CSV file for analysis or record-keeping.
+            Pile[Branch]: A pile of validated branches.
 
-        Args:
-                filename (str): The name of the output CSV file. Defaults to 'log.csv'.
-                dir_exist_ok (bool): If True, will not raise an error if the directory already exists. Defaults to True.
-                timestamp (bool): If True, appends a timestamp to the filename for uniqueness. Defaults to True.
-                time_prefix (bool): If True, adds a timestamp prefix to the filename. Defaults to False.
-                verbose (bool): If True, prints a success message upon completion. Defaults to True.
-                clear (bool): If True, clears the log after saving. Defaults to True.
-                **kwargs: Additional keyword arguments for `DataFrame.to_csv()`.
-
-        Examples:
-                >>> branch.log_to_csv("branch_log.csv")
-                >>> branch.log_to_csv("detailed_branch_log.csv", timestamp=True, verbose=True)
-        """
-        for name, branch in self.branches.items():
-            f_name = f"{name}_{filename}"
-            branch.log_to_csv(
-                filename=f_name,
-                dir_exist_ok=dir_exist_ok,
-                timestamp=timestamp,
-                time_prefix=time_prefix,
-                verbose=verbose,
-                clear=clear,
-                **kwargs,
-            )
-
-    def log_to_json(
-        self,
-        filename: str = "log.json",
-        dir_exist_ok: bool = True,
-        timestamp: bool = True,
-        time_prefix: bool = False,
-        verbose: bool = True,
-        clear: bool = True,
-        **kwargs,
-    ):
-        """
-        Saves the branch's log data to a JSON file.
-
-        Useful for exporting log data in JSON format, allowing for easy integration with web applications
-        and services that consume JSON.
-
-        Args:
-                filename (str): The name of the output JSON file. Defaults to 'log.json'.
-                dir_exist_ok (bool): If directory existence should not raise an error. Defaults to True.
-                timestamp (bool): If True, appends a timestamp to the filename. Defaults to True.
-                time_prefix (bool): If True, adds a timestamp prefix to the filename. Defaults to False.
-                verbose (bool): If True, prints a success message upon completion. Defaults to True.
-                clear (bool): If True, clears the log after saving. Defaults to True.
-                **kwargs: Additional keyword arguments for `DataFrame.to_json()`.
-
-        Examples:
-                >>> branch.log_to_json("branch_log.json")
-                >>> branch.log_to_json("detailed_branch_log.json", verbose=True, timestamp=True)
-        """
-        for name, branch in self.branches.items():
-            f_name = f"{name}_{filename}"
-            branch.log_to_json(
-                filename=f_name,
-                dir_exist_ok=dir_exist_ok,
-                timestamp=timestamp,
-                time_prefix=time_prefix,
-                verbose=verbose,
-                clear=clear,
-                **kwargs,
-            )
-
-    @property
-    def all_messages(self) -> dataframe.ln_DataFrame:
-        """
-        return all messages across branches
+        Raises:
+            ValueError: If the input value contains non-Branch objects.
         """
-        dfs = deque()
-        for _, v in self.branches.items():
-            dfs.append(convert.to_df(v.messages))
-        return convert.to_df(convert.to_list(dfs, flatten=True, dropna=True))
+        if isinstance(value, Pile):
+            for branch in value:
+                if not isinstance(branch, Branch):
+                    raise ValueError("The branches pile contains non-Branch object")
+            return value
+        else:
+            try:
+                value = pile(items=value, item_type=Branch)
+                return value
+            except Exception as e:
+                raise ValueError(f"Invalid branches value. Error:{e}")
 
-    # ----- chatflow ----#
-    async def call_chatcompletion(
+    # ---- branch manipulation ---- #
+    def new_branch(
         self,
-        branch: Branch | str | None = None,
-        sender: str | None = None,
-        with_sender=False,
-        **kwargs,
+        system: System | None = None,
+        system_sender: str | None = None,
+        user: str | None = None,
+        messages: Pile = None,
+        progress: Progression = None,
+        tool_manager: ToolManager = None,
+        tools: Any = None,
+        imodel=None,
     ):
         """
-        Asynchronously calls the chat completion service with the current message queue.
-
-        This method prepares the messages for chat completion, sends the request to the configured service, and handles the response. The method supports additional keyword arguments that are passed directly to the service.
+        Creates a new branch and adds it to the session.
 
         Args:
-                sender (str | None): The name of the sender to be included in the chat completion request. Defaults to None.
-                with_sender (bool): If True, includes the sender's name in the messages. Defaults to False.
-                **kwargs: Arbitrary keyword arguments passed directly to the chat completion service.
-
-        Examples:
-                >>> await branch.call_chatcompletion()
-        """
-        branch = self.get_branch(branch)
-        await branch.call_chatcompletion(
-            sender=sender,
-            with_sender=with_sender,
-            **kwargs,
-        )
-
-    async def chat(
-        self,
-        instruction: dict | list | Instruction | str,
-        branch: Branch | str | None = None,
-        context: dict | list | str = None,
-        sender: str | None = None,
-        system: dict | list | System | None = None,
-        tools: TOOL_TYPE = False,
-        out: bool = True,
-        invoke: bool = True,
-        output_fields=None,
-        **kwargs,
-    ) -> str | None:
-        """
-        a chat conversation with LLM, processing instructions and system messages, optionally invoking tools.
+            system (System, optional): The system message for the branch.
+            system_sender (str, optional): The sender of the system message.
+            user (str, optional): The user associated with the branch.
+            messages (Pile, optional): The pile of messages for the branch.
+            progress (Progression, optional): The progression of messages.
+            tool_manager (ToolManager, optional): The tool manager for the branch.
+            tools (Any, optional): The tools to register with the tool manager.
+            imodel (iModel, optional): The model associated with the branch.
 
-        Args:
-                branch: The Branch instance to perform chat operations.
-                instruction (dict | list | Instruction | str): The instruction for the chat.
-                context (Optional[Any]): Additional context for the chat.
-                sender (str | None): The sender of the chat message.
-                system (Optional[Union[System, str, dict[str, Any]]]): System message to be processed.
-                tools (Union[bool, Tool, List[Tool], str, List[str]]): Specifies tools to be invoked.
-                out (bool): If True, outputs the chat response.
-                invoke (bool): If True, invokes tools as part of the chat.
-                **kwargs: Arbitrary keyword arguments for chat completion.
-
-        Examples:
-                >>> await ChatFlow.chat(branch, "Ask about user preferences")
+        Returns:
+            Branch: The created branch.
         """
-
-        branch = self.get_branch(branch)
-        return await branch.chat(
-            instruction=instruction,
-            context=context,
-            sender=sender,
+        if system is None:
+            system = self.system.clone()
+            system.sender = self.ln_id
+            system_sender = self.ln_id
+        branch = Branch(
             system=system,
+            system_sender=system_sender,
+            user=user,
+            messages=messages,
+            progress=progress,
+            tool_manager=tool_manager,
             tools=tools,
-            out=out,
-            invoke=invoke,
-            output_fields=output_fields,
-            **kwargs,
+            imodel=imodel or self.imodel,
         )
+        self.branches.append(branch)
+        self.mail_manager.add_sources(branch)
+        if self.default_branch is None:
+            self.default_branch = branch
+        return branch
 
-    async def ReAct(
-        self,
-        instruction: dict | list | Instruction | str,
-        branch: Branch | str | None = None,
-        context=None,
-        sender=None,
-        system=None,
-        tools=None,
-        auto=False,
-        num_rounds: int = 1,
-        reason_prompt=None,
-        action_prompt=None,
-        output_prompt=None,
-        **kwargs,
-    ):
+    def delete_branch(self, branch):
         """
-        Performs a reason-tool cycle with optional tool invocation over multiple rounds.
+        Deletes a branch from the session.
 
         Args:
-                branch: The Branch instance to perform ReAct operations.
-                instruction (dict | list | Instruction | str): Initial instruction for the cycle.
-                context: Context relevant to the instruction.
-                sender (str | None): Identifier for the message sender.
-                system: Initial system message or configuration.
-                tools: Tools to be registered or used during the cycle.
-                num_rounds (int): Number of reason-tool cycles to perform.
-                **kwargs: Additional keyword arguments for customization.
-
-        Examples:
-                >>> await ChatFlow.ReAct(branch, "Analyze user feedback", num_rounds=2)
+            branch (Branch | str): The branch or its ID to delete.
         """
-        branch = self.get_branch(branch)
+        branch_id = get_lion_id(branch)
+        self.branches.pop(branch_id)
+        self.mail_manager.delete_source(branch_id)
 
-        return await branch.ReAct(
-            instruction=instruction,
-            context=context,
-            sender=sender,
-            system=system,
-            tools=tools,
-            auto=auto,
-            num_rounds=num_rounds,
-            reason_prompt=reason_prompt,
-            action_prompt=action_prompt,
-            output_prompt=output_prompt,
-            **kwargs,
-        )
+        if self.default_branch == branch:
+            if self.branches.size() == 0:
+                self.default_branch = None
+            else:
+                self.default_branch = self.branches[0]
 
-    async def followup(
-        self,
-        instruction: dict | list | Instruction | str,
-        branch=None,
-        context=None,
-        sender=None,
-        system=None,
-        tools=None,
-        max_followup: int = 1,
-        auto=False,
-        followup_prompt=None,
-        output_prompt=None,
-        out=True,
-        **kwargs,
-    ):
+    def split_branch(self, branch):
         """
-        Automatically performs follow-up tools based on chat intertools and tool invocations.
+        Splits a branch, creating a new branch with the same messages and tools.
 
         Args:
-                branch: The Branch instance to perform follow-up operations.
-                instruction (dict | list | Instruction | str): The initial instruction for follow-up.
-                context: Context relevant to the instruction.
-                sender (str | None): Identifier for the message sender.
-                system: Initial system message or configuration.
-                tools: Specifies tools to be considered for follow-up tools.
-                max_followup (int): Maximum number of follow-up chats allowed.
-                out (bool): If True, outputs the result of the follow-up tool.
-                **kwargs: Additional keyword arguments for follow-up customization.
-
-        Examples:
-                >>> await ChatFlow.auto_followup(branch, "Finalize report", max_followup=2)
-        """
-        branch = self.get_branch(branch)
-        return await branch.followup(
-            instruction=instruction,
-            context=context,
-            sender=sender,
-            system=system,
-            tools=tools,
-            max_followup=max_followup,
-            auto=auto,
-            followup_prompt=followup_prompt,
-            output_prompt=output_prompt,
-            out=out,
-            **kwargs,
-        )
+            branch (Branch | str): The branch or its ID to split.
 
-    async def parallel_chat(
-        self,
-        instruction: Instruction | str,
-        num_instances=1,
-        context=None,
-        sender=None,
-        branch_system=None,
-        messages=None,
-        tools=False,
-        out=True,
-        invoke: bool = True,
-        output_fields=None,
-        persist_path=None,
-        branch_config=None,
-        explode=False,
-        include_mapping=False,
-        **kwargs,
-    ):
-        """
-        parallel chat
+        Returns:
+            Branch: The newly created branch.
         """
-
-        if branch_config is None:
-            branch_config = {}
-        flow = PolyChat(self)
-
-        return await flow.parallel_chat(
-            instruction,
-            num_instances=num_instances,
-            context=context,
-            sender=sender,
-            branch_system=branch_system,
-            messages=messages,
-            tools=tools,
-            out=out,
-            invoke=invoke,
-            output_fields=output_fields,
-            persist_path=persist_path,
-            branch_config=branch_config,
-            explode=explode,
-            include_mapping=include_mapping,
-            **kwargs,
+        branch = self.branches[branch]
+        system = branch.system.clone() if branch.system else None
+        if system:
+            system.sender = branch.ln_id
+        progress = progression()
+        messages = pile()
+
+        for id_ in branch.progress:
+            clone_message = branch.messages[id_].clone()
+            progress.append(clone_message.ln_id)
+            messages.append(clone_message)
+
+        tools = (
+            list(branch.tool_manager.registry.values())
+            if branch.tool_manager.registry
+            else None
         )
-
-    # ---- branch manipulation ---- #
-    def new_branch(
-        self,
-        branch_name: str,
-        system: dict | list | System | None = None,
-        sender: str | None = None,
-        messages: dataframe.ln_DataFrame | None = None,
-        tool_manager=None,
-        service=None,
-        llmconfig=None,
-        tools: TOOL_TYPE = False,
-    ) -> None:
-        """Create a new branch with the specified configurations.
-
-        Args:
-                branch_name (str | None): Name of the new branch.
-                system (Optional[Union[System, str]]): System or context identifier for the new branch.
-                sender (str | None): Default sender identifier for the new branch.
-                messages (Optional[dataframe.ln_DataFrame]): Initial set of messages for the new branch.
-                instruction_sets (Optional[Any]): Instruction sets for the new branch.
-                tool_manager (Optional[Any]): Tool manager for handling tools in the new branch.
-                service (BaseService]): External service instance for the ne | None branch.
-                llmconfig (dict[str, Any] | None): Configuration for language learning models in the new branch.
-                tools (TOOL_TYPE | None): List of tools available for the new branch.
-
-        Raises:
-                ValueError: If the branch name already exists.
-
-        Examples:
-                >>> session.new_branch("new_branch_name")
-        """
-        if branch_name in self.branches.keys():
-            raise ValueError(
-                f"Invalid new branch name {branch_name}. Branch already existed."
-            )
-        if isinstance(tools, Tool):
-            tools = [tools]
-        new_branch = Branch(
-            name=branch_name,
+        branch_clone = Branch(
+            system=system,
+            system_sender=branch.ln_id,
+            user=branch.user,
+            progress=progress,
             messages=messages,
-            tool_manager=tool_manager,
-            service=service,
-            llmconfig=llmconfig,
             tools=tools,
         )
-        if system:
-            new_branch.add_message(system=system, sender=sender)
-
-        self.branches[branch_name] = new_branch
-        self.mail_manager.sources[branch_name] = new_branch
-        self.mail_manager.mails[branch_name] = {}
-
-    def get_branch(
-        self, branch: Branch | str | None = None, get_name: bool = False
-    ) -> Branch | Tuple[Branch, str]:
-        """
-        Retrieve a branch by name or instance.
-
-        Args:
-                branch (Optional[Branch | str]): The branch name or instance to retrieve.
-                get_name (bool): If True, returns a tuple of the branch instance and its name.
-
-        Returns:
-                Union[Branch, Tuple[Branch, str]]: The branch instance or a tuple of the branch instance and its name.
-
-        Raises:
-                ValueError: If the branch name does not exist or the branch input is invalid.
+        for message in branch_clone.messages:
+            message.sender = branch.ln_id
+            message.recipient = branch_clone.ln_id
+        self.branches.append(branch_clone)
+        self.mail_manager.add_sources(branch_clone)
+        return branch_clone
 
-        Examples:
-                >>> branch_instance = session.get_branch("existing_branch_name")
-                >>> branch_instance, branch_name = session.get_branch("existing_branch_name", get_name=True)
+    def change_default_branch(self, branch):
         """
-        if isinstance(branch, str):
-            if branch not in self.branches.keys():
-                raise ValueError(f"Invalid branch name {branch}. Not exist.")
-            return (
-                (self.branches[branch], branch) if get_name else self.branches[branch]
-            )
-        elif isinstance(branch, Branch) and branch in self.branches.values():
-            return (branch, branch.name) if get_name else branch
-        elif branch is None:
-            if get_name:
-                return self.default_branch, self.default_branch_name
-            return self.default_branch
-
-        else:
-            raise ValueError(f"Invalid branch input {branch}.")
-
-    def change_default_branch(self, branch: str | Branch) -> None:
-        """Change the default branch of the session.
-
-        Args:
-                branch (str | Branch): The branch name or instance to set as the new default.
-
-        Examples:
-                >>> session.change_default_branch("new_default_branch")
-        """
-        branch_, name_ = self.get_branch(branch, get_name=True)
-        self.default_branch = branch_
-        self.default_branch_name = name_
-
-    def delete_branch(self, branch: Branch | str, verbose: bool = True) -> bool:
-        """Delete a branch from the session.
+        Changes the default branch of the session.
 
         Args:
-                branch (Branch | str): The branch name or instance to delete.
-                verbose (bool): If True, prints a message upon deletion.
-
-        Returns:
-                bool: True if the branch was successfully deleted.
-
-        Raises:
-                ValueError: If attempting to delete the current default branch.
-
-        Examples:
-                >>> session.delete_branch("branch_to_delete")
+            branch (Branch | str): The branch or its ID to set as the default.
         """
-        _, branch_name = self.get_branch(branch, get_name=True)
-
-        if branch_name == self.default_branch_name:
-            raise ValueError(
-                f"{branch_name} is the current default branch, please switch to another branch before delete it."
-            )
-        self.branches.pop(branch_name)
-        # self.mail_manager.sources.pop(branch_name)
-        self.mail_manager.mails.pop(branch_name)
-        if verbose:
-            print(f"Branch {branch_name} is deleted.")
-        return True
-
-    def merge_branch(
-        self,
-        from_: str | Branch,
-        to_branch: str | Branch,
-        update: bool = True,
-        del_: bool = False,
-    ) -> None:
-        """Merge messages and settings from one branch to another.
-
-        Args:
-                from_ (str | Branch): The source branch name or instance.
-                to_branch (str | Branch): The target branch name or instance where the merge will happen.
-                update (bool): If True, updates the target branch with the source branch's settings.
-                del_ (bool): If True, deletes the source branch after merging.
+        branch = self.branches[branch]
+        self.default_branch = branch
 
-        Examples:
-                >>> session.merge_branch("source_branch", "target_branch", del_=True)
-        """
-        from_ = self.get_branch(branch=from_)
-        to_branch, to_name = self.get_branch(branch=to_branch, get_name=True)
-        to_branch.merge_branch(from_, update=update)
-
-        if del_:
-            if from_ == self.default_branch:
-                self.default_branch_name = to_name
-                self.default_branch = to_branch
-            self.delete_branch(from_, verbose=False)
-
-    def take_branch(self, branch):
-        self.branches[branch.branch_name] = branch
-        self.mail_manager.sources[branch.id_] = branch
-        self.mail_manager.mails[branch.id_] = {}
-
-    def collect(self, from_: str | Branch | list[str | Branch] | None = None):
+    def collect(self, from_: Branch | str | Pile[Branch] | None = None):
         """
-        Collects requests from specified branches or from all branches if none are specified.
-
-        This method is intended to aggregate data or requests from one or more branches for processing or analysis.
+        Collects mail from specified branches.
 
         Args:
-                from_ (Optional[Union[str, Branch, List[str | Branch]]]): The branch(es) from which to collect requests.
-                        Can be a single branch name, a single branch instance, a list of branch names, a list of branch instances, or None.
-                        If None, requests are collected from all branches.
-
-        Examples:
-                >>> session.collect("branch_name")
-                >>> session.collect([branch_instance_1, "branch_name_2"])
-                >>> session.collect()  # Collects from all branches
+            from_ (Branch | str | Pile[Branch], optional): The branches to collect mail from.
+                If None, collects mail from all branches.
         """
         if from_ is None:
-            for branch in self.branches.values():
-                self.mail_manager.collect(branch.id_)
+            self.mail_manager.collect_all()
         else:
-            if not isinstance(from_, list):
-                from_ = [from_]
-            for branch in from_:
-                if isinstance(branch, str):
-                    branch = self.branches[branch]
-                    self.mail_manager.collect(branch.id_)
-                elif isinstance(branch, Branch):
-                    self.mail_manager.collect(branch.id_)
-
-    def send(self, to_: str | Branch | list[str | Branch] | None = None):
-        """
-        Sends requests or data to specified branches or to all branches if none are specified.
+            try:
+                sources = to_list_type(from_)
+                for source in sources:
+                    self.mail_manager.collect(get_lion_id(source))
+            except Exception as e:
+                raise ValueError(f"Failed to collect mail. Error: {e}")
 
-        This method facilitates the distribution of data or requests to one or more branches, potentially for further tool or processing.
+    def send(self, to_: Branch | str | Pile[Branch] | None = None):
+        """
+        Sends mail to specified branches.
 
         Args:
-                to_ (Optional[Union[str, Branch, List[str | Branch]]]): The target branch(es) to which to send requests.
-                        Can be a single branch name, a single branch instance, a list of branch names, a list of branch instances, or None.
-                        If None, requests are sent to all branches.
-
-        Examples:
-                >>> session.send("target_branch")
-                >>> session.send([branch_instance_1, "target_branch_2"])
-                >>> session.send()  # Sends to all branches
+            to_ (Branch | str | Pile[Branch], optional): The branches to send mail to.
+                If None, sends mail to all branches.
         """
         if to_ is None:
-            for branch in self.branches.values():
-                self.mail_manager.send(branch.id_)
+            self.mail_manager.send_all()
         else:
-            if not isinstance(to_, list):
-                to_ = [to_]
-            for branch in to_:
-                if isinstance(branch, str):
-                    branch = self.branches[branch]
-                    self.mail_manager.send(branch.id_)
-                if isinstance(branch, Branch):
-                    self.mail_manager.send(branch.id_)
+            try:
+                sources = to_list_type(to_)
+                for source in sources:
+                    self.mail_manager.send(get_lion_id(source))
+            except Exception as e:
+                raise ValueError(f"Failed to send mail. Error: {e}")
 
     def collect_send_all(self, receive_all=False):
         """
-        Collects and sends requests across all branches, optionally invoking a receive operation on each branch.
-
-        This method is a convenience function for performing a full cycle of collect and send operations across all branches,
-        useful in scenarios where data or requests need to be aggregated and then distributed uniformly.
+        Collects and sends mail for all branches, optionally receiving all mail.
 
         Args:
-                receive_all (bool): If True, triggers a `receive_all` method on each branch after sending requests,
-                        which can be used to process or acknowledge the received data.
-
-        Examples:
-                >>> session.collect_send_all()
-                >>> session.collect_send_all(receive_all=True)
+            receive_all (bool, optional): Whether to receive all mail for all branches.
         """
         self.collect()
         self.send()
         if receive_all:
-            for branch in self.branches.values():
+            for branch in self.branches:
                 branch.receive_all()
 
-    def setup_default_branch(self, **kwargs):
-        self._setup_default_branch(**kwargs)
-        self._verify_default_branch()
+    async def chat(self, *args, branch=None, **kwargs):
+        """
+        Initiates a chat interaction with a branch.
 
-    def _verify_default_branch(self):
-        if self.branches:
-            if self.default_branch_name not in self.branches.keys():
-                raise ValueError("default branch name is not in imported branches")
-            if self.default_branch is not self.branches[self.default_branch_name]:
-                raise ValueError(
-                    f"default branch does not match Branch object under {self.default_branch_name}"
-                )
+        Args:
+            *args: Positional arguments to pass to the chat method.
+            branch (Branch, optional): The branch to chat with. Defaults to the default branch.
+            **kwargs: Keyword arguments to pass to the chat method.
 
-        if not self.branches:
-            self.branches[self.default_branch_name] = self.default_branch
+        Returns:
+            Any: The result of the chat interaction.
 
-    def _setup_default_branch(
-        self,
-        system,
-        sender,
-        default_branch,
-        default_branch_name,
-        messages,  # instruction_sets,
-        tool_manager,
-        service,
-        llmconfig,
-        tools,
-        persist_path,
-        datalogger,
-    ):
+        Raises:
+            ValueError: If the specified branch is not found in the session branches.
+        """
+        if branch is None:
+            branch = self.default_branch
+        if branch not in self.branches:
+            raise ValueError("Branch not found in session branches")
+        return await self.branches[branch].chat(*args, **kwargs)
 
-        branch = default_branch or Branch(
-            name=default_branch_name,
-            service=service,
-            llmconfig=llmconfig,
-            tools=tools,
-            tool_manager=tool_manager,
-            # instruction_sets=instruction_sets,
-            messages=messages,
-            persist_path=persist_path,
-            datalogger=datalogger,
-        )
+    async def direct(self, *args, branch=None, **kwargs):
+        """
+        Initiates a direct interaction with a branch.
 
-        self.default_branch = branch
-        self.default_branch_name = default_branch_name or "main"
-        if system:
-            self.default_branch.add_message(system=system, sender=sender)
+        Args:
+            *args: Positional arguments to pass to the direct method.
+            branch (Branch, optional): The branch to interact with. Defaults to the default branch.
+            **kwargs: Keyword arguments to pass to the direct method.
 
-        self.llmconfig = self.default_branch.llmconfig
+        Returns:
+            Any: The result of the direct interaction.
+
+        Raises:
+            ValueError: If the specified branch is not found in the session branches.
+        """
+        if branch is None:
+            branch = self.default_branch
+        if branch not in self.branches:
+            raise ValueError("Branch not found in session branches")
+        return await self.branches[branch].direct(*args, **kwargs)