unique_toolkit 0.0.2__py3-none-any.whl → 0.5.1__py3-none-any.whl

unique_toolkit/chat/service.py

@@ -0,0 +1,380 @@
+import logging
+import re
+from typing import Optional
+
+import unique_sdk
+
+from unique_toolkit.app.performance.async_wrapper import async_warning, to_async
+from unique_toolkit.chat.schemas import ChatMessage, ChatMessageRole
+from unique_toolkit.chat.state import ChatState
+from unique_toolkit.content.schemas import ContentReference
+from unique_toolkit.content.utils import count_tokens
+
+
+class ChatService:
+    """
+    Provides all functionalities to manage the chat session.
+
+    Attributes:
+        state (ChatState): The chat state.
+        logger (Optional[logging.Logger]): The logger. Defaults to None.
+    """
+
+    def __init__(self, state: ChatState, logger: Optional[logging.Logger] = None):
+        self.state = state
+        self.logger = logger or logging.getLogger(__name__)
+
+    def modify_assistant_message(
+        self,
+        content: str,
+        references: list[ContentReference] = [],
+        debug_info: dict = {},
+        message_id: Optional[str] = None,
+    ) -> ChatMessage:
+        """
+        Modifies a message in the chat session synchronously.
+
+        Args:
+            content (str): The new content for the message.
+            references (list[ContentReference]): list of ContentReference objects. Defaults to [].
+            debug_info (dict[str, Any]]]): Debug information. Defaults to {}.
+            message_id (Optional[str]): The message ID. Defaults to None.
+
+        Returns:
+            ChatMessage: The modified message.
+
+        Raises:
+            Exception: If the modification fails.
+        """
+        return self._trigger_modify_assistant_message(
+            content=content,
+            message_id=message_id,
+            references=references,
+            debug_info=debug_info,
+        )
+
+    @to_async
+    @async_warning
+    def async_modify_assistant_message(
+        self,
+        content: str,
+        references: list[ContentReference] = [],
+        debug_info: dict = {},
+        message_id: Optional[str] = None,
+    ) -> ChatMessage:
+        """
+        Modifies a message in the chat session asynchronously.
+
+        Args:
+            content (str): The new 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.
+
+        Returns:
+            ChatMessage: The modified message.
+
+        Raises:
+            Exception: If the modification fails.
+        """
+        return self._trigger_modify_assistant_message(
+            content,
+            message_id,
+            references,
+            debug_info,
+        )
+
+    def get_full_history(self) -> list[ChatMessage]:
+        """
+        Loads the full chat history for the chat session synchronously.
+
+        Returns:
+            list[ChatMessage]: The full chat history.
+
+        Raises:
+            Exception: If the loading fails.
+        """
+        return self._get_full_history()
+
+    @to_async
+    @async_warning
+    def async_get_full_history(self) -> list[ChatMessage]:
+        """
+        Loads the full chat history for the chat session asynchronously.
+
+        Returns:
+            list[ChatMessage]: The full chat history.
+
+        Raises:
+            Exception: If the loading fails.
+        """
+        return self._get_full_history()
+
+    def get_full_and_selected_history(
+        self,
+        token_limit: int,
+        percent_of_max_tokens: float,
+        max_messages: int,
+    ) -> tuple[list[ChatMessage], list[ChatMessage]]:
+        """
+        Loads the chat history for the chat session synchronously.
+
+        Args:
+            token_limit (int): The maximum number of tokens to load.
+            percent_of_max_tokens (float): The percentage of the maximum tokens to load.
+            max_messages (int): The maximum number of messages to load.
+
+        Returns:
+            tuple[list[ChatMessage], list[ChatMessage]]: The selected and full chat history.
+
+        Raises:
+            Exception: If the loading fails.
+        """
+        return self._get_full_and_selected_history(
+            token_limit=token_limit,
+            percent_of_max_tokens=percent_of_max_tokens,
+            max_messages=max_messages,
+        )
+
+    @to_async
+    @async_warning
+    def async_get_full_and_selected_history(
+        self,
+        token_limit: int,
+        percent_of_max_tokens: float,
+        max_messages: int,
+    ) -> tuple[list[ChatMessage], list[ChatMessage]]:
+        """
+        Loads the chat history for the chat session asynchronously.
+
+        Args:
+            token_limit (int): The maximum number of tokens to load.
+            percent_of_max_tokens (float): The percentage of the maximum tokens to load.
+            max_messages (int): The maximum number of messages to load.
+
+        Returns:
+            tuple[list[ChatMessage], list[ChatMessage]]: The selected and full chat history.
+
+        Raises:
+            Exception: If the loading fails.
+        """
+        return self._get_full_and_selected_history(
+            token_limit=token_limit,
+            percent_of_max_tokens=percent_of_max_tokens,
+            max_messages=max_messages,
+        )
+
+    def create_assistant_message(
+        self,
+        content: str,
+        references: list[ContentReference] = [],
+        debug_info: dict = {},
+    ):
+        """
+        Creates a message in the chat session synchronously.
+
+        Args:
+            content (str): The content for the message.
+            references (list[ContentReference]): list of ContentReference objects. Defaults to None.
+            debug_info (dict[str, Any]]): Debug information. Defaults to None.
+
+        Returns:
+            ChatMessage: The created message.
+
+        Raises:
+            Exception: If the creation fails.
+        """
+        return self._trigger_create_assistant_message(
+            content=content,
+            references=references,
+            debug_info=debug_info,
+        )
+
+    @to_async
+    @async_warning
+    def async_create_assistant_message(
+        self,
+        content: str,
+        references: list[ContentReference] = [],
+        debug_info: dict = {},
+    ):
+        """
+        Creates a message in the chat session asynchronously.
+
+        Args:
+            content (str): The content for the message.
+            references (list[ContentReference]): list of references. Defaults to None.
+            debug_info (dict[str, Any]]): Debug information. Defaults to None.
+
+        Returns:
+            ChatMessage: The created message.
+
+        Raises:
+            Exception: If the creation fails.
+        """
+
+        return self._trigger_create_assistant_message(
+            content=content,
+            references=references,
+            debug_info=debug_info,
+        )
+
+    def _trigger_modify_assistant_message(
+        self,
+        content: str,
+        message_id: Optional[str],
+        references: list[ContentReference],
+        debug_info: dict,
+    ) -> ChatMessage:
+        message_id = message_id or self.state.assistant_message_id
+
+        try:
+            message = unique_sdk.Message.modify(
+                user_id=self.state.user_id,
+                company_id=self.state.company_id,
+                id=message_id,  # type: ignore
+                chatId=self.state.chat_id,
+                text=content,
+                references=self._map_references(references),  # type: ignore
+                debugInfo=debug_info or {},
+            )
+        except Exception as e:
+            self.logger.error(f"Failed to modify assistant message: {e}")
+            raise e
+        return ChatMessage(**message)
+
+    def _trigger_create_assistant_message(
+        self,
+        content: str,
+        references: list[ContentReference],
+        debug_info: dict,
+    ) -> ChatMessage:
+        try:
+            message = unique_sdk.Message.create(
+                user_id=self.state.user_id,
+                company_id=self.state.company_id,
+                chatId=self.state.chat_id,
+                assistantId=self.state.assistant_id,
+                text=content,
+                role=ChatMessageRole.ASSISTANT.name,
+                references=self._map_references(references),  # type: ignore
+                debugInfo=debug_info,
+            )
+        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_and_selected_history(
+        self,
+        token_limit,
+        percent_of_max_tokens=0.15,
+        max_messages=4,
+    ):
+        full_history = self._get_full_history()
+        selected_history = self._get_selection_from_history(
+            full_history,
+            int(round(token_limit * percent_of_max_tokens)),
+            max_messages,
+        )
+
+        return full_history, selected_history
+
+    def _get_full_history(self):
+        SYSTEM_MESSAGE_PREFIX = "[SYSTEM] "
+
+        messages = self._trigger_list_messages(self.state.chat_id)
+
+        # 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 self._map_to_chat_messages(filtered_messages)
+
+    def _trigger_list_messages(self, chat_id: str):
+        try:
+            messages = unique_sdk.Message.list(
+                user_id=self.state.user_id,
+                company_id=self.state.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:]

unique_toolkit/chat/state.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from typing import Self
+
+from unique_toolkit.app.schemas import Event
+
+
+@dataclass
+class ChatState:
+    """
+    Represents the state of the chat session.
+
+    Attributes:
+        company_id (str): The company ID.
+        user_id (str): The user ID.
+        chat_id (str): The chat ID.
+        scope_ids (list[str] | None): The scope IDs.
+        chat_only (bool): The chat only flag.
+        user_message_text (str): The user message text.
+        user_message_id (str): The user message ID.
+        assistant_message_id (str): The assistant message ID.
+    """
+
+    company_id: str
+    user_id: str
+    assistant_id: str
+    chat_id: str
+    scope_ids: list[str] | None = None
+    chat_only: bool = False
+    user_message_text: str | None = None
+    user_message_id: str | None = None
+    assistant_message_id: str | None = None
+    module_name: str | None = None
+
+    @classmethod
+    def from_event(cls, event: Event) -> Self:
+        """
+        Creates a ChatState instance from the Event.
+
+        Args:
+            event (Event): The Event object.
+
+        Returns:
+            ChatManager: The ChatManager instance.
+        """
+        config = event.payload.configuration
+
+        scope_ids = config.get("scopeIds") or None
+        chat_only = config.get("scopeToChatOnUpload", False)
+        return cls(
+            user_id=event.user_id,
+            chat_id=event.payload.chat_id,
+            company_id=event.company_id,
+            assistant_id=event.payload.assistant_id,
+            scope_ids=scope_ids,
+            chat_only=chat_only,
+            user_message_text=event.payload.user_message.text,
+            user_message_id=event.payload.user_message.id,
+            assistant_message_id=event.payload.assistant_message.id,
+            module_name=event.payload.name,
+        )

unique_toolkit/chat/utils.py

@@ -0,0 +1,25 @@
+from unique_toolkit.chat.schemas import ChatMessage
+from unique_toolkit.content.utils import count_tokens
+
+
+def convert_chat_history_to_injectable_string(
+    history: list[ChatMessage],
+) -> tuple[list[str], int]:
+    """
+    Converts chat history to a string that can be injected into the model.
+
+    Args:
+        history (list[ChatMessage]): The chat history.
+
+    Returns:
+        tuple[list[str], int]: The chat history and the token length of the chat context.
+    """
+    chatHistory = []
+    for msg in history:
+        if msg.role.value == "assistant":
+            chatHistory.append(f"previous_answer: {msg.content}")
+        else:
+            chatHistory.append(f"previous_question: {msg.content}")
+    chatContext = "\n".join(chatHistory)
+    chatContextTokenLength = count_tokens(chatContext)
+    return chatHistory, chatContextTokenLength

unique_toolkit/content/schemas.py

@@ -0,0 +1,90 @@
+from datetime import datetime
+from enum import StrEnum
+from typing import Optional
+
+from humps import camelize
+from pydantic import BaseModel, ConfigDict
+
+# set config to convert camelCase to snake_case
+model_config = ConfigDict(
+    alias_generator=camelize,
+    populate_by_name=True,
+    arbitrary_types_allowed=True,
+)
+
+
+class ContentMetadata(BaseModel):
+    model_config = model_config
+    key: str
+    mime_type: str
+
+
+class ContentChunk(BaseModel):
+    model_config = model_config
+    id: str
+    text: str
+    order: int
+    key: str | None = None
+    chunk_id: str | None = None
+    url: str | None = None
+    title: str | None = None
+    start_page: int | None = None
+    end_page: int | None = None
+
+    object: str | None = None
+    metadata: ContentMetadata | None = None
+    internally_stored_at: datetime | None = None
+    created_at: datetime | None = None
+    updated_at: datetime | None = None
+
+
+class Content(BaseModel):
+    model_config = model_config
+    id: str
+    key: str
+    title: str | None = None
+    url: str | None = None
+    chunks: list[ContentChunk] = []
+    write_url: str | None = None
+    read_url: str | None = None
+
+
+class ContentReference(BaseModel):
+    model_config = model_config
+    id: str
+    message_id: str
+    name: str
+    sequence_number: int
+    source: str
+    source_id: str
+    url: str
+
+
+class ContentSearchType(StrEnum):
+    COMBINED = "COMBINED"
+    VECTOR = "VECTOR"
+
+
+class ContentSearchResult(BaseModel):
+    """Schema corresponding to unique_sdk.SearchResult"""
+
+    id: str
+    text: str
+    order: int
+    chunkId: str | None = None
+    key: str | None = None
+    title: str | None = None
+    url: str | None = None
+    startPage: int | None = None
+    endPage: int | None = None
+    object: str | None = None
+
+
+class ContentUploadInput(BaseModel):
+    key: str
+    title: str
+    mime_type: str
+
+    owner_type: Optional[str] = None
+    owner_id: Optional[str] = None
+    byte_size: Optional[int] = None