unique_toolkit 0.0.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

unique_toolkit/chat/service.py

@@ -1,9 +1,14 @@
-import unique_sdk
-from unique_sdk.utils.chat_history import load_history
+import logging
+import re
+from typing import Optional
 
-from unique_toolkit.performance.async_wrapper import to_async
+import unique_sdk
 
-from .state import ChatState
+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:
@@ -12,171 +17,364 @@ class ChatService:
 
     Attributes:
         state (ChatState): The chat state.
+        logger (Optional[logging.Logger]): The logger. Defaults to None.
     """
 
-    def __init__(self, state: ChatState):
+    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,
-        *,
-        text: str,
-        references: list = [],
+        content: str,
+        references: list[ContentReference] = [],
         debug_info: dict = {},
-    ) -> None:
+        message_id: Optional[str] = None,
+    ) -> ChatMessage:
         """
         Modifies a message in the chat session synchronously.
 
         Args:
-            text (str): The new text for the message.
-            references (Optional[list[dict[str, Any]]], optional): list of references. Defaults to None.
-            debug_info (Optional[dict[str, Any]]], optional): Debug information. Defaults to None.
+            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.
         """
-        self._trigger_modify_assistant_message(text, references, debug_info)
+        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,
-        *,
-        text: str,
-        references: list = [],
+        content: str,
+        references: list[ContentReference] = [],
         debug_info: dict = {},
-    ) -> None:
+        message_id: Optional[str] = None,
+    ) -> ChatMessage:
         """
         Modifies a message in the chat session asynchronously.
 
         Args:
-            text (str): The new text for the message.
-            references (Optional[list[dict[str, Any]]], optional): list of references. Defaults to None.
+            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.
         """
-        self._trigger_modify_assistant_message(text, references, debug_info)
+        return self._get_full_history()
 
-    def get_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,
-        *,
-        max_tokens: int,
+        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:
-            max_tokens (int): The maximum number of tokens to load.
+            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:
-            list[dict[str, Any]]: The chat history.
+            tuple[list[ChatMessage], list[ChatMessage]]: The selected and full chat history.
+
+        Raises:
+            Exception: If the loading fails.
         """
-        return self._trigger_load_history(
-            max_tokens, percent_of_max_tokens, max_messages
+        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
-    def async_get_history(
+    @async_warning
+    def async_get_full_and_selected_history(
         self,
-        *,
-        max_tokens: int,
+        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:
-            max_tokens (int): The maximum number of tokens to load.
+            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:
-            list[dict[str, Any]]: The chat history.
+            tuple[list[ChatMessage], list[ChatMessage]]: The selected and full chat history.
+
+        Raises:
+            Exception: If the loading fails.
         """
-        return self._trigger_load_history(
-            max_tokens, percent_of_max_tokens, max_messages
+        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,
-        *,
-        text: str,
-        references: list = [],
+        content: str,
+        references: list[ContentReference] = [],
         debug_info: dict = {},
     ):
         """
         Creates a message in the chat session synchronously.
 
         Args:
-            text (str): The text for the message.
-            references (Optional[list[dict[str, Any]]], optional): list of references. Defaults to None.
-            debug_info (Optional[dict[str, Any]]], optional): Debug information. Defaults to None.
+            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(text, references, debug_info)
+        return self._trigger_create_assistant_message(
+            content=content,
+            references=references,
+            debug_info=debug_info,
+        )
 
     @to_async
+    @async_warning
     def async_create_assistant_message(
         self,
-        *,
-        text: str,
-        references: list = [],
+        content: str,
+        references: list[ContentReference] = [],
         debug_info: dict = {},
     ):
         """
         Creates a message in the chat session asynchronously.
 
         Args:
-            text (str): The text for the message.
-            references (Optional[list[dict[str, Any]]], optional): list of references. Defaults to None.
-            debug_info (Optional[dict[str, Any]]], optional): Debug information. Defaults to None.
+            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(text, references, debug_info)
+
+        return self._trigger_create_assistant_message(
+            content=content,
+            references=references,
+            debug_info=debug_info,
+        )
 
     def _trigger_modify_assistant_message(
         self,
-        text: str,
-        references: list = [],
-        debug_info: dict = {},
-    ) -> None:
-        unique_sdk.Message.modify(
-            user_id=self.state.user_id,
-            company_id=self.state.company_id,
-            id=self.state.assistant_message_id,
-            chatId=self.state.chat_id,
-            text=text,
-            references=references or [],
-            debugInfo=debug_info or {},
-        )
+        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_load_history(
+    def _trigger_create_assistant_message(
         self,
-        max_tokens: int,
-        percent_of_max_tokens: float,
-        max_messages: int,
+        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,
     ):
-        return load_history(
-            self.state.user_id,
-            self.state.company_id,
-            self.state.chat_id,
-            max_tokens,
-            percent_of_max_tokens,
+        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,
         )
 
-    # TODO throws error at the moment
-    def _trigger_create_assistant_message(
+        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,
-        text: str,
-        references: list = [],
-        debug_info: dict = {},
+        full_history: list[ChatMessage],
+        max_tokens: int,
+        max_messages=4,
     ):
-        return 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=text,
-            role="ASSISTANT",
-            references=references,
-            debugInfo=debug_info,
+        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

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 from typing import Self
 
-from unique_toolkit.event.schema import Event
+from unique_toolkit.app.schemas import Event
 
 
 @dataclass
@@ -13,6 +13,8 @@ class ChatState:
         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.
@@ -20,29 +22,39 @@ class ChatState:
 
     company_id: str
     user_id: str
-    assistant_id: str = ""
-    chat_id: str = ""
-    user_message_text: str = ""
-    user_message_id: str = ""
-    assistant_message_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 an event dictionary.
+        Creates a ChatState instance from the Event.
 
         Args:
-            event (dict): The event dictionary.
+            event (Event): The Event object.
 
         Returns:
-            ChatState: The ChatState instance.
+            ChatManager: The ChatManager instance.
         """
-        # Chat only flag
+        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