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

lionagi/core/execute/branch_executor.py

@@ -0,0 +1,296 @@
+import contextlib
+from lionagi.libs import convert, AsyncUtil, ParseUtil
+from lionagi.core.generic import ActionNode, Edge
+from lionagi.core.mail.schema import BaseMail
+from lionagi.core.messages.schema import System, Instruction
+
+from lionagi.core.branch.branch import Branch
+from lionagi.core.execute.base_executor import BaseExecutor
+
+
+class BranchExecutor(Branch, BaseExecutor):
+
+    async def forward(self) -> None:
+        """
+        Forwards the execution by processing all pending incoming mails in each branch. Depending on the category of the mail,
+        it processes starts, nodes, node lists, conditions, or ends, accordingly executing different functions.
+        """
+        for key in list(self.pending_ins.keys()):
+            while self.pending_ins[key]:
+                mail = self.pending_ins[key].popleft()
+                if mail.category == "start":
+                    self._process_start(mail)
+                elif mail.category == "node":
+                    await self._process_node(mail)
+                elif mail.category == "node_list":
+                    self._process_node_list(mail)
+                elif mail.category == "condition":
+                    self._process_condition(mail)
+                elif mail.category == "end":
+                    self._process_end(mail)
+
+    async def execute(self, refresh_time=1) -> None:
+        """
+        Executes the forward process repeatedly at specified time intervals until execution is instructed to stop.
+
+        Args:
+            refresh_time (int): The interval, in seconds, at which the forward method is called repeatedly.
+        """
+        while not self.execute_stop:
+            await self.forward()
+            await AsyncUtil.sleep(refresh_time)
+
+    async def _process_node(self, mail: BaseMail):
+        """
+        Processes a single node based on the node type specified in the mail's package. It handles different types of nodes such as System,
+        Instruction, ActionNode, and generic nodes through separate processes.
+
+        Args:
+            mail (BaseMail): The mail containing the node to be processed along with associated details.
+
+        Raises:
+            ValueError: If an invalid mail is encountered or the process encounters errors.
+        """
+        if isinstance(mail.package["package"], System):
+            self._system_process(mail.package["package"], verbose=self.verbose)
+            self.send(
+                mail.sender_id,
+                "node_id",
+                {"request_source": self.id_, "package": mail.package["package"].id_},
+            )
+
+        elif isinstance(mail.package["package"], Instruction):
+            await self._instruction_process(
+                mail.package["package"], verbose=self.verbose
+            )
+            self.send(
+                mail.sender_id,
+                "node_id",
+                {"request_source": self.id_, "package": mail.package["package"].id_},
+            )
+
+        elif isinstance(mail.package["package"], ActionNode):
+            await self._action_process(mail.package["package"], verbose=self.verbose)
+            self.send(
+                mail.sender_id,
+                "node_id",
+                {
+                    "request_source": self.id_,
+                    "package": mail.package["package"].instruction.id_,
+                },
+            )
+        else:
+            try:
+                await self._agent_process(mail.package["package"], verbose=self.verbose)
+                self.send(
+                    mail.sender_id,
+                    "node_id",
+                    {
+                        "request_source": self.id_,
+                        "package": mail.package["package"].id_,
+                    },
+                )
+            except:
+                raise ValueError(f"Invalid mail to process. Mail:{mail}")
+
+    def _process_node_list(self, mail: BaseMail):
+        """
+        Processes a list of nodes provided in the mail, but currently only sends an end signal as multiple path selection is not supported.
+
+        Args:
+            mail (BaseMail): The mail containing a list of nodes to be processed.
+
+        Raises:
+            ValueError: When trying to process multiple paths which is currently unsupported.
+        """
+        self.send(mail.sender_id, "end", {"request_source": self.id_, "package": "end"})
+        self.execute_stop = True
+        raise ValueError("Multiple path selection is currently not supported")
+
+    def _process_condition(self, mail: BaseMail):
+        """
+        Processes a condition associated with an edge based on the mail's package, setting up the result of the condition check.
+
+        Args:
+            mail (BaseMail): The mail containing the condition to be processed.
+        """
+        relationship: Edge = mail.package["package"]
+        check_result = relationship.condition(self)
+        back_mail = {
+            "from": self.id_,
+            "edge_id": mail.package["package"].id_,
+            "check_result": check_result,
+        }
+        self.send(
+            mail.sender_id,
+            "condition",
+            {"request_source": self.id_, "package": back_mail},
+        )
+
+    def _system_process(self, system: System, verbose=True, context_verbose=False):
+        """
+        Processes a system node, possibly displaying its content and context if verbose is enabled.
+
+        Args:
+            system (System): The system node to process.
+            verbose (bool): Flag to enable verbose output.
+            context_verbose (bool): Flag to enable verbose output specifically for context.
+        """
+        from lionagi.libs import SysUtil
+
+        SysUtil.check_import("IPython")
+        from IPython.display import Markdown, display
+
+        if verbose:
+            print(f"------------------Welcome: {system.sender}--------------------")
+            with contextlib.suppress(Exception):
+                system.content = ParseUtil.fuzzy_parse_json(system.content)
+            display(Markdown(f"system: {convert.to_str(system.system_info)}"))
+            if self.context and context_verbose:
+                display(Markdown(f"context: {convert.to_str(self.context)}"))
+
+        self.add_message(system=system)
+
+    async def _instruction_process(
+        self, instruction: Instruction, verbose=True, **kwargs
+    ):
+        """
+        Processes an instruction node, possibly displaying its content if verbose is enabled, and handling any additional keyword arguments.
+
+        Args:
+            instruction (Instruction): The instruction node to process.
+            verbose (bool): Flag to enable verbose output.
+            **kwargs: Additional keyword arguments that might affect how instructions are processed.
+        """
+        from lionagi.libs import SysUtil
+
+        SysUtil.check_import("IPython")
+        from IPython.display import Markdown, display
+
+        if verbose:
+            with contextlib.suppress(Exception):
+                instruction.content = ParseUtil.fuzzy_parse_json(instruction.content)
+            display(
+                Markdown(
+                    f"{instruction.sender}: {convert.to_str(instruction.instruct)}"
+                )
+            )
+
+        if self.context:
+            instruction.content.update({"context": self.context})
+            self.context_log.append(self.context)
+            self.context = None
+
+        result = await self.chat(instruction, **kwargs)
+        with contextlib.suppress(Exception):
+            result = ParseUtil.fuzzy_parse_json(result)
+            if "response" in result.keys():
+                result = result["response"]
+        if verbose and len(self.assistant_responses) != 0:
+            display(
+                Markdown(
+                    f"{self.last_assistant_response.sender}: {convert.to_str(result)}"
+                )
+            )
+            print("-----------------------------------------------------")
+
+        self.execution_responses.append(result)
+
+    async def _action_process(self, action: ActionNode, verbose=True):
+        """
+        Processes an action node, executing the defined action along with any tools specified within the node.
+
+        Args:
+            action (ActionNode): The action node to process.
+            verbose (bool): Flag to enable verbose output of the action results.
+        """
+        from lionagi.libs import SysUtil
+
+        SysUtil.check_import("IPython")
+        from IPython.display import Markdown, display
+
+        try:
+            func = getattr(self, action.action)
+        except:
+            raise ValueError(f"{action.action} is not a valid action")
+
+        if verbose:
+            display(
+                Markdown(
+                    f"{action.instruction.sender}: {convert.to_str(action.instruction.instruct)}"
+                )
+            )
+
+        if action.tools:
+            self.register_tools(action.tools)
+        if self.context:
+            result = await func(
+                action.instruction.content["instruction"],
+                context=self.context,
+                tools=action.tools,
+                **action.action_kwargs,
+            )
+            self.context = None
+        else:
+            result = await func(
+                action.instruction.content, tools=action.tools, **action.action_kwargs
+            )
+
+        if verbose and len(self.assistant_responses) != 0:
+            display(
+                Markdown(
+                    f"{self.last_assistant_response.sender}: {convert.to_str(result)}"
+                )
+            )
+            print("-----------------------------------------------------")
+
+        self.execution_responses.append(result)
+
+    async def _agent_process(self, agent, verbose=True):
+        """
+        Processes an agent.
+
+        Args:
+            agent: The agent to process.
+            verbose (bool): A flag indicating whether to provide verbose output (default: True).
+        """
+        context = list(self.messages["content"])
+        if verbose:
+            print("*****************************************************")
+        result = await agent.execute(context)
+
+        if verbose:
+            print("*****************************************************")
+
+        from pandas import DataFrame
+
+        if isinstance(result, DataFrame):
+            self.context = list(result["content"])
+        else:
+            self.context = result
+        self.execution_responses.append(result)
+
+    def _process_start(self, mail):
+        """
+        Processes a start mail.
+
+        Args:
+            mail (BaseMail): The start mail to process.
+        """
+        start_mail_content = mail.package
+        self.context = start_mail_content["context"]
+        self.send(
+            start_mail_content["structure_id"],
+            "start",
+            {"request_source": self.id_, "package": "start"},
+        )
+
+    def _process_end(self, mail: BaseMail):
+        """
+        Processes an end mail.
+
+        Args:
+            mail (BaseMail): The end mail to process.
+        """
+        self.execute_stop = True
+        self.send(mail.sender_id, "end", {"request_source": self.id_, "package": "end"})

lionagi/core/execute/instruction_map_executor.py

@@ -0,0 +1,179 @@
+import asyncio
+from pydantic import Field
+
+from lionagi.libs import AsyncUtil
+
+from lionagi.core.mail.schema import BaseMail, MailTransfer
+from lionagi.core.mail.mail_manager import MailManager
+from lionagi.core.execute.base_executor import BaseExecutor
+from lionagi.core.execute.branch_executor import BranchExecutor
+
+
+class InstructionMapExecutor(BaseExecutor):
+    """
+    Manages the execution of a mapped set of instructions across multiple branches within an executable structure.
+
+    Attributes:
+        branches (dict[str, BranchExecutor]): A dictionary of branch executors managing individual instruction flows.
+        structure_id (str): The identifier for the structure within which these branches operate.
+        mail_transfer (MailTransfer): Handles the transfer of mail between branches and other components.
+        branch_kwargs (dict): Keyword arguments used for initializing branches.
+        num_end_branches (int): Tracks the number of branches that have completed execution.
+        mail_manager (MailManager): Manages the distribution and collection of mails across branches.
+    """
+
+    branches: dict[str, BranchExecutor] = Field(
+        default_factory=dict, description="The branches of the instruction mapping."
+    )
+    structure_id: str = Field("", description="The ID of the executable structure.")
+    mail_transfer: MailTransfer = Field(
+        default_factory=MailTransfer, description="The mail transfer."
+    )
+    branch_kwargs: dict = Field(
+        default_factory=dict,
+        description="The keyword arguments for the initializing the branches.",
+    )
+    num_end_branches: int = Field(0, description="The number of end branches.")
+    mail_manager: MailManager = Field(
+        default_factory=MailManager, description="The mail manager."
+    )
+
+    def __init__(self, **kwargs):
+        """
+        Initializes an InstructionMapExecutor with the given parameters.
+
+        Args:
+            **kwargs: Arbitrary keyword arguments passed to the base executor and used for initializing branch executors.
+        """
+        super().__init__(**kwargs)
+        self.mail_manager = MailManager([self.mail_transfer])
+
+    def transfer_ins(self):
+        """
+        Processes incoming mails, directing them appropriately based on their categories, and handles the initial setup
+        of branches or the routing of node and condition mails.
+        """
+        for key in list(self.pending_ins.keys()):
+            while self.pending_ins[key]:
+                mail: BaseMail = self.pending_ins[key].popleft()
+                if mail.category == "start":
+                    self._process_start(mail)
+                elif mail.category == "node_list":
+                    self._process_node_list(mail)
+                elif (
+                    (mail.category == "node")
+                    or (mail.category == "condition")
+                    or (mail.category == "end")
+                ):
+                    mail.sender_id = self.mail_transfer.id_
+                    mail.recipient_id = mail.package["request_source"]
+                    self.mail_transfer.pending_outs.append(mail)
+
+    def transfer_outs(self):
+        """
+        Processes outgoing mails from the central mail transfer, handling end-of-execution notifications and routing
+        other mails to appropriate recipients.
+        """
+        for key in list(self.mail_transfer.pending_ins.keys()):
+            while self.mail_transfer.pending_ins[key]:
+                mail: BaseMail = self.mail_transfer.pending_ins[key].popleft()
+                if mail.category == "end":
+                    self.num_end_branches += 1
+                    if self.num_end_branches == len(
+                        self.branches
+                    ):  # tell when structure should stop
+                        mail.sender_id = self.id_
+                        mail.recipient_id = self.structure_id
+                        self.pending_outs.append(mail)
+                        self.execute_stop = True
+                else:
+                    mail.sender_id = self.id_
+                    mail.recipient_id = self.structure_id
+                    self.pending_outs.append(mail)
+
+    def _process_start(self, start_mail: BaseMail):
+        """
+        Processes a start mail to initialize a new branch executor and configures it based on the mail's package content.
+
+        Args:
+            start_mail (BaseMail): The mail initiating the start of a new branch execution.
+        """
+        branch = BranchExecutor(verbose=self.verbose, **self.branch_kwargs)
+        branch.context = start_mail.package["context"]
+        self.branches[branch.id_] = branch
+        self.mail_manager.add_sources([branch])
+        self.structure_id = start_mail.package["structure_id"]
+        mail = BaseMail(
+            sender_id=self.id_,
+            recipient_id=self.structure_id,
+            category="start",
+            package={"request_source": branch.id_, "package": "start"},
+        )
+        self.pending_outs.append(mail)
+
+    def _process_node_list(self, nl_mail: BaseMail):
+        """
+        Processes a node list mail, setting up new branches or propagating the execution context based on the node list
+        provided in the mail.
+
+        Args:
+            nl_mail (BaseMail): The mail containing a list of nodes to be processed in subsequent branches.
+        """
+        source_branch_id = nl_mail.package["request_source"]
+        node_list = nl_mail.package["package"]
+        shared_context = self.branches[source_branch_id].context
+        shared_context_log = self.branches[source_branch_id].context_log
+        base_branch = self.branches[source_branch_id]
+
+        first_node_mail = BaseMail(
+            sender_id=self.mail_transfer.id_,
+            recipient_id=source_branch_id,
+            category="node",
+            package={"request_source": source_branch_id, "package": node_list[0]},
+        )
+        self.mail_transfer.pending_outs.append(first_node_mail)
+
+        for i in range(1, len(node_list)):
+            branch = BranchExecutor(
+                verbose=self.verbose,
+                messages=base_branch.messages.copy(),
+                service=base_branch.service,
+                llmconfig=base_branch.llmconfig,
+                datalogger=base_branch.datalogger,
+            )
+            branch.context = shared_context
+            branch.context_log = shared_context_log
+            self.branches[branch.id_] = branch
+            self.mail_manager.add_sources([branch])
+            node_mail = BaseMail(
+                sender_id=self.mail_transfer.id_,
+                recipient_id=branch.id_,
+                category="node",
+                package={"request_source": source_branch_id, "package": node_list[i]},
+            )
+            self.mail_transfer.pending_outs.append(node_mail)
+
+    async def forward(self):
+        """
+        Forwards the execution by processing all incoming and outgoing mails and advancing the state of all active branches.
+        """
+        self.transfer_ins()
+        self.transfer_outs()
+        self.mail_manager.collect_all()
+        self.mail_manager.send_all()
+        tasks = [
+            branch.forward() for branch in self.branches.values() if branch.pending_ins
+        ]
+        await asyncio.gather(*tasks)
+        return
+
+    async def execute(self, refresh_time=1):
+        """
+        Continuously executes the forward process at specified intervals until instructed to stop.
+
+        Args:
+            refresh_time (int): The time in seconds between execution cycles.
+        """
+        while not self.execute_stop:
+            await self.forward()
+            await asyncio.sleep(refresh_time)