unique_toolkit 0.5.55__py3-none-any.whl → 0.6.0__py3-none-any.whl

unique_toolkit/chat/service.py

@@ -1,39 +1,76 @@
 import logging
-import re
 from typing import Optional
 
-import unique_sdk
-from unique_sdk._list_object import ListObject
+from typing_extensions import deprecated
 
-from unique_toolkit._common import _time_utils
-from unique_toolkit._common._base_service import BaseService
-from unique_toolkit.app.schemas import Event
+from unique_toolkit.app.schemas import ChatEvent, Event
+from unique_toolkit.chat.constants import (
+    DEFAULT_MAX_MESSAGES,
+    DEFAULT_PERCENT_OF_MAX_TOKENS,
+    DOMAIN_NAME,
+)
+from unique_toolkit.chat.functions import (
+    create_message,
+    create_message_assessment,
+    create_message_assessment_async,
+    create_message_async,
+    get_full_history,
+    get_full_history_async,
+    get_selection_from_history,
+    modify_message,
+    modify_message_assessment,
+    modify_message_assessment_async,
+    modify_message_async,
+)
 from unique_toolkit.chat.schemas import (
     ChatMessage,
+    ChatMessageAssessment,
+    ChatMessageAssessmentLabel,
+    ChatMessageAssessmentStatus,
+    ChatMessageAssessmentType,
     ChatMessageRole,
-    MessageAssessment,
-    MessageAssessmentLabel,
-    MessageAssessmentStatus,
-    MessageAssessmentType,
 )
 from unique_toolkit.content.schemas import ContentReference
-from unique_toolkit.content.utils import count_tokens
 
+logger = logging.getLogger(f"toolkit.{DOMAIN_NAME}.{__name__}")
 
-class ChatService(BaseService):
+
+class ChatService:
     """
     Provides all functionalities to manage the chat session.
 
     Attributes:
-        event (Event): The Event object.
-        logger (Optional[logging.Logger]): The logger. Defaults to None.
+        company_id (str | None): The company ID.
+        user_id (str | None): The user ID.
+        assistant_message_id (str | None): The assistant message ID.
+        user_message_id (str | None): The user message ID.
+        chat_id (str | None): The chat ID.
+        assistant_id (str | None): The assistant ID.
+        user_message_text (str | None): The user message text.
     """
 
-    def __init__(self, event: Event, logger: Optional[logging.Logger] = None):
-        super().__init__(event, logger)
+    def __init__(self, event: ChatEvent | Event):
+        self._event = event
+        self.company_id = event.company_id
+        self.user_id = event.user_id
+        self.assistant_message_id = event.payload.assistant_message.id
+        self.user_message_id = event.payload.user_message.id
+        self.chat_id = event.payload.chat_id
+        self.assistant_id = event.payload.assistant_id
+        self.user_message_text = event.payload.user_message.text
 
-    DEFAULT_PERCENT_OF_MAX_TOKENS = 0.15
-    DEFAULT_MAX_MESSAGES = 4
+    @property
+    @deprecated(
+        "The event property is deprecated and will be removed in a future version."
+    )
+    def event(self) -> Event | ChatEvent | None:
+        """
+        Get the event object (deprecated).
+
+        Returns:
+            Event | BaseEvent | None: The event object.
+        """
+        return self._event
 
     async def update_debug_info_async(self, debug_info: dict):
         """
@@ -42,14 +79,17 @@ class ChatService(BaseService):
         Args:
             debug_info (dict): The new debug information.
         """
-        params = self._construct_message_modify_params(
-            assistant=False, debug_info=debug_info
+
+        return await modify_message_async(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=self.assistant_message_id,
+            chat_id=self.chat_id,
+            user_message_id=self.user_message_id,
+            user_message_text=self.user_message_text,
+            assistant=False,
+            debug_info=debug_info,
         )
-        try:
-            await unique_sdk.Message.modify_async(**params)
-        except Exception as e:
-            self.logger.error(f"Failed to update debug info: {e}")
-            raise e
 
     def update_debug_info(self, debug_info: dict):
         """
@@ -58,25 +98,28 @@ class ChatService(BaseService):
         Args:
             debug_info (dict): The new debug information.
         """
-        params = self._construct_message_modify_params(
-            assistant=False, debug_info=debug_info
+
+        return modify_message(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=self.assistant_message_id,
+            chat_id=self.chat_id,
+            user_message_id=self.user_message_id,
+            user_message_text=self.user_message_text,
+            assistant=False,
+            debug_info=debug_info,
         )
-        try:
-            unique_sdk.Message.modify(**params)
-        except Exception as e:
-            self.logger.error(f"Failed to update debug info: {e}")
-            raise e
 
     def modify_user_message(
         self,
         content: str,
-        references: Optional[list[ContentReference]] = None,
-        debug_info: Optional[dict] = None,
-        message_id: Optional[str] = None,
-        set_completed_at: Optional[bool] = False,
+        references: list[ContentReference] = [],
+        debug_info: dict = {},
+        message_id: str | None = None,
+        set_completed_at: bool | None = False,
     ) -> ChatMessage:
         """
-        Modifies a message in the chat session synchronously.
+        Modifies a user message in the chat session synchronously.
 
         Args:
             content (str): The new content for the message.
@@ -91,28 +134,28 @@ class ChatService(BaseService):
         Raises:
             Exception: If the modification fails.
         """
-        try:
-            params = self._construct_message_modify_params(
-                assistant=False,
-                content=content,
-                references=references,
-                debug_info=debug_info,
-                message_id=message_id,
-                set_completed_at=set_completed_at,
-            )
-            message = unique_sdk.Message.modify(**params)
-        except Exception as e:
-            self.logger.error(f"Failed to modify user message: {e}")
-            raise e
-        return ChatMessage(**message)
+        return modify_message(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=self.assistant_message_id,
+            chat_id=self.chat_id,
+            user_message_id=self.user_message_id,
+            user_message_text=self.user_message_text,
+            assistant=False,
+            content=content,
+            references=references,
+            debug_info=debug_info,
+            message_id=message_id,
+            set_completed_at=set_completed_at or False,
+        )
 
     async def modify_user_message_async(
         self,
         content: str,
         references: list[ContentReference] = [],
         debug_info: dict = {},
-        message_id: Optional[str] = None,
-        set_completed_at: Optional[bool] = False,
+        message_id: str | None = None,
+        set_completed_at: bool | None = False,
     ) -> ChatMessage:
         """
         Modifies a message in the chat session asynchronously.
@@ -121,8 +164,8 @@ class ChatService(BaseService):
             content (str): The new content for the message.
             message_id (str, optional): The message ID. Defaults to None, then the ChatState user message id is used.
             references (list[ContentReference]): list of ContentReference objects. Defaults to None.
-            debug_info (Optional[dict[str, Any]]], optional): Debug information. Defaults to None.
-            set_completed_at (Optional[bool]): Whether to set the completedAt field with the current date time. Defaults to False.
+            debug_info (dict[str, Any]]]): Debug information. Defaults to {}.
+            set_completed_at (bool, optional): Whether to set the completedAt field with the current date time. Defaults to False.
 
         Returns:
             ChatMessage: The modified message.
@@ -130,20 +173,20 @@ class ChatService(BaseService):
         Raises:
             Exception: If the modification fails.
         """
-        try:
-            params = self._construct_message_modify_params(
-                assistant=False,
-                content=content,
-                references=references,
-                debug_info=debug_info,
-                message_id=message_id,
-                set_completed_at=set_completed_at,
-            )
-            message = await unique_sdk.Message.modify_async(**params)
-        except Exception as e:
-            self.logger.error(f"Failed to modify user message: {e}")
-            raise e
-        return ChatMessage(**message)
+        return await modify_message_async(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=self.assistant_message_id,
+            chat_id=self.chat_id,
+            user_message_id=self.user_message_id,
+            user_message_text=self.user_message_text,
+            assistant=False,
+            content=content,
+            references=references,
+            debug_info=debug_info,
+            message_id=message_id,
+            set_completed_at=set_completed_at or False,
+        )
 
     def modify_assistant_message(
         self,
@@ -151,8 +194,8 @@ class ChatService(BaseService):
         original_content: str | None = None,
         references: list[ContentReference] = [],
         debug_info: dict = {},
-        message_id: Optional[str] = None,
-        set_completed_at: Optional[bool] = False,
+        message_id: str | None = None,
+        set_completed_at: bool | None = False,
     ) -> ChatMessage:
         """
         Modifies a message in the chat session synchronously.
@@ -171,21 +214,21 @@ class ChatService(BaseService):
         Raises:
             Exception: If the modification fails.
         """
-        try:
-            params = self._construct_message_modify_params(
-                assistant=True,
-                content=content,
-                original_content=original_content,
-                references=references,
-                debug_info=debug_info,
-                message_id=message_id,
-                set_completed_at=set_completed_at,
-            )
-            message = unique_sdk.Message.modify(**params)
-        except Exception as e:
-            self.logger.error(f"Failed to modify assistant message: {e}")
-            raise e
-        return ChatMessage(**message)
+        return modify_message(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=self.assistant_message_id,
+            chat_id=self.chat_id,
+            user_message_id=self.user_message_id,
+            user_message_text=self.user_message_text,
+            assistant=True,
+            content=content,
+            original_content=original_content,
+            references=references,
+            debug_info=debug_info,
+            message_id=message_id,
+            set_completed_at=set_completed_at or False,
+        )
 
     async def modify_assistant_message_async(
         self,
@@ -193,8 +236,8 @@ class ChatService(BaseService):
         original_content: str | None = None,
         references: list[ContentReference] = [],
         debug_info: dict = {},
-        message_id: Optional[str] = None,
-        set_completed_at: Optional[bool] = False,
+        message_id: str | None = None,
+        set_completed_at: bool | None = False,
     ) -> ChatMessage:
         """
         Modifies a message in the chat session asynchronously.
@@ -204,8 +247,8 @@ class ChatService(BaseService):
             original_content (str, optional): The original content for the message.
             message_id (str, optional): The message ID. Defaults to None, then the ChatState assistant message id is used.
             references (list[ContentReference]): list of ContentReference objects. Defaults to None.
-            debug_info (Optional[dict[str, Any]]], optional): Debug information. Defaults to None.
-            set_completed_at (Optional[bool]): Whether to set the completedAt field with the current date time. Defaults to False.
+            debug_info (dict[str, Any]], optional): Debug information. Defaults to None.
+            set_completed_at (bool, optional): Whether to set the completedAt field with the current date time. Defaults to False.
 
         Returns:
             ChatMessage: The modified message.
@@ -213,21 +256,22 @@ class ChatService(BaseService):
         Raises:
             Exception: If the modification fails.
         """
-        try:
-            params = self._construct_message_modify_params(
-                assistant=True,
-                content=content,
-                original_content=original_content,
-                references=references,
-                debug_info=debug_info,
-                message_id=message_id,
-                set_completed_at=set_completed_at,
-            )
-            message = await unique_sdk.Message.modify_async(**params)
-        except Exception as e:
-            self.logger.error(f"Failed to modify assistant message: {e}")
-            raise e
-        return ChatMessage(**message)
+
+        return await modify_message_async(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=self.assistant_message_id,
+            chat_id=self.chat_id,
+            user_message_id=self.user_message_id,
+            user_message_text=self.user_message_text,
+            assistant=True,
+            content=content,
+            original_content=original_content,
+            references=references,
+            debug_info=debug_info,
+            message_id=message_id,
+            set_completed_at=set_completed_at or False,
+        )
 
     def get_full_history(self) -> list[ChatMessage]:
         """
@@ -239,7 +283,11 @@ class ChatService(BaseService):
         Raises:
             Exception: If the loading fails.
         """
-        return self._get_full_history()
+        return get_full_history(
+            event_user_id=self.user_id,
+            event_company_id=self.company_id,
+            event_payload_chat_id=self.chat_id,
+        )
 
     async def get_full_history_async(self) -> list[ChatMessage]:
         """
@@ -251,7 +299,11 @@ class ChatService(BaseService):
         Raises:
             Exception: If the loading fails.
         """
-        return await self._get_full_history_async()
+        return await get_full_history_async(
+            event_user_id=self.user_id,
+            event_company_id=self.company_id,
+            event_payload_chat_id=self.chat_id,
+        )
 
     def get_full_and_selected_history(
         self,
@@ -273,8 +325,12 @@ class ChatService(BaseService):
         Raises:
             Exception: If the loading fails.
         """
-        full_history = self._get_full_history()
-        selected_history = self._get_selection_from_history(
+        full_history = get_full_history(
+            event_user_id=self.user_id,
+            event_company_id=self.company_id,
+            event_payload_chat_id=self.chat_id,
+        )
+        selected_history = get_selection_from_history(
             full_history=full_history,
             max_tokens=int(round(token_limit * percent_of_max_tokens)),
             max_messages=max_messages,
@@ -302,8 +358,12 @@ class ChatService(BaseService):
         Raises:
             Exception: If the loading fails.
         """
-        full_history = await self._get_full_history_async()
-        selected_history = self._get_selection_from_history(
+        full_history = await get_full_history_async(
+            event_user_id=self.user_id,
+            event_company_id=self.company_id,
+            event_payload_chat_id=self.chat_id,
+        )
+        selected_history = get_selection_from_history(
             full_history=full_history,
             max_tokens=int(round(token_limit * percent_of_max_tokens)),
             max_messages=max_messages,
@@ -335,23 +395,18 @@ class ChatService(BaseService):
         Raises:
             Exception: If the creation fails.
         """
-        if original_content is None:
-            original_content = content
-
-        try:
-            params = self._construct_message_create_params(
-                content=content,
-                original_content=original_content,
-                references=references,
-                debug_info=debug_info,
-                set_completed_at=set_completed_at,
-            )
-
-            message = unique_sdk.Message.create(**params)
-        except Exception as e:
-            self.logger.error(f"Failed to create assistant message: {e}")
-            raise e
-        return ChatMessage(**message)
+        return create_message(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            chat_id=self.chat_id,
+            assistant_id=self.assistant_id,
+            role=ChatMessageRole.ASSISTANT,
+            content=content,
+            original_content=original_content,
+            references=references,
+            debug_info=debug_info,
+            set_completed_at=set_completed_at,
+        )
 
     async def create_assistant_message_async(
         self,
@@ -377,233 +432,30 @@ class ChatService(BaseService):
         Raises:
             Exception: If the creation fails.
         """
-        if original_content is None:
-            original_content = content
-
-        try:
-            params = self._construct_message_create_params(
-                content=content,
-                original_content=original_content,
-                references=references,
-                debug_info=debug_info,
-                set_completed_at=set_completed_at,
-            )
-
-            message = await unique_sdk.Message.create_async(**params)
-        except Exception as e:
-            self.logger.error(f"Failed to create assistant message: {e}")
-            raise e
-        return ChatMessage(**message)
-
-    @staticmethod
-    def _map_references(references: list[ContentReference]):
-        return [
-            {
-                "name": ref.name,
-                "url": ref.url,
-                "sequenceNumber": ref.sequence_number,
-                "sourceId": ref.source_id,
-                "source": ref.source,
-            }
-            for ref in references
-        ]
-
-    def _get_full_history(self):
-        messages = self._trigger_list_messages(self.event.payload.chat_id)
-        messages = self._filter_valid_messages(messages)
-
-        return self._map_to_chat_messages(messages)
-
-    async def _get_full_history_async(self):
-        messages = await self._trigger_list_messages_async(self.event.payload.chat_id)
-        messages = self._filter_valid_messages(messages)
-
-        return self._map_to_chat_messages(messages)
-
-    @staticmethod
-    def _filter_valid_messages(messages: ListObject[unique_sdk.Message]):
-        SYSTEM_MESSAGE_PREFIX = "[SYSTEM] "
-
-        # Remove the last two messages
-        messages = messages["data"][:-2]  # type: ignore
-        filtered_messages = []
-        for message in messages:
-            if message["text"] is None:
-                continue
-            elif SYSTEM_MESSAGE_PREFIX in message["text"]:
-                continue
-            else:
-                filtered_messages.append(message)
-
-        return filtered_messages
-
-    def _trigger_list_messages(self, chat_id: str):
-        try:
-            messages = unique_sdk.Message.list(
-                user_id=self.event.user_id,
-                company_id=self.event.company_id,
-                chatId=chat_id,
-            )
-            return messages
-        except Exception as e:
-            self.logger.error(f"Failed to list chat history: {e}")
-            raise e
-
-    async def _trigger_list_messages_async(self, chat_id: str):
-        try:
-            messages = await unique_sdk.Message.list_async(
-                user_id=self.event.user_id,
-                company_id=self.event.company_id,
-                chatId=chat_id,
-            )
-            return messages
-        except Exception as e:
-            self.logger.error(f"Failed to list chat history: {e}")
-            raise e
-
-    @staticmethod
-    def _map_to_chat_messages(messages: list[dict]):
-        return [ChatMessage(**msg) for msg in messages]
-
-    def _get_selection_from_history(
-        self,
-        full_history: list[ChatMessage],
-        max_tokens: int,
-        max_messages=4,
-    ):
-        messages = full_history[-max_messages:]
-        filtered_messages = [m for m in messages if m.content]
-        mapped_messages = []
-
-        for m in filtered_messages:
-            m.content = re.sub(r"<sup>\d+</sup>", "", m.content)
-            m.role = (
-                ChatMessageRole.ASSISTANT
-                if m.role == ChatMessageRole.ASSISTANT
-                else ChatMessageRole.USER
-            )
-            mapped_messages.append(m)
-
-        return self._pick_messages_in_reverse_for_token_window(
-            messages=mapped_messages,
-            limit=max_tokens,
-        )
 
-    def _pick_messages_in_reverse_for_token_window(
-        self,
-        messages: list[ChatMessage],
-        limit: int,
-    ) -> list[ChatMessage]:
-        if len(messages) < 1 or limit < 1:
-            return []
-
-        last_index = len(messages) - 1
-        token_count = count_tokens(messages[last_index].content)
-        while token_count > limit:
-            self.logger.debug(
-                f"Limit too low for the initial message. Last message TokenCount {token_count} available tokens {limit} - cutting message in half until it fits"
-            )
-            content = messages[last_index].content
-            messages[last_index].content = content[: len(content) // 2] + "..."
-            token_count = count_tokens(messages[last_index].content)
-
-        while token_count <= limit and last_index > 0:
-            token_count = count_tokens(
-                "".join([msg.content for msg in messages[:last_index]])
-            )
-            if token_count <= limit:
-                last_index -= 1
-
-        last_index = max(0, last_index)
-        return messages[last_index:]
-
-    def _construct_message_modify_params(
-        self,
-        assistant: bool = True,
-        content: Optional[str] = None,
-        original_content: Optional[str] = None,
-        references: Optional[list[ContentReference]] = None,
-        debug_info: Optional[dict] = None,
-        message_id: Optional[str] = None,
-        set_completed_at: Optional[bool] = False,
-    ):
-        if message_id:
-            # Message ID specified. No need to guess
-            message_id = message_id
-        elif assistant:
-            # Assistant message ID
-            message_id = self.event.payload.assistant_message.id
-        else:
-            # User message ID
-            message_id = self.event.payload.user_message.id
-            if content is None:
-                content = self.event.payload.user_message.text
-
-        if set_completed_at:
-            completed_at_datetime = _time_utils.get_datetime_now()
-        else:
-            completed_at_datetime = None
-
-        params = {
-            "user_id": self.event.user_id,
-            "company_id": self.event.company_id,
-            "id": message_id,
-            "chatId": self.event.payload.chat_id,
-            "text": content,
-            "originalText": original_content,
-            "references": self._map_references(references) if references else [],
-            "debugInfo": debug_info,
-            "completedAt": completed_at_datetime,
-        }
-        return params
-
-    def _construct_message_create_params(
-        self,
-        role: ChatMessageRole = ChatMessageRole.ASSISTANT,
-        content: Optional[str] = None,
-        original_content: Optional[str] = None,
-        references: Optional[list[ContentReference]] = None,
-        debug_info: Optional[dict] = None,
-        assistantId: Optional[str] = None,
-        set_completed_at: Optional[bool] = False,
-    ):
-        if assistantId:
-            # Assistant ID specified. No need to guess
-            assistantId = assistantId
-        else:
-            assistantId = self.event.payload.assistant_id
-
-        if set_completed_at:
-            completed_at_datetime = _time_utils.get_datetime_now()
-        else:
-            completed_at_datetime = None
-
-        if original_content is None:
-            original_content = content
-
-        params = {
-            "user_id": self.event.user_id,
-            "company_id": self.event.company_id,
-            "assistantId": assistantId,
-            "role": role.value.upper(),
-            "chatId": self.event.payload.chat_id,
-            "text": content,
-            "originalText": original_content,
-            "references": self._map_references(references) if references else [],
-            "debugInfo": debug_info,
-            "completedAt": completed_at_datetime,
-        }
-        return params
+        return await create_message_async(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            chat_id=self.chat_id,
+            assistant_id=self.assistant_id,
+            role=ChatMessageRole.ASSISTANT,
+            content=content,
+            original_content=original_content,
+            references=references,
+            debug_info=debug_info,
+            set_completed_at=set_completed_at,
+        )
 
     def create_message_assessment(
         self,
         assistant_message_id: str,
-        status: MessageAssessmentStatus,
-        type: MessageAssessmentType,
+        status: ChatMessageAssessmentStatus,
+        type: ChatMessageAssessmentType,
+        title: str | None = None,
         explanation: str | None = None,
-        label: MessageAssessmentLabel | None = None,
+        label: ChatMessageAssessmentLabel | None = None,
         is_visible: bool = True,
-    ) -> MessageAssessment:
+    ) -> ChatMessageAssessment:
         """
         Creates a message assessment for an assistant message synchronously.
 
@@ -611,91 +463,88 @@ class ChatService(BaseService):
             assistant_message_id (str): The ID of the assistant message to assess
             status (MessageAssessmentStatus): The status of the assessment (e.g. "DONE")
             type (MessageAssessmentType): The type of assessment (e.g. "HALLUCINATION")
+            title (str | None): The title of the assessment
             explanation (str | None): Explanation of the assessment
-            label (MessageAssessmentLabel | None): The assessment label (e.g. "NEGATIVE")
+            label (MessageAssessmentLabel | None): The assessment label (e.g. "RED")
             is_visible (bool): Whether the assessment is visible to users. Defaults to True.
 
         Returns:
-            MessageAssessment: The created message assessment
+            ChatMessageAssessment: The created message assessment
 
         Raises:
             Exception: If the creation fails
         """
-        try:
-            assessment = unique_sdk.MessageAssessment.create(
-                user_id=self.event.user_id,
-                company_id=self.event.company_id,
-                messageId=assistant_message_id,
-                status=status.name,
-                explanation=explanation,
-                label=label.name if label else None,
-                type=type.name if type else None,
-                isVisible=is_visible,
-            )
-            return MessageAssessment(**assessment)
-        except Exception as e:
-            self.logger.error(f"Failed to create message assessment: {e}")
-            raise e
+        return create_message_assessment(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=assistant_message_id,
+            status=status,
+            type=type,
+            title=title,
+            explanation=explanation,
+            label=label,
+            is_visible=is_visible,
+        )
 
     async def create_message_assessment_async(
         self,
         assistant_message_id: str,
-        status: MessageAssessmentStatus,
-        type: MessageAssessmentType,
+        status: ChatMessageAssessmentStatus,
+        type: ChatMessageAssessmentType,
+        title: str | None = None,
         explanation: str | None = None,
-        label: MessageAssessmentLabel | None = None,
+        label: ChatMessageAssessmentLabel | None = None,
         is_visible: bool = True,
-    ) -> MessageAssessment:
+    ) -> ChatMessageAssessment:
         """
         Creates a message assessment for an assistant message asynchronously.
 
         Args:
             assistant_message_id (str): The ID of the assistant message to assess
-            status (MessageAssessmentStatus): The status of the assessment (e.g. "DONE")
-            type (MessageAssessmentType): The type of assessment (e.g. "HALLUCINATION")
+            status (ChatMessageAssessmentStatus): The status of the assessment (e.g. "DONE")
+            type (ChatMessageAssessmentType): The type of assessment (e.g. "HALLUCINATION")
+            title (str | None): The title of the assessment
             explanation (str | None): Explanation of the assessment
-            label (MessageAssessmentLabel | None): The assessment label (e.g. "NEGATIVE")
+            label (ChatMessageAssessmentLabel | None): The assessment label (e.g. "RED")
             is_visible (bool): Whether the assessment is visible to users. Defaults to True.
 
         Returns:
-            MessageAssessment: The created message assessment
+            ChatMessageAssessment: The created message assessment
 
         Raises:
             Exception: If the creation fails
         """
-        try:
-            assessment = await unique_sdk.MessageAssessment.create_async(
-                user_id=self.event.user_id,
-                company_id=self.event.company_id,
-                messageId=assistant_message_id,
-                status=status.name,
-                explanation=explanation,
-                label=label.name if label else None,
-                type=type.name if type else None,
-                isVisible=is_visible,
-            )
-            return MessageAssessment(**assessment)
-        except Exception as e:
-            self.logger.error(f"Failed to create message assessment: {e}")
-            raise e
+        return await create_message_assessment_async(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=assistant_message_id,
+            status=status,
+            type=type,
+            title=title,
+            explanation=explanation,
+            label=label,
+            is_visible=is_visible,
+        )
 
     def modify_message_assessment(
         self,
         assistant_message_id: str,
-        status: MessageAssessmentStatus,
-        type: MessageAssessmentType,
+        status: ChatMessageAssessmentStatus,
+        type: ChatMessageAssessmentType,
+        title: str | None = None,
         explanation: str | None = None,
-        label: MessageAssessmentLabel | None = None,
-    ) -> MessageAssessment:
+        label: ChatMessageAssessmentLabel | None = None,
+    ) -> ChatMessageAssessment:
         """
         Modifies a message assessment for an assistant message synchronously.
 
         Args:
             assistant_message_id (str): The ID of the assistant message to assess
             status (MessageAssessmentStatus): The status of the assessment (e.g. "DONE")
+            title (str | None): The title of the assessment
             explanation (str | None): Explanation of the assessment
-            label (MessageAssessmentLabel | None): The assessment label (e.g. "NEGATIVE")
-            type (MessageAssessmentType): The type of assessment (e.g. "HALLUCINATION")
+            label (ChatMessageAssessmentLabel | None): The assessment label (e.g. "RED")
+            type (ChatMessageAssessmentType): The type of assessment (e.g. "HALLUCINATION")
 
         Returns:
             dict: The modified message assessment
@@ -703,56 +552,50 @@ class ChatService(BaseService):
         Raises:
             Exception: If the modification fails
         """
-        try:
-            assessment = unique_sdk.MessageAssessment.modify(
-                user_id=self.event.user_id,
-                company_id=self.event.company_id,
-                messageId=assistant_message_id,
-                status=status.name,
-                explanation=explanation,
-                label=label.name if label else None,
-                type=type.name,
-            )
-            return MessageAssessment(**assessment)
-        except Exception as e:
-            self.logger.error(f"Failed to modify message assessment: {e}")
-            raise e
+        return modify_message_assessment(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=assistant_message_id,
+            status=status,
+            type=type,
+            title=title,
+            explanation=explanation,
+            label=label,
+        )
 
     async def modify_message_assessment_async(
         self,
         assistant_message_id: str,
-        type: MessageAssessmentType,
-        status: MessageAssessmentStatus | None = None,
+        type: ChatMessageAssessmentType,
+        title: str | None = None,
+        status: ChatMessageAssessmentStatus | None = None,
         explanation: str | None = None,
-        label: MessageAssessmentLabel | None = None,
-    ) -> MessageAssessment:
+        label: ChatMessageAssessmentLabel | None = None,
+    ) -> ChatMessageAssessment:
         """
         Modifies a message assessment for an assistant message asynchronously.
 
         Args:
             assistant_message_id (str): The ID of the assistant message to assess
-            status (MessageAssessmentStatus): The status of the assessment (e.g. "DONE")
+            status (ChatMessageAssessmentStatus): The status of the assessment (e.g. "DONE")
+            title (str | None): The title of the assessment
             explanation (str | None): Explanation of the assessment
-            label (MessageAssessmentLabel | None): The assessment label (e.g. "NEGATIVE")
-            type (MessageAssessmentType): The type of assessment (e.g. "HALLUCINATION")
+            label (ChatMessageAssessmentLabel | None): The assessment label (e.g. "RED")
+            type (ChatMessageAssessmentType): The type of assessment (e.g. "HALLUCINATION")
 
         Returns:
-            MessageAssessment: The modified message assessment
+            ChatMessageAssessment: The modified message assessment
 
         Raises:
             Exception: If the modification fails
         """
-        try:
-            assessment = await unique_sdk.MessageAssessment.modify_async(
-                user_id=self.event.user_id,
-                company_id=self.event.company_id,
-                messageId=assistant_message_id,
-                status=status.name if status else None,
-                explanation=explanation,
-                label=label.name if label else None,
-                type=type.name,
-            )
-            return MessageAssessment(**assessment)
-        except Exception as e:
-            self.logger.error(f"Failed to modify message assessment: {e}")
-            raise e
+        return await modify_message_assessment_async(
+            user_id=self.user_id,
+            company_id=self.company_id,
+            assistant_message_id=assistant_message_id,
+            status=status,
+            type=type,
+            title=title,
+            explanation=explanation,
+            label=label,
+        )