lionagi 0.0.312__py3-none-any.whl → 0.2.1__py3-none-any.whl

lionagi/core/branch/base_branch.py

@@ -1,654 +0,0 @@
-from abc import ABC
-from typing import Any
-
-from lionagi.libs.sys_util import PATH_TYPE
-from lionagi.libs import convert, dataframe, SysUtil
-
-from ..schema.base_node import BaseRelatableNode
-from ..schema.data_logger import DataLogger, DLog
-from ..messages.schema import (
-    BranchColumns,
-    System,
-    Response,
-    Instruction,
-    BaseMessage,
-)
-from .util import MessageUtil
-
-
-class BaseBranch(BaseRelatableNode, ABC):
-    """
-    Base class for managing branches of conversation, incorporating messages
-    and logging functionality.
-
-    Attributes:
-            messages (dataframe.ln_DataFrame): Holds the messages in the branch.
-            datalogger (DataLogger): Logs data related to the branch's operation.
-            persist_path (PATH_TYPE): Filesystem path for data persistence.
-    """
-
-    _columns: list[str] = BranchColumns.COLUMNS.value
-
-    def __init__(
-        self,
-        messages: dataframe.ln_DataFrame | None = None,
-        datalogger: DataLogger | None = None,
-        persist_path: PATH_TYPE | None = None,
-        name=None,
-        **kwargs,
-    ) -> None:
-        super().__init__(**kwargs)
-        if isinstance(messages, dataframe.ln_DataFrame):
-            if MessageUtil.validate_messages(messages):
-                self.messages = messages
-            else:
-                raise ValueError("Invalid messages format")
-        else:
-            self.messages = dataframe.ln_DataFrame(columns=self._columns)
-
-        self.datalogger = datalogger or DataLogger(persist_path=persist_path)
-        self.name = name
-
-    def add_message(
-        self,
-        system: dict | list | System | None = None,
-        instruction: dict | list | Instruction | None = None,
-        context: str | dict[str, Any] | None = None,
-        response: dict | list | BaseMessage | None = None,
-        output_fields=None,
-        recipient=None,
-        **kwargs,
-    ) -> None:
-        """
-        Adds a message to the branch.
-
-        Args:
-                system: Information for creating a System message.
-                instruction: Information for creating an Instruction message.
-                context: Context information for the message.
-                response: Response data for creating a message.
-                **kwargs: Additional keyword arguments for message creation.
-        """
-        _msg = MessageUtil.create_message(
-            system=system,
-            instruction=instruction,
-            context=context,
-            response=response,
-            output_fields=output_fields,
-            recipient=recipient,
-            **kwargs,
-        )
-
-        if isinstance(_msg, System):
-            self.system_node = _msg
-
-        # sourcery skip: merge-nested-ifs
-        if isinstance(_msg, Instruction):
-            if recipient is None and self.name is not None:
-                _msg.recipient = self.name
-
-        if isinstance(_msg, Response):
-            if "action_response" in _msg.content.keys():
-                if recipient is None and self.name is not None:
-                    _msg.recipient = self.name
-                if recipient is not None and self.name is None:
-                    _msg.recipient = recipient
-            if "response" in _msg.content.keys():
-                if self.name is not None:
-                    _msg.sender = self.name
-
-        _msg.content = _msg.msg_content
-        self.messages.loc[len(self.messages)] = _msg.to_pd_series()
-
-    def _to_chatcompletion_message(
-        self, with_sender: bool = False
-    ) -> list[dict[str, Any]]:
-        """
-        Converts messages to a list of dictionaries formatted for chat completion,
-        optionally including sender information.
-
-        Args:
-                with_sender: Flag to include sender information in the output.
-
-        Returns:
-                A list of message dictionaries, each with 'role' and 'content' keys,
-                and optionally prefixed by 'Sender' if with_sender is True.
-        """
-
-        message = []
-
-        for _, row in self.messages.iterrows():
-            content_ = row["content"]
-            if content_.startswith("Sender"):
-                content_ = content_.split(":", 1)[1]
-
-            # if isinstance(content_, str):
-            #     try:
-            #         content_ = json.dumps(to_dict(content_))
-            #     except Exception as e:
-            #         raise ValueError(
-            #             f"Error in serializing, {row['node_id']} {content_}: {e}"
-            #         )
-
-            out = {"role": row["role"], "content": content_}
-            if with_sender:
-                out["content"] = f"Sender {row['sender']}: {content_}"
-
-            message.append(out)
-        return message
-
-    @property
-    def chat_messages(self) -> list[dict[str, Any]]:
-        """
-        Retrieves all chat messages without sender information.
-
-        Returns:
-                A list of dictionaries representing chat messages.
-        """
-
-        return self._to_chatcompletion_message()
-
-    @property
-    def chat_messages_with_sender(self) -> list[dict[str, Any]]:
-        """
-        Retrieves all chat messages, including sender information.
-
-        Returns:
-                A list of dictionaries representing chat messages, each prefixed with its sender.
-        """
-
-        return self._to_chatcompletion_message(with_sender=True)
-
-    @property
-    def last_message(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves the last message from the branch as a pandas Series.
-
-        Returns:
-                A pandas Series representing the last message in the branch.
-        """
-
-        return MessageUtil.get_message_rows(self.messages, n=1, from_="last")
-
-    @property
-    def last_message_content(self) -> dict[str, Any]:
-        """
-        Extracts the content of the last message in the branch.
-
-        Returns:
-                A dictionary representing the content of the last message.
-        """
-
-        return convert.to_dict(self.messages.content.iloc[-1])
-
-    @property
-    def first_system(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves the first message marked with the 'system' role.
-
-        Returns:
-                A pandas Series representing the first 'system' message in the branch.
-        """
-
-        return MessageUtil.get_message_rows(
-            self.messages, role="system", n=1, from_="front"
-        )
-
-    @property
-    def last_response(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves the last message marked with the 'assistant' role.
-
-        Returns:
-                A pandas Series representing the last 'assistant' (response) message in the branch.
-        """
-
-        return MessageUtil.get_message_rows(
-            self.messages, role="assistant", n=1, from_="last"
-        )
-
-    @property
-    def last_response_content(self) -> dict[str, Any]:
-        """
-        Extracts the content of the last 'assistant' (response) message.
-
-        Returns:
-                A dictionary representing the content of the last 'assistant' message.
-        """
-
-        return convert.to_dict(self.last_response.content.iloc[-1])
-
-    @property
-    def action_request(self) -> dataframe.ln_DataFrame:
-        """
-        Filters and retrieves all messages sent by 'action_request'.
-
-        Returns:
-                A pandas DataFrame containing all 'action_request' messages.
-        """
-
-        return convert.to_df(self.messages[self.messages.sender == "action_request"])
-
-    @property
-    def action_response(self) -> dataframe.ln_DataFrame:
-        """
-        Filters and retrieves all messages sent by 'action_response'.
-
-        Returns:
-                A pandas DataFrame containing all 'action_response' messages.
-        """
-
-        return convert.to_df(self.messages[self.messages.sender == "action_response"])
-
-    @property
-    def responses(self) -> dataframe.ln_DataFrame:
-        """
-        Retrieves all messages marked with the 'assistant' role.
-
-        Returns:
-                A pandas DataFrame containing all messages with an 'assistant' role.
-        """
-
-        return convert.to_df(self.messages[self.messages.role == "assistant"])
-
-    @property
-    def assistant_responses(self) -> dataframe.ln_DataFrame:
-        """
-        Filters 'assistant' role messages excluding 'action_request' and 'action_response'.
-
-        Returns:
-                A pandas DataFrame of 'assistant' messages excluding action requests/responses.
-        """
-
-        a_responses = self.responses[self.responses.sender != "action_response"]
-        a_responses = a_responses[a_responses.sender != "action_request"]
-        return convert.to_df(a_responses)
-
-    @property
-    def last_assistant_response(self):
-        return self.assistant_responses.iloc[-1]
-
-    @property
-    def info(self) -> dict[str, Any]:
-        """
-        Summarizes branch information, including message counts by role.
-
-        Returns:
-                A dictionary containing counts of messages categorized by their role.
-        """
-
-        return self._info()
-
-    @property
-    def sender_info(self) -> dict[str, int]:
-        """
-        Provides a summary of message counts categorized by sender.
-
-        Returns:
-                A dictionary with senders as keys and counts of their messages as values.
-        """
-
-        return self._info(use_sender=True)
-
-    @property
-    def describe(self) -> dict[str, Any]:
-        """
-        Provides a detailed description of the branch, including a summary of messages.
-
-        Returns:
-                A dictionary with a summary of total messages, a breakdown by role, and
-                a preview of the first five messages.
-        """
-
-        return {
-            "total_messages": len(self.messages),
-            "summary_by_role": self._info(),
-            "messages": [msg.to_dict() for _, msg in self.messages.iterrows()][
-                : len(self.messages) - 1 if len(self.messages) < 5 else 5
-            ],
-        }
-
-    @classmethod
-    def _from_csv(cls, filename: str, read_kwargs=None, **kwargs) -> "BaseBranch":
-        read_kwargs = {} if read_kwargs is None else read_kwargs
-        messages = dataframe.read_csv(filename, **read_kwargs)
-        return cls(messages=messages, **kwargs)
-
-    @classmethod
-    def from_csv(cls, **kwargs) -> "BaseBranch":
-
-        return cls._from_csv(**kwargs)
-
-    @classmethod
-    def from_json_string(cls, **kwargs) -> "BaseBranch":
-
-        return cls._from_json(**kwargs)
-
-    @classmethod
-    def _from_json(cls, filename: str, read_kwargs=None, **kwargs) -> "BaseBranch":
-        read_kwargs = {} if read_kwargs is None else read_kwargs
-        messages = dataframe.read_json(filename, **read_kwargs)
-        return cls(messages=messages, **kwargs)
-
-    def to_csv_file(
-        self,
-        filename: PATH_TYPE = "messages.csv",
-        dir_exist_ok: bool = True,
-        timestamp: bool = True,
-        time_prefix: bool = False,
-        verbose: bool = True,
-        clear: bool = True,
-        **kwargs,
-    ) -> None:
-        """
-        Exports the branch messages to a CSV file.
-
-        Args:
-                filepath: Destination path for the CSV file. Defaults to 'messages.csv'.
-                dir_exist_ok: If False, an error is raised if the directory exists. Defaults to True.
-                timestamp: If True, appends a timestamp to the filename. Defaults to True.
-                time_prefix: If True, prefixes the filename with a timestamp. Defaults to False.
-                verbose: If True, prints a message upon successful export. Defaults to True.
-                clear: If True, clears the messages after exporting. Defaults to True.
-                **kwargs: Additional keyword arguments for pandas.DataFrame.to_csv().
-        """
-
-        if not filename.endswith(".csv"):
-            filename += ".csv"
-
-        filename = SysUtil.create_path(
-            self.datalogger.persist_path,
-            filename,
-            timestamp=timestamp,
-            dir_exist_ok=dir_exist_ok,
-            time_prefix=time_prefix,
-        )
-
-        try:
-            self.messages.to_csv(filename, **kwargs)
-            if verbose:
-                print(f"{len(self.messages)} messages saved to {filename}")
-            if clear:
-                self.clear_messages()
-        except Exception as e:
-            raise ValueError(f"Error in saving to csv: {e}") from e
-
-    def to_json_file(
-        self,
-        filename: PATH_TYPE = "messages.json",
-        dir_exist_ok: bool = True,
-        timestamp: bool = True,
-        time_prefix: bool = False,
-        verbose: bool = True,
-        clear: bool = True,
-        **kwargs,
-    ) -> None:
-        """
-        Exports the branch messages to a JSON file.
-
-        Args:
-                filename: Destination path for the JSON file. Defaults to 'messages.json'.
-                dir_exist_ok: If False, an error is raised if the dirctory exists. Defaults to True.
-                timestamp: If True, appends a timestamp to the filename. Defaults to True.
-                time_prefix: If True, prefixes the filename with a timestamp. Defaults to False.
-                verbose: If True, prints a message upon successful export. Defaults to True.
-                clear: If True, clears the messages after exporting. Defaults to True.
-                **kwargs: Additional keyword arguments for pandas.DataFrame.to_json().
-        """
-
-        if not filename.endswith(".json"):
-            filename += ".json"
-
-        filename = SysUtil.create_path(
-            self.datalogger.persist_path,
-            filename,
-            timestamp=timestamp,
-            dir_exist_ok=dir_exist_ok,
-            time_prefix=time_prefix,
-        )
-
-        try:
-            self.messages.to_json(
-                filename, orient="records", lines=True, date_format="iso", **kwargs
-            )
-            if verbose:
-                print(f"{len(self.messages)} messages saved to {filename}")
-            if clear:
-                self.clear_messages()
-        except Exception as e:
-            raise ValueError(f"Error in saving to json: {e}") from e
-
-    def log_to_csv(
-        self,
-        filename: PATH_TYPE = "log.csv",
-        dir_exist_ok: bool = True,
-        timestamp: bool = True,
-        time_prefix: bool = False,
-        verbose: bool = True,
-        clear: bool = True,
-        flatten_=True,
-        sep="[^_^]",
-        **kwargs,
-    ) -> None:
-        """
-        Exports the data logger contents to a CSV file.
-
-        Args:
-                filename: Destination path for the CSV file. Defaults to 'log.csv'.
-                dir_exist_ok: If False, an error is raised if the directory exists. Defaults to True.
-                timestamp: If True, appends a timestamp to the filename. Defaults to True.
-                time_prefix: If True, prefixes the filename with a timestamp. Defaults to False.
-                verbose: If True, prints a message upon successful export. Defaults to True.
-                clear: If True, clears the logger after exporting. Defaults to True.
-                **kwargs: Additional keyword arguments for pandas.DataFrame.to_csv().
-        """
-        self.datalogger.to_csv_file(
-            filename=filename,
-            dir_exist_ok=dir_exist_ok,
-            timestamp=timestamp,
-            time_prefix=time_prefix,
-            verbose=verbose,
-            clear=clear,
-            flatten_=flatten_,
-            sep=sep,
-            **kwargs,
-        )
-
-    def log_to_json(
-        self,
-        filename: PATH_TYPE = "log.json",
-        dir_exist_ok: bool = True,
-        timestamp: bool = True,
-        time_prefix: bool = False,
-        verbose: bool = True,
-        clear: bool = True,
-        flatten_=True,
-        sep="[^_^]",
-        **kwargs,
-    ) -> None:
-        """
-        Exports the data logger contents to a JSON file.
-
-        Args:
-                filename: Destination path for the JSON file. Defaults to 'log.json'.
-                dir_exist_ok: If False, an error is raised if the directory exists. Defaults to True.
-                timestamp: If True, appends a timestamp to the filename. Defaults to True.
-                time_prefix: If True, prefixes the filename with a timestamp. Defaults to False.
-                verbose: If True, prints a message upon successful export. Defaults to True.
-                clear: If True, clears the logger after exporting. Defaults to True.
-                **kwargs: Additional keyword arguments for pandas.DataFrame.to_json().
-        """
-
-        self.datalogger.to_json_file(
-            filename=filename,
-            dir_exist_ok=dir_exist_ok,
-            timestamp=timestamp,
-            time_prefix=time_prefix,
-            verbose=verbose,
-            clear=clear,
-            flatten_=flatten_,
-            sep=sep,
-            **kwargs,
-        )
-
-    def load_log(self, filename, flattened=True, sep="[^_^]", verbose=True, **kwargs):
-        df = ""
-        try:
-            if filename.endswith(".csv"):
-                df = dataframe.read_csv(filename, **kwargs)
-
-            elif filename.endswith(".json"):
-                df = dataframe.read_json(filename, **kwargs)
-
-            for _, row in df.iterrows():
-                self.datalogger.log.append(
-                    DLog.deserialize(
-                        input_str=row.input_data,
-                        output_str=row.output_data,
-                        unflatten_=flattened,
-                        sep=sep,
-                    )
-                )
-
-            if verbose:
-                print(f"Loaded {len(df)} logs from {filename}")
-        except Exception as e:
-            raise ValueError(f"Error in loading log: {e}") from e
-
-    def remove_message(self, node_id: str) -> None:
-        """
-        Removes a message from the branch based on its node ID.
-
-        Args:
-                node_id: The unique identifier of the message to be removed.
-        """
-        MessageUtil.remove_message(self.messages, node_id)
-
-    def update_message(self, node_id: str, column: str, value: Any) -> bool:
-        """
-        Updates a specific column of a message identified by node_id with a new value.
-
-        Args:
-                value: The new value to update the message with.
-                node_id: The unique identifier of the message to update.
-                column: The column of the message to update.
-        """
-
-        index = self.messages[self.messages["node_id"] == node_id].index[0]
-
-        return dataframe.update_row(
-            self.messages, row=index, column=column, value=value
-        )
-
-    def change_first_system_message(
-        self, system: str | dict[str, Any] | System, sender: str | None = None
-    ) -> None:
-        """
-        Updates the first system message with new content and/or sender.
-
-        Args:
-                system: The new system message content or a System object.
-                sender: The identifier of the sender for the system message.
-        """
-
-        if len(self.messages[self.messages["role"] == "system"]) == 0:
-            raise ValueError("There is no system message in the messages.")
-
-        if not isinstance(system, (str, dict, System)):
-            raise ValueError("Input cannot be converted into a system message.")
-
-        if isinstance(system, (str, dict)):
-            system = System(system, sender=sender)
-
-        if isinstance(system, System):
-            system.timestamp = SysUtil.get_timestamp()
-            sys_index = self.messages[self.messages.role == "system"].index
-            self.messages.loc[sys_index[0]] = system.to_pd_series()
-
-    def rollback(self, steps: int) -> None:
-        """
-        Removes the last 'n' messages from the branch.
-
-        Args:
-                steps: The number of messages to remove from the end.
-        """
-
-        self.messages = dataframe.remove_last_n_rows(self.messages, steps)
-
-    def clear_messages(self) -> None:
-        """
-        Clears all messages from the branch.
-        """
-        self.messages = dataframe.ln_DataFrame(columns=self._columns)
-
-    def replace_keyword(
-        self,
-        keyword: str,
-        replacement: str,
-        column: str = "content",
-        case_sensitive: bool = False,
-    ) -> None:
-
-        dataframe.replace_keyword(
-            self.messages,
-            keyword,
-            replacement,
-            column=column,
-            case_sensitive=case_sensitive,
-        )
-
-    def search_keywords(
-        self,
-        keywords: str | list[str],
-        case_sensitive: bool = False,
-        reset_index: bool = False,
-        dropna: bool = False,
-    ) -> dataframe.ln_DataFrame:
-        return dataframe.search_keywords(
-            self.messages,
-            keywords,
-            case_sensitive=case_sensitive,
-            reset_index=reset_index,
-            dropna=dropna,
-        )
-
-    def extend(self, messages: dataframe.ln_DataFrame, **kwargs) -> None:
-
-        self.messages = MessageUtil.extend(self.messages, messages, **kwargs)
-
-    def filter_by(
-        self,
-        role: str | None = None,
-        sender: str | None = None,
-        start_time=None,
-        end_time=None,
-        content_keywords: str | list[str] | None = None,
-        case_sensitive: bool = False,
-    ) -> dataframe.ln_DataFrame:
-
-        return MessageUtil.filter_messages_by(
-            self.messages,
-            role=role,
-            sender=sender,
-            start_time=start_time,
-            end_time=end_time,
-            content_keywords=content_keywords,
-            case_sensitive=case_sensitive,
-        )
-
-    # noinspection PyTestUnpassedFixture
-    def _info(self, use_sender: bool = False) -> dict[str, int]:
-        """
-        Helper method to generate summaries of messages either by role or sender.
-
-        Args:
-                use_sender: If True, summary is categorized by sender. Otherwise, by role.
-
-        Returns:
-                A dictionary summarizing the count of messages either by role or sender.
-        """
-
-        messages = self.messages["sender"] if use_sender else self.messages["role"]
-        result = messages.value_counts().to_dict()
-        result["total"] = len(self.messages)
-        return result