unique_toolkit 0.7.7__py3-none-any.whl → 1.23.0__py3-none-any.whl

unique_toolkit/agentic/history_manager/history_construction_with_contents.py

@@ -0,0 +1,297 @@
+import base64
+import mimetypes
+from datetime import datetime
+from enum import StrEnum
+
+import numpy as np
+import tiktoken
+from pydantic import RootModel
+
+from unique_toolkit._common.token.token_counting import (
+    num_tokens_per_language_model_message,
+)
+from unique_toolkit.app import ChatEventUserMessage
+from unique_toolkit.chat.schemas import ChatMessage
+from unique_toolkit.chat.schemas import ChatMessageRole as ChatRole
+from unique_toolkit.chat.service import ChatService
+from unique_toolkit.content.schemas import Content
+from unique_toolkit.content.service import ContentService
+from unique_toolkit.language_model import LanguageModelMessageRole as LLMRole
+from unique_toolkit.language_model.infos import EncoderName
+from unique_toolkit.language_model.schemas import LanguageModelMessages
+
+# TODO: Test this once it moves into the unique toolkit
+
+map_chat_llm_message_role = {
+    ChatRole.USER: LLMRole.USER,
+    ChatRole.ASSISTANT: LLMRole.ASSISTANT,
+}
+
+
+class ImageMimeType(StrEnum):
+    JPEG = "image/jpeg"
+    PNG = "image/png"
+    GIF = "image/gif"
+    BMP = "image/bmp"
+    WEBP = "image/webp"
+    TIFF = "image/tiff"
+    SVG = "image/svg+xml"
+
+
+class FileMimeType(StrEnum):
+    PDF = "application/pdf"
+    DOCX = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
+    DOC = "application/msword"
+    XLSX = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
+    XLS = "application/vnd.ms-excel"
+    PPTX = "application/vnd.openxmlformats-officedocument.presentationml.presentation"
+    CSV = "text/csv"
+    HTML = "text/html"
+    MD = "text/markdown"
+    TXT = "text/plain"
+
+
+class ChatMessageWithContents(ChatMessage):
+    contents: list[Content] = []
+
+
+class ChatHistoryWithContent(RootModel):
+    root: list[ChatMessageWithContents]
+
+    @classmethod
+    def from_chat_history_and_contents(
+        cls,
+        chat_history: list[ChatMessage],
+        chat_contents: list[Content],
+    ):
+        combined = chat_contents + chat_history
+        combined.sort(key=lambda x: x.created_at or datetime.min)
+
+        grouped_elements = []
+        content_container = []
+
+        # Content is collected and added to the next chat message
+        for c in combined:
+            if isinstance(c, ChatMessage):
+                grouped_elements.append(
+                    ChatMessageWithContents(
+                        contents=content_container.copy(),
+                        **c.model_dump(),
+                    ),
+                )
+                content_container.clear()
+            else:
+                content_container.append(c)
+
+        return cls(root=grouped_elements)
+
+    def __iter__(self):
+        return iter(self.root)
+
+    def __getitem__(self, item):
+        return self.root[item]
+
+
+def is_image_content(filename: str) -> bool:
+    mimetype, _ = mimetypes.guess_type(filename)
+
+    if not mimetype:
+        return False
+
+    return mimetype in ImageMimeType.__members__.values()
+
+
+def is_file_content(filename: str) -> bool:
+    mimetype, _ = mimetypes.guess_type(filename)
+
+    if not mimetype:
+        return False
+
+    return mimetype in FileMimeType.__members__.values()
+
+
+def get_chat_history_with_contents(
+    user_message: ChatEventUserMessage,
+    chat_id: str,
+    chat_history: list[ChatMessage],
+    content_service: ContentService,
+) -> ChatHistoryWithContent:
+    last_user_message = ChatMessage(
+        id=user_message.id,
+        chat_id=chat_id,
+        text=user_message.text,
+        originalText=user_message.original_text,
+        role=ChatRole.USER,
+        gpt_request=None,
+        created_at=datetime.fromisoformat(user_message.created_at),
+    )
+    if len(chat_history) > 0 and last_user_message.id == chat_history[-1].id:
+        pass
+    else:
+        chat_history.append(last_user_message)
+
+    chat_contents = content_service.search_contents(
+        where={
+            "ownerId": {
+                "equals": chat_id,
+            },
+        },
+    )
+
+    return ChatHistoryWithContent.from_chat_history_and_contents(
+        chat_history,
+        chat_contents,
+    )
+
+
+def download_encoded_images(
+    contents: list[Content],
+    content_service: ContentService,
+    chat_id: str,
+) -> list[str]:
+    base64_encoded_images = []
+    for im in contents:
+        if is_image_content(im.key):
+            try:
+                file_bytes = content_service.download_content_to_bytes(
+                    content_id=im.id,
+                    chat_id=chat_id,
+                )
+
+                mime_type, _ = mimetypes.guess_type(im.key)
+                encoded_string = base64.b64encode(file_bytes).decode("utf-8")
+                image_string = f"data:{mime_type};base64," + encoded_string
+                base64_encoded_images.append(image_string)
+            except Exception as e:
+                print(e)
+    return base64_encoded_images
+
+
+class FileContentSerialization(StrEnum):
+    NONE = "none"
+    FILE_NAME = "file_name"
+
+
+class ImageContentInclusion(StrEnum):
+    NONE = "none"
+    ALL = "all"
+
+
+def file_content_serialization(
+    file_contents: list[Content],
+    file_content_serialization: FileContentSerialization,
+) -> str:
+    match file_content_serialization:
+        case FileContentSerialization.NONE:
+            return ""
+        case FileContentSerialization.FILE_NAME:
+            file_names = [
+                f"- Uploaded file: {f.key} at {f.created_at}" for f in file_contents
+            ]
+            return "\n".join(
+                [
+                    "Files Uploaded to Chat can be accessed by internal search tool if available:\n",
+                ]
+                + file_names,
+            )
+
+
+def get_full_history_with_contents(
+    user_message: ChatEventUserMessage,
+    chat_id: str,
+    chat_service: ChatService,
+    content_service: ContentService,
+    include_images: ImageContentInclusion = ImageContentInclusion.ALL,
+    file_content_serialization_type: FileContentSerialization = FileContentSerialization.FILE_NAME,
+) -> LanguageModelMessages:
+    grouped_elements = get_chat_history_with_contents(
+        user_message=user_message,
+        chat_id=chat_id,
+        chat_history=chat_service.get_full_history(),
+        content_service=content_service,
+    )
+
+    builder = LanguageModelMessages([]).builder()
+    for c in grouped_elements:
+        # LanguageModelUserMessage has not field original content
+        text = c.original_content if c.original_content else c.content
+        if text is None:
+            if c.role == ChatRole.USER:
+                raise ValueError(
+                    "Content or original_content of LanguageModelMessages should exist.",
+                )
+            text = ""
+
+        if len(c.contents) > 0:
+            file_contents = [co for co in c.contents if is_file_content(co.key)]
+            image_contents = [co for co in c.contents if is_image_content(co.key)]
+
+            content = (
+                text
+                + "\n\n"
+                + file_content_serialization(
+                    file_contents,
+                    file_content_serialization_type,
+                )
+            )
+            content = content.strip()
+
+            if include_images and len(image_contents) > 0:
+                builder.image_message_append(
+                    content=content,
+                    images=download_encoded_images(
+                        contents=image_contents,
+                        content_service=content_service,
+                        chat_id=chat_id,
+                    ),
+                    role=map_chat_llm_message_role[c.role],
+                )
+            else:
+                builder.message_append(
+                    role=map_chat_llm_message_role[c.role],
+                    content=content,
+                )
+        else:
+            builder.message_append(
+                role=map_chat_llm_message_role[c.role],
+                content=text,
+            )
+    return builder.build()
+
+
+def get_full_history_as_llm_messages(
+    chat_service: ChatService,
+) -> LanguageModelMessages:
+    chat_history = chat_service.get_full_history()
+
+    map_chat_llm_message_role = {
+        ChatRole.USER: LLMRole.USER,
+        ChatRole.ASSISTANT: LLMRole.ASSISTANT,
+    }
+
+    builder = LanguageModelMessages([]).builder()
+    for c in chat_history:
+        builder.message_append(
+            role=map_chat_llm_message_role[c.role],
+            content=c.content or "",
+        )
+    return builder.build()
+
+
+def limit_to_token_window(
+    messages: LanguageModelMessages,
+    token_limit: int,
+    encoding_name: EncoderName = EncoderName.O200K_BASE,
+) -> LanguageModelMessages:
+    encoder = tiktoken.get_encoding(encoding_name)
+    token_per_message_reversed = num_tokens_per_language_model_message(
+        messages,
+        encode=encoder.encode,
+    )
+
+    to_take: list[bool] = (np.cumsum(token_per_message_reversed) < token_limit).tolist()
+    to_take.reverse()
+
+    return LanguageModelMessages(
+        root=[m for m, tt in zip(messages, to_take, strict=False) if tt],
+    )

unique_toolkit/agentic/history_manager/history_manager.py

@@ -0,0 +1,242 @@
+from logging import Logger
+from typing import Annotated, Awaitable, Callable
+
+from pydantic import BaseModel, Field
+
+from unique_toolkit._common.feature_flags.schema import (
+    FeatureExtendedSourceSerialization,
+)
+from unique_toolkit._common.validators import LMI
+from unique_toolkit.agentic.history_manager.loop_token_reducer import LoopTokenReducer
+from unique_toolkit.agentic.history_manager.utils import transform_chunks_to_string
+from unique_toolkit.agentic.reference_manager.reference_manager import ReferenceManager
+from unique_toolkit.agentic.tools.config import get_configuration_dict
+from unique_toolkit.agentic.tools.schemas import ToolCallResponse
+from unique_toolkit.app.schemas import ChatEvent
+from unique_toolkit.language_model.default_language_model import DEFAULT_GPT_4o
+from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.schemas import (
+    LanguageModelAssistantMessage,
+    LanguageModelFunction,
+    LanguageModelMessage,
+    LanguageModelMessages,
+    LanguageModelToolMessage,
+)
+
+DeactivatedNone = Annotated[
+    None,
+    Field(title="Deactivated", description="None"),
+]
+
+
+class UploadedContentConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    user_context_window_limit_warning: str = Field(
+        default="The uploaded content is too large to fit into the ai model. "
+        "Unique AI will search for relevant sections in the material and if needed combine the data with knowledge base content",
+        description="Message to show when using the Internal Search instead of upload and chat tool due to context window limit. Jinja template.",
+    )
+    percent_for_uploaded_content: float = Field(
+        default=0.6,
+        ge=0.0,
+        le=1.0,
+        description="The fraction of the max input tokens that will be reserved for the uploaded content.",
+    )
+
+
+class ExperimentalFeatures(FeatureExtendedSourceSerialization): ...
+
+
+class HistoryManagerConfig(BaseModel):
+    experimental_features: ExperimentalFeatures = Field(
+        default=ExperimentalFeatures(),
+        description="Experimental features for the history manager.",
+    )
+
+    percent_of_max_tokens_for_history: float = Field(
+        default=0.2,
+        ge=0.0,
+        lt=1.0,
+        description="The fraction of the max input tokens that will be reserved for the history.",
+    )
+
+    language_model: LMI = LanguageModelInfo.from_name(DEFAULT_GPT_4o)
+
+    @property
+    def max_history_tokens(self) -> int:
+        return int(
+            self.language_model.token_limits.token_limit_input
+            * self.percent_of_max_tokens_for_history,
+        )
+
+    uploaded_content_config: (
+        Annotated[
+            UploadedContentConfig,
+            Field(title="Active"),
+        ]
+        | DeactivatedNone
+    ) = UploadedContentConfig()
+
+
+class HistoryManager:
+    """
+    Manages the history of tool calls and conversation loops.
+
+    This class is responsible for:
+    - Storing and maintaining the history of tool call results and conversation messages.
+    - Merging uploaded content with the conversation history for a unified view.
+    - Limiting the history to fit within a configurable token window for efficient processing.
+    - Providing methods to retrieve, manipulate, and append to the conversation history.
+    - Handling post-processing steps to clean or modify the history as needed.
+
+    Key Features:
+    - Tool Call History: Tracks the results of tool calls and appends them to the conversation history.
+    - Loop History: Maintains a record of conversation loops, including assistant and user messages.
+    - History Merging: Combines uploaded files and chat messages into a cohesive history.
+    - Token Window Management: Ensures the history stays within a specified token limit for optimal performance.
+    - Post-Processing Support: Allows for custom transformations or cleanup of the conversation history.
+
+    The HistoryManager serves as the backbone for managing and retrieving conversation history in a structured and efficient manner.
+    """
+
+    def __init__(
+        self,
+        logger: Logger,
+        event: ChatEvent,
+        config: HistoryManagerConfig,
+        language_model: LMI,
+        reference_manager: ReferenceManager,
+    ):
+        self._config = config
+        self._logger = logger
+        self._language_model = language_model
+        self._token_reducer = LoopTokenReducer(
+            logger=self._logger,
+            event=event,
+            max_history_tokens=self._config.max_history_tokens,
+            has_uploaded_content_config=bool(self._config.uploaded_content_config),
+            language_model=self._language_model,
+            reference_manager=reference_manager,
+        )
+        self._tool_call_result_history: list[ToolCallResponse] = []
+        self._tool_calls: list[LanguageModelFunction] = []
+        self._loop_history: list[LanguageModelMessage] = []
+        self._source_enumerator = 0
+
+    def add_tool_call(self, tool_call: LanguageModelFunction) -> None:
+        self._tool_calls.append(tool_call)
+
+    def get_tool_calls(self) -> list[LanguageModelFunction]:
+        return self._tool_calls
+
+    def has_no_loop_messages(self) -> bool:
+        return len(self._loop_history) == 0
+
+    def add_tool_call_results(self, tool_call_results: list[ToolCallResponse]):
+        for tool_response in tool_call_results:
+            if not tool_response.successful:
+                self._loop_history.append(
+                    LanguageModelToolMessage(
+                        name=tool_response.name,
+                        tool_call_id=tool_response.id,
+                        content=f"Tool call {tool_response.name} failed with error: {tool_response.error_message}",
+                    )
+                )
+                continue
+            self._append_tool_call_result_to_history(tool_response)
+
+    def _append_tool_call_result_to_history(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> None:
+        tool_call_result_for_history = self._get_tool_call_result_for_loop_history(
+            tool_response=tool_response
+        )
+        self._loop_history.append(tool_call_result_for_history)
+
+    def _get_tool_call_result_for_loop_history(
+        self,
+        tool_response: ToolCallResponse,
+    ) -> LanguageModelMessage:
+        self._logger.debug(
+            f"Appending tool call result to history: {tool_response.name}"
+        )
+
+        if tool_response.content != "":
+            return LanguageModelToolMessage(
+                content=tool_response.content,
+                tool_call_id=tool_response.id,  # type: ignore
+                name=tool_response.name,
+            )
+
+        content_chunks = (
+            tool_response.content_chunks or []
+        )  # it can be that the tool response does not have content chunks
+
+        # Transform content chunks into sources to be appended to tool result
+        stringified_sources, sources = transform_chunks_to_string(
+            content_chunks,
+            self._source_enumerator,
+        )
+
+        self._source_enumerator += len(
+            sources
+        )  # To make sure all sources have unique source numbers
+
+        # Append the result to the history
+        return LanguageModelToolMessage(
+            content=stringified_sources,
+            tool_call_id=tool_response.id,  # type: ignore
+            name=tool_response.name,
+        )
+
+    def _append_tool_calls_to_history(
+        self, tool_calls: list[LanguageModelFunction]
+    ) -> None:
+        self._loop_history.append(
+            LanguageModelAssistantMessage.from_functions(tool_calls=tool_calls)
+        )
+
+    def add_assistant_message(self, message: LanguageModelAssistantMessage) -> None:
+        self._loop_history.append(message)
+
+    async def get_history_for_model_call(
+        self,
+        original_user_message: str,
+        rendered_user_message_string: str,
+        rendered_system_message_string: str,
+        remove_from_text: Callable[[str], Awaitable[str]],
+    ) -> LanguageModelMessages:
+        self._logger.info("Getting history for model call -> ")
+
+        messages = await self._token_reducer.get_history_for_model_call(
+            original_user_message=original_user_message,
+            rendered_user_message_string=rendered_user_message_string,
+            rendered_system_message_string=rendered_system_message_string,
+            loop_history=self._loop_history,
+            remove_from_text=remove_from_text,
+        )
+        return messages
+
+    async def get_user_visible_chat_history(
+        self,
+        assistant_message_text: str | None = None,
+        remove_from_text: Callable[[str], Awaitable[str]] | None = None,
+    ) -> LanguageModelMessages:
+        """Get the user visible chat history.
+
+        Args:
+            assistant_message_text (str | None): The latest assistant message to append to the history, as this is not extracted from the history.
+            If None, the history will be returned without the latest assistant message.
+            remove_from_text (Callable[[str], Awaitable[str]] | None): A function to remove text from the history before returning it.
+
+        Returns:
+            LanguageModelMessages: The user visible chat history.
+        """
+        history = await self._token_reducer.get_history_from_db(remove_from_text)
+        if assistant_message_text:
+            history.append(
+                LanguageModelAssistantMessage(content=assistant_message_text)
+            )
+        return LanguageModelMessages(history)