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

lionagi/protocols/messages/manager.py

@@ -0,0 +1,429 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Literal
+
+from jinja2 import Template
+from pydantic import BaseModel, JsonValue
+
+from .._concepts import Manager
+from ..generic.log import Log
+from ..generic.pile import Pile
+from ..generic.progression import Progression
+from .action_request import ActionRequest
+from .action_response import ActionResponse
+from .assistant_response import AssistantResponse
+from .instruction import Instruction
+from .message import RoledMessage, SenderRecipient
+from .system import System
+
+DEFAULT_SYSTEM = "You are a helpful AI assistant. Let's think step by step."
+
+
+class MessageManager(Manager):
+
+    def __init__(
+        self,
+        messages: list[RoledMessage] | None = None,
+        progression: Progression | None = None,
+        system: System | None = None,
+    ):
+        super().__init__()
+        m_ = []
+        if isinstance(messages, list):
+            for i in messages:
+                if isinstance(i, dict):
+                    i = RoledMessage.from_dict(i)
+                if isinstance(i, RoledMessage):
+                    m_.append(i)
+        if isinstance(messages, dict):
+            self.messages = Pile.from_dict(messages)
+        else:
+            self.messages: Pile[RoledMessage] = Pile(
+                collections=m_,
+                item_type={RoledMessage},
+                strict_type=False,
+                progression=progression,
+            )
+        if system and not isinstance(system, System):
+            raise ValueError("System message must be a System instance.")
+        self.system = system  # system must be the first message
+        if self.system:
+            self.add_message(system=self.system)
+
+    @property
+    def progression(self) -> Progression:
+        return self.messages.progression
+
+    def set_system(self, system: System) -> None:
+        """
+        Set or update the system message.
+
+        Args:
+            system: The new system message to set
+        """
+        if not self.system:
+            self.system = system
+            self.messages.insert(0, self.system)
+        else:
+            old_system = self.system
+            self.system = system
+            self.messages.insert(0, self.system)
+            self.messages.exclude(old_system)
+
+    async def aclear_messages(self):
+        """Asynchronously clear all messages except system message."""
+        async with self.messages:
+            self.clear_messages()
+
+    async def a_add_message(self, **kwargs):
+        """
+        Asynchronously add a message.
+
+        Args:
+            **kwargs: Message creation parameters
+        """
+        async with self.messages:
+            return self.add_message(**kwargs)
+
+    @staticmethod
+    def create_instruction(
+        *,
+        instruction: JsonValue = None,
+        context: JsonValue = None,
+        guidance: JsonValue = None,
+        images: list = None,
+        request_fields: JsonValue = None,
+        plain_content: JsonValue = None,
+        image_detail: Literal["low", "high", "auto"] = None,
+        request_model: BaseModel | type[BaseModel] = None,
+        response_format: BaseModel | type[BaseModel] = None,
+        tool_schemas: dict = None,
+        sender: SenderRecipient = None,
+        recipient: SenderRecipient = None,
+    ) -> Instruction:
+
+        params = {k: v for k, v in locals().items() if v is not None}
+
+        if isinstance(instruction, Instruction):
+            params.pop("instruction")
+            instruction.update(**params)
+            return instruction
+        else:
+            return Instruction.create(**params)
+
+    @staticmethod
+    def create_assistant_response(
+        *,
+        sender: Any = None,
+        recipient: Any = None,
+        assistant_response: AssistantResponse | Any = None,
+        template: Template | str = None,
+        template_context: dict[str, Any] = None,
+    ) -> AssistantResponse:
+
+        params = {k: v for k, v in locals().items() if v is not None}
+        template_context = params.pop("template_context", {})
+        params.update(template_context)
+
+        if isinstance(assistant_response, AssistantResponse):
+            params.pop("assistant_response")
+            assistant_response.update(**params)
+            return assistant_response
+
+        return AssistantResponse.create(**params)
+
+    @staticmethod
+    def create_action_request(
+        *,
+        sender: SenderRecipient = None,
+        recipient: SenderRecipient = None,
+        function: str = None,
+        arguments: dict[str, Any] = None,
+        action_request: ActionRequest | None = None,
+        template: Template | str = None,
+        template_context: dict[str, Any] = None,
+    ) -> ActionRequest:
+
+        params = {
+            "sender": sender,
+            "recipient": recipient,
+            "function": function,
+            "arguments": arguments,
+            "template": template,
+            **(template_context or {}),
+        }
+
+        if isinstance(action_request, ActionRequest):
+            action_request.update(**params)
+            return action_request
+
+        return ActionRequest.create(**params)
+
+    @staticmethod
+    def create_action_response(
+        *,
+        action_request: ActionRequest,
+        action_output: Any = None,
+        action_response: ActionResponse | Any = None,
+        sender: SenderRecipient = None,
+        recipient: SenderRecipient = None,
+    ) -> ActionResponse:
+
+        if not isinstance(action_request, ActionRequest):
+            raise ValueError(
+                "Error: please provide a corresponding action request for an "
+                "action response."
+            )
+
+        params = {
+            "action_request": action_request,
+            "output": action_output,
+            "sender": sender,
+            "recipient": recipient,
+        }
+
+        if isinstance(action_response, ActionResponse):
+            action_response.update(**params)
+            return action_response
+
+        return ActionResponse.create(**params)
+
+    @staticmethod
+    def create_system(
+        *,
+        system: Any = None,
+        system_datetime: bool | str = None,
+        sender: Any = None,
+        recipient: Any = None,
+        template: Template | str = None,
+        template_context: dict[str, Any] = None,
+    ) -> System:
+        params = {
+            "system_datetime": system_datetime,
+            "sender": sender,
+            "recipient": recipient,
+            "template": template,
+            **(template_context or {}),
+        }
+        if isinstance(system, System):
+            system.update(**params)
+            return system
+
+        if system:
+            params["system_message"] = system
+
+        return System.create(**params)
+
+    def add_message(
+        self,
+        *,
+        # common
+        sender: SenderRecipient = None,
+        recipient: SenderRecipient = None,
+        template: Template | str = None,
+        template_context: dict[str, Any] = None,
+        metadata: dict[str, Any] = None,
+        # instruction
+        instruction: JsonValue = None,
+        context: JsonValue = None,
+        guidance: JsonValue = None,
+        request_fields: JsonValue = None,
+        plain_content: JsonValue = None,
+        request_model: BaseModel | type[BaseModel] = None,
+        response_format: BaseModel | type[BaseModel] = None,
+        images: list = None,
+        image_detail: Literal["low", "high", "auto"] = None,
+        tool_schemas: dict = None,
+        # system
+        system: Any = None,
+        system_datetime: bool | str = None,
+        # assistant_response
+        assistant_response: AssistantResponse | Any = None,
+        # actions
+        action_function: str = None,
+        action_arguments: dict[str, Any] = None,
+        action_output: Any = None,
+        action_request: ActionRequest | None = None,
+        action_response: ActionResponse | Any = None,
+    ) -> tuple[RoledMessage, Log]:
+
+        _msg = None
+        if (
+            sum(
+                bool(x)
+                for x in (
+                    instruction,
+                    assistant_response,
+                    system,
+                    action_request,
+                )
+            )
+            > 1
+        ):
+            raise ValueError("Only one message type can be added at a time.")
+
+        if system:
+            _msg = self.create_system(
+                system=system,
+                system_datetime=system_datetime,
+                sender=sender,
+                recipient=recipient,
+                template=template,
+                template_context=template_context,
+            )
+            self.set_system(_msg)
+
+        elif action_output:
+            _msg = self.create_action_response(
+                action_request=action_request,
+                action_output=action_output,
+                action_response=action_response,
+                sender=sender,
+                recipient=recipient,
+            )
+
+        elif action_request or (action_function and action_arguments):
+            _msg = self.create_action_request(
+                sender=sender,
+                recipient=recipient,
+                function=action_function,
+                arguments=action_arguments,
+                action_request=action_request,
+                template=template,
+                template_context=template_context,
+            )
+
+        elif assistant_response:
+            _msg = self.create_assistant_response(
+                sender=sender,
+                recipient=recipient,
+                assistant_response=assistant_response,
+                template=template,
+                template_context=template_context,
+            )
+
+        else:
+            _msg = self.create_instruction(
+                instruction=instruction,
+                context=context,
+                guidance=guidance,
+                images=images,
+                request_fields=request_fields,
+                plain_content=plain_content,
+                image_detail=image_detail,
+                request_model=request_model,
+                response_format=response_format,
+                tool_schemas=tool_schemas,
+                sender=sender,
+                recipient=recipient,
+            )
+
+        if metadata:
+            _msg.metadata.setdefault("extra", {})
+            _msg.metadata["extra"].update(metadata)
+
+        if _msg in self.messages:
+            idx = self.messages.progression.index(_msg.id)
+            self.messages.exclude(_msg.id)
+            self.messages.insert(idx, _msg)
+        else:
+            self.messages.include(_msg)
+
+        return _msg
+
+    def clear_messages(self):
+        self.messages.clear()
+        if self.system:
+            self.messages.insert(0, self.system)
+
+    @property
+    def last_response(self) -> AssistantResponse | None:
+        """Get the last assistant response message."""
+        for i in reversed(self.messages.progression):
+            if isinstance(self.messages[i], AssistantResponse):
+                return self.messages[i]
+
+    @property
+    def last_instruction(self) -> Instruction | None:
+        """Get the last instruction message."""
+        for i in reversed(self.messages.progression):
+            if isinstance(self.messages[i], Instruction):
+                return self.messages[i]
+
+    @property
+    def assistant_responses(self) -> Pile[AssistantResponse]:
+        """Get all assistant response messages."""
+        return Pile(
+            collections=[
+                self.messages[i]
+                for i in self.messages.progression
+                if isinstance(self.messages[i], AssistantResponse)
+            ]
+        )
+
+    @property
+    def action_requests(self) -> Pile[ActionRequest]:
+        """Get all action request messages."""
+        return Pile(
+            collections=[
+                self.messages[i]
+                for i in self.messages.progression
+                if isinstance(self.messages[i], ActionRequest)
+            ]
+        )
+
+    @property
+    def action_responses(self) -> Pile[ActionResponse]:
+        """Get all action response messages."""
+        return Pile(
+            collections=[
+                self.messages[i]
+                for i in self.messages.progression
+                if isinstance(self.messages[i], ActionResponse)
+            ]
+        )
+
+    @property
+    def instructions(self) -> Pile[Instruction]:
+        """Get all instruction messages."""
+        return Pile(
+            collections=[
+                self.messages[i]
+                for i in self.messages.progression
+                if isinstance(self.messages[i], Instruction)
+            ]
+        )
+
+    def remove_last_instruction_tool_schemas(self) -> None:
+        id_ = self.last_instruction.id
+        self.messages[id_].tool_schemas = None
+
+    def concat_recent_action_responses_to_instruction(
+        self, instruction: Instruction
+    ) -> None:
+        for i in reversed(self.messages.progression):
+            if isinstance(self.messages[i], ActionResponse):
+                instruction.context.append(self.messages[i].content)
+            else:
+                break
+
+    def to_chat_msgs(self, progression=None) -> list[dict]:
+        if progression == []:
+            return []
+        try:
+            return [
+                self.messages[i].chat_msg
+                for i in (progression or self.progression)
+            ]
+        except Exception as e:
+            raise ValueError(
+                "Invalid progress, not all requested messages are in the message pile"
+            ) from e
+
+    def __bool__(self):
+        return bool(self.messages)
+
+    def __contains__(self, message: RoledMessage) -> bool:
+        return message in self.messages

lionagi/protocols/messages/message.py

@@ -0,0 +1,216 @@
+# Copyright (c) 2023 - 2024, HaiyangLi <quantocean.li at gmail dot com>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import json
+from pathlib import Path
+from typing import Any
+
+from jinja2 import Environment, FileSystemLoader, Template
+from pydantic import Field, PrivateAttr, field_serializer
+
+from .._concepts import Sendable
+from ..generic.element import Element, IDType
+from ..generic.log import Log
+from ..graph.node import Node
+from .base import (
+    MessageFlag,
+    MessageRole,
+    SenderRecipient,
+    validate_sender_recipient,
+)
+
+template_path = Path(__file__).parent / "templates"
+jinja_env = Environment(loader=FileSystemLoader(template_path))
+
+__all__ = ("RoledMessage",)
+
+
+class RoledMessage(Node, Sendable):
+
+    content: dict = Field(
+        default_factory=dict,
+        description="The content of the message.",
+    )
+
+    role: MessageRole | None = Field(
+        None,
+        description="The role of the message in the conversation.",
+    )
+
+    _flag: MessageFlag | None = PrivateAttr(None)
+
+    template: str | Template | None = None
+
+    sender: SenderRecipient | None = Field(
+        default=MessageRole.UNSET,
+        title="Sender",
+        description="The ID of the sender node or a role.",
+    )
+
+    recipient: SenderRecipient | None = Field(
+        default=MessageRole.UNSET,
+        title="Recipient",
+        description="The ID of the recipient node or a role.",
+    )
+
+    @field_serializer("sender", "recipient")
+    def _serialize_sender_recipient(self, value: SenderRecipient) -> str:
+        if isinstance(value, MessageRole | MessageFlag):
+            return value.value
+        if isinstance(value, str):
+            return value
+        if isinstance(value, Element):
+            return str(value.id)
+        if isinstance(value, IDType):
+            return str(value)
+        return str(value)
+
+    @property
+    def image_content(self) -> list[dict[str, Any]] | None:
+        """
+        Return image content if present in the message.
+
+        Returns:
+            Optional[List[Dict[str, Any]]]: A list of image content
+                dictionaries, or None if no images are present.
+        """
+        msg_ = self.chat_msg
+        if isinstance(msg_, dict) and isinstance(msg_["content"], list):
+            return [i for i in msg_["content"] if i["type"] == "image_url"]
+        return None
+
+    @property
+    def chat_msg(self) -> dict[str, Any] | None:
+        """
+        Return message in chat representation format.
+
+        Returns:
+            Optional[Dict[str, Any]]: The message formatted for chat use,
+                or None if formatting fails.
+        """
+        try:
+            return {"role": str(self.role), "content": self.rendered}
+        except Exception:
+            return None
+
+    @property
+    def rendered(self) -> str:
+        try:
+            if isinstance(self.template, str):
+                return self.template.format(**self.content)
+            if isinstance(self.template, Template):
+                return self.template.render(**self.content)
+        except Exception:
+            return json.dumps(self.content)
+
+    @classmethod
+    def create(cls, **kwargs):
+        raise NotImplementedError(
+            "create method must be implemented in subclass"
+        )
+
+    @classmethod
+    def from_dict(cls, dict_: dict):
+        try:
+            self: RoledMessage = super().from_dict(
+                {k: v for k, v in dict_.items() if v}
+            )
+            self._flag = MessageFlag.MESSAGE_LOAD
+            return self
+        except Exception as e:
+            raise ValueError(f"Invalid RoledMessage data: {e}")
+
+    def is_clone(self) -> bool:
+        """
+        Check if the message is a clone of another message.
+
+        Returns:
+            bool: True if the message is a clone, False otherwise.
+        """
+        return self._flag == MessageFlag.MESSAGE_CLONE
+
+    def clone(self, keep_role: bool = True) -> "RoledMessage":
+        instance = self.__class__(
+            content=self.content,
+            role=self.role if keep_role else MessageRole.UNSET,
+            metadata={"clone_from": self},
+        )
+        instance._flag = MessageFlag.MESSAGE_CLONE
+        return instance
+
+    def to_log(self) -> Log:
+        """
+        Convert the message to a Log object.
+
+        Creates a Log instance containing the message content and additional
+        information as loginfo.
+
+        Returns:
+            Log: A Log object representing the message.
+        """
+        return Log.create(self)
+
+    @field_serializer("role")
+    def _serialize_role(self, value: MessageRole):
+        """
+        Serialize the role for storage or transmission.
+
+        Args:
+            value: The MessageRole to serialize
+
+        Returns:
+            str: The string value of the role
+        """
+        if isinstance(value, MessageRole):
+            return value.value
+        return str(value)
+
+    @field_serializer("metadata")
+    def _serialize_metadata(self, value: dict):
+        if "clone_from" in value:
+            origin_obj: RoledMessage = value.pop("clone_from")
+            origin_info = origin_obj.to_dict()
+            value["clone_from_info"] = {
+                "clone_from_info": {
+                    "original_id": origin_info["id"],
+                    "original_created_at": origin_info["created_at"],
+                    "original_sender": origin_info["sender"],
+                    "original_recipient": origin_info["recipient"],
+                    "original_lion_class": origin_info["metadata"][
+                        "lion_class"
+                    ],
+                    "original_role": origin_info["role"],
+                }
+            }
+        return value
+
+    @field_serializer("template")
+    def _serialize_template(self, value: Template | str):
+        if isinstance(value, Template):
+            return None
+        return value
+
+    def update(self, sender, recipient, template, **kwargs):
+        if sender:
+            self.sender = validate_sender_recipient(sender)
+        if recipient:
+            self.recipient = validate_sender_recipient(recipient)
+        if kwargs:
+            self.content.update(kwargs)
+        if template:
+            if not isinstance(template, Template | str):
+                raise ValueError("Template must be a Jinja2 Template or str")
+            self.template = template
+
+    def __str__(self) -> str:
+
+        content_preview = (
+            f"{str(self.content)[:75]}..."
+            if len(str(self.content)) > 75
+            else str(self.content)
+        )
+        return (
+            f"Message(role={self.role}, sender={self.sender}, "
+            f"content='{content_preview}')"
+        )